OptiProxy-NAS: Differentiable Neural Architecture Search

• The paper introduces an optimization proxy that converts the discrete NAS problem into a continuous, differentiable search space for direct gradient-based optimization.
• It employs a proxy model with relaxed architecture parameters using Gumbel-softmax and Binary Concrete methods to integrate metrics such as accuracy and latency.
• Empirical results across multiple benchmarks show competitive performance with reduced search cost and improved flexibility for diverse domains.

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" (Lyu et al., 6 Sep 2025), 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 (Lyu et al., 6 Sep 2025).

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}$ (Lyu et al., 6 Sep 2025).

Within this setting, the paper distinguishes two prevalent acceleration strategies. Predictor-based NAS trains a surrogate $$\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}$$0 to approximate $$\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}$$1, 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 (Lyu et al., 6 Sep 2025).

2. Proxy-space formulation and architecture parameterization

The central construction is a proxy space $$\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}$$2 defined by continuous parameters over operations and topology. Architectures are parameterized by operation features $$\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}$$3 and topology $$\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}$$4, and OptiProxy-NAS introduces continuous logits $$\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}$$5 for operations and $$\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}$$6 for edges. These induce distributions

$$\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}$$7

and

$$\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}$$8

The resulting proxy space is

$$\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}$$9

which is described as continuous, differentiable, and unconstrained (Lyu et al., 6 Sep 2025).

This yields a continuous reformulation of predictor-based NAS. The paper writes the upper-level search objective as

$\mathcal{A}$0

where $\mathcal{A}$1 is the proxy model, $\mathcal{A}$2 is the discrete sampling function, and $\mathcal{A}$3 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 $\mathcal{A}$4 together with $\mathcal{A}$5 constitutes the optimization proxy (Lyu et al., 6 Sep 2025).

3. Differentiable relaxation and proxy gradients

Because the discrete sampler $\mathcal{A}$6 is non-differentiable, OptiProxy-NAS introduces a proxy sampling function $\mathcal{A}$7 based on Gumbel-softmax for operation selection and Binary Concrete for topology. For operation $\mathcal{A}$8 at node $\mathcal{A}$9,

$\mathbb{S}$0

and for topology,

$\mathbb{S}$1

where $\mathbb{S}$2 are Gumbel random variables and $\mathbb{S}$3 is a temperature parameter (Lyu et al., 6 Sep 2025).

The objective at fixed $\mathbb{S}$4 becomes

$\mathbb{S}$5

and the non-differentiable gradients with respect to $\mathbb{S}$6 are replaced by proxy gradients

$\mathbb{S}$7

As $\mathbb{S}$8, the relaxation is stated to be unbiased: $\mathbb{S}$9 converges in distribution to the categorical distribution defined by $w(\mathcal{A})$0, and $w(\mathcal{A})$1 converges to the Bernoulli distribution defined by $w(\mathcal{A})$2. The framework therefore preserves a stochastic link to the original discrete search space while supporting ordinary gradient-based optimizers such as Adam (Lyu et al., 6 Sep 2025).

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 (Lyu et al., 6 Sep 2025).

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 $w(\mathcal{A})$3 and adjacency matrix $w(\mathcal{A})$4, and one update is written as

$w(\mathcal{A})$5

Training uses a combination of mean squared error and ranking loss,

$w(\mathcal{A})$6

with pairwise ranking loss

$w(\mathcal{A})$7

where $w(\mathcal{A})$8, $w(\mathcal{A})$9 is a margin, and $\mathbb{D}_\mathtt{train}$0 is a weighting factor (Lyu et al., 6 Sep 2025).

Search is organized as a sequential model-based optimization loop. Initialization samples architectures and evaluates them under the true metric $\mathbb{D}_\mathtt{train}$1. Each SMBO round then fits $\mathbb{D}_\mathtt{train}$2, runs gradient-based search in proxy space for several groups of $\mathbb{D}_\mathtt{train}$3, 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 $\mathbb{D}_\mathtt{train}$4. The paper also uses multiple $\mathbb{D}_\mathtt{train}$5 groups initialized by Latin Hypercube Sampling, Gumbel temperature decay for exploration–exploitation scheduling, and a verify ratio $\mathbb{D}_\mathtt{train}$6 for proxy-based preselection before expensive evaluation (Lyu et al., 6 Sep 2025).

