CoLLM-NAS: Collaborative Large Language Models for Efficient Knowledge-Guided Neural Architecture Search

TL;DR

CoLLM-NAS tackles the inefficiency and validity issues of neural architecture search by introducing a collaborative two-stage framework that pairs a stateful Navigator LLM with a stateless Generator LLM, coordinated by a management module and guided by inherited supernet weights. This knowledge-guided search leverages LLM priors and progressive trajectory information to rapidly steer the search toward high-performing architectures across multiple spaces, including ImageNet-scale tasks. Empirical results demonstrate state-of-the-art or competitive performance on ImageNet and NAS-Bench-201, with substantial reductions in search cost (3–10x) and improved robustness across ablations, prompt variations, and LLM choices. The framework also generalizes across different NAS baselines (OFA, SPOS, AutoFormer) and search spaces (MobileNet, ShuffleNet, AutoFormer), suggesting broad applicability of trajectory-to-strategy-to-solution search in neural architecture design.

Abstract

The integration of Large Language Models (LLMs) with Neural Architecture Search (NAS) has introduced new possibilities for automating the design of neural architectures. However, most existing methods face critical limitations, including architectural invalidity, computational inefficiency, and inferior performance compared to traditional NAS. In this work, we present Collaborative LLM-based NAS (CoLLM-NAS), a two-stage NAS framework with knowledge-guided search driven by two complementary LLMs. Specifically, we propose a Navigator LLM to guide search direction and a Generator LLM to synthesize high-quality candidates, with a dedicated Coordinator module to manage their interaction. CoLLM-NAS efficiently guides the search process by combining LLMs' inherent knowledge of structured neural architectures with progressive knowledge from iterative feedback and historical trajectory. Experimental results on ImageNet and NAS-Bench-201 show that CoLLM-NAS surpasses existing NAS methods and conventional search algorithms, achieving new state-of-the-art results. Furthermore, CoLLM-NAS consistently enhances the performance and efficiency of various two-stage NAS methods (e.g., OFA, SPOS, and AutoFormer) across diverse search spaces (e.g., MobileNet, ShuffleNet, and AutoFormer), demonstrating its excellent generalization.

Figures (5)

• Figure 1: Consistency heatmap between LLM-predicted and ground-truth rankings on CIFAR-10 and CIFAR-100 within NAS-Bench-201 search space.
• Figure 2: Pipeline of CoLLM-NAS. The search starts with the Navigator LLM generating an initial exploration strategy based on target accuracy and resource constraint (e.g., FLOPs, Params). The Coordinator then transmits this strategy to the Generator LLM, which synthesizes high-quality candidates accordingly. After Coordinator validation and evaluation, results are fed back to the Navigator LLM for strategy refinement. This loop iterates until either achieving the target accuracy or reaching the iteration limit. Orange regions indicate more promising areas in the search space.
• Figure 3: T-SNE visualization of CoLLM-NAS's search dynamics on ImageNet-16-120 within NAS-Bench-201 search space. Architectural performance rankings are represented through color-coding, with unexplored ones displayed in gray.
• Figure 4: Ablation on main mechanisms: (a) Comparison of iterative performance between CoLLM-NAS and SiLLM-NAS on NAS-Bench-201. (b) Impact of memory retention settings on different datasets. N$\checkmark$/$\times$ and G$\checkmark$/$\times$ denote whether the memory of Navigator/Generator LLM is retained.
• Figure 5: Performance comparison of different temperature settings on CIFAR-100 within NAS-Bench-201 search space.