---
title: 'OptiProxy-NAS: Differentiable Neural Architecture Search'
url: https://www.emergentmind.com/topics/optiproxy-nas
type: topic
---

# OptiProxy-NAS: Differentiable Neural Architecture Search

OptiProxy-NAS is a neural architecture search framework that introduces an “optimization proxy” to turn the inherently discrete NAS problem into a continuous, differentiable, and smooth optimization problem, thereby enabling end-to-end gradient-based search without supernets and with direct integration of real-world objectives such as accuracy and latency. In the formulation introduced by "OptiProxy-NAS: Optimization Proxy based End-to-End Neural Architecture Search" [2509.05656], the proxy is not a reduced dataset, a zero-cost score, or a proxy device, but a continuous proxy representation of architecture distributions over operations and topology. The reported evaluation covers 12 NAS tasks of 4 search spaces across three different domains including computer vision, natural language processing, and resource-constrained NAS, and further experiments on low-fidelity scenarios are used to verify flexibility [2509.05656].

## 1. Discrete NAS and the motivation for an optimization proxy

Standard NAS is posed as a bilevel optimization problem over a discrete architecture space \(\mathbb{S}\):
$$
\begin{gathered}
\max_{\mathcal{A} \in \mathbb{S}} \operatorname{Eval}(\langle\mathcal{A}, w^{*}(\mathcal{A})\rangle; \mathbb{D}_\mathtt{val}) \\
\text{s.t.} \quad w^{*}(\mathcal{A})=\arg \min_{w} \mathcal{L}(w(\mathcal{A}); \mathbb{D}_\mathtt{train}),
\end{gathered}
$$
where \(\mathcal{A}\) is an architecture, \(\mathbb{S}\) the architecture space, \(w(\mathcal{A})\) its weights, \(\mathbb{D}_\mathtt{train}\) and \(\mathbb{D}_\mathtt{val}\) are train and validation sets, and \(\operatorname{Eval}\) is usually the validation accuracy. The search space is characterized as discrete, vast, and spiky or isolated: NAS-Bench-201 has \(6^5 = 15{,}625\) cells, and NAS-Bench-301 has cardinality up to \(\sim 10^{53}\) [2509.05656].

Within this setting, the paper distinguishes two prevalent acceleration strategies. Predictor-based NAS trains a surrogate \(\hat{f}_\theta(\mathcal{A})\) to approximate \(f(\mathcal{A})\), but the sampler still operates in the discrete space and has no gradient access through the surrogate. Differentiable NAS via supernets relaxes architecture parameters inside a supernetwork and optimizes both weights and architecture by gradient descent on the supernet loss, but it requires training a huge supernet, is often costly in memory and computation, and ties the search objective to differentiable quantities such as cross-entropy. OptiProxy-NAS is proposed for the gap between these regimes: it seeks gradient-based optimization over architecture choices without training a supernet, while directly optimizing real-world metrics, including non-differentiable ones, through a learned proxy model [2509.05656].

## 2. Proxy-space formulation and architecture parameterization

The central construction is a proxy space \(\mathcal{X}\) defined by continuous parameters over operations and topology. Architectures are parameterized by operation features \(\mathcal{A}_o\) and topology \(\mathcal{A}_t\), and OptiProxy-NAS introduces continuous logits
\(\alpha \in \mathbb{R}^{N \times M}\) for operations and \(\beta \in \mathbb{R}^{N \times N}\) for edges. These induce distributions
$$
\mathcal{A}_o^i \sim \mathrm{Cat}(\text{softmax}(\alpha^i)),
$$
and
$$
\mathcal{A}_t^{(h,k)} \sim \mathrm{Bernoulli}(\sigma(\beta^{(h,k)})).
$$
The resulting proxy space is
$$
\mathcal{X} := \{(\alpha,\beta) : \alpha \in \mathbb{R}^{N \times M}, \beta \in \mathbb{R}^{N \times N}\},
$$
which is described as continuous, differentiable, and unconstrained [2509.05656].

This yields a continuous reformulation of predictor-based NAS. The paper writes the upper-level search objective as
$$
\begin{aligned}
& \max_{(\alpha,\beta)\in\mathcal{X}} \hat{f}\big(s(\alpha,\beta), \theta^*\big) \\
\text{s.t.}\quad & \mathcal{D} = \{(\mathcal{A}_b, f(\mathcal{A}_b))\}_{b=1}^{B},\quad \mathcal{A}_b \sim s(\alpha,\beta),
\end{aligned}
$$
where \(\hat{f}\) is the proxy model, \(s\) is the discrete sampling function, and \(\theta^*\) denotes fitted proxy parameters. This changes the role of the proxy from post hoc ranking to a differentiable object that mediates the optimization itself. In the authors’ terminology, the pair \((\hat{f}, \hat{s})\) together with \((\alpha,\beta)\) constitutes the optimization proxy [2509.05656].

