RARTS: An Efficient First-Order Relaxed Architecture Search Method

DARTS training relies on an iterative algorithm to solve a bilevel optimization problem [32, 14] which involves two loss functions computed via data splitting (splitting the dataset into two halves, i.e. training data and validation data):

Here $w$ denotes the network weights, $\alpha$ is the architecture parameter, $L_{train}$ and $L_{val}$ are the loss functions computed on the training data $D_{train}$ and the validation data $D_{val}$. Since many common datasets like CIFAR do not include the validation data, $D_{train}$ and $D_{val}$ are usually two non-overlapping halves of the original training data. We denote $L_{train}$ and $L_{val}$ by $L_{t}$ and $L_{v}$ to avoid any confusions with the meaning of the subscripts. $D_{t}$ and $D_{v}$ are defined similarly. DARTS has adopted data splitting because it is believed that joint training of both $\alpha$ and $w$ via gradient descent on the whole dataset by minimizing the overall loss function:

can lead to overfitting [14, 15]. Therefore, DARTS searches for the architectures through a two-step differentiable algorithm which updates the network weights and the architecture parameters in an alternating way:

• •

  update weight $w$ by descending along $\nabla_{w}L_{t}(w,\alpha)$

• •

  update architecture parameter $\alpha$ by descending along:

  $\displaystyle\nabla_{\alpha}\,L_{v}(w-\xi\,\nabla_{w}L_{t}(w,\alpha),\alpha)$

where $\xi=0$ ($\xi>0$ ) gives the first or second-order approximation. The bilevel optimization problem also arises in hyperparameter optimization and meta-learning, where a second-order algorithm and a convergence theorem on minimizers have been proposed in previous work [33] (Theorem 3.2), under the assumption that the $\alpha$-minimization is solved exactly, and $w^{t}(\alpha)$ converges uniformly to $w(\alpha)$. However, the $\alpha$-minimization of DARTS is approximated by gradient methods only, and hence the convergence of DARTS algorithm remains unknown theoretically.

We are aware of the fact that the first-order DARTS updates the architecture parameters on $D_{v}$ by descending along $\nabla_{\alpha}\,L_{v}(w,\alpha)$, which means it merely uses half of the data to train $\alpha$ and might cause some convergence issues (see Fig. 1). MiLeNAS has developed a mixed-level solution, where the architecture parameters can be learned on $D=D_{t}\cup D_{v}$ via a first-order descending algorithm [15]:

We shall see that MiLeNAS is actually a constrained case of RARTS when our two network splits become identical. However, we point out that computing $L_{v}$ using an identical network makes MiLeNAS still suffer from the same convergence issue in a later example (Section III-D). The second-order DARTS is observed to approximate the optimum better than first-order DARTS in a solvable model and through experiments, yet it requires computing the mixed derivative $\nabla^{2}_{\alpha,w}L_{t}$, at a considerable overhead. Searching by DARTS can also lead to the architecture collapse issue, meaning the selected architecture contains too many skip-connections. Typically such a bias in operation selection degrades the model performance. SNAS [27], FBNet [12], and GDAS [18] use differentiable Gumbel-Softmax to mimic one-hot encoding which implies exclusive competition and risks of unfair advantages [19]. This unfair dominance of skip-connections in DARTS has also been noted by FairDARTS [19], which has proposed a collaborative competition approach by making the architecture parameters independent, through replacing the softmax with sigmoid. They have further penalized the operations in the search space with probability close to $\frac{1}{2}$, i.e. a neutral and ambiguous selection. As these methods focus on replacing some operations or the loss function, it would be worthwhile to explore other solutions such as replacing the gradient-based DARTS search algorithm.

In addition to DARTS, many other differentiable methods for architecture search have been proposed, considering various aspects such as the search space, selection criterion, and training tricks. SNAS [27] has discussed it from a statistical perspective with however a minor to moderate performance improvement. The search efficiency has also been improved by sampling a portion of the search space during each update in training. A perturbation-based selection scheme has been proposed in [34], as the magnitude of architecture parameters is believed to be inadequate as a selection criterion. P-DARTS [35] has adopted operation dropout and regularization on skip-connections. From the procedure side to delay a quick short cut aggregation, it has also divided the search stage into multiple stages and progressively adds more depth than DARTS. PC-DARTS [36] samples a proportion of channels to reduce the bias of operation selection and enlarge the batch size as well. GDAS [18] searches the architecture with one operation sampled at a time. Other approaches apply differentiable methods on much larger search spaces with sampling techniques to save memory and avoid model transfer [7, 12]. We will see that these variants of the differentiable architecture search method are actually complementary to our approach that advances DARTS on the purely algorithmic side by mobilizing weights. Moreover, many works [7, 37, 12, 38, 16, 39] manage to balance latency with the performance of the model to enhance the efficiency of the model. Despite the broad use of differentiable methods in the works we have mentioned, one may wonder how DARTS and its variants beat random search. A detailed comparison in [40] has elaborated the advantage of DARTS in accuracy and efficiency compared with random search.

