Searching for A Robust Neural Architecture in Four GPU Hours
Xuanyi Dong, Yi Yang
Introduction
Designing an efficient and effective neural architecture requires substantial human effort and takes a long time . Since the birth of AlexNet in 2012, human experts have conducted a huge number of experiments, and consequently devised several useful structures, such as attention and residual connection . However, the infinite possible choices of network architecture make the manual search unfeasible . Recently, neural architecture search (NAS) has increasingly attracted the interest of researchers . These approaches learn to automatically discover good architectures. They can thus reduce the labour of human experts and find better neural architectures than the human-invented architectures. Therefore, NAS is an important research topic in machine learning.
Most NAS approaches apply evolutionary algorithms (EA) or reinforcement learning (RL) to design neural architectures automatically. In both RL-based and EA-based approaches, their searching procedures require the validation accuracy of numerous architecture candidates, which is computationally expensive . For example, the typical RL-based method utilizes the validation accuracy as a reward to optimize the architecture generator . An EA-based method leverages the validation accuracy to decide whether a model will be removed from the population or not . These approaches use a large amount of computational resources, which is inefficient and unaffordable. This motivates researchers to reduce the computational cost.
In this paper, we propose a Gradient-based searching approach using Differentiable Architecture Sampling (GDAS). It can search for a robust neural architecture in four hours with a single V100 GPU. GDAS significantly improves efficiency compared to the previous methods. We start by searching for a robust neural “cell” instead of a neural network . A neural cell contains multiple functions to transform features, and a neural network consists of many copies of the discovered neural cell . Fig. 1 illustrates our searching procedure in detail. We represent the search space of a cell by a DAG. Every grey square node indicates a feature tensor, numbered by the computation order. Different colored arrows indicate different kinds of operations, which transform one node into its intermediate features. Meanwhile, each node is the sum of the intermediate features transformed from the previous nodes. During training, the proposed GDAS samples a sub-graph from the whole DAG, indicated by solid connections in Fig. 1. In this sub-graph, each node only receives one intermediate feature from every previous node. Specifically, among the intermediate features between every two nodes, GDAS samples one feature in a differentiable way. In this way, GDAS can be trained by gradient descent to discover a robust neural cell in an end-to-end fashion.
The fast searching ability of GDAS is mainly due to the sampling behavior. A DAG contains hundreds of parametric operations with millions of parameters. Directly optimizing this DAG instead of sampling a sub-graph leads to two disadvantages. First, it costs a lot of time to update numerous parameters in one training iteration, increasing the overall training time to more than one day . Second, optimizing different operations together could make them compete with each other. For example, different operations could generate opposite values. The sum of these opposite values tends to vanish, breaking the information flow between the two connected nodes and destabilizing the optimization procedure. To solve these two problems, the proposed GDAS samples a sub-graph at one training iteration. As a result, we only need to optimize a part of the DAG at one iteration, which accelerates the training procedure. Moreover, the inappropriate competition is avoided, which makes the optimization effective.
In summary, GDAS has the following benefits:
1. Compared to previous RL-based and EA-based methods, GDAS makes the searching procedure differentiable, which allows us to end-to-end learn a robust searching rule by gradient descent. For RL-based and EA-based methods, feedback (reward) is obtained after a prolonged training trajectory, while feedback (loss) in our gradient-based method is instant and is given in every iteration. As a result, the optimization of GDAS is potentially more efficient.
2. Instead of using the whole DAG, GDAS samples one sub-graph at one training iteration, accelerating the searching procedure. Besides, the sampling in GDAS is learnable and contributes to finding a better cell.
3. GDAS delivers a strong empirical performances while using fewer GPU resources. On CIFAR-10, GDAS can finish one searching procedure in several GPU hours and discover a robust neural network with a test error of 2.82%. On PTB, GDAS discovers a RNN model with a test perplexity of 57.5. Moreover, the networks discovered on CIFAR and PTB can be successfully transferred to ImageNet and WT2.
Related Work
Recently, researchers have made significant progress in automatically discovering good architectures . Most NAS approaches can be categorized in two modalities: macro search and micro search.
Macro search algorithms aim to directly discover the entire neural networks . To search convolutional neural networks (CNNs) , typical approaches apply RL to optimize the searching policy to discover architectures . Baker et al. trained a learning agent by Q-learning to sequentially choose CNN layers. Zoph and Le utilized long short-term memory (LSTM) as a controller to configure each convolutional layer, such as the filter shape and the number of filters. In these macro search algorithms , the number of possible networks is exponential to the depth of a network, e.g., a depth of 12 can result in more than 1029 possible networks . It is difficult and ineffective to search networks in such a large search space, and, therefore, these macro search methods usually limit the CNN models to be shallow, e.g., a depth is less than 12. Since macro-discovered networks are shallower than deep CNNs , their accuracies are limited. In contrast, our GDAS allows the network to be much deeper by stacking tens of discovered cells and thus can achieve a better accuracy.
Micro search algorithms aim to discover neural cells and design a neural architecture by stacking many copies of the discovered cells . A typical micro search approach is NASNet , which extends the approach of to search neural cells in the proposed “NASNet search space”. Following NASNet , many researchers propose their methods based on the NASNet search space . For example, Real et al. applied EA algorithm with a simple regularization technique to search neural cells. Liu et al. proposed a progressive approach to search cells from shallow to deep gradually. These micro search algorithms usually take more than 100 GPU days . Even though some of them reduce the searching cost, they still take more than one GPU day . Our GDAS is a also micro search algorithm, focusing on search cost reduction. In experiments, we can find a robust network within fewer GPU hours, which is 1000 less than the standard NAS approach .
Improving Efficiency. Since NAS algorithms usually require expensive computational resources , an increasing number of researchers focus on improving the architecture search speed . A variety of techniques have been proposed, such as progressive-complexity search stages , accuracy prediction , HyperNet , Net2Net transformation , and parameter sharing . For instance, Cai et al. reused weights of previously discovered networks to amortize the training cost. Pham et al. shared parameters between different child networks to improve the efficiency of the searching procedure. Brock et al. utilized a network to generate model parameters given a discovered network, avoiding fully training from scratch. Liu et al. relaxed the search space to be continuous, so that they can use gradient descent to effectively search cells. Though these approaches successfully accelerate the architecture search procedure, several GPU days are still required . Our GDAS samples individual architecture in a differentiable way to effectively discover architecture. As a result, GDAS can finish the search procedure in several GPU hours on CIFAR-10, which is much faster than these efficient methods.
Contemporary to this work, Xie et al. applied a similar technique to relax the discrete candidate sampling as ours. They focus on fixing the inconsistency between the loss of attention-based NAS and their objective. In contrast, we focus on making the sampling procedure differentiable and accelerating the searching procedure.
Methodology
We search for the neural cell in the search space and stack this cell in series to compose the whole neural network. For CNN, a cell is a fully convolutional network that takes output tensors of previous cells as inputs and generates another feature tensor. For recurrent neural network (RNN), a cell takes the feature vector of the current step and the hidden state of the previous step as inputs, and generates the current hidden state. For simplification, we take CNN as an example for the following description.
We represent the cell in CNN as a DAG consisting of an ordered sequence of computational nodes. Each computational node represents one feature tensor, which is transformed from two previous feature tensors. This procedure can be formulated as shown in Eq. (1) following .
2 Searching by Differentiable Model Sampling
Formally, we denote a neural architecture as and the weights of this neural architecture as . The goal of NAS is to find an architecture , which can achieve the minimum validation loss after being trained by minimizeing the training loss, as shown in Eq. (3.2).
where is sampled from and is its associated weight. The discrete probability distribution is characterized by a learnable probability mass function as in Eq. (4):
Training. Reviewing the objective of NAS in Eq. (3.2), the main challenge is learning to find architecture . By utilizing Eq. (7), we can make the sampling procedure differentiable and learn a distribution of neural cells (representing architectures). However, it is still intractable to directly solve Eq. (3.2), because the nested formulation in Eq. (3.2) needs to calculate high order derivatives. In practice, to avoid calculating high order derivatives, we apply the alternative optimization strategy to update the sampling distribution and the weights of all functions in an iterative way. Given one data sample and its associated label , we calculate the loss as:
One benefit of this acceleration trick is that it allows us to directly search on the large-scale dataset (e.g., ImageNet) due to the saved GPU memory. We did some experiments to directly search on ImageNet using the same hyper-parameters as on the small datasets, however, failed to obtain a good performance. Searching on a large-scale dataset might require different hyper-parameters and needs careful tuning. We will explore this in our future work.
3 Discussion on the Reduction Cell
Revisiting state-of-the-art architectures designed by human experts, AlexNet and VGGNet use the max pooling to reduce the spatial dimension; ResNet uses a convolutional layer with stride of 2; and DenseNet uses a 1 by 1 convolutional layer followed by average pooling to reduce dimension. These human-designed reduction cells are simple and effective. The automatically discovered reduction cells are also usually similar and simple . For example, the reduction cell discovered by only has max pooling and identity operations.
Most human-designed and automatically discovered reduction cells are simple and can achieve a high accuracy. Moreover, compared to searching one normal cell, jointly searching a normal cell and a reduction cell will greatly increase the search space and make the optimization difficult. We hope to find a better network by fixing the reduction cell. Inspired by , we design a fixed reduction cell as shown in Fig. 3. In the experiments, with this human-designed reduction cell, GDAS finds a better architecture, yielding fewer parameters and higher accuracy.
Experimental Study
CIFAR-10 and CIFAR-100 consist of 50K training images and 10K test images. CIFAR-10 categorizes images into 10 classes, while CIFAR-100 has 100 classes.
ImageNet is a large-scale and well-known benchmark for image classification. It contains 1K classes, 1.28 million images for training, and 50K images for validation.
Penn Treebank (PTB) is a corpus consisting of over 4.5 million words of American English words. We pre-process PTB following .
WikiText-2 (WT2) is a collection of 2 million tokens from the set of verified Good and Featured articles on Wikipedia. The training set contains 600 articles with 2,088,628 tokens. The validation set contains 60 articles with 217,646 tokens. The test set contains 60 articles with 245,569 tokens.
2 Search for CNN
Clarifications on the searching cost (GPU days) of different methods. The searching costs listed in Tab. 1 and Tab. 2 are not normalized across different GPU devices. Different algorithms might run on different machines, and we simply refer the searching costs reported in their papers.It is difficult for us to run all algorithms on the same GPU. If we use other GPU devices, the searching cost of “GDAS (FRC)” could be a different number. For example, if we use Titan 1080Ti, the search cost will increase to about seven GPU hours.
Results on CIFAR. After the searching procedure, we use C=36, B=4, and N=6 to form a CNN. Following the previous works , we train the network by 600 epochs in total. We start the learning rate of 0.025 and reduce it to 0 with the cosine learning rate scheduler. We set the probability of path dropout as 0.2 and the auxiliary tower with the weight of 0.4 . We use the standard pre-processing and data augmentation, i.e., randomly cropping, horizontally flipping, normalization, and CutOut .
We compare the models discovered by our approach with other state-of-the-art models in Tab. 1. The models discovered by the macro search algorithms obtain a higher error than the models discovered by the micro search algorithms. Using GDAS, we discover a model with 3.3M parameters, which achieves 2.93% error on CIFAR-10. Using GDAS (FRC), we discover a model with only 2.5M parameters, which achieves 2.82% error on CIFAR-10. NASNet-A achieves a lower error rate than ours, but it contains more than 80% of the parameters than the model discovered by GDAS (FRC). Notably, our GDAS discovers a comparable model with the state-of-the-art, whereas the searching cost of our approach is much less than the others. For example, GDAS (FRC) takes less than 4 hours on a single V100 GPU, which is about 0.17 GPU days. It is faster than NASNet by almost 104 times. ENAS is a recent work that focuses on accelerating the searching procedure. ENAS is very efficient, whereas our GDAS (FRC) is three times faster than ENAS.
Results on ImageNet. Following , we use the ImageNet-mobile setting, in which the input size is 224224 and the number of multiply-add operations is restricted to be less than 600M. We train models by SGD with 250 epochs and use the batch size of 128. We initialize the learning rate of 0.1 and reduce it by 0.97 after each epoch.
We compare our results on ImageNet with the other methods in Tab. 2. Most algorithms in Tab. 2 take more than 1000 GPU days to discover a good CNN cell. DARTS uses minimum resources among the compared algorithms, whereas ours is even faster than DARTS by more than 10 times. For GDAS (FRC), we use C=52 and N=4 to construct the model following the setting in . For GDAS, if we use C=52 and N=4, the number of multiply-add operations will be larger than 600 MB, and thus we use C=50 to restrict it to be less than 600MB. Our model, GDAS (FRC) [C=52,N=4], costs about 20% less multiply-add operations than but obtains the same top-5 error. AmoebaNet-A and Progressive NAS achieve a slightly lower test error than ours. However, their methods cost a prohibitive amount of GPU resources. The results in Tab. 2 show the discovered cell on CIFAR-10 can be successfully transferred to ImageNet and achieve competitive performance.
3 Search for RNN
Results on PTB. We evaluate the RNN model formed by the discovered recurrent cell on PTB. We use a batch size of 64 and a hidden size of 850. We train the model using the A-SGD by 2000 epochs. The learning rate is fixed as 20 and the weight decay is 8e-7. DARTS and ENAS greatly reduce the search cost compared to previous methods. Our GDAS incurs a lower search cost than all the previous methods. Note that our code is not heavily optimized and the theoretical search cost should be less than the one reported in Tab. 3.
We compare different RNN models in Tab. 3. The model discovered by GDAS achieves a validation perplexity of 59.8 and a test perplexity of 57.5. The performance of our discovered RNN is on par with the state-of-the-art models in Tab. 3. LSTM + SE obtains better results than ours, but it is an ensemble method using mixture of softmax. By applying the SE technique , GDAS can achieve the lower perplexity without doubt. LSTM is an extensively tuned model, whereas our automatically discovered model is superior to it. Compared to other efficient approaches, the search cost of GDAS is the lowest.
Results on WT2. To train the model on WT2, we use the same experiment settings as PTB, but we use a hidden size of 700 and a weight decay of 5e-7. We train the model in 3000 epochs in total. Tab. 4 compares different RNN models on WT2. Our approach achieves competitive results among all automatically searching approaches. GDAS is worse than “LSTM + SC” . Since our model is searched on a small dataset PTB, and the transferable ability of the discovered model might be a little bit weak. If we directly search the RNN model on WT2, we could obtain a better model and improve the transferable ability.
4 Discussion
We visualize the discovered cells in Fig. 4. These automatically discovered cells are complex and hard to be designed manually. Moreover, networks with these discovered cells can achieve more superior performance than hand-crafted networks. This demonstrates that automated neural architecture search is the future of architecture design.
Revisiting Sec. 3.3, we propose a new reduction cell as a replacement for automated reduction cell. With this reduction cell, we can more effectively search neural cells. For further analysis, we use the normal cell found by GDAS and the proposed reduction cell to construct a new CNN, denoted as “GDAS-N + FIX-R” in Tab. 5. The accuracy of this network on CIFAR-10 is similar to “GDAS-N + GDAS-R” and “FRC-N + FIX-R” in Tab. 5. This result implies that the reduction cell might have a negligible effect on the performance of networks and the hand-crafted reduction cell could be on par with the automatically discovered one.
Most recent NAS approaches search neural networks on the small-scale datasets, such as CIFAR, and then transfer the discovered networks to the large-scale datasets, such as ImageNet. The obstacle of directly searching on ImageNet is the huge computational cost. GDAS is an efficient NAS algorithm and gives us an opportunity to search on ImageNet. We will explore this research direction in our future work.
Conclusion
In this paper, we propose a Gradient-based neural architecture search approach using Differentiable Architecture Sampler (GDAS). Our approach is efficient and reduces the search cost of the standard NAS approach by about 104 times. Moreover, both CNN and RNN models discovered by our GDAS can achieve competitive performance compared to state-of-the-art models.