DASS: Differentiable Architecture Search for Sparse Neural Networks

Neural Architecture Search (NAS) has recently attracted remarkable attention by relieving human experts from the laborious effort of designing neural networks. Early NAS methods mainly utilized evolutionary-based [9, 36, 37, 38] or reinforcement-learning-based methods [39, 40, 41]. Despite the efficiencies of the architecture designed by evolutionary-based and reinforcement-learning-based methods, they require tremendous computing resources. For example, the proposed method in [39] evaluates 20,000 neural candidates in 500 NVIDIA^® P100 GPUs over four days. One-shot architecture search methods [42, 43, 44] have been proposed to identify optimal neural architectures in a few GPU days (\(\gt\)1 GPU day [45]). In particular, Differentiable Architecture Search (DARTS) [27, 28, 29] is a variation of one-shot NAS methods that relaxes the search space to be continuous and differentiable. The detailed description of DARTS can be found in Section 3.1. Despite the broad successes of DARTS in advancing NAS applicability, achieving optimal results remains a challenge for real-world problems. Many subsequent works investigate some of these challenges by focusing on (i) increasing search speed [46, 47], (ii) improving generalization performance [35, 48], (iii) addressing robustness issues [49, 50, 51], (iv) reducing quantization error [14, 16], and (v) designing hardware-aware architectures [52, 53, 54]. On the other hand, few works attempt to prune the search space by removing inferior network operations. [55, 56, 57, 58, 59, 60]. These works utilized the pruning mechanism to progressively remove some operations from the search space. Unlike them, our method aims to extend the search space to improve the performance of the sparse network by searching for the best operations with sparse weight structures. Technically, our method extends the search space by adding the parametric sparse version of convolution and linear operations to find the best sparse architecture. Therefore, there is a lack of research on sparse weight parameters when designing neural architectures. Our proposed method (DASS) searches for the operations that are most effective for sparse weight parameters in order to achieve higher generalizing performance.

Network pruning is an effective method for reducing the size of DNNs, enabling them to be effectively deployed on devices with limited resource capacity. Prior works on network pruning can be classified into two categories: structured and unstructured pruning methods. The purpose of structured pruning is to remove redundant channels or filters to preserve the entire structure of weight tensors with dimension reduction [22, 23, 24, 25, 61, 62]. While structured pruning is famous for hardware acceleration, it sacrifices a certain degree of flexibility as well as weight sparsity [63].

On the other hand, unstructured pruning methods offer superior flexibility and compression rate by removing parameters with the least impact on the network accuracy of the weight tensors [10, 11, 12, 20, 22, 63, 64, 65, 66]. In general, unstructured pruning entails three stages to make a sparse network, including (i) pre-training, (ii) pruning, and (iii) fine-tuning. Prior unstructured pruning methods used various criteria to select the lowest pruning weight parameters. [67, 68] pruned weight parameters based on the second-derivative values of the loss function. Several studies proposed to remove the weight parameters below a fixed pruning threshold, regardless of the training objective [64, 65, 66, 69, 70, 71]. To address the limitation of fixed thresholding methods, [20, 72] proposed layer-wise trainable thresholds to determine the optimal value for each layer separately. The lottery-ticket hypothesis [66, 73, 74] is a different line of the method that identifies the pruning mask for an initialized CNN and trains the resulting sparse model from scratch without changing the pruning mask. HYDRA [10] formulate the pruning objective as empirical risk minimization and integrate it with the training objective. Unlike other methods, optimization-based pruning criteria improve the performance of sparse networks in comparison to other metrics. Despite the success of optimization-based pruning in achieving a significant compression rate, the classification accuracy is compromised, notably when the pruning ratio is extremely high (up to 99%). We show that the main reason for this issue is due to the non-optimal backbone architecture. We extend the search space of DASS by parametric sparse operations and formulate pruning as an empirical risk minimization problem and integrate it into the bi-level optimization problem to find the best sparse network.

