AdvantageNAS: Efficient Neural Architecture Search with Credit Assignment

The probability distribution for $\bm{a}$ is modeled by an independent categorical distribution. Each $\bm{a}_{i}$ for $i\in\mathcal{I}$ independently follows a categorical distribution $p_{\theta_{i}}$ parameterized by $\theta_{i}$. The probability mass function for $\bm{a}$ is written as the product $p_{\theta}(\bm{a})=\prod_{i\in\mathcal{I}}p_{\theta_{i}}(\bm{a}_{i})$, where $\theta=(\theta_{1},\dots,\theta_{\lvert\mathcal{I}\rvert})$. The domain of $\theta_{i}$ is $\mathbb{R}^{\lvert\mathcal{J}_{i}\rvert}$ and the probability of $\bm{a}_{i}$ being $1$ in the $j$th category is $\mu_{i}^{j}(\theta_{i})$ defined by (2).

The objective function of AdvantageNAS is the reward function expected over the architecture distribution, i.e., $J(\theta,\bm{w})=\mathbb{E}_{\bm{a}\sim p_{\theta}}[R(\bm{a},\bm{w})]$, similar to REINFORCE. It is easy to see that the maximization of $J(\theta,\bm{w})$ leads to the maximization of $R(\bm{a},\bm{w})$ since $J(\theta,\bm{w})\leqslant\max_{\bm{a}}R(\bm{a},\bm{w})$ and $\sup_{\theta}J(\theta,\bm{w})=\max_{\bm{a}}R(\bm{a},\bm{w})$, where $p_{\theta}$ represents the Dirac delta function concentrated at $\operatornamewithlimits{argmax}_{\bm{a}}R(\bm{a},\bm{w})$.

AdvantageNAS update $\bm{w}$ and $\theta$ alternately. The weights $\bm{w}$ is updated by estimating the gradient $\nabla_{\bm{w}}J(\theta,\bm{w})$ using the Monte Carlo method with an architecture sample $\bm{a}\sim p_{\theta}$ and a minibatch reward $r(\bm{a},\bm{w};D)=-\frac{1}{\lvert D\rvert}\sum_{d\in D}\mathcal{L}(\bm{a},\bm{w};d)$, where $D\subseteq\mathcal{D}$. Then, $\theta$ is updated by estimating the policy gradient $\nabla_{\theta}J(\theta,\bm{w})$ using the Monte Carlo method with another architecture sample $\bm{a}\sim p_{\theta}$ and a minibatch reward $r(\bm{a},\bm{w};V)$, where $V\subseteq\mathcal{V}$. Next, we discuss how to estimate the policy gradient to improve the convergence speed while maintaining the time and space complexities at levels similar to those in previous studies. The algorithm of AdvantageNAS is described in Algorithm 1 using pseudocode.

The REINFORCE algorithm estimates the gradient using a simple Monte Carlo method,

where $\bm{a}$ is an architecture sampled from $p_{\theta}$. In general, this gradient is called a policy gradient, where $p_{\theta}$ is regarded as a policy that takes no input (i.e., no state) and outputs the probability of selecting each architecture (as an action). For our choice of $p_{\theta}$, the gradient of the log-likelihood is $\nabla_{\theta}\log(p_{\theta}(\bm{a}))=\bm{a}-\mu(\theta)$. In other words, the probability vector $\mu(\theta)$ is updated to approach $\bm{a}$ if the minibatch reward $r(\bm{a},\bm{w};V)$ is positive. A higher reward results in a greater shift of $\mu(\theta)$ towards current sample $\bm{a}$.

Intuitively, a high estimation variance will lead to slow convergence. This is indeed theoretically supported by the subsequent result, whose proof follows (Akimoto et al. 2012); the details are presented in Section A.1. The lower the estimation variance $\operatorname{Var}_{\delta}\left[\delta\right]$ of policy gradient $\delta$, the greater is the lower bound of the expected improvement in one step. Therefore, reducing the estimation variance is key to speed up the architecture update.

