TREX: Automating the Full LLM Training Pipeline with Multi-Agent Tree Search

Paper: 2604.14116 Authors: Zerun Ma, Guoqiang Wang, Xinchen Xie, Yicheng Chen, He Du, Bowen Li, Yanan Sun, Wenran Liu, Kai Chen, Yining Li Categories: cs.AI, cs.CL

The Gap

LLM agents can now handle isolated scientific tasks—write code, summarize papers, run experiments. But real research workflows aren’t isolated. Training an LLM for a new domain requires chaining together literature review, dataset curation, hyperparameter tuning, and iterative refinement. Prior work (AutoML systems, neural architecture search, prompt optimization frameworks) automates pieces of this pipeline but assumes humans handle the glue: deciding what papers matter, which datasets to mix, when to pivot strategies. The boundary is clear: agents do tasks, humans do workflows.

This paper asks: can an agent system own the entire training lifecycle, including the meta-decisions about how to explore the solution space?

Problem: LLM training = multi-stage workflow with exploration
    |
    v
Assumption: Model it as tree search over experimental paths
    |
    v
Method: Two agents (Researcher + Executor) + tree structure
    |
    +---> Researcher: plans experiments, reviews results
    +---> Executor: runs training, collects metrics
    +---> Tree: tracks history, enables backtracking
    |
    v
Evidence: FT-Bench (10 real-world tasks), consistent improvement
    |
    v
Conclusion: Automated training is feasible with structured exploration

The Increment

One sentence: Before TREX, you automated training runs but manually designed the training strategy; after TREX, the system explores strategies autonomously and learns from failed experiments.

Core Mechanism

TREX has two agents working in a loop. The Researcher agent reads papers, analyzes task requirements, and proposes training strategies (which datasets to use, what hyperparameters to try). The Executor agent takes those proposals, prepares data recipes, runs training jobs, and reports back metrics. Each experimental round becomes a node in a tree: the root is the initial task, branches are different strategies, and leaves are trained models with evaluation scores.

The tree structure does three things. First, it prevents redundant work—if you’ve already tried a particular dataset mix, the system reuses those results instead of retraining. Second, it enables backtracking—if a strategy hits a dead end (model doesn’t converge, performance plateaus), the system can jump back to an earlier node and try a different branch. Third, it accumulates knowledge—after several rounds, the Researcher can look at the tree and extract patterns (e.g., “adding math data helps reasoning tasks but hurts creative writing”).

User Task --> Researcher Agent
                  |
                  v (proposes strategy)
              Tree Node
                  |
                  v
              Executor Agent
                  |
                  v (trains model)
              Evaluation Metrics
                  |
                  v (feedback)
              Researcher Agent
                  |
                  v (next strategy or refine)
              New Tree Node ...

Think of TREX like a chess engine exploring moves. The Researcher is the evaluation function—it looks at the board (current model performance) and suggests promising moves (training strategies). The Executor is the move executor—it actually makes the move and reports the new board state. The tree is the game tree—it remembers all positions explored so far. When a line of play looks bad, the engine doesn’t start over; it backtracks to a promising position and tries a different continuation. Over time, the engine learns which opening strategies lead to strong positions, just as TREX learns which data mixes lead to better models.

Key Concepts

• Tree-based exploration: Most AutoML systems treat experiments as independent trials—run 100 random configurations, pick the best. TREX treats them as a connected search space. Each experiment is a node, and edges represent refinements (e.g., “same strategy but with more data”). This structure lets the system reason about relationships: if Strategy A worked well and Strategy B is a small modification of A, try B next. It’s the difference between random search and guided search. The tree also makes failure informative—a dead-end branch tells you what not to try, which is knowledge you can reuse.

• Agent role separation: The Researcher and Executor aren’t just two functions; they’re two cognitive modes. The Researcher operates in “planning space”—it reads abstracts, compares approaches, makes high-level decisions. The Executor operates in “execution space”—it writes data processing scripts, monitors GPU usage, debugs training crashes. Separating these modes prevents mode confusion: the Researcher never gets bogged down in implementation details, and the Executor never second-guesses the strategy mid-training. It’s like having a coach and an athlete—the coach studies opponents and designs plays, the athlete executes them without overthinking.

Framework Shift

Before (mainstream approach):        After (this paper):

Human designs strategy               Agent explores strategy space
    |                                    |
    v                                    v
AutoML tunes hyperparams         Tree tracks all experiments
    |                                    |
    v                                    v
Train model once                 Iterative refinement with memory
    |                                    |
    v                                    v
Evaluate and repeat              Distill insights from tree
(no memory of past trials)       (reuse and learn from history)

From isolated optimization to cumulative exploration, the core shift is treating training as a search problem with memory.

Expert Assessment

Problem choice: Real gap. Anyone who’s fine-tuned models knows the pain: you try a dataset, it doesn’t work, you try another, you forget what you tried before, you waste GPU hours retraining the same thing. The paper targets a genuine bottleneck in the research loop. It sits at the intersection of AutoML (which automates tuning) and research agents (which automate literature review)—combining them is natural but underexplored.

Method maturity: The tree structure is clever but not novel (it’s Monte Carlo Tree Search applied to training). The real contribution is the agent orchestration—getting two LLMs to collaborate reliably over multiple rounds without diverging or getting stuck. The paper doesn’t deeply explore failure modes (what if the Researcher proposes an impossible strategy? what if the Executor misinterprets instructions?). It feels like a strong prototype, not a production system.

Experimental integrity: FT-Bench is well-designed—10 tasks spanning different domains, clear metrics. But the baselines are weak: they compare against single-round training and basic AutoML, not against a human expert doing iterative refinement. The paper shows TREX improves over naive approaches, but doesn’t show it matches or beats a skilled practitioner. The ablation studies are solid (removing the tree hurts performance, confirming its value).

Writing quality: The introduction oversells (“automating complex, real-world workflows”)—this automates one workflow (LLM training), not workflows in general. The method section is dense; the tree formalism could be explained with a concrete example first, then generalized. The related work section is thorough but reads like a checklist. Rewriting the results section to lead with insights (e.g., “we found that tree-based exploration reduces redundant experiments by 40%”) instead of tables would make the contribution clearer.

Verdict: weak accept — Solid execution on a real problem, but the evaluation doesn’t prove it’s ready for practitioners to use, and the method isn’t surprising enough to be a conceptual breakthrough.

Takeaways

The tree structure is the transferable idea. Anytime you’re doing iterative experimentation (hyperparameter tuning, prompt engineering, architecture search), representing your trials as a tree instead of a flat list gives you three things: (1) you can reuse partial results, (2) you can backtrack intelligently, (3) you can extract patterns across branches. The specific agents here are LLM-based, but the tree framework works with any search process. If you’re building an AutoML system, steal the tree; if you’re doing manual experimentation, draw your own tree on paper to avoid repeating dead ends.