## 3. Differentiable relaxation and proxy gradients

Because the discrete sampler \(s(\alpha,\beta)\) is non-differentiable, OptiProxy-NAS introduces a proxy sampling function \(\hat{s}\) based on Gumbel-softmax for operation selection and Binary Concrete for topology. For operation \(j\) at node \(i\),
$$
\hat{s}_{o}^{(i,j)}
= \frac{\exp \big( (\alpha^{(i,j)} + g^{(i,j)}) / \tau \big)}
       {\sum_{m=0}^{M-1} \exp \big( (\alpha^{(i,m)} + g^{(i,m)}) / \tau \big)},
$$
and for topology,
$$
\hat{s}_{t}^{(h,k)} = \sigma \left( (\beta^{(h,k)} + g^{(h,k)}) / \tau \right),
$$
where \(g\) are Gumbel random variables and \(\tau > 0\) is a temperature parameter [2509.05656].

The objective at fixed \(\theta^*\) becomes
$$
\psi(\alpha,\beta) = \hat{f}(\hat{s}(\alpha,\beta), \theta^*),
$$
and the non-differentiable gradients with respect to \(s\) are replaced by proxy gradients
$$
\nabla_\alpha \psi \approx \frac{\partial \hat{f}}{\partial \hat{s}} \cdot \frac{\partial \hat{s}}{\partial \alpha},
\qquad
\nabla_\beta \psi \approx \frac{\partial \hat{f}}{\partial \hat{s}} \cdot \frac{\partial \hat{s}}{\partial \beta}.
$$
As \(\tau \to 0\), the relaxation is stated to be unbiased: \(\hat{s}_o\) converges in distribution to the categorical distribution defined by \(\alpha\), and \(\hat{s}_t\) converges to the Bernoulli distribution defined by \(\beta\). The framework therefore preserves a stochastic link to the original discrete search space while supporting ordinary gradient-based optimizers such as Adam [2509.05656].

A practical consequence is that non-differentiable deployment metrics can be handled indirectly through the proxy model. The paper emphasizes that this differs from supernet-based differentiable NAS, where the architecture optimization is tied to differentiable losses of the supernetwork, and from black-box predictor-based NAS, where the predictor does not provide end-to-end gradients to the sampler [2509.05656].

## 4. Proxy model, SMBO loop, and reference implementation

The reference OptiProxy-NAS implementation uses a graph convolutional network as the proxy model. The GCN consumes a DAG representation with node features \(V_0\) and adjacency matrix \(A\), and one update is written as
$$
V_{l+1} = \mathcal{LN}\left(
\mathrm{Linear}\big(
\mathrm{concat}(
\mathrm{GeLU}(A V_l W_l) +
\mathrm{GeLU}(A^\top V_l W_l)
)\big)\right).
$$
Training uses a combination of mean squared error and ranking loss,
$$
\mathcal{L} = \mathcal{L}_{\mathrm{mse}} + \lambda \mathcal{L}_{\mathrm{rank}},
$$
with pairwise ranking loss
$$
\mathcal{L}_{\mathrm{rank}}(\{(\mathcal{A}_j,y_j)\})
= \sum_{j} \sum_{i: y_i>y_j} \max\big(0, m - (\hat{f}(\mathcal{A}_i) - \hat{f}(\mathcal{A}_j))\big),
$$
where \(y_i=f(\mathcal{A}_i)\), \(m\) is a margin, and \(\lambda\) is a weighting factor [2509.05656].

Search is organized as a sequential model-based optimization loop. Initialization samples architectures and evaluates them under the true metric \(f\). Each SMBO round then fits \(\hat{f}_\theta\), runs gradient-based search in proxy space for several groups of \((\alpha,\beta)\), samples discrete candidates from the optimized distributions, selects the most promising candidates using the proxy, evaluates them under the true objective, and augments the dataset \(\mathcal{D}\). The paper also uses multiple \((\alpha_i,\beta_i)\) groups initialized by Latin Hypercube Sampling, Gumbel temperature decay for exploration–exploitation scheduling, and a verify ratio \(Q/B\) for proxy-based preselection before expensive evaluation [2509.05656].

The reported hyperparameters make the proxy deliberately lightweight. The reference GCN uses `gcn_hidden = 144`, `gcn_layers = 2`, `linear_size = 144`, batch size \(7\), AdamW with learning rate \(0.001\), and \(100\) training epochs per update. Search uses `search_epochs = 300`, `base_temp = 0.7`, `min_temp = 0.2`, `lr_alpha = 0.02`, `lr_beta = 0.001`, `parallel_batch = 5`, `num_sample = 10`, and `verify_ratio = 20`. The GCN proxy has about \(139.7\text{K}\) parameters and about \(1.3\) MFLOPs. On NAS-Bench-201 CIFAR-10, training \(100\) architectures is reported to take more than \(15\) GPU-hours and \(4\)GB RAM, whereas the extra cost of the proxy is about \(20\) seconds and \(10\)MB, corresponding to \(0.037\%\) runtime and \(0.25\%\) memory [2509.05656].

