Probeable DARTS with Application to Computational Pathology

The goal of DARTS [23] is to search for two types of cells (namely normal and reduction) as building blocks, which are stacked to form a full network. Each cell is represented as a directed acyclic graph with $N$ nodes, including two input nodes, intermediate nodes and one output node. Every node $x_{i}$ is a latent representation (e.g., feature map in CNN) and every edge $(i,j)$ is a mixture of weighted candidate operations in a pre-defined operation search space $\mathcal{O}$ (e.g., convolution, skip-connection). The output $\bar{o}_{i,j}$ of an edge $(i,j)$ is then a weighted sum of candidate operations [23]:

where $\alpha^{o}_{i,j}$ is an architecture parameter for weighting operation $o\left(x_{i}\right)$. The output of an intermediate node $x_{j}$ is the sum of all input edges, i.e., $x_{j}=\sum_{i<j}\bar{o}_{i,j}\left(x_{j}\right)$. The output node of a cell is the concatenation of all intermediate nodes. Normal cells keep the input resolution while reduction cells decrease resolutions with stride 2 in all candidate operations.

In the searching procedure, the network weight $w$ and architecture parameter $\alpha$ are jointly learned via bi-level optimization [23]:

where $\mathcal{L}_{val}$ and $\mathcal{L}_{train}$ denote the validation and training datasets, respectively. Using gradient descent, $w$ and $\alpha$ can be updated alternatively during each training iteration.

When the searching is finished, the discrete cell architecture is obtained by replacing each edge by the operation with the largest architecture weight, then selecting the two strongest input edges for each intermediate node. Fig. 2 (b) illustrates the evolution of cell structure during searching. The discrete network is then retrained from scratch for final evaluation. The whole process is shown in Fig. 2 (a).

To monitor the searching process of DARTS, we adopt the explainability metric stable rank to probe the intermediate convolutional layers in different cells and quantify their learning quality as explained in [12]. Given a convolutional weight matrix, we first decompose it by low-rank factorization. This factors out the perturbation noise in the layer while keeping the most useful information in the low-rank component. The stable rank $\mathcal{S}$ is the normalized sum of the singular values of the low-rank matrix. It measures the norm energy of the convolutional weights and encodes the low-rank structure’s space span of the output mapping. A higher value indicates better propagation of information through a convolutional layer [12].

Using this probing metric, we monitor the searching phase of the existing DARTS framework applied on the CIFAR100 dataset [19]. The left column of Fig. 3 shows an example of the stable rank evolution (top row) of the layers in one cell as well as the architecture weights evolution (bottom 2 rows) in two edges. We can see from the stable rank evolution that the convolutional layers are not learning well in the original DARTS with stochastic gradient descent (DARTS+SGD) and default initial learning rate 0.025 [23] – most layers generate zero stable rank through all training epochs. Recall the edge structure introduced in Sec. 3.1, each candidate operation is multiplied by a weight, which is between $0$ and $1$. This makes their gradients small during backpropagation. Therefore the convolutional layers are learning slowly. The bottom two images show that skip-connections (green curves) are preferred. The same phenomenon is reported in [21] that when searched on CV datasets, the original DARTS tends to select too many skip-connections, which is a kind of overfitting, resulting in a shallow network with poor representation ability.

In light of this, we increase the initial learning rate from $0.025$ to $0.05$ and $0.175$. The stable rank evolution of layers in the same cell, as well as the architecture weights in the same edges, are shown in Fig. 3 (b) and (c). With the increase in initial learning rates, more layers generate a higher stable rank, which means they’re learning better. In the meantime, the architecture weights evolution reveals that the preference for skip-connections is suppressed. With the help of the probing metric, we know that the existing DARTS framework lacks proper tuning in hyperparameters and how the searching can be improved.

