Differentiable Architecture Search
From discrete search to continuous relaxation
Classic Neural Architecture Search treats architecture choice as a discrete combinatorial problem — sample a candidate, train it, score it, repeat — which is why it needed thousands of GPU-hours. Differentiable Architecture Search (DARTS), introduced in 2018, reframes the problem so it can be solved with ordinary gradient descent instead. Rather than choosing one discrete operation per connection (e.g. 'this edge is a 3x3 convolution'), DARTS keeps *all* candidate operations active simultaneously, each weighted by a learnable continuous parameter, and optimizes those weights jointly with the network's normal weights.
The continuous relaxation, concretely
DARTS represents the search space as a directed acyclic graph where each edge between nodes could be any operation from a candidate set $\mathcal{O}$ (e.g. 3x3 conv, 5x5 conv, max pooling, skip connection, zero/none). Instead of picking one, the output of an edge is a softmax-weighted mixture: $\bar{o}(x) = \sum_{o \in \mathcal{O}} \frac{\exp(\alpha_o)}{\sum_{o'} \exp(\alpha_{o'})} \, o(x)$, where $\alpha_o$ are learnable architecture parameters, one per operation per edge. At the end of search, the final discrete architecture is recovered by keeping, on each edge, only the operation with the largest $\alpha_o$ — collapsing the soft mixture back into a hard choice.
import torch
import torch.nn as nn
import torch.nn.functional as F
class MixedOp(nn.Module):
"""One edge in the DARTS search graph: a weighted mix of candidate ops."""
def __init__(self, candidate_ops: list[nn.Module]):
super().__init__()
self.ops = nn.ModuleList(candidate_ops)
self.alpha = nn.Parameter(torch.randn(len(candidate_ops)) * 1e-3)
def forward(self, x):
weights = F.softmax(self.alpha, dim=0)
return sum(w * op(x) for w, op in zip(weights, self.ops))
def discretize(self):
"""After search: collapse to the single best operation."""
best_idx = torch.argmax(self.alpha).item()
return self.ops[best_idx]Bilevel optimization: two sets of parameters
DARTS optimizes two distinct parameter groups: the network's ordinary weights $w$ (conv filters, etc.) and the architecture parameters $\alpha$. Formally this is a bilevel optimization problem — find $\alpha^*$ that minimizes validation loss, where $w^*(\alpha)$ is itself the result of minimizing training loss given $\alpha$. Solving this exactly would require fully retraining $w$ for every update to $\alpha$, which is intractable. DARTS instead approximates it: alternate a gradient step on $w$ using the training set with a gradient step on $\alpha$ using the validation set, treating $w$ as roughly fixed during the $\alpha$ update. This approximation is what makes the whole method a single, relatively cheap training run instead of a search loop around many runs.
The discretization gap
Discrete NAS vs differentiable NAS
| Property | Discrete NAS (RL/evolutionary) | Differentiable NAS (DARTS) |
|---|---|---|
| Search cost | Thousands of training runs | ~1 training run on a relaxed super-network |
| Optimization method | Black-box (reward, fitness) | Gradient descent (bilevel) |
| Candidate evaluation | Train-and-score each candidate | All candidates co-exist, weighted |
| Known failure mode | Prohibitive compute cost | Discretization gap, skip-connection collapse |
What's Next?
DARTS treats the architecture itself as something to optimize with gradients; the same impulse — replacing manual tuning with a principled search procedure — applies just as naturally to the hyperparameters that sit around an architecture, like learning rates and regularization strength, which is where Bayesian optimization-style approaches take over.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.