---
title: Architecture Gradient in Differentiable NAS
url: https://www.emergentmind.com/topics/architecture-gradient
type: topic
---

# Architecture Gradient in Differentiable NAS

Architecture gradient denotes the derivative of an architecture-selection objective with respect to continuous architecture parameters, typically written as $\nabla_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha)$ in differentiable neural architecture search (NAS). It arises when a discrete search space of operations, edges, blocks, or macro-level design choices is relaxed into a continuous parameterization so that architecture choices can be updated by gradient-based optimization jointly with network weights. In the literature, the term encompasses classical bilevel derivatives used in DARTS-style methods, geometry-aware and natural-gradient updates on relaxed architecture manifolds, gradient-based zero-cost proxies evaluated at initialization, and extensions to weight-entangled and domain-specific search spaces such as end-to-end automatic speech recognition (ASR) [2104.02868] [2103.12529] [2201.09785].

## 1. Formal definition and objective structure

The canonical formulation treats network weights $w$ and architecture parameters $\alpha$ asymmetrically. In DARTS-style bilevel NAS, the lower level fits weights on training loss and the upper level chooses architectures by validation loss:
$$
w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad
\alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).
$$
Under this formulation, the architecture gradient is the derivative of the validation-level objective through the implicit dependence of $w^*(\alpha)$ on $\alpha$ [2104.02868].

An exact expression follows from the chain rule and implicit differentiation. With $w^*(\alpha)$ satisfying $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$, the derivative can be written as
$$
\nabla_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha)
=
\partial_\alpha L_{\mathrm{val}}
-
\partial_w L_{\mathrm{val}}
\,[\partial^2_{w,w}L_{\mathrm{train}}]^{-1}
\partial^2_{w,\alpha}L_{\mathrm{train}}.
$$
Because the Hessian inverse is expensive, differentiable NAS typically replaces it with a one-step approximation $w' = w-\xi \nabla_w L_{\mathrm{train}}(w,\alpha)$ and computes an approximate upper-level gradient instead [2103.12529].

A first-order simplification sets $\xi=0$, so that $\nabla_\alpha L_{\mathrm{val}}(w',\alpha)$ is used directly for updating $\alpha$. Darts-Conformer reports using exactly this first-order variant in ASR, together with a warm-up period in which $\alpha$ is frozen and only $w$ is updated to avoid early collapse on speech data [2104.02868].

A different line of analysis argues that weight-sharing NAS can be studied as a single-level empirical risk minimization problem over both shared weights and continuous architecture parameters. In this view, one optimizes
$$
L(w,\alpha) := \frac{1}{|T|}\sum_{(x,y)\in T}\ell(h_{w,\alpha}(x),y),
\qquad
\min_{w\in\mathbb R^d,\alpha\in\Theta}L(w,\alpha),
$$
where $\Theta$ is a continuous relaxation of a discrete search space, often a Cartesian product of simplices. This perspective shifts attention from only approximating the bilevel derivative to designing optimizers and regularizers appropriate for the geometry of $\alpha$ [2004.07802].

## 2. Continuous relaxations and computable parameterizations

Architecture gradients become computable only after a discrete search space is embedded into a differentiable parameterization. In the most common mixture relaxation, each directed edge chooses a convex combination of candidate operations. For an edge $(i,j)$ with operation set $O$, the relaxed weight of operation $o$ is
$$
\bar w_{o}^{(i,j)}
=
\frac{\exp(\alpha_{o}^{(i,j)})}{\sum_{o'\in O}\exp(\alpha_{o'}^{(i,j)})},
$$
and the destination node aggregates all incoming mixed operations. Back-propagation through this linear combination yields $\partial L/\partial \alpha_{e,o}$ directly [2103.12529] [2004.07802].

A stochastic relaxation replaces deterministic mixtures by a distribution $p_\alpha$ over architectures. In that setting,
$$
\nabla_\alpha \mathbb E_{a\sim p_\alpha}[\ell(h_{w,a},y)]
=
\mathbb E_a[\ell(h_{w,a},y)\nabla_\alpha \log p_\alpha(a)],
$$
possibly combined with a straight-through estimator or Gumbel-softmax. This makes architecture gradients applicable to arbitrary search spaces rather than only simplex-valued mixtures [2004.07802].

