Closed-Loop LLM Discovery of Non-Standard Channel Priors in Vision Models

TL;DR

This work introduces a closed-loop NAS framework where a fine-tuned large language model optimizes vision-network channel configurations directly in executable code. It overcomes data scarcity by bootstrapping with AST mutations to create a large, valid corpus (LEMER) that teaches the LLM architectural priors, enabling conditional code generation to improve performance. On CIFAR-100, the approach yields a 24.1% relative improvement over the initial distribution and reveals non-standard channel widths and late-stage expansion as effective priors, along with Pareto-efficient, parameter-saving configurations. The results demonstrate that language models can acquire domain-specific architectural knowledge from synthetic data and performance feedback, offering a scalable, architecture-agnostic path for vision NAS with practical impact on design efficiency and hardware-aware optimization.

Abstract

Channel configuration search the optimization of layer specifications such as layer widths in deep neural networks presents a complex combinatorial challenge constrained by tensor shape compatibility and computational budgets. We posit that Large Language Models (LLMs) offer a transformative approach to Neural Architecture Search (NAS), capable of reasoning about architectural code structure in ways that traditional heuristics cannot. In this paper, we investigate the application of an LLM-driven NAS framework to the problem of channel configuration. We formulate the search as a sequence of conditional code generation tasks, where an LLM refines architectural specifications based on performance telemetry. Crucially, we address the data scarcity problem by generating a vast corpus of valid, shape-consistent architectures via Abstract Syntax Tree (AST) mutations. While these mutated networks are not necessarily high-performing, they provide the critical volume of structural data required for the LLM to learn the latent relationship between channel configurations and model performance. This allows the LLM to internalize complex design patterns and apply them to optimize feature extraction strategies. Experimental results on CIFAR-100 validate the efficacy of this approach, demonstrating that the model yields statistically significant improvements in accuracy. Our analysis confirms that the LLM successfully acquires domain-specific architectural priors, distinguishing this method from random search and highlighting the immense potential of language-driven design in deep learning.

Figures (5)

• Figure 1: The general overview of the experimental cycle. Neural network codes are retrieved from LEMUR database and went through mutations in abstract syntax tree bootstrapper to generate valid models and exemplify the problem specification to LLM. Iterative inference-training cycle according structured prompts, including accuracy metric signal, self improve the channel predictions done by the LLM.
• Figure 2: The AST-based bootstrapping pipeline. It illustrates the process of AST parsing (extracting layer definitions and offsets), dependency analysis using the TorchFX symbolic graph, the constraint-aware editing and repairing phase, and the verification protocol using dummy tensor checks to generate valid seed models. Whenever a random mutation is chosen and applied, the network must be corrected. The invalid layers are shown in red. By examining the input and outputs to Conv2D layers, the network is corrected later on, denoted by downstream green layers.
• Figure 3: CIFAR-100 validation accuracy (one training epoch) of LLM-generated image classification models across LLM fine-tuning epochs. Left: mean accuracy of all valid models; fluctuations reflect the exploration of diverse configurations. Right: rolling average of validation-set classification accuracy (window size = 3) for image classification models generated by an LLM across fine-tuning epochs.
• Figure 4: Search dynamics. Left: generation success rate (valid/total); low generation success stability here indicates the LLM maintains metric focused improvement during exploration rather than memorizing the specifics of the evaluation constraints. Right: best-so-far trajectory (high-water mark), where step-function jumps highlight discrete breakthroughs in the generative process for the obtained max accuracies.
• Figure 5: Efficiency and flow analysis. Left: accuracy vs. parameter count. The red dashed line indicates the Pareto frontier of efficient models. The LLM discovers models that maximize accuracy for a given parameter budget. The color and the size of the points remark the epoch that the models were generated. Right: architectural flow (parallel coordinates) for all models, colored by accuracy. High-performing models (yellow/green) consistently exhibit a specific trajectory: narrow early layers followed by a massive expansion in the final layer.