We have shown in the previous section that the probing metric can be used to tune proper initial learning rates for DARTS. Then why not utilize an optimizer that incorporates such metric for learning rates adjustment during searching? To verify this, we adopt the Adas optimizer [12], which adaptively adjusts the learning rate for each layer based on their stable rank evolution. At the end of each searching epoch, it first computes the difference in stable rank over consecutive epochs for each convolutional layer, and then adds (with a weight) result to the learning rate momentum. A hyperparameter scheduler beta is used for weighting this term. This process is illustrated in Fig. 2 (c). The Adas optimizer is aware of the learning quality of each layer and therefore tunes their learning rates accordingly.

Fig. 4 shows the training and validation errors when searching with different optimizers and initial learning rates. We can see that SGD with a 0.175 initial learning rate leads to overfitting during searching. The final gap between training and validation error is around $40\%$. While for Adas, the overfitting problem is reduced. The resulting final gap is around $20\%$. This indicates that Adas improves the searching process of DARTS with better generalization ability.

We investigate the trade-off between network size and test performance by searching for the optimal architectures with different numbers of cells and intermediate nodes. This trade-off is crucial for high-throughput CPath applications in real-time. On the ADP dataset, DARTS+Adas obtains the best-performing architecture. It consists of four cells, and each cell contains three intermediate nodes. Fig. 6 (c) and (d) show the snapshots of the cell structures.

On the CIFAR dataset, we search for the optimum optimizer and number of intermediate nodes. During searching, eight cells are stacked as in DARTS, while the number of nodes in each cell is tuned between 4, 5, and 6. During evaluation, to prevent the network size from being too large with more nodes, we also change the number of cells accordingly, i.e., 20 cells for 4 nodes, 17 cells for 5 nodes, and 14 cells for 6 nodes. Table. 3 and Table. 4 show the test performance of architectures searched on CIFAR-10 and CIFAR-100. We can see that in each setting, DARTS+Adas outperforms the default DARTS+SGD in terms of test accuracy, while a cost is paid in parameter size. This is because the original DARTS+SGD tends to select skip-connections in the final architecture as described above in Sec. 3.2. The optimum number of nodes is 4.

To find the optimum architecture on the ADP dataset, we run the architecture search in all different parameter settings, i.e., different numbers of cells and nodes, different choices of optimizer. Note that in CIFAR experiments we search for a shallower network (with few cells) during searching but train a deeper one (with more cells) for evaluation due to the complexity of CV datasets. In ADP, however, we keep the number of cells the same in two stages. This is because ADP is a simpler dataset so we don’t need to increase the model complexity during evaluation. This also brings more consistency to the search-evaluation pipeline.

Optimum optimizer and number of cells. We first search for the optimum optimizer and the number of cells. The test results of the searched architectures after final evaluation are shown in Table. 5. We also plot the accuracy versus parameter size in Fig. 5 (a), where each dot represents a different choice of cell number. We can see that as the number of cells increases the accuracy of DARTS+SGD drops while its parameter size increases. When using DARTS+Adas, the test accuracy remains the highest across different numbers of cells, and the parameter size remains the smallest. The highest accuracy is achieved with 6 cells, while for 4 cells, the searched architecture has the smallest size but still obtains the second highest accuracy.

Optimum number of intermediate nodes. We then fix the number of cells as four and search for the optimum number of intermediate nodes. The test performance are shown in Tabel.6 and in Fig. 5 (b). We can see that DARTS+Adas achieves higher accuracy than DARTS+SGD with fewer nodes, hence less computation complexity. The highest accuracy $94.46\%$ is achieved with 3 nodes, leading to 0.31M parameters and 0.27G MAC operations. Fig. 6 shows the cell architectures searched with 2, 3, and 4 nodes using DARTS+Adas.

We select the searched architectures with 2, 3 and 4 nodes (namely DARTS-ADP-N2, -N3 and -N4), and train them on three more datasets: BCSS [1], BACH [2] and Osteosarcoma [20]. Details including data augmentations can be found in Section Datasets of the Supplementary Material. The goal is to evaluate how the searched architectures perform when transferred to different CPath datasets that cover different variations of single- vs multi-labels, multiclass problems, data samples, and organs. We also train several mobile-friendly architectures on them for comparison, including a 4-cell DARTS [23]. All networks are trained for 600 epochs with batch size 96, using the SGD optimizer with a 0.025 initial learning rate and cosine annealing scheduler.