Credit assignment in reinforcement learning tasks is a well-known cause of slow convergence of policy gradient methods (Minsky 1961). The reward is the result of a set of consecutive actions (in our case, the choice $\bm{a}_{i}$ at each edge $i\in\mathcal{I}$). A positive reward for $\bm{a}=(\bm{a}_{1},\dots,\bm{a}_{\lvert I\rvert})$ does not necessarily mean that all the actions $\bm{a}_{i}$ are promising. However, since (4) updates $\theta$ based on a single reward, it cannot recognize the contribution of each edge. To understand this problem, consider a situation in which we have two edges with two candidate operations for each edge. The reward is assumed to be the sum of the contributions of the choices on the two edges, $r_{i}^{j}$ for $\bm{a}_{i}=\bm{e}_{j}$ for $i=1$ and $2$ and $j=1$ and $2$. If $r_{1}^{1}\gg r_{2}^{1}>0$ and $r_{1}^{2}\ll r_{2}^{2}<0$, i.e., the contribution of the first edge dominates, the sign of the reward is determined by $\bm{a}_{1}$. The update direction of $\mu_{2}(\theta_{2})$ is determined by $\bm{a}_{1}$; it is not based on whether $\bm{a}_{2}$ is a better choice. This results in a high variance for the $\mu_{2}(\theta_{2})$ update, leading to slow convergence.

We introduce the notion of advantage, used in policy gradient methods, to address the issue of credit assignment (Minsky 1961). We regard $\bm{s}_{i}=(\bm{a}_{1},\dots,\bm{a}_{i-1},\bm{a}_{i+1},\dots,\bm{a}_{\lvert I\rvert})$ as a state and $\bm{a}_{i}$ as an action. By introducing the advantage, we evaluate the goodness of the choice $\bm{a}_{i}$ at each edge $i$ rather than evaluate the goodness of successive actions $\bm{a}$ by a scalar reward $r$. We expect to reduce the estimation variance of the policy gradient by the advantage.

We define the advantage of $\bm{a}_{i}$ as (by letting $\bm{a}=(\bm{s}_{i},\bm{a}_{i})$)

It is understood as follows. Given the state $\bm{s}_{i}$, the advantage of the action $\bm{a}_{i}$ is the rewards improvement over $\bm{a}=(\bm{s}_{i},\bm{0})$, where the output of the $i$th edge is zero.⁴⁴4When a zero operation (returns zero consistently) is included in the search space, we replace the advantage of the zero operation with the worst expected advantage of the other operations estimated by a moving average. The rationale for this modification is as follows. First, (6) for the zero operation is always zero, resulting in no architecture update, which is not preferred. Second, the zero operation can mostly be represented by other operations with specific weights. Hence, there is no reason to prefer it unless we consider a latency constraint. See Appendix B for further details. We replace the reward with the advantage in the policy gradient estimation in (4). If we replace the second term on the right-hand side of (6) with a single scalar independent of $\bm{s}_{i}$, the resulting policy gradient is known to be an unbiased estimator (Evans and Swartz 2000). However, a state-dependent advantage such as (6) is not unbiased in general. The following result states that the policy gradient estimate using this advantage is an unbiased estimate of (4). Its proof is presented in Section A.2.

Thus, we can use $\lvert\mathcal{I}\rvert$ different advantage values rather than a single reward value. A shortcoming of the above approach is that it requires the computation of an additional reward to compute $r((\bm{s}_{i},\bm{0}),\bm{w};V)$ for each $i\in\mathcal{I}$; hence, it requires $\lvert\mathcal{I}\rvert$ additional forward propagation to estimate (4). As our objective is to reduce the time complexity by improving the convergence speed of SP-NAS, we aim to compute an advantage with ideally no additional computational cost.

To reduce the computational cost, we approximate the advantage function (6) as (dropping $\bm{w}$ and $V$ for short)