The reported hyperparameters make the proxy deliberately lightweight. The reference GCN uses gcn_hidden = 144, gcn_layers = 2, linear_size = 144, batch size $\mathbb{D}_\mathtt{train}$7, AdamW with learning rate $\mathbb{D}_\mathtt{train}$8, and $\mathbb{D}_\mathtt{train}$9 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 $\mathbb{D}_\mathtt{val}$0 parameters and about $\mathbb{D}_\mathtt{val}$1 MFLOPs. On NAS-Bench-201 CIFAR-10, training $\mathbb{D}_\mathtt{val}$2 architectures is reported to take more than $\mathbb{D}_\mathtt{val}$3 GPU-hours and $\mathbb{D}_\mathtt{val}$4GB RAM, whereas the extra cost of the proxy is about $\mathbb{D}_\mathtt{val}$5 seconds and $\mathbb{D}_\mathtt{val}$6MB, corresponding to $\mathbb{D}_\mathtt{val}$7 runtime and $\mathbb{D}_\mathtt{val}$8 memory (Lyu et al., 6 Sep 2025).

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 (Lyu et al., 6 Sep 2025). On NAS-Bench-201, OptiProxy-NAS reaches the global optimum on CIFAR-10 and CIFAR-100 with $\mathbb{D}_\mathtt{val}$9 queries, and on ImageNet16-120 it reaches the optimum validation accuracy $\operatorname{Eval}$0 with $\operatorname{Eval}$1 queries. On NAS-Bench-101, the reported best result is $\operatorname{Eval}$2 validation accuracy with $\operatorname{Eval}$3 queries. On NAS-Bench-301, the method achieves $\operatorname{Eval}$4 validation accuracy with $\operatorname{Eval}$5 queries, and on NAS-Bench-NLP it reports $\operatorname{Eval}$6 with $\operatorname{Eval}$7 queries while also reaching $\operatorname{Eval}$8 with only $\operatorname{Eval}$9 queries (Lyu et al., 6 Sep 2025).

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 $6^5 = 15{,}625$0 queries, the method is reported to consistently match or surpass the published best results, and with only $6^5 = 15{,}625$1 queries it often still outperforms AG-Net and DiNAS. One representative result is EdgeGPU with constraint $6^5 = 15{,}625$2, where OptiProxy-NAS reports $6^5 = 15{,}625$3 with $6^5 = 15{,}625$4 queries and $6^5 = 15{,}625$5 with $6^5 = 15{,}625$6 queries (Lyu et al., 6 Sep 2025).

The paper also evaluates low-fidelity scenarios by replacing the true metric $6^5 = 15{,}625$7 during search with short-training validation accuracy, such as $6^5 = 15{,}625$8-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 $6^5 = 15{,}625$9 to $\sim 10^{53}$0 training epochs on CIFAR-10 for a mean validation accuracy of $\sim 10^{53}$1, from $\sim 10^{53}$2 to $\sim 10^{53}$3 on CIFAR-100 for $\sim 10^{53}$4, and from $\sim 10^{53}$5 to $\sim 10^{53}$6 on ImageNet16-120 for $\sim 10^{53}$7. 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 (Lyu et al., 6 Sep 2025).

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 (Park, 2019), a computationally reduced training setting over channels, resolution, sample ratio, and epochs (Zhou et al., 2020), a representative subset of target data for differentiable NAS (Na et al., 2021), a zero-cost score or learned zero-cost performance predictor (Krishnakumar et al., 2022, Qiao et al., 2024, Cortês et al., 2024), a robustness-oriented zero-cost proxy (Ha et al., 2023, Lukasik et al., 2023), a hardware proxy device for latency ranking (Lu et al., 2021), or a stabilized share-parameter proxy in one-shot topological NAS (Guo et al., 2020). 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 (Cai et al., 2018). 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 (Lyu et al., 6 Sep 2025).

This difference has methodological consequences. Earlier data-proxy work focused on preserving the relative ranking of architectures under smaller datasets (Park, 2019, Na et al., 2021), economical proxy settings emphasized high rank consistency under low-fidelity training (Zhou et al., 2020), and zero-cost proxy research showed that multiple proxies capture complementary information and can substantially improve surrogate quality when combined (Krishnakumar et al., 2022). 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 (Lyu et al., 6 Sep 2025).

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 $\sim 10^{53}$8, and the authors identify future work in parallelization, multi-fidelity evaluation, learning curve extrapolation, weight-sharing, and zero-cost metrics (Lyu et al., 6 Sep 2025). 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.