Differentiable Architecture Search (DARTS) [27] is a NAS method that significantly reduces the search cost by relaxing the search space to be continuous and differentiable. DARTS cell template is represented by a Directed Acyclic Graph (DAG) containing \(N\) intra-nodes. The edge \((i,j)\) between two nodes is associated with an operation \(o^{(i,j)}\) (e.g., skip connection or \(3\times 3\) max-pooling) within \(\mathcal {O}\) search space. Equation (1) computes the output of intermediate nodes.where \(\mathcal {O}\) and \(\alpha _{o}^{(i,j)}\) denote the set of all candidate operations and the selection probability of \(o\), respectively. The output node in the cell is the concatenation of all intermediate nodes. DARTS optimizes architecture parameters (\(\alpha\)) and network weights (\(\theta\)) with the following bi-level objective function:whereandThe operation with the largest \(\alpha _o\) is selected for each edge. \(X_{train}\) and \(Y_{train}\) represent the training dataset and corresponding labels, respectively. Similarly, the validation dataset and labels are indicated by \(X_{val}\) and \(Y_{val}\), respectively. After the search process has been completed, the final architecture is re-trained from scratch to obtain maximum accuracy.

Pruning is considered unstructured if it removes low-importance parameters from weight tensors and makes sparse [63]. This paper uses the optimization-based unstructured network pruning method to provide greater flexibility and an extreme compression rate compared to structured pruning methods. The pruning method includes three main optimization stages: (i) pre-training: training the network on the target dataset, (ii) pruning: pruning unimportant weights from the pre-trained network, and (iii) fine-tuning: the sparse network is re-trained to recover its original accuracy. For the pruning stage, we consider an optimization-based method with the following steps: First, we define the pruning parameters that show the importance of each weight of the network (\(s^0\)) and initialize them according to Equation (3).where \(\theta _{pre,i}\) denotes the weight of \(i_{th}\) layer in the pre-trained network. Next, to learn the pruning parameters \((\hat{s})\), we formulate the optimization problem as Equation (4), which is then solved by stochastic gradient descent (SGD) [75].

\(\theta _{pre}\) and \(\mathbb {E}\) refer to the pre-trained network parameters and mathematical expectation, respectively. By solving this optimization problem, we are able to determine the effect of each weight parameter on the loss function and, consequently, the accuracy of the network. Finally, we convert the floating values of the pruning parameters to a binary mask based on selecting top-\(k\) weights with the highest magnitude of pruning parameters.

We propose DASS, a differentiable architecture search method for sparse neural networks. DASS at first extends the search space of the NAS with parametric sparse operations. Then it modifies the bi-level optimization problem to learn the architecture, weights, and pruning parameters. DASS employs a three-step approach to solve the complicated bi-level optimization problem, which consists of (1) Pre-training: find the best dense architecture (pruning parameters equal to zero) from the search space and pre-train it (2) Pruning and sparse architecture design: find the best pruning mask (optimizing pruning parameters) and update the architecture parameters based on the sparse weights and finally (3) Fine-tuning: re-train the sparse architecture to achieve the maximum classification performance.

To support sparse operations, DASS proposes the parametric sparse version of convolution and linear operations called SparseConv and SparseLinear, respectively. These operations have a sparsity mask (\(m\)) to remove redundant weight parameters from the network. Figure 4 illustrates the functionality of these two operations. In addition, Table 1 summarizes the operations of the DASS search space.

To empirically investigate the efficiency of the proposed sparse search space, we compare the similarity of the feature maps of high-performance dense architecture (with a large number of parameters) with the sparse architecture discovered by DASS and the architecture designed from the original search space; DARTS\(_{spasre}\); methods. We use Kendall’s \(\tau\) [76] metric to measure the similarity between output feature maps. The \(\tau\) correlation coefficient returns a value between \(-\)1 and 1. To present the outcome more clearly, we scale up these values between \(-\)100 and 100. Closer values to 100 indicate stronger positive similarity between the feature maps. Figure 5 summarizes the results. Our observations reveal a similarity between DASS feature maps and dense architecture (up to 16%). On the other hand, the correlation between DARTS\(_{sparse}\) and dense architecture is insignificant. Therefore, it shows that the architecture designed by DASS based on new search space can extract features more similar to high-performance dense architecture while DARTS\(_{sparse}\) that use dense search space lost important features after pruning. The level of similarity is not very high because DASS is a sparse network with a pruning ratio of 99%. However, it can demonstrate that DASS retrieves useful features.

DASS aims to search for optimal architecture parameters (\(\alpha ^\star\)) to minimize the loss of validation of the sparse network. Thus, to achieve a consistent search objective with the proposed sparse search space, we formulate the entire search objective as a complex bi-level optimization problem:

Here \(m\) denotes the parameters of the binary pruning mask. This formulation learns the architecture parameters based on the sparse weight parameters. However, Equation (5) is not a straightforward bi-level optimization problem because the lower-level problem consists of two optimization problems. To overcome this challenge, we break the search objective into three distinct steps. Thus, the problem is transformed into two bi-level optimization problems to determine the optimal architecture parameters for dense and sparse weights and an optimization problem to fine-tune the weight parameters. In addition, the lower-level optimization problem consists of a discrete optimization problem for pruning masks.