Darts-Conformer illustrates how a domain architecture can be recast into a differentiable search space. A standard Conformer layer contains a half-step feed-forward module, multi-head self-attention, a convolution module, and a final feed-forward layer. The search method turns each sub-module into a node in a directed acyclic graph and inserts an Op-mutator on each allowable edge that softly chooses among the original Conformer operation, a zero operation, and a skip connection. The resulting Darts-Conformer cell learns both which paths to use and which concrete operation on each path, such as convolution kernel size in $\{15,23,31\}$ or number of heads $2$ versus $4$ [2104.02868].

Weight entanglement extends architecture gradients to macro-level spaces where storing an independent tensor for every candidate is infeasible. For cell $j$, only the largest operator weights $W_{\max_j}$ are stored, and sub-operations are realized by slicing or zero-padding that tensor. The effective superposed weight is
$$
W_{\mathrm{sup}_j}
=
\sum_{k=1}^N f(\alpha_{j,k})\cdot \mathrm{Slice}(W_{\max_j};\text{size of }op_k),
$$
so the gradient with respect to $\alpha_{j,k}$ is obtained by back-propagating through the superposed tensor and collecting the slice-specific gradient:
$$
\frac{\partial L_{\mathrm{val}}}{\partial \alpha_{j,k}}
=
f'(\alpha_{j,k})\cdot
\left\langle
\frac{\partial L_{\mathrm{val}}}{\partial W_{\mathrm{sup}_j}},
S_{j,k}
\right\rangle .
$$
This construction adapts gradient-based NAS to kernel sizes, channel widths, expansion ratios, transformer heads, and depths while preserving the memory efficiency of entangled supernets [2312.10440].

A still broader stochastic relaxation is used by the adaptive stochastic natural gradient method, which introduces a parameterized family of distributions $\{P_\theta:c\sim p_\theta(c)\}$ over architecture choices and optimizes
$$
J(x,\theta)=\mathbb E_{c\sim p_\theta}[f(x,c)].
$$
Its architecture gradient is the score-function derivative
$$
\nabla_\theta J(x,\theta)
=
\mathbb E_{p_\theta}[f(x,c)\nabla_\theta \ln p_\theta(c)],
$$
which explicitly treats architectural search as optimization over a distribution family rather than only over relaxed edge weights [1905.08537].

## 3. Optimization geometries for architecture updates

Once architecture gradients are available, the central design choice is how to update $\alpha$. Euclidean gradient descent or Adam is the most direct option and remains common in bilevel NAS. Darts-Conformer, for example, uses Adam for both $w$ and $\alpha$, with learning rates $\eta_w=3\times 10^{-4}$ and $\eta_\alpha=2\times 10^{-4}$, a batch size of $48$, and a three-epoch warm-up before joint bilevel updates [2104.02868].

Mirror descent replaces Euclidean geometry by a geometry adapted to the feasible set of architecture parameters. For a distance-generating function $\phi$, the generic update is
$$
\alpha^{t+1}
=
\arg\min_{u\in\Theta}
\left\langle \nabla_\alpha L(w^t,\alpha^t), u-\alpha^t\right\rangle
+\frac{1}{\eta}D_\phi(u\|\alpha^t).
$$
When $\Theta$ is a simplex and $\phi(u)=\sum_i u_i\log u_i$, the update becomes exponentiated gradient:
$$
\tilde \alpha_i=\alpha_i^t \exp\!\bigl(-\eta[\nabla_\alpha L(w^t,\alpha^t)]_i\bigr),
\qquad
\alpha_i^{t+1}=\frac{\tilde\alpha_i}{\sum_j \tilde\alpha_j}.
$$
This KL-geometry encourages sparsity in $\alpha$ and reduces the relaxation-to-discrete gap when pruning at the end of search [2004.07802].

Natural gradient methods go further by respecting the geometry of a distribution family over architectures. The stochastic natural gradient is
$$
\tilde\nabla_\theta J = F(\theta)^{-1}\nabla_\theta J,
$$
where $F(\theta)$ is the Fisher information. In practice, the method estimates the gradient from sampled architectures and uses an adaptive step-size mechanism based on a target signal-to-noise ratio, so that the trust-region radius is adjusted without manual tuning. The paper emphasizes robustness to initial step size and search space, addressing a recurrent difficulty of NAS methods that are highly sensitive to such inputs [1905.08537].

