LLM-Guided Neural Architecture Search for Robust Co-Design of Physical Neural Networks

AI summary

Overview

• Research area: Neural architecture search (NAS), automated machine learning, and unconventional/photonic computing hardware co-design.
• Technical level: Intermediate. Familiarity with NAS concepts (search spaces, Pareto fronts, proxies) helps, but the paper's framing is accessible.
• Scope: The paper introduces UH-NAS, a hardware-agnostic, LLM-guided NAS framework that co-optimizes MNIST classification accuracy and inference energy across CPU, GPU, and MZI-mesh optical backends, with an emphasis on robustness to optical hardware non-idealities.

What This Paper Is About

Deploying neural networks on unconventional hardware (such as silicon-photonic MZI meshes) requires architectures that trade off task accuracy against platform-specific costs like energy, physical noise, and numerical precision. Existing NAS methods are typically built for one hardware family and rely on hand-designed objectives, which makes cross-platform comparison and generalization difficult. This paper proposes UH-NAS, which treats hardware as a swappable backend and uses a large language model as an evolutionary operator to search for architectures that are both accurate and robust under hardware non-idealities.

Key Contributions

1. A hardware-agnostic NAS framework (UH-NAS). Hardware is abstracted as a swappable backend containing per-operation energy models, physical constraints, non-ideality simulators, and quantization settings, so the same search algorithm can run on CPU, GPU, and MZI-based optical hardware without modification.
2. LLMs as evolutionary operators plus a knowledge base. The LLM acts as mutation/crossover proposal generator and as a synthesizer of hardware-specific design heuristics that are updated each generation, integrated into an NSGA-II multi-objective loop with MAP-Elites-style diversity sampling.
3. Full noisy training instead of zero-cost proxies. The paper reports that 13 zero-cost proxies show negative Spearman correlations with validation accuracy under worst-case optical non-idealities, so UH-NAS scores every candidate with full training (20 epochs, Adam, quantization-aware training).
4. Physically interpretable design principles for MZI optical networks, derived from the architectures the search discovers, including preference for shallow wide MLPs, avoidance of convolutions and skip connections under high crosstalk, and the role of batch normalization as a noise-recalibration mechanism.

Main Findings

• Best accuracy under worst-case optical non-idealities: UH-NAS with GPT-4.1 reached 97.00% ± 0.30% top-1 MNIST accuracy, versus 88.69% ± 6.85% for RZ-NAS and 66.03% ± 14.51% for LLMatic under the same LLM backbone; LLMatic trained with non-ideality accuracy as its fitness criterion reached 91.81% ± 1.83%.
• LLM backbone matters: With GPT-4.1 Nano, UH-NAS scored 91.96% ± 0.43%; with Gemini 3.1 Flash-Lite, 96.95% ± 0.03%. The paper attributes weaker results with small models to weak or incorrect generated prompts, and notes Gemini 3.1 Flash-Lite tends to bias toward convolutional layers, which degrade catastrophically under non-idealities.
• Removing the LLM collapses search diversity: without LLM-guided generation, the search produced 26 unique designs out of 250 evaluations; with the LLM it produced 203 unique architectures, giving broader Pareto coverage and lower validation error at comparable energy. UH-NAS without the LLM scored 92.74%, 92.46%, and 92.61% across the three backbones.
• Zero-cost proxies break down under non-idealities: on 200 randomly drawn NAS-Bench-201/CIFAR-10 architectures, most of the 13 ZC proxies had positive correlations with accuracy under clean conditions but negative correlations under worst-case MZI conditions (σφ = 0.05, σγ = 0.005, crosstalk = 0.15), largely driven by convolutional layers.
• Robustness generalizes from worst-case training: Pareto fronts optimized under worst-case noise remained well-structured when re-evaluated at none, realistic, medium, and worst-case noise levels, and worst-case-trained architectures stayed competitive at realistic and medium noise levels (the inverse did not necessarily hold).
• System prompt ablation: removing the hardware-informed knowledge base reduced GPT-4.1 from 97.00% ± 0.30% to 96.02% ± 1.59% and Gemini 3.1 Flash-Lite from 96.95% ± 0.03% to 96.55% ± 0.42%, while GPT-4.1 Nano improved from 91.96% ± 0.43% to 93.21% ± 2.34%.
• Convergence speed: GPT-4.1 and Gemini 3.1 Flash-Lite converged to approximately 97% accuracy within 10–15 of the 30 NAS generations (averaged over 3 runs), while GPT-4.1-nano showed limited improvement.
• Hardware-dependent architecture motifs: CPU and GPU Pareto-optimal architectures were convolutional-linear hybrids with 4–6 layers (described as akin to LeNet-5), whereas the optical Pareto front was dominated by shallow, wide MLPs with convolutions largely absent outside the energy-efficient end, normalization appearing at high-energy regimes, and no skip connections.

Methodology in Plain English

The researchers framed the task as a multi-objective problem: for each hardware backend, find architectures that are non-dominated in accuracy and inference energy. They restrict the objectives to maximizing classification accuracy and minimizing inference energy, with energy approximated by platform-specific analytic cost models.

