Architecture Gradient in Differentiable NAS

• Architecture Gradient is the derivative of a NAS objective with respect to continuous architecture parameters, transforming discrete choices into a differentiable space.
• It employs various techniques such as first-order approximations, mirror descent, and natural-gradient methods to update architecture parameters efficiently.
• Applied in methods like DARTS and Darts-Conformer, it guides operation selection and routing in networks, balancing training speed with generalization.

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) (Shi et al., 2021, Zhang et al., 2021, Shu et al., 2022).

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$ (Shi et al., 2021).

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 (Zhang et al., 2021).

A first-order simplification sets $w$0, so that $w$1 is used directly for updating $w$2. Darts-Conformer reports using exactly this first-order variant in ASR, together with a warm-up period in which $w$3 is frozen and only $w$4 is updated to avoid early collapse on speech data (Shi et al., 2021).

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

$w$5

where $w$6 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 $w$7 (Li et al., 2020).

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 $w$8 with operation set $w$9, the relaxed weight of operation $\alpha$0 is

$\alpha$1

and the destination node aggregates all incoming mixed operations. Back-propagation through this linear combination yields $\alpha$2 directly (Zhang et al., 2021, Li et al., 2020).

A stochastic relaxation replaces deterministic mixtures by a distribution $\alpha$3 over architectures. In that setting,

$\alpha$4

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 (Li et al., 2020).

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 $\alpha$5 or number of heads $\alpha$6 versus $\alpha$7 (Shi et al., 2021).

Weight entanglement extends architecture gradients to macro-level spaces where storing an independent tensor for every candidate is infeasible. For cell $\alpha$8, only the largest operator weights $\alpha$9 are stored, and sub-operations are realized by slicing or zero-padding that tensor. The effective superposed weight is

$$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$0

so the gradient with respect to $$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$1 is obtained by back-propagating through the superposed tensor and collecting the slice-specific gradient:

$$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$2

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 (Sukthanker et al., 2023).

A still broader stochastic relaxation is used by the adaptive stochastic natural gradient method, which introduces a parameterized family of distributions $$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$3 over architecture choices and optimizes

$$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$4

Its architecture gradient is the score-function derivative

$$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$5

which explicitly treats architectural search as optimization over a distribution family rather than only over relaxed edge weights (Akimoto et al., 2019).

3. Optimization geometries for architecture updates

Once architecture gradients are available, the central design choice is how to update $$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$6. 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^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$7 and $$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$8, with learning rates $$w^*(\alpha)=\arg\min_w L_{\mathrm{train}}(w,\alpha), \qquad \alpha^*=\arg\min_\alpha L_{\mathrm{val}}(w^*(\alpha),\alpha).$$9 and $w^*(\alpha)$0, a batch size of $w^*(\alpha)$1, and a three-epoch warm-up before joint bilevel updates (Shi et al., 2021).

Mirror descent replaces Euclidean geometry by a geometry adapted to the feasible set of architecture parameters. For a distance-generating function $w^*(\alpha)$2, the generic update is

$w^*(\alpha)$3

When $w^*(\alpha)$4 is a simplex and $w^*(\alpha)$5, the update becomes exponentiated gradient:

$w^*(\alpha)$6

This KL-geometry encourages sparsity in $w^*(\alpha)$7 and reduces the relaxation-to-discrete gap when pruning at the end of search (Li et al., 2020).

Natural gradient methods go further by respecting the geometry of a distribution family over architectures. The stochastic natural gradient is

$w^*(\alpha)$8

where $w^*(\alpha)$9 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 (Akimoto et al., 2019).

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, $\alpha$0 is optimized while $\alpha$1 remains frozen. During exploration, an Improved Competitive Swarm Optimizer updates candidate architectures in continuous architecture space, while gradient descent concurrently updates $\alpha$2; the best swarm member $\alpha$3 then pulls the shared architecture parameters via

$\alpha$4

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 (Du et al., 1 Apr 2026).

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$5.

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 $\alpha$6 and initialization $\alpha$7, representative metrics include the gradient norm

$\alpha$8

SNIP,

$\alpha$9

GraSP, which introduces a Hessian term, and NTK-based quantities such as the trace norm of the finite-width neural tangent kernel (Shu et al., 2022).

The unifying theory shows that these metrics all depend on $w^*(\alpha)$0. Under a $w^*(\alpha)$1-Lipschitz and $w^*(\alpha)$2-smooth loss, Theorem 1 proves that there exist constants $w^*(\alpha)$3 such that, with high probability,

$w^*(\alpha)$4

The result interprets apparently different zero-cost signals as correlated views of a single underlying NTK complexity (Shu et al., 2022).

The same work derives a generalization bound for networks trained by gradient descent:

$w^*(\alpha)$5

where $w^*(\alpha)$6 is the condition number of the NTK and $w^*(\alpha)$7 can be any of the gradient-based zero-cost metrics. In the realizable regime this yields the prescription to maximize $w^*(\alpha)$8; in the non-realizable regime it yields a trade-off between training speed and generalization, motivating the unified objective

$w^*(\alpha)$9

Hybrid NAS (HNAS) implements this objective by tuning only the two scalars $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$0 with a Gaussian-process surrogate while selecting architectures from zero-cost scores and validating them with short training runs (Shu et al., 2022).

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 $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$1; NodeMHA chooses two from $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$2; NodeCNN chooses two from $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$3; NodeFF takes CNN as its single predecessor (Shi et al., 2021).

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 $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$4 epochs, then a joint bilevel search runs until epoch $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$5, updating $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$6 on training loss and $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$7 on validation loss for each minibatch of size $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$8. 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 $\nabla_w L_{\mathrm{train}}(w^*(\alpha),\alpha)=0$9 epochs on AISHELL-1 and $$\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}}.$$0 epochs on Sogou-2k (Shi et al., 2021).

The compute budget is explicitly modest: single-GPU Titan Xp search time is $$\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}}.$$1 hours, approximately $$\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}}.$$2 GPU-days. After discretization, the searched cell preserves the original Conformer ordering but adds learned skip pathways: two parallel CNN paths with kernel size $$\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}}.$$3, 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 (Shi et al., 2021).

Empirically, the baseline six-layer Conformer achieves $$\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}}.$$4 CER on AISHELL-1 test, while the six-layer Darts-Conformer achieves $$\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}}.$$5, a $$\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}}.$$6 relative reduction in CER. On Sogou-2k, the Conformer baseline reports $$\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}}.$$7 CER on clean, noisy, and far-field subsets, whereas Darts-Conformer reports $$\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}}.$$8. The paper identifies this as the first successful attempt to apply gradient-based architecture search in the attention-based encoder-decoder ASR model (Shi et al., 2021).

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 (Du et al., 1 Apr 2026). 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 (Akimoto et al., 2019).

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 (Sukthanker et al., 2023).

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, $$\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}}.$$9, over the same operator and target language, and shows that across over $w' = w-\xi \nabla_w L_{\mathrm{train}}(w,\alpha)$0 training runs one structure recovers a target at $w' = w-\xi \nabla_w L_{\mathrm{train}}(w,\alpha)$1 while another scores $w' = w-\xi \nabla_w L_{\mathrm{train}}(w,\alpha)$2, 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 (Gupta, 25 Apr 2026).

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.