Differentiable search method has contributed to a wide range of tasks other than topological architecture search. TAS [20] searches for the width of each layer, i. e. number of channels, by learning the optimal one from the aggregation of several candidate feature maps via a differentiable method and sampling. FasterSeg [16] searches for the cell operations and layer width, as well as the multi-resolution network path over the semantic segmentation task. These works of searching width are closely related to channel pruning, which means pruning redundant channels from the convolutional layers. Among numerous methods to prune redundant channels [41, 24, 42, 43, 44, 45], a classical approach is to apply group LASSO [46] on the weights to identify unimportant channels. The weights in each channel form one group, and the magnitude of each group is measured by $\ell_{2}$ norm of its weights. The network is trained by minimizing a loss function penalized by the $\ell_{1}$ norm of these magnitudes from all groups. The channels are pruned based on thresholding their norms. Selecting good thresholds as hyperparameters for different channels can be laborious for deep networks. On the other hand, channel selection is intrinsically a network architecture issue. It is debatable if thresholding by weight magnitudes is always meaningful [47].

Another approach of channel pruning [21, 22] involves assigning a channel scaling (scoring) factor to each channel, which is a learnable parameter independent of the weights. In the training process, the factors and the weights are learned jointly, and the channels with low scaling factors are pruned. After that, the optimal weights of the pruned network are adjusted by one more stage of fine-tuning. In terms of channel scaling factors, the channel pruning problem becomes a special case of neural architecture search. Besides this formulation, there are several pruning methods based on NAS. AMC [37] has defined a reward function and pruned the channels via reinforcement learning. MetaPruning [48] generates the best pruned model and weights from a meta network.

As pointed out in DARTS [14] and MiLeNAS [15], when learning the architecture parameter $\alpha$, splitting training and validation data should be taken into account to avoid overfitting. However, we have discussed that the bilevel formulation (1) and training algorithm of DARTS may lead to several issues: unknown convergence, low efficiency and the unfair selection of operations. Therefore, we follow the routine of train-validation data splitting, but want to formulate a single level problem, in contrast to DARTS and MiLeNAS. First, if we use $(w,\alpha)$, the pair of weight and architecture parameters in Eq. (2) to represent a network, what we propose to do is to further relax the network weights $w$ via splitting a network copy denoted by $(y,\alpha)$. We call $(y,\alpha)$ and $(w,\alpha)$ the primary and auxiliary networks, which share the same architecture $\alpha$ and the same dimensions as weight tensors, but can have different weight initialization.

Next, a primary loss $L_{v}(y,\alpha)$ is computed with parameters $(y,\alpha)$ fed on data $D_{v}$, while an auxiliary loss $L_{t}(w,\alpha)$ is computed with parameters $(w,\alpha)$ fed on data $D_{t}$. Note that the computation of the auxiliary loss $L_{t}(w,\alpha)$ is the same as that of DARTS. The difference is that the primary loss is computed on the primary network $(y,\alpha)$, instead of $(w,\alpha)$. Now we present the single level objective of our relaxed architecture search (RARTS) framework. With a $\ell_{2}$ penalty on the distance between $w$ and $y$, the two loss functions are combined through the following relaxed Lagrangian $L=L(y,w,\alpha)$ of Eq. (2):

where $\lambda$ and $\beta$ are hyperparameters controlling the penalty scale and the learning process. We will see in the search algorithm that the penalty term enables the two networks to exchange information and cooperate to search the architecture which they share together. This technique of splitting $w$ and $y$ is called network splitting, which is also inspired by some previous work [49]. In their work, splitting of variables is able to approximate a non-smooth minimization problem via an algorithm of combined closed-form solutions and gradient descent.

Since various NAS approaches discover architectures of inconsistent sizes or FLOPS, it has made the comparison through different methods unfair, because larger models are likely to have better performance but low efficiency. Many NAS methods have adopted latency as a model constraint [7, 16]. To control model size, we follow the technique of approximating the model latency with the sum of latency from all the operations [16], and add the approximated latency to the loss function as a penalty. Since each component of the latency tensor (denoted by Lat) is the latency amount associated with a candidate operation, the dimension of Lat is the same as that of $\alpha$. Therefore, we provide an alternative objective which is penalized by the latency of the model:

where the bracket is inner product.

We minimize the relaxed Lagrangian $L(y,w,\alpha)$ in (4) by iteration on the three variables in an alternating way to allow individual and flexible learning schedules for the three variables. Similar to Gauss-Seidel method in numerical linear algebra [50], we use updated variables immediately in each step and obtain the following three-step iteration:

With explicit gradient $\nabla_{w,y}\|y-w\|^{2}_{2}$, we have:

To minimize the Lagrangian (III-A), the first two steps are the same as Eq. (7) since the latency only depends on $\alpha$. The third step becomes:

