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

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 LLMs 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.

• The paper introduces UH-NAS, a physics-aware neural architecture search framework that leverages LLMs to co-optimize task accuracy and inference energy across heterogeneous hardware backends.
• It integrates multi-objective NSGA-II selection and full noisy training evaluations to address hardware non-idealities, moving beyond unreliable zero-cost proxies.
• The method achieves competitive MNIST accuracy (~97%) and discovers diverse, robust architectures tailored for both optical and digital systems.

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

Framework Overview

The paper introduces UH-NAS (Unconventional Hardware Neural Architecture Search), a hardware-agnostic NAS framework integrating LLMs as evolutionary operators to co-optimize task accuracy and inference energy across heterogeneous hardware backends. Hardware is abstracted as a swappable backend, exposing per-platform energy models, physical constraints, and non-ideality simulators. This enables fair system-level comparison and robust architecture discovery without modifying the search algorithm, addressing the deficiencies of existing NAS methods which are limited to single hardware families and fail to generalize across unconventional platforms.

UH-NAS operates in iterative multi-objective evolutionary cycles. At each generation, an LLM-driven operator generates new candidate architectures informed by hardware-specific system prompts and a dynamically updated knowledge base. NSGA-II is used for multi-objective selection, optimizing accuracy (↑) and inference energy (↓) to construct rank-0 Pareto fronts for each backend. Unlike prevailing NAS approaches, every candidate is evaluated with full noisy training under hardware constraints and non-idealities—bypassing zero-cost (ZC) proxies, whose breakdown under hardware noise is empirically demonstrated.

Figure 1: Pareto fronts for MNIST validation error vs energy per inference across GPU, CPU, and MZI mesh optical systems; optical platforms exhibit orders-of-magnitude lower energy but higher error, with UH-NAS narrowing this gap via robust architecture selection.

Hardware-Aware Multi-Objective Optimization and Backend Abstractions

UH-NAS's hardware abstraction exposes each backend via operation-level energy cost, physical constraints, and detailed non-ideality models calibrated from silicon-photonics and system documentation. Energy per MAC operation differs drastically: optical hardware achieves ~0.02 pJ/MAC, GPU ~0.89 pJ/MAC, CPU ~91.7 pJ/MAC, while ReLU activations are orders-of-magnitude more expensive in optical regimes due to ADC/DAC overhead. Non-idealities for MZI meshes include phase error (σφ), thermal crosstalk (ε), and gamma noise (σγ), evaluated under realistic and worst-case bounds extracted from the literature.

These backend abstractions inform candidate scoring and population augmentation. In optical systems, architectural motifs—such as skip connections or convolutional layers—are heavily penalized or even prohibited due to prohibitive overhead and catastrophic noise amplification. UH-NAS leverages quantization-aware training for 8-bit precision and injects realistic non-idealities during evaluation.

LLM-Guided Search: Diversity and Robustness

The integration of LLMs acts at two levels: (1) as mutation/crossover operators, proposing architectural modifications via hardware-specific prompts and (2) as knowledge synthesizers, refining search heuristics every generation based on system-level performance and robustness degradations. This enables non-local exploration and circumvents search space degeneracy induced by hardware constraints.

Figure 2: UH-NAS with LLM achieves broader search diversity (203 unique designs in 250 evaluations), lower validation error, and broader Pareto coverage than the ablation without LLM (26 unique designs, degenerate front).

UH-NAS with strong LLM backbones (e.g., GPT-4.1, Gemini 3.1 Flash-Lite) reaches MNIST accuracy of ~97% within 10–15 NAS generations, while smaller LLMs collapse to poor architectural prompts, leading to stagnated search and suboptimal performance. System prompts incorporating physical context are essential; ablations removing them show consistent accuracy degradation.

Figure 3: Strong LLMs (GPT-4.1, Gemini 3.1 Flash-Lite) enable rapid convergence to ~97% accuracy, while weaker models plateau with limited improvement due to poor system prompt synthesis.

Robustness Across Non-Ideality Regimes

UH-NAS discovers architectures that generalize robustness across non-ideal hardware regimes. Pareto fronts optimized under worst-case noise retain competitive performance when re-evaluated at moderate and realistic noise levels; inversely, architectures found under idealized conditions fail to generalize to noisy regimes. This robustness is preserved across the front structure, as revealed by cross-regime re-evaluations.

Figure 4: Robustness of UH-NAS architectures across non-idealities; worst-case-trained models remain competitive as noise increases, indicating robustness generalization.

Furthermore, zero-cost proxies—which are widely used for rapid NAS evaluation—are shown to be unreliable. Spearman correlations between ZC proxies and validation accuracy collapse under hardware non-idealities, reversing sign and indicating destructive performance, especially for convolutional layers in optical MZI systems.

Figure 5: Spearman rank correlation between ZC proxies and validation accuracy; proxies are consistent under clean conditions but negatively correlated under worst-case non-idealities, invalidating ZC for NAS in non-ideal optical hardware.

Physically Interpretable Design Principles

UH-NAS empirically establishes that for MZI-based optical neural networks, optimal architectures diverge from conventional CPU/GPU design paradigms. Shallow, wide MLPs dominate over convolutional-linear hybrids, due to non-idealities compounding through depth and convolutional kernel weight reuse amplifying crosstalk failures. Batch normalization surfaces as the principal strategy for noise recalibration, and architectures eschew skip connections entirely due to ADC/DAC conversion costs.

In digital (CPU/GPU) backends, Pareto-optimal architectures mirror traditional design patterns, maximizing depth and convolutional front-ends. For optical hardware, depth sensitivity from phase error and catastrophic convolutional layer failures from crosstalk motivate architecture shifts to shallow, flat linear networks with aggressive batch normalization, validating the hardware-aware LLM-guided search methodology.

Energy Validation

UH-NAS employs closed-form analytics for energy estimation, calibrating against real RAPL measurements on Xeon CPUs. Although absolute values are underestimated, relative ranking fidelity remains high ($r^2 = 0.828$). Cross-platform comparisons remain theoretically valid since analytical models capture ranking and relative energy scales.

Figure 6: Log-log correlation between analytical energy estimates and real RAPL validation; high ranking fidelity despite absolute offset validates theoretical cross-platform comparison.

Conclusion

UH-NAS demonstrates a robust, physics-aware LLM-guided NAS framework for unconventional hardware platforms, enabling architecture–hardware co-design via swappable backend abstractions and full-noisy evaluation. Empirically, UH-NAS discovers more diverse, robust architectures than existing NAS baselines, achieving competitive task accuracy on photonic MZI hardware even under worst-case non-idealities. The work uncovers physically interpretable architectural motifs suitable for emerging hardware regimes and invalidates ZC proxies for NAS evaluation under noise.

Implications are significant: the framework provides actionable architectural design principles for optical neural networks, such as favoring shallow/wide MLPs and batch normalization, and sets a methodological precedent for hardware-aware evolutionary search across novel computing paradigms (memristor, quantum, neuromorphic). Extensions to larger datasets, broader hardware classes, and integration with early-stopping are natural future directions.

See (2606.10294) for full bibliographic details and related works.