Hybrid search schemes combine architectural gradients with population-based exploration. G-ICSO-NAS separates the search into a Warm-up Phase, an Exploration Phase, and a Stability Phase. During warm-up, $w$ is optimized while $\alpha$ remains frozen. During exploration, an Improved Competitive Swarm Optimizer updates candidate architectures in continuous architecture space, while gradient descent concurrently updates $w$; the best swarm member $x^*$ then pulls the shared architecture parameters via
$$
\alpha_{t+1}=\alpha_t+\eta_\alpha (x^*-\alpha_t).
$$
The final Stability Phase returns to pure gradient-based bilevel steps with a greatly reduced architecture learning rate and early stopping based on Hoeffding’s test over architecture-step norms [2604.00703].

These variants show that “architecture gradient” does not refer to a single optimizer. It refers to a derivative signal defined on a relaxed architecture space, while the actual update may be Euclidean, mirror-descent, natural-gradient, or hybrid swarm-gradient. This suggests that the computational meaning of the gradient is inseparable from the geometry imposed on $\alpha$.

## 4. Training-free architecture gradients and zero-cost theory

A common misconception is that architecture gradients require supernet training. Training-free NAS demonstrates a different regime: one computes gradient-based proxies from a network at random initialization using a single forward and backward pass. For a dataset $\mathcal S=\{(x_i,y_i)\}_{i=1}^m$ and initialization $\theta_0$, representative metrics include the gradient norm
$$
G_{\mathrm{Grad}}
=
\left\|
\frac1m\sum_{i=1}^m
\nabla_\theta \ell(f(x_i;\theta_0),y_i)
\right\|_2,
$$
SNIP,
$$
G_{\mathrm{SNIP}}
=
\left|
\frac1m\sum_{i=1}^m
\theta_0^\top
\nabla_\theta \ell(f(x_i;\theta_0),y_i)
\right|,
$$
GraSP, which introduces a Hessian term, and NTK-based quantities such as the trace norm of the finite-width neural tangent kernel [2201.09785].

The unifying theory shows that these metrics all depend on $\nabla_\theta f(x_i;\theta_0)$. Under a $\beta$-Lipschitz and $\gamma$-smooth loss, Theorem 1 proves that there exist constants $C_1,C_2,C_3>0$ such that, with high probability,
$$
G_{\mathrm{Grad}}\le C_1 G_{\mathrm{Trace}},\qquad
G_{\mathrm{SNIP}}\le C_2 G_{\mathrm{Trace}},\qquad
G_{\mathrm{GraSP}}\le C_3 G_{\mathrm{Trace}}.
$$
The result interprets apparently different zero-cost signals as correlated views of a single underlying NTK complexity [2201.09785].

The same work derives a generalization bound for networks trained by gradient descent:
$$
L_{\mathcal D}(f_t)\le L_{\mathcal S}(f_t)+\frac{\kappa}{G},
$$
where $\kappa=\lambda_{\max}(\Theta_0)/\lambda_{\min}(\Theta_0)$ is the condition number of the NTK and $G$ can be any of the gradient-based zero-cost metrics. In the realizable regime this yields the prescription to maximize $G$; in the non-realizable regime it yields a trade-off between training speed and generalization, motivating the unified objective
$$
\min_\alpha \frac{\kappa(\alpha)}{G(\alpha)}+\mu\,F(G(\alpha)^2-\nu).
$$
Hybrid NAS (HNAS) implements this objective by tuning only the two scalars $(\mu,\nu)$ with a Gaussian-process surrogate while selecting architectures from zero-cost scores and validating them with short training runs [2201.09785].

The training-free perspective broadens the meaning of architecture gradient. Here the “gradient” is not the derivative of validation loss through a trained supernet, but a differentiable statistic of the initialized network that correlates with future trainability and generalization. This suggests that architecture gradient can function either as an update direction or as a ranking signal.

## 5. Domain-specific instantiation: end-to-end ASR