Note that the update of $\alpha$ in Eq. (7) involves both $L_{t}$ and $L_{v}$, which is similar to the second-order DARTS but without the mixed second derivatives. The first-order DARTS uses $\nabla_{\alpha}L_{v}$ only in this step. In the previous section, we have discussed the architecture collapse issue of DARTS, i.e., selecting to many skip-connections. A possible reason why DARTS may lead to architecture collapse is that its architecture parameters converge more quickly than the weights in the convolutional layers. That means, when DARTS selects architecture parameters, it tends to select skip-connection operations, since the convolutional layers are not trained well. The fact that first-order DARTS has only used one of the two data splits to train the weights, makes the training of convolutional layers worse. For RARTS, we make use of both $L_{t}$ and $L_{v}$ to update the weight parameters $w$ and $y$ in the first two steps of Eq. (7). In the third step of Eq. (7), both $L_{t}$ and $L_{v}$ are also used to update the shared architecture $\alpha$. In this way, the architecture is learned better, as more data are involved during training. If $y=w$ is enforced in Eq. (7) e.g. through a multiplier, RARTS essentially reduces to first-order MiLeNAS [15]. However, relaxing to $y\not=w$ has its advantages of having more generality and robustness as it is optimized on two networks with different but related weights. In contrast, MiLeNAS trains the network weigts on the training data $D_{t}$ only, and suffers from the same convergence issue as first-order DARTS (Section III-D). We summarize the RARTS algorithm in Algorithm 1.

Suppose that $L_{t}$ and $L_{v}$ both satisfy Lipschitz gradient property, or there exist positive constants $L_{1}$ and $L_{2}$ such that ($z=(y,\alpha)$, $z^{\prime}=(y^{\prime},\alpha^{\prime})$):

which implies:

for any $(z,z^{\prime})$; similarly ($\zeta=(w,\alpha)$, $\zeta^{\prime}=(w^{\prime},\alpha^{\prime})$):

which implies:

for any $(\zeta,\zeta^{\prime})$.

We compare a few differentibale methods through an example [14] which has an analytical solution. Regardless of the latency penalty, we consider quadratic functions $L_{v}=\alpha\,w-2\alpha+1$, $L_{t}=w^{2}-2\,\alpha\,w+\alpha^{2}$ for the bilevel problem (1). Therefore, the solution to the inner level problem is:

Then $L_{v}(w^{\ast}(\alpha),\alpha)=\alpha^{2}-2\alpha+1$, and the global minimizer of this bilevel problem is $(w^{\ast},\alpha^{\ast})=(1,1)$. However, the equilibrium equations of first-order DARTS is:

which gives a spurious equilibrium $(\bar{w},\bar{\alpha})=(2,2)$. The equilibrium equations of first-order MiLeNAS is:

which also results in the spurious equilibrium $(\bar{w},\bar{\alpha})=(2,2)$.

On the other hand, RARTS can approximate the correct minimizer $(w^{\ast},\alpha^{\ast})=(1,1)$ better. Note that both $L_{v}$ and $L_{t}$ satisfy the Lipschitz gradient property, which implies the descent of Lagrangian $L$ by the proof of Theorem 1. If $\lambda>1/2$, $\beta>3/2$, $L$ is bounded and coercive, which follows from an eigenvalue analysis of linear system (7) and is observed in computation. Hence, Theorem 1 can be applied to this example, and the equilibrium system (9) reads:

Adding (13) and (14) gives: $\bar{w}=\frac{2\lambda-1}{2\lambda}\bar{\alpha}$, which together with (15) determines $(\bar{\alpha},\bar{w},\bar{y})$ uniquely: $(\bar{\alpha},\bar{w},\bar{y})=(\frac{4\beta\lambda}{4\beta\lambda-\beta-2\lambda},\frac{4\beta\lambda-2\beta}{4\beta\lambda-\beta-2\lambda},\frac{4\beta\lambda-2\beta-4\lambda}{4\beta\lambda-\beta-2\lambda})$, if $4\beta\lambda-\beta-2\lambda\neq 0$. At $\lambda=\beta=15$, $(\bar{\alpha},\bar{w},\bar{y})\approx(1.053,1.018,0.947)$ where global convergence holds for the whole RARTS sequence. The learning dynamics starting from $(\alpha_{0},w_{0},y_{0})=(2,-2,y_{0})$, is reproduced in Fig. 1, along with three learning curves from RARTS as the parameters $(\lambda,\beta)$ and the initial value $y_{0}$ vary. In Fig. 1a, $\beta=10$, $y_{0}=0$. In Fig. 1b, $\lambda=10$, $y_{0}=0$. In Fig. 1c, $\lambda=\beta=10$. In all experiments, the learning rates are fixed at $0.01$. For a range of $(\lambda,\beta)$ and $y_{0}$, we see that our learning curves enter a small circle around $(1,1)$, while first-order DARTS converges to the spurious point.