Section 5.4 proposes a multi-step optimization algorithm to solve the optimization problem and handle the discrete optimization problem by converting it to a continuous optimization problem.

As shown in Figure 6, the optimization algorithm consists of three main steps, which are described as follows.

(1) DATASET: to evaluate DASS, we use CIFAR-10 [77] and ImageNet [78] public classification datasets. For the search process, we split the CIFAR-10 dataset into 30k data points for training and 30k for validation. We transfer the best-learned cells on CIFAR-10 to ImageNet [27] and re-train the final sparse network from scratch.

(2) Details on Searching Networks: we create a network with 16 initial channels and eight cells. Each cell consists of seven nodes equipped with a depth-wise concatenation operation as the output node. The SparseConv operations follow the ReLU+SpasreConv+Batch Normalization order. We train the network using SGD for 50 epochs with a batch size of 64 in the DASS pre-train step. Then, we update the values of pruning and architecture parameters for 20 epochs in the DASS pruning step. Finally, we fine-tune the network for 200 epochs. The initial learning rate for the DASS in pre-train, pruning, and fine-tuning steps is 0.025, 0.1, and 0.01, respectively. In our experiments, we use the cosine annealing learning rate [79]. We use weight decay = 3 \(\times\) 10^-4 and momentum = 0.9 in all steps. The search process for CIFAR-10 takes \(\approx\)3 GPU-days on a single NVIDIA^® RTX A4000 that produces 4.35 Kg \(CO_2\). We compare the sparse architecture designed by our method, DASS, with other dense and sparse networks. NAS-Bench-101 [80] and NAS-Bench-201 [81] are examples of NAS algorithm evaluation benchmarks. They consist of numerous dense designs and their respective performance. Due to the fact that they do not support sparse architectures, we cannot evaluate DASS using these benchmarks. Creating sparse benchmarks for evaluating NAS algorithms is a suggestion for future work.

(3) DASS variants and Hardware Configuration: Table 2 provides the configuration details of the DASS variants. Each variation is built by stacking a different number of DASS cells and the output channels of the first layer to generate networks for various resource budgets. Table 3 presents specifications of hardware devices utilized for evaluating the performance of DASS at inference time.

Table 4 compares the performance of DASS against the state-of-the-art and the state-of-the-practice DNNs. We select the architecture with the highest accuracy, DrNAS [35], as the baseline for comparing compression rates. In comparison with DrNAS [35], DASS-Large provides 37.73\(\times\) and 29.23\(\times\) higher network compression rates while delivering a comparable accuracy (less than 2.5% accuracy loss) on the CIFAR-10 and ImageNet datasets, respectively. Compared to the best handcrafted designed network [33] on the CIFAR-10 (CCT-6/3x1), DASS-Large significantly decreases the parameters of the network by 29.9\(\times\) with providing slightly higher accuracy.

As we focus on improving the accuracy of sparse networks at extremely high pruning ratios, we compare DASS with other sparse networks with the unstructured pruning method at 99% pruning ratio (Table 5). In comparison with DARTS\(_{sparse}\), DASS-Small yields 7.81% and 7.81% higher top-1 accuracies with 1.23\(\times\) and 1.05\(\times\) reduction in network size on the CIFAR-10 and ImageNet datasets, respectively. It indicates that the network design based on new search space and sparse objective function finds better sparse architecture. In comparison with ResNet-18\(_{sparse}\) on the CIFAR-10 dataset, we provide 1.56% and 4.7% higher accuracy with 2.08\(\times\) and 1.05\(\times\) network size reduction for DASS-Medium and DASS-Large, respectively. Compared to ResNet-18\(_{sparse}\) on the ImageNet dataset, DASS-Medium provides 0.76% higher accuracy with 1.42\(\times\) network size reduction. MCUNET [8] is a lightweight neural network for microcontrollers. It is designed by a tiny neural architecture search mechanism. Compared to MCUNET on the ImageNet dataset, DASS-Large provides 1% higher accuracy with 2.89\(\times\) network size reduction. This result shows that only optimizing the size of the filters without considering the sparsity can not generate the best architecture. DASS directly search for the best operations in sparse version to design high-performance lightweight network. We can conclude that DASS increases sparse networks’ accuracy at high pruning ratios compared to NAS-based and handcrafted networks.

