TREX: Automating LLM Fine-tuning via Agent-Driven Tree-based Exploration

1. Problem Statement

While Large Language Models (LLMs) have successfully automated isolated scientific tasks (e.g., code generation, literature review), automating the entire lifecycle of LLM training and fine-tuning remains a significant challenge. Existing automated research agents struggle with LLM fine-tuning due to:

• Open-ended Complexity: Unlike optimizing fixed parameters or code patches, designing a training scheme involves a vast, open-ended space of variables including data distribution, algorithm selection, and hyperparameters.
• Context Limitations: Training data is massive and cannot be fully placed within an agent's context window, making it difficult for agents to reason about data directly.
• High Computational Cost: LLM training and evaluation are time-consuming and resource-intensive, rendering traditional evolutionary search methods (which rely on massive batch proposals) inefficient or infeasible.
• Lack of Specialized Benchmarks: Existing benchmarks focus on code synthesis or general ML engineering, lacking a systematic evaluation for end-to-end LLM fine-tuning.

2. Methodology: The TREX Framework

TREX is a multi-agent system designed to automate the LLM fine-tuning lifecycle through a dual-loop workflow orchestrated by two core agents: the Researcher and the Executor.

A. Agent Architecture

1. Researcher Agent:
   • Role: Responsible for high-level experimental design, literature review, and result analysis.
   • Function: It interprets task goals, retrieves relevant literature (via arXiv/Hugging Face), and formulates experimental plans. It uses a coarse-to-fine approach, first determining high-level strategies (e.g., "augment data") and then refining them into concrete configurations (e.g., specific data mixing ratios, learning rates).
   • Tools: Equipped with academic search engines and the AIDP (AI Data Processor) library.
2. Executor Agent:
   • Role: Implements the experimental plans on GPU clusters.
   • Function: Translates Researcher instructions into executable code, manages cluster resources, performs data processing, model training, and evaluation.
   • Safety: Operates within a sandboxed environment to prevent data corruption.

B. Core Mechanisms

• Tree-Based Exploration (MCTS):
  • The experimental process is modeled as a Monte Carlo Tree Search (MCTS). Each node represents an experimental iteration.
  • Strategy: The system balances exploration (trying novel strategies) and exploitation (refining high-performing paths) using the Upper Confidence Bound for Trees (UCT) algorithm.
  • Memory Management: To prevent context overflow, the agent uses a condensed memory context ($MC(v)$) that integrates the trajectory from the root to the current node, sibling nodes (for diversity), and critical historical nodes (for global insights).
• AIDP Library: A modular library built on HuggingFace Datasets that provides high-performance, deterministic operators for data ingestion, filtering, synthesis, and selection. This allows the agent to construct complex data pipelines autonomously.
• Fine-Grained Diagnostics: Instead of relying solely on final scores, the agent performs bad-case analysis and attribution studies on validation sets to understand why a model failed, feeding these insights back into the next experimental design.

3. Key Contributions

1. TREX System: The first automated research agent dedicated specifically to the full lifecycle of LLM fine-tuning, capable of handling open-ended tasks from data curation to model evaluation.
2. FT-Bench (Fine-Tuning Benchmark): A new benchmark comprising 10 diverse tasks derived from real-world scenarios (e.g., molecule generation, legal reasoning, financial QA). It evaluates agents on their ability to autonomously navigate the complete fine-tuning pipeline, filling a gap in existing benchmarks like MLE-bench or RE-bench.
3. Tree-Driven Strategy for High-Cost Tasks: Demonstrates that MCTS-based exploration is effective for LLM training, where computational budgets are tight, by efficiently reusing historical results and distilling insights rather than relying on brute-force sampling.
4. AIDP Library: Introduces a specialized toolkit for LLM agents to manage data engineering tasks reliably and reproducibly.

4. Experimental Results

The authors evaluated TREX on FT-Bench using Qwen3-1.7B as the base model for fine-tuning, with the Researcher powered by either Qwen3-Next-80B-Thinking or Gemini 3 Pro.

• Performance Gains: TREX consistently improved model performance across all 10 tasks.
  • Notable Successes: On the ACI-Bench (clinical notes), TREX achieved a 849% relative gain over the base model (using Gemini 3 Pro). On TOMG-Bench (molecule generation), it achieved a 108% gain, surpassing human-expert fine-tuned baselines (OpenMolIns-Large).
  • Comparison with Humans: In several tasks (e.g., TOMG-Bench, OpenFinData), TREX's automated optimization yielded performance gains comparable to or exceeding those achieved by human experts using complex, multi-stage pipelines (e.g., CPT-SFT-RL).
• Ablation Studies:
  • MCTS vs. Baselines: MCTS significantly outperformed Greedy Best-First Search (GBFS) and Sequential Expansion Search (SES) in stability and final performance.
  • AIDP Impact: Removing the AIDP library led to frequent data-processing failures and lower performance, highlighting its necessity for robust pipeline construction.
  • Bad-Case Analysis: Enabling the agent to analyze failure cases improved final scores, proving the value of fine-grained diagnostics.
• Strategy Diversity: The system successfully explored diverse strategies, including data pipeline refinement, synthetic data construction, and training scheme adjustments (e.g., mixing ratios, LoRA ranks).

5. Significance

• Scalability in AI Research: TREX demonstrates that AI agents can autonomously manage complex, resource-heavy scientific workflows (LLM training) that were previously too costly for automated systems.
• Democratization of Fine-Tuning: By automating the search for optimal data recipes and training strategies, TREX lowers the barrier for high-performance model adaptation, potentially reducing the reliance on manual expert tuning.
• New Paradigm for Evaluation: The introduction of FT-Bench sets a new standard for evaluating autonomous research agents, moving beyond code generation to the core of model capability enhancement.
• Future of Automated Science: The paper validates the "Agent-Driven Tree-based Exploration" paradigm as a viable path for closing the loop in scientific discovery, where agents not only generate hypotheses but also execute and refine the physical/digital experiments required to validate them.

In conclusion, TREX represents a significant leap forward in AI-for-AI, proving that multi-agent systems can autonomously optimize LLMs through iterative, tree-guided exploration, often matching or surpassing human expertise in specific domains.