where $\mathcal{O}_{i}$ and $\mathcal{O}_{i}^{j}$ are as in (1) and $j\in\mathcal{J}_{i}$ is the index of the selected operation at edge $i$, i.e., $\bm{a}_{i}$ is the one-hot vector with $1$ at $j$th. Also note that here we consider the situation where $\mathcal{O}_{i}$ and $\mathcal{O}_{i}^{j}$ are column vectors and the batch size equals 1 for simplicity. The above approximation follows the first-order Taylor approximation of the reward around $\bm{a}$. Compared to a naive policy gradient computation that uses only a single reward $r(\bm{a},\bm{w};V)$, the policy gradient estimate with the above approximate advantage additionally requires only one reward gradient computation, $\nabla_{\mathcal{O}_{i}}r(\bm{a},\bm{w};V)^{\mathrm{T}}\mathcal{O}_{i}^{j}(x_{Q(i)};w_{i}^{j})$. The output $\mathcal{O}_{i}^{j}(x_{Q(i)};w_{i}^{j})$ of the selected operation is computed in the forward process and $\nabla_{\mathcal{O}_{i}}r(\bm{a},\bm{w};V)$ is computed in the backward process. Therefore, our advantage involves no additional computational cost.

AdvantageNAS uses the gradient information of the reward rather than a scalar reward. Existing works on SP-NAS also use the gradient (Cai, Zhu, and Han 2019; Dong and Yang 2019). AdvantageNAS aims to reduce the estimation variance while keeping the policy gradient unbiased, whereas the previous methods simply try to incorporate the gradient information of the reward and they do not consider the estimation bias.

We know from Theorem 4.2 that the advantage (6) provides an unbiased estimator of the gradient $\nabla_{\theta}J(\theta,\bm{w})$ and from Theorem 4.1 that reducing the variance in an unbiased estimator of the gradient contributes toward improving the lower bound of the expected improvement. Since (7) is an approximation of (6), we expect that our advantage contributes toward reducing the estimation variance, resulting in faster convergence.

To understand the variance reduction effect of the proposed advantage, we consider a simple situation in which the reward is given by $\bm{r}^{\mathrm{T}}\bm{a}=\sum_{i=1}^{\lvert\mathcal{I}\rvert}\bm{r}_{i}^{\mathrm{T}}\bm{a}_{i}$ for some vector $\bm{r}=(\bm{r}_{1},\dots,\bm{r}_{\mathcal{I}})$. In this example, neither $\mathcal{V}$ nor $\bm{w}$ is introduced. In this linear objective situation, the advantages (6) and (7) are equivalent and $A_{i}(\bm{a})=\bm{r}_{i}^{\mathrm{T}}\bm{a}_{i}$. The square norm of the gradient component $\nabla_{\theta_{i}}J(\theta)$ is $\sum_{j\in\mathcal{J}_{i}}(r_{i}^{j})^{2}(\mu_{i}^{j})^{2}(1-\mu_{i}^{j})^{2}$. The estimation variance of the policy gradient $\delta_{\mathrm{adv}}$ using the advantage (6) or (7) is $\operatorname{Var}[\delta_{\mathrm{adv}}]=\sum_{j}\mu_{i}^{j}(r_{i}^{j})^{2}(1-2\mu_{i}^{j}+\lVert\mu_{i}\rVert^{2})$. Meanwhile, the estimation variance of the policy gradient $\delta_{\mathrm{base}}$ using a reward value (i.e., (4)) is $\operatorname{Var}[\delta_{\mathrm{base}}]=\operatorname{Var}[\delta_{\mathrm{adv}}]+(1-\lVert\mu_{i}\rVert^{2})\sum_{k\neq i}\sum{j}(r_{k}^{j})^{2}\mu_{k}^{j}$. In other words, the variance reduction effect for each gradient component typically scales linearly with $\lvert\mathcal{I}\rvert$ (because $\sum_{k\neq i}$), whereas $\lVert\nabla_{\theta_{i}}J(\theta)\rVert^{2}$ and $\operatorname{Var}[\delta_{\mathrm{adv}}]$ do not. In light of Theorem 4.1, the proposed advantage potentially contributes toward improving the search performance, especially when the number of edges is $\lvert\mathcal{I}\rvert\gg 1$.