## 5. Empirical behavior across search spaces, domains, and fidelities

The empirical study covers NAS-Bench-201, NAS-Bench-101, NAS-Bench-301, NAS-Bench-NLP, and HW-NAS-Bench, summarized in the paper as 12 NAS tasks of 4 search spaces across computer vision, natural language processing, and resource-constrained NAS [2509.05656]. On NAS-Bench-201, OptiProxy-NAS reaches the global optimum on CIFAR-10 and CIFAR-100 with \(100\) queries, and on ImageNet16-120 it reaches the optimum validation accuracy \(46.77\%\) with \(280\) queries. On NAS-Bench-101, the reported best result is \(94.98 \pm 0.16\) validation accuracy with \(142\) queries. On NAS-Bench-301, the method achieves \(94.97 \pm 0.05\) validation accuracy with \(100\) queries, and on NAS-Bench-NLP it reports \(96.18 \pm 0.20\) with \(150\) queries while also reaching \(96.08 \pm 0.19\) with only \(75\) queries [2509.05656].

The hardware-aware experiments use HW-NAS-Bench over devices and latency constraints including EdgeGPU, Raspi4, EdgeTPU, Pixel3, Eyeris, and FPGA. Under a budget of \(200\) queries, the method is reported to consistently match or surpass the published best results, and with only \(100\) queries it often still outperforms AG-Net and DiNAS. One representative result is EdgeGPU with constraint \(4\), where OptiProxy-NAS reports \(44.56 \pm 0.23\) with \(200\) queries and \(44.15 \pm 0.63\) with \(100\) queries [2509.05656].

The paper also evaluates low-fidelity scenarios by replacing the true metric \(f\) during search with short-training validation accuracy, such as \(5\)-epoch training on NAS-Bench-201, and then fully training the selected architectures afterward. To reach similar final validation accuracy, the reported total search cost drops from \(6000\) to \(3200\) training epochs on CIFAR-10 for a mean validation accuracy of \(91.46\%\), from \(20000\) to \(3500\) on CIFAR-100 for \(73.49\%\), and from \(20000\) to \(7250\) on ImageNet16-120 for \(46.65\%\). The authors use these results to argue that the framework is not tied to a single fidelity level and can absorb low-fidelity objectives without changing its basic structure [2509.05656].

## 6. Relation to earlier proxy paradigms and recognized limitations

The term “proxy” in NAS predates OptiProxy-NAS and has been used for several distinct mechanisms. Prior work treated a proxy as a reduced dataset that preserves relative architecture ranking [1911.09322], a computationally reduced training setting over channels, resolution, sample ratio, and epochs [2001.01233], a representative subset of target data for differentiable NAS [2106.04784], a zero-cost score or learned zero-cost performance predictor [2210.03230; 2404.00271; 2411.15290], a robustness-oriented zero-cost proxy [2306.05031; 2307.09365], a hardware proxy device for latency ranking [2111.01203], or a stabilized share-parameter proxy in one-shot topological NAS [2005.10511]. ProxylessNAS is an especially important contrast case: despite its name, it was designed to remove proxy tasks and search directly on the target task and hardware [1812.00332]. Within that broader lineage, OptiProxy-NAS gives “proxy” a narrower and more formal meaning: a continuous proxy representation that makes the search itself differentiable and smooth [2509.05656].

This difference has methodological consequences. Earlier data-proxy work focused on preserving the relative ranking of architectures under smaller datasets [1911.09322; 2106.04784], economical proxy settings emphasized high rank consistency under low-fidelity training [2001.01233], and zero-cost proxy research showed that multiple proxies capture complementary information and can substantially improve surrogate quality when combined [2210.03230]. OptiProxy-NAS instead keeps a learned predictor but internalizes it into the optimization loop, so gradients flow through the proxy model and the relaxed architecture variables rather than through a supernet [2509.05656].

The paper nevertheless recognizes several limitations. It assumes that the proxy model can approximate the architecture–performance landscape well enough; if the mapping is extremely noisy or irregular, the proxy may be misleading. There remains a nontrivial projection from soft architectures to discrete ones, and validity constraints such as NAS-Bench-101 graph validity and isomorphism require filtering and deduplication. The current framework uses purely gradient-based optimization over \((\alpha,\beta)\), and the authors identify future work in parallelization, multi-fidelity evaluation, learning curve extrapolation, weight-sharing, and zero-cost metrics [2509.05656]. These constraints place OptiProxy-NAS within the broader movement toward proxy-aware NAS rather than outside it: the framework changes what a proxy is allowed to be, but not the continuing requirement that the proxy remain informative enough to guide search.

Source: https://www.emergentmind.com/topics/optiproxy-nas