🤖 machine learning

LLM as a Tool, Not an Agent: Code-Mined Tree Transformations for Neural Architecture Search

This paper introduces LLMasTool, a hierarchical tree-based Neural Architecture Search framework that treats LLMs as tools for resolving specific degrees of freedom within a reliable algorithmic search process, thereby enabling stable, open-ended model evolution that outperforms existing methods on CIFAR and ImageNet benchmarks.

Original authors: Masakazu Yoshimura, Zitang Sun, Yuiko Sakuma, Junji Otsuka, Atsushi Irie, Takeshi Ohashi

Published 2026-04-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Masakazu Yoshimura, Zitang Sun, Yuiko Sakuma, Junji Otsuka, Atsushi Irie, Takeshi Ohashi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build the world's best LEGO castle.

For a long time, scientists trying to build the best AI brains (Neural Networks) have used two main approaches:

The Strict Architect: They draw a very specific, limited blueprint. They can only use certain LEGO bricks and can only connect them in specific ways. This is safe and reliable, but it limits creativity. You'll never build a castle with a dragon's head because the blueprint didn't allow it.
The Magic Genie (LLM): You ask a super-smart AI (a Large Language Model) to just "write the code for a better castle." The problem? The AI often gets confused, hallucinates (makes up bricks that don't exist), or just copies the same old castles it saw in its training books. It's like asking a genie to build a castle, and it hands you a blueprint that falls apart when you try to build it.

The authors of this paper, "LLMasTool," decided to try a third way. They realized: Don't let the AI be the architect. Let the AI be the helpful assistant.

Here is how their system works, using a simple analogy:

1. The "Module Library" (The LEGO Box)

Instead of asking the AI to invent new bricks from scratch, the system first scans thousands of existing, working codebases (like PyTorch). It automatically finds all the valid, working "LEGO bricks" (modules) and puts them in a giant, organized library.

Why this matters: The AI doesn't have to guess what a "Convolutional Block" is. It just grabs a real, working one from the library. This guarantees that every piece it uses actually exists and works.

2. The "Tree" vs. The "Code"

Most AI systems try to write a whole new program (a long, messy paragraph of text) to describe a new architecture. If the AI makes one typo, the whole thing breaks.

The Paper's Trick: They represent the AI brain as a Tree.
- The trunk is the whole model.
- The big branches are major parts (like the "Backbone" or "Neck").
- The small twigs are specific settings (like "use 32 channels" or "kernel size 3").
The Benefit: It's much easier to chop off a branch and replace it with a different one from the library than to rewrite a whole paragraph of code.

3. The "Foreman" and the "Specialist" (Coarse-to-Fine)

This is the core innovation. They split the work between a Rule-Based Algorithm (The Foreman) and the LLM (The Specialist).

The Foreman (The Algorithm): This is the boss. It looks at the Tree and decides the big picture strategy.
- Example: "Okay, we need to swap out the 'Backbone' branch. Let's try removing a layer here, or maybe adding a new branch there."
- The Foreman uses math (Bayesian statistics) to decide where to look next, ensuring they don't just keep doing the same thing over and over. It forces the system to explore new, weird directions.
The Specialist (The LLM): Once the Foreman says, "Hey, let's swap this branch for a 'ResBlock'," the LLM steps in.
- The Job: The LLM doesn't have to invent the whole castle. It just has to fill in the small details. "Okay, if we use a ResBlock here, what should the settings be? Should the bias be on or off? What's the best number?"
- The Magic: Because the LLM is only doing small, specific tasks, it rarely makes mistakes. It acts like a tool, not a boss.

4. The "Evolution" (Trial and Error)

The system works like natural selection:

Propose: The Foreman picks a spot on the Tree and asks the LLM to make a small change.
Test: They build the new "castle" (train the AI model) and see how well it performs.
Learn: If it gets better, they remember that change. If it gets worse, they forget it.
Repeat: They do this thousands of times, slowly evolving the architecture into something incredibly efficient and powerful.

Why is this a big deal?

No More Hallucinations: Because the AI is just picking from a library of real code and filling in small blanks, it almost never creates a "broken" model.
True Creativity: The "Foreman" forces the system to try weird combinations the AI might never have thought of on its own, breaking out of the AI's "training bias."
Better Results: In their tests, this method beat the previous state-of-the-art methods by a significant margin on standard image recognition tests (like CIFAR and ImageNet).

In summary:
Instead of asking a genius AI to write a novel from scratch (which often fails), they gave it a dictionary of real words (the module library) and a strict editor (the algorithm) to tell it exactly which words to use. The AI then just fills in the grammar. The result? A masterpiece that is both creative and perfectly grammatical.

1. Problem Statement

Neural Architecture Search (NAS) aims to automate the discovery of high-performing deep neural network (DNN) architectures. However, existing approaches face significant limitations:

Traditional Algorithm-Driven NAS: Relies on carefully hand-crafted search spaces (e.g., specific cells or graph primitives) to ensure executability. This restricts open-ended exploration and biases discoveries toward the designer's prior knowledge.
LLM-Driven (Agentic) NAS: Recent approaches use Large Language Models (LLMs) as autonomous agents to generate code or architectures directly. These methods struggle with:
- Reliability: Difficulty in generating complex, valid, and executable architectures without syntax errors or logical flaws.
- Bias: LLMs tend to reproduce patterns seen in their training data, hindering truly novel or open-ended exploration.
- Controllability: Lack of strict constraints makes it hard to guarantee that generated models meet specific resource constraints (FLOPs, parameters) or structural requirements.

The core challenge is bridging the gap between the reliability of algorithmic search and the generative power of LLMs without falling into the traps of bias or invalidity.

2. Methodology: LLMasTool

The authors propose LLMasTool, a hybrid NAS framework that treats the LLM as a structured tool rather than an autonomous agent. The framework operates through three main components:

A. Algorithm-Based Module Mining

Instead of relying on LLMs to understand code or memorize architectures, the system automatically mines reusable modules from existing open-source codebases (e.g., PyTorch, MMPretrain, timm).

Process: A rule-based algorithm parses source code into Abstract Syntax Trees (AST) to identify classes inheriting from torch.nn.Module.
Output: A Module Database ( $D_M$ ) containing modules, their metadata (constructor arguments, input/output shapes), and source code. This ensures 100% correctness in the search space definition and allows the system to leverage existing, proven knowledge assets without hallucination.

B. Deploy-Friendly Hierarchical Tree Representation

The paper moves away from the standard data-flow graph representation used in many NAS methods.

Structure: Architectures are represented as hierarchical trees.
- Nodes: Represent macro-level modules (e.g., Backbone, Neck), components, and leaf-level hyperparameters (e.g., kernel size, channel width).
- Advantage: This structure mirrors how models are defined in frameworks like PyTorch (nested modules), making it easier to map existing source code directly into the search space. It facilitates top-down design and simplifies the integration of complex human-written code blocks.

C. Coarse-to-Fine Evolution with LLM as a Tool

The evolution process transforms the architecture tree rather than generating raw code. It employs a hierarchical coarse-to-fine decision-making strategy:

Coarse-Level Planning (Algorithm-Driven):
- A diversity-guided evolutionary algorithm determines the type of operation (e.g., Swap, Insert, Delete, Change Hyperparameter) and the general direction.
- Bayesian Modeling: Uses Hierarchical Thompson Sampling on a Beta-Bernoulli model to balance exploration and exploitation. It tracks historical success rates of different operations and prompt categories to guide future choices efficiently.
- Diversity Control: The algorithm explicitly enumerates possible modification directions to ensure the search is not constrained by the LLM's prior bias.
Fine-Grained Decision (LLM-Driven):
- Once the coarse direction is set (e.g., "Swap a module at location X"), the LLM acts as a tool to resolve the remaining degrees of freedom.
- Prompting: The LLM receives specific prompts (templates) containing the current architecture, the module database, and historical feedback. It is asked to select specific modules, determine hyperparameters, or generate the necessary configuration text to ensure the transformation is executable and valid.
- Categories: The system uses different prompt categories ranging from "Rely LLM" (high reliance on prior knowledge) to "Minimum LLM" (algorithm decides location/module, LLM only fills hyperparameters) to ensure diverse exploration.

3. Key Contributions

LLMasTool Framework: A novel NAS framework that integrates algorithmic reliability with LLM flexibility by representing architectures as hierarchical trees and evolution as tree transformations.
Code-Mined Search Space: An automated module mining technique that extracts reusable building blocks from arbitrary source code, creating a reliable and diverse search space without manual engineering.
Coarse-to-Fine Decision Making: A strategy where algorithms handle high-level planning (ensuring diversity and avoiding LLM bias) while LLMs handle low-level refinement (ensuring validity and executability).
Hierarchical Thompson Sampling: A Bayesian optimization method applied to the decision hierarchy to efficiently navigate the vast open-ended search space using historical feedback.

4. Experimental Results

The method was evaluated on NAS-Bench-201 (CIFAR-10, CIFAR-100, ImageNet16-120) and ImageNet-100.

Performance Gains:
- CIFAR-10: Improved by 0.69 points over the previous SOTA (NADER).
- CIFAR-100: Improved by 1.83 points.
- ImageNet16-120: Improved by 2.68 points.
- ImageNet-100: Achieved state-of-the-art results, significantly outperforming methods like NADER, LLMatic, and LeMo-NADe.
Ablation Studies:
- Tree Transformation vs. Code Generation: Tree-based evolution significantly outperformed direct code generation or config modification, proving the importance of structured representation.
- Module Mining: Using mined modules from external libraries yielded much higher accuracy than using only primitive modules.
- LLM Role: The "Minimum LLM" approach (where algorithms decide what to change and LLM decides how) often outperformed "Rely LLM" approaches, confirming that reducing LLM bias leads to better open-ended exploration.
- Hierarchical Sampling: The proposed Hierarchical Thompson Sampling was crucial for efficiency; uniform sampling or non-hierarchical sampling led to stagnation.

5. Significance

This paper represents a paradigm shift in how LLMs are utilized for scientific discovery and engineering tasks:

From Agent to Tool: It argues that for complex, constrained tasks like NAS, LLMs should not be autonomous agents but rather specialized tools operating within a rigorous algorithmic framework. This mitigates the risks of hallucination and bias.
Open-Ended Exploration: By decoupling the search direction (algorithm) from the implementation details (LLM), the method can explore architectural spaces that are fundamentally different from the LLM's training data, leading to genuinely novel discoveries.
Reusability of Knowledge: The module mining approach demonstrates how to systematically harvest and reuse existing software engineering assets (code) to accelerate AI research, bridging the gap between traditional software engineering and automated machine learning.

In summary, LLMasTool demonstrates that combining the structural rigor of algorithmic search with the semantic understanding of LLMs, while strictly controlling the LLM's role, yields superior results in discovering high-performance neural architectures.