GraphIR: Architecture-Level Search States for LLM-Guided Neural Architecture Evolution

2026-08-03Machine Learning

Machine Learning
AI summary

The authors address a challenge in using large language models (LLMs) to design neural networks by improving how these models understand and change network code. They introduce GraphIR, a new way to represent network designs that clearly shows how data flows, what parts can be changed, and the rules parts must follow. To test their approach, they created a set of questions to see how well GraphIR helps with understanding dependencies in the network. Their results show that GraphIR helps LLMs better analyze and modify networks, leading to improved performance on several benchmark tests without increasing model size or reducing efficiency.

Large Language ModelsNeural Architecture SearchGraphIRIntermediate RepresentationTensor FlowMutation SurfaceDependency ReasoningExecutable Neural ProgramsOpenEvolveBenchmarking
Authors
Zhen Liu, Wanqi Zhou, Shuanghao Bai, Yuhan Liu, Jinjun Wang, Jingwen Fu
Abstract
Large language models (LLMs) enable neural architecture search (NAS) directly over executable neural network programs. However, code-level flexibility does not provide the architecture state needed for effective mutation: LLMs must infer tensor dependencies, editable components, and compatibility constraints from implementation details. To address this representation mismatch, we propose GraphIR, an architecture-aware intermediate representation that supplements executable programs with a mutation-aligned candidate state. GraphIR organizes each candidate through three complementary views: a computation skeleton describing tensor flow, a mutation surface exposing editable modules and operations, and a validity envelope capturing interface contracts, propagated shapes, and downstream dependencies. To evaluate our method, we construct NAS-Dependency, a 120-question benchmark covering six complementary dependency-reasoning dimensions. The diagnostic shows that GraphIR is particularly effective at identifying exact producer occurrences, tracing dependency propagation, and diagnosing interface and failure risks. Across six downstream benchmarks including CLRS, GraphIR achieves the best overall search performance while maintaining comparable model size and favorable end-to-end NAS efficiency when integrated into OpenEvolve. These results show that a mutation-oriented architecture state provides an effective interface between executable neural programs and LLM-guided architecture evolution.