The search loop works like this. It starts from 10 hand-designed seed architectures spanning pure MLP, convolutional-linear hybrid, and varying-depth configurations. Each generation, the LLM reads the current Pareto front, recently evaluated candidates, and measured accuracy degradation under noise, then updates a structured list of design heuristics (the knowledge base). Parents are chosen from an NSGA-II archive via binary tournament, and the LLM proposes new candidates through mutation or crossover conditioned on that archive and knowledge base, with a target region in a MAP-Elites diversity grid sampled using a 70/30 exploit/explore ratio. Candidates are trained for 20 epochs and evaluated to produce accuracy, energy, and a degradation score measured as the accuracy difference between the trained model and the same model evaluated with injected hardware noise. Results are inserted back into the archive via NSGA-II, which ranks Pareto fronts and uses crowding distance for spread.

The search ran for 30 generations with 8 candidates generated per generation and an NSGA-II population size of 20. Each candidate was trained for 20 epochs using Adam with a learning rate of 10⁻³ and cross-entropy loss, with quantization-aware training to approximate hardware effects. Baselines were adapted for fairness: LLMatic used the same seed architectures, the same 30 generations and 8 candidates per generation, the same training conditions, and the same MAP-Elites archive and prompts, but with its fitness switched to non-ideality accuracy; RZ-NAS used 1500 generations with population size 100 and SynFlow as the zero-cost fitness function.

Hardware modeling covers CPU (Xeon 8380, FP32), GPU (Blackwell B200, FP16), and an 8-bit MZI-based optical neural network simulated with pytorch-onn. Per-operation energy figures are 91.7 pJ per MAC on CPU, 0.89 pJ on GPU, and 0.02 pJ on optics; ReLU costs 3 pJ per element on CPU and GPU and 10 pJ on optics; DAC costs 2 pJ per element and ADC 4 pJ per element on the optical path. The optical backend additionally models phase error, thermal crosstalk, gamma noise, and quantization noise at realistic, medium, and worst-case levels. MZI-specific architectural constraints include block-diagonal decomposition of weight matrices into 4×4 sub-matrices implemented as a Clements-topology mesh, a soft width cap of ≤ 512 for linear layers, and heavy penalties on residual connections because they require extra ADC/DAC conversions.

Why This Matters

Impact on research. The paper argues that hardware does not merely rescale a model's score but changes which architectural motifs are preferred, and it provides evidence that zero-cost proxies—widely used to accelerate NAS—can invert their sign under hardware non-idealities. It also offers a template for comparing heterogeneous platforms under their own constraints rather than reusing one hardware family's search procedure.

Real-world applications:

• Designing networks for silicon-photonic accelerators, where the paper reports optical MACs roughly four orders of magnitude cheaper than CPU MACs but with DAC/ADC conversion costs that partially offset the gain.
• Energy-constrained inference on edge and embedded devices, where inference energy per model matters as much as accuracy.
• Analog in-memory and memristor-style crossbar computing, a paradigm the authors list as a natural extension of the backend abstraction.
• Automated co-design pipelines where a new hardware backend can be plugged in without rewriting the search algorithm.

Industry relevance. The framework targets the practical problem of porting machine learning onto specialized and emerging accelerators, where designers currently need per-platform expertise. The paper's claim is that LLM-guided search can absorb hardware context through prompts and knowledge bases, which could lower the barrier for evaluating whether a new hardware platform is worth deploying for a given workload.

Future Directions

• Extending beyond feedforward architectures composed of standard layers, which the authors note may not transfer to fundamentally different hardware such as spiking neuromorphic processors or quantum computers.
• Scaling to harder datasets, since the current demonstration is on MNIST classification.
• Expanding to additional unconventional computing paradigms, with memristors named explicitly.
• Reducing computational cost through early-stopping strategies such as Hyperband-style successive halving, since UH-NAS trains every candidate rather than using proxies.
• Reducing the reliance on hardware-specific inputs, since the framework is algorithmically hardware-agnostic but still requires seed architectures, energy models, and non-ideality simulators per platform.

Target Audience

Researchers and practitioners in neural architecture search, automated machine learning, and hardware–software co-design, particularly those working on photonic or other unconventional accelerators. It is also relevant to ML systems engineers evaluating energy-accuracy trade-offs across deployment platforms, and to readers interested in how LLMs can serve as search operators rather than as end-task models. Readers without background in NAS terminology will need to consult the referenced background literature.

Authors’ abstract

Deploying neural networks on unconventional hardware demands architectures that co-optimize task accuracy and platform-specific constraints such as energy cost, physical non-idealities, and numerical precision. Existing neural architecture search (NAS) methods are typically tailored to a single hardware family, limiting cross-platform comparison and generalization. We introduce Unconventional Hardware Neural Architecture Search (UH-NAS), a hardware-agnostic, LLM-guided NAS framework that integrates language models as evolutionary operators to co-optimize accuracy and inference energy. By exposing hardware as a swappable backend with per-platform energy models, physical constraints, and non-ideality simulators, UH-NAS enables fair system-level comparisons across various backends without modifying the search algorithm. Tested on optical MZI hardware, UH-NAS discovers more diverse, robust architectures than conventional baselines while outperforming existing LLM-to-NAS approaches. Additional ablations on architecture robustness under non-idealities and the role of system prompts highlight the importance of architecture-hardware co-design for emerging computing platforms.