The most explicit domain transfer of architecture gradients in the supplied literature is Darts-Conformer for end-to-end ASR. The method fuses a DARTS mutator with Conformer blocks and searches directly in an attention-based encoder-decoder ASR model. The search space includes node-specific input choices: NodeMAC chooses one input from $\{\text{raw-input-0},\text{raw-input-1}\}$; NodeMHA chooses two from $\{\text{raw-0},\text{raw-1},\text{MAC}\}$; NodeCNN chooses two from $\{\text{raw-1},\text{MAC},\text{MHA}\}$; NodeFF takes CNN as its single predecessor [2104.02868].

The search-and-train procedure has two stages. During search, a supernet is built by stacking Darts-Conformer cells with Op-mutators. The architecture parameters are frozen for $e_w=3$ epochs, then a joint bilevel search runs until epoch $30$, updating $w$ on training loss and $\alpha$ on validation loss for each minibatch of size $48$. The loss used in both search and training is a mixture of Connectionist Temporal Classification and cross-entropy; the final training stage fixes the discovered cell, stacks it six times in the encoder, and trains the full encoder-decoder model from scratch for $65$ epochs on AISHELL-1 and $30$ epochs on Sogou-2k [2104.02868].

The compute budget is explicitly modest: single-GPU Titan Xp search time is $16$ hours, approximately $0.7$ GPU-days. After discretization, the searched cell preserves the original Conformer ordering but adds learned skip pathways: two parallel CNN paths with kernel size $15$, one in the original CNN slot and one from the raw input; two parallel MHAs with four heads each, one in the usual MHA slot and one attending directly to the raw input; and a final FFN that consumes the sum of the CNN outputs [2104.02868].

Empirically, the baseline six-layer Conformer achieves $6.4\%$ CER on AISHELL-1 test, while the six-layer Darts-Conformer achieves $6.1\%$, a $4.7\%$ relative reduction in CER. On Sogou-2k, the Conformer baseline reports $10.9\% / 7.5\% / 14.0\%$ CER on clean, noisy, and far-field subsets, whereas Darts-Conformer reports $10.6\% / 7.2\% / 14.0\%$. The paper identifies this as the first successful attempt to apply gradient-based architecture search in the attention-based encoder-decoder ASR model [2104.02868].

This case clarifies what architecture gradients contribute in practice: they do not merely prune operations, but can alter the internal routing pattern of a hand-designed module while retaining the module’s global inductive bias.

## 6. Failure modes, misconceptions, and landscape effects

Several recurring failure modes are attached to architecture gradients. Gradient-based NAS in the DARTS family is described as efficient but prone to premature convergence and performance collapse, motivating hybrid methods such as G-ICSO-NAS that inject diversity-aware swarm exploration before returning to fine-grained gradient search [2604.00703]. A related concern is sensitivity to learning rate and search space, which the adaptive stochastic natural gradient method identifies as a practical obstacle to applying NAS out-of-the-box [1905.08537].

Another misconception is that gradient-based NAS is restricted to classical cell-based weight-sharing spaces. Weight entanglement shows otherwise. TangleNAS adapts gradient-based methods such as DrNAS and DARTS to macro search spaces by representing all options in each cell with a single largest tensor and differentiating through slices. The empirical claim is that this preserves the memory efficiency of weight-entangled spaces while retaining the benefits of gradient-based methods, including enhanced performance, improved supernet training properties, and superior any-time performance [2312.10440].

A more fundamental misconception is that a more expressive architecture space should be easier for gradient descent to exploit. The symbolic-regression study directly contradicts that expectation. It compares three depth-3 tree architectures, $F_{V16}\subset F_{\mathrm{Hybrid}}\subset F_{\mathrm{Eq6}}$, over the same operator and target language, and shows that across over $12{,}700$ training runs one structure recovers a target at $100\%$ while another scores $0\%$, with the ranking reversed on a different target. The most expressive structure fails on targets that a restricted alternative solves reliably; balanced tree shapes are never recovered; and reversing the operator’s gradient profile collapses recovery entirely [2604.23256].

These results support an objective interpretation of architecture gradient: it is not only a derivative of a relaxed architecture parameterization, but also a descriptor of the optimization landscape induced by that parameterization. Expressiveness guarantees that a solution exists in the search space, but not that gradient descent finds it. In the supplied literature, this principle recurs across bilevel NAS, zero-cost ranking, macro-level weight entanglement, ASR cell search, and symbolic regression.

Source: https://www.emergentmind.com/topics/architecture-gradient