As shown in Tabel.7, across all datasets, the group of DARTS-ADP networks achieves higher or comparable test accuracy than state-of-the-art networks, but with smaller parameter sizes. The 2-node version contains only 0.24M parameters but still ranks high in test accuracy. As for computation complexity, though MobileNetV2 0.35 [28], MobileNetV3-small [13], and MNASNet-small [33] achieve fewer MAC operations, their accuracies are two percent lower. This shows the superiority of DARTS-ADP in the speed-accuracy trade-off, which is desirable for CPath applications of high-throughput image analysis.

Another way to meet the needs of high-throughput applications is to decrease the image resolution. To evaluate the robustness of the architectures against downscaled inputs, we retrain several models with different resolutions (272, 136, 68, 34) in three datasets. Fig. 7 shows the performance of different network selection. Each line represents a network and each dot represents a specific resolution. The DARTS-based networks consistently achieve the highest accuracy with the lowest computation complexity in all resolutions and across all datasets. As the resolution decreases, their test accuracy exhibits a much less drop compared to ResNet18 [8] and MobileNetV2 [28], which shows the robustness of the DARTS-based networks. Such robustness is also illustrated in the standard deviation (denoted by shades). Compared to 4-cell DARTS [23], the two DARTS-ADP networks obtain higher or comparable test accuracy with lower computation complexity, which again shows their superiority in the speed-accuracy trade-off.

To better understand and interpret the performance of feature representation of different networks, we apply Grad-CAM [30] on the last convolutional layers of different networks to obtain the heatmaps for predicting ground truth labels. This visualization technique allows us to evaluate the reliability of different networks in label prediction. We randomly select five patches that contain different ground truth labels from the test set of ADP, BCSS, and BACH, and feed them into three networks for comparison. The selected networks are DARTS-ADP-N4, MobileNetV3-large [13], and MNASNet-A1 [33]. Results are shown in Fig. 8, where the heatmap indicates pixel-level confidence of pertinent labels of the image patch.

According to pathologists’ assessment, the overall performance of DARTS-ADP-N4 is the best. Examples are shown in the first column of Fig. 8(a), where DARTS-ADP-N4 successfully highlights the region of Erythrocytes. Either MobileNetV3-large or MNASNet-A1 discovers incomplete or false regions. This demonstrates the superiority of DARTS-ADP in extracting label-pertinent features from image patches, and hence more reliable predictions.

In CIFAR experiments, we train the network for 50 epochs with batch size 64 and initial channels 16. We test two optimizers for optimizing model weights, which are the original SGD [23] and Adas [12]. For DARTS+SGD, we follow [23] to use initial learning rate 0.025, cosine annealing scheduler, momentum 0.9 and weight decay $3\times 10^{-4}$. For DARTS+Adas, we use initial learning rate 0.175, scheduler beta 0.98, momentum 0.9, and weight decay $3\times 10^{-4}$. As for architecture parameter optimization, we follow [23] to use Adam [18] optimizer with initial learning rate $3\times 10^{-4}$, momentum $(0.5,0.999)$, and weight decay $10^{-3}$.

In ADP experiments, most hyperparameters are the same except that we use batch size 32 due to computation overhead. We also increase the initial learning rate of DARTS+SGD to 0.175 for model weights optimization.

In both CIFAR and CPath experiments, we follow [23] to train the network for 600 epochs with batch size 96 and initial channels 36. We use SGD optimizer with an initial learning rate of 0.025, cosine annealing scheduler, momentum 0.9, and weight decay $3\times 10^{-4}$. Additional enhancements include cutout and auxiliary towers as in [23]. Note that we disable auxiliary towers in training when we compare the performance of the searched architectures with the state-of-the-art networks.