Table 6 compares DASS and the DARTS\(_{sparse}\) method with three different pruning ratios including 90%, 95%, and 99% on the CIFAR-10 dataset. DASS achieves 1.57%, 1.04%, and 7.8% higher accuracies with 7%, 6.9%, and 23% network size reduction compared to the DARTS\(_{sparse}\) at 90%, 95%, and 99% pruning ratios, respectively. Thus, DASS is significantly more effective at extremely higher pruning ratios (99%) than lower pruning ratios (90%).

Table 7 compares DASS with state-of-the-art pruning algorithms. The results indicate that DASS outperforms other pruning algorithms with different backbone architectures on CIFAR-10 and ImageNet datasets. On CIFAR-10, DASS-Large shows a 1.6% higher accuracy and 3.8\(\times\) reduction in the network size compared to the most accurate results provided by TAS\(_{Pruning}\) [87]. DASS-Large also provides 4.68% accuracy improvement with 38.14\(\times\) reduction in the network size over TAS\(_{Pruning}\) [87] on ImageNet. In light of DASS’ higher efficiency compared to other pruning methods, we can conclude that the pruning method was not the only reason for the DASS’s effectiveness and it is independent of the pruning algorithm.

Network quantization emerged as a promising research direction to reduce the computation of neural networks. Recently, [14, 15, 16] proposed to integrate the quantization mechanism into the differentiable NAS procedure to improve the performance of quantized networks. Table 8 compares DASS with the best results of NAS-based quantized networks. The compression rate is calculated as \(\frac{\sum _{l=1}^{L} \#W_l \times 32}{\sum _{l=1}^{L} \#W_l^t \times q}\) where \(\#W_l\) and \(\#W_l^t\) are the number of weights in layer \(l\) for full-precision (32-bit) and quantized network with \(q-\)bit resolution [14]. DASS-Medium yields 0.24% and 3.24% higher accuracies and significantly higher compression rate by 2.7\(\times\) and 4.24\(\times\) compared to TAS [14] as the most accurate quantized network on the CIFAR-10 and ImageNet datasets, respectively.

We extensively study the effectiveness of DASS in the context of hardware efficiency by computing the inference time (latency) of various state-of-the-art sparse networks for a wide range of resource-constrained edge devices on the CIFAR-10 dataset (Figure 9). The batch size is equal to 1 for all experiments. It is worth noting that we did not utilize any simplification techniques, such as [90], to compact the sparse filters by fusing weight parameters. Our results reveal that the Pareto-frontier of DASS consistently outperforms all other counterparts by a significant margin, especially on CPUs that have very limited parallelism. DASS-Tiny as the fastest network improves the accuracy from MobileNet-v2’s 73.44% to 81.35% (+7.91% improvement) and accelerates the inference by up to 3.87\(\times\). More importantly, DASS-Tiny runs much faster than DARTS\(_{sparse}\) by 1.67–4.74\(\times\) with slightly better accuracy. Compared to ResNet-18\(_{sparse}\) as the closest network to DASS in terms of accuracy, DASS-Medium provides 1.46% accuracy improvement and up to 1.94\(\times\) acceleration on hardware.

We use the t-distributed stochastic neighbor embedding (t-SNE) method [91] for visualizing decision boundaries of dense high-performance architecture designed by DARTS, DARTS\(_{sparse}\) (sparse dense DART architecture with pruning), and DASS (our sparse architecture) on the CIFAR-10 dataset. Figure 10 illustrates the decision boundaries of classification for each network. According to the results, DASS has a higher discrimination power than DARTS\(_{sparse}\), and DASS with a 99% pruning ratio behaves very similarly to the dense and high-performance DARTS architecture.

Figure 11 shows the best cells searched by DASS-Small. An interesting finding is that, for the normal cell, DASS-Small tends to select SparseConv operation with larger kernel sizes (\(5 \times 5\)), providing more pruning candidates to optimize the pruning mask. DASS-Small tends to leverage max-pooling operations in the reduction cell instead of avg-pooling operations. This is because the max-pooling operation has a higher feature extraction capability with sparse filters [92].

To verify the reproducibility of results, the DASS-Small search procedure was run five times with different random seeds. Figure 12 plots the average of accuracy and loss variations as well as the shades to indicate the confidence intervals. Results show that, while the confidence interval is wide at first, the average of multiple runs converges to neural architectures with similar performance with an average standard deviation (STDEV) of 2.22%.