PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zerun Ma, Guoqiang Wang, Xinchen Xie, Yicheng Chen, He Du, Bowen Li, Yanan Sun, Wenran Liu, Kai Chen, Yining Li

arXiv:2604.14116v2cs.AIcs.CL

TL;DR

Automating open-ended LLM fine-tuning remains difficult because training requires coordinating data, algorithms, hyperparameters, and expensive experiments. TREX addresses this with Researcher–Executor collaboration, tree-based exploration, and historical-memory reuse, evaluated on the 10-task FT-Bench benchmark. TREX consistently improves performance across tasks, surpasses expert-designed recipes on several tasks, and remains effective with an open-source reasoning backend.

  • Problem

    Automating LLM fine-tuning is challenging because schemes combine data, algorithms, and hyperparameters, while existing evaluation leaves end-to-end LLM training insufficiently covered.

  • Method

    TREX orchestrates Researcher–Executor collaboration for research, data construction, training, and evaluation, while MCTS organizes iterative exploration and condensed memory reuses experimental history.

  • Results

    TREX consistently enhances base-model performance across FT-Bench tasks and surpasses expert-designed fine-tuning recipes on several tasks.

  • Takeaways & Limitations

    Automated agents can optimize open-ended LLM fine-tuning strategies and data across diverse real-world tasks under limited resource and time budgets.

Abstract

from arXiv · show

While Large Language Models (LLMs) have empowered AI research agents to perform isolated scientific tasks, automating complex, real-world workflows, such as LLM training, remains a significant challenge. In this paper, we introduce TREX, a multi-agent system that automates the entire LLM training life-cycle. By orchestrating collaboration between two core modules-the Researcher and the Executor-the system seamlessly performs requirement analysis, open-domain literature and data research, formulation of training strategies, preparation of data recipes, and model training and evaluation. The multi-round experimental process is modeled as a search tree, enabling the system to efficiently plan exploration paths, reuse historical results, and distill high-level insights from iterative trials. To evaluate the capability of automated LLM training, we construct FT-Bench, a benchmark comprising 10 tasks derived from real-world scenarios, ranging from optimizing fundamental model capabilities to enhancing performance on domain-specific tasks. Experimental results demonstrate that the TREX agent consistently optimizes model performance on target tasks.

1. Introduction

TREX addresses the difficulty of automating open-ended LLM fine-tuning by coordinating research, data preparation, experimentation, and evaluation. It uses iterative tree-based exploration and reports improvements across diverse fine-tuning tasks, including gains beyond expert-designed recipes on several tasks.

  • Motivation: Automating LLM training is difficult because fine-tuning requires jointly designing data distributions, algorithms, and hyperparameters, while training data is too large for direct agent-context inclusion.The challenge remains even when model architecture and initialization weights are fixed.
  • TREX: TREX coordinates a Researcher that analyzes requirements, literature, and strategies with an Executor that constructs data, trains models, evaluates them, and returns results.This collaboration closes the loop between experimental execution and research analysis.
  • TREX: MCTS-based exploration selects and optimizes promising paths from historical experiments to generate training plans without extensive rollouts.The approach is inspired by researchers retaining empirically validated strategies while exploring alternatives.
  • Evaluation: FT-Bench contains 10 evaluation datasets spanning general capability enhancement and vertical domain adaptation for assessing automated LLM fine-tuning.The benchmark measures whether agents improve model performance through fine-tuning.
  • Results: TREX continuously optimizes fine-tuning strategies and data under limited resource and time budgets, enhancing performance on the evaluation tasks.The reported improvements cover the benchmark tasks rather than a single isolated setting.
  • Results: On several tasks, TREX achieves performance gains that surpass expert-designed fine-tuning recipes, while remaining effective with the open-source Qwen3-Next-80B backend.The strongest results are obtained with Gemini 3 Pro or Claude 4.5 Sonnet as proprietary reasoning backends.

2. Related Work

Prior AI research systems automate parts of discovery, writing, implementation, and evaluation, but existing approaches generally constrain search or isolate components. TREX instead integrates data preparation into a holistic automatic loop for open-ended LLM training.

  • Automated AI research: AI systems have demonstrated autonomous idea generation and experiment implementation, although human oversight remains necessary.Examples include Dolphin, InternAgent, AI-Scientist, and AI-Researcher.
  • Automated AI research: End-to-end research agents cover workflows such as idea generation, experimentation, literature review, algorithm optimization, and result analysis.The cited systems differ in how many research functions they integrate.
  • Search-based methods: Evolutionary and tree-search methods support diverse research ideas but depend on large-scale sampling, limiting their practicality for computationally expensive LLM fine-tuning.The cited limitation concerns tasks with high computational overhead.
  • Benchmarking: Existing benchmarks evaluate code generation, machine-learning engineering, and related research tasks, but do not fully cover automated LLM training workflows.These limitations motivate dedicated evaluation for LLM fine-tuning.
  • AutoML: Traditional AutoML automates model selection and hyperparameter configuration, while newer LLM methods remain constrained by predefined search spaces or isolated components.TREX targets a more open-ended setting by automating the entire LLM training lifecycle.
  • AI for data construction: LLM data-construction research addresses synthesis, refinement, and filtering, but typically uses LLMs as discrete tools within predetermined protocols.TREX integrates data preparation into an autonomous optimization loop.

3. TREX

TREX uses a dual-loop, multi-agent workflow to design, execute, diagnose, and iteratively improve LLM fine-tuning experiments. Its tree search, data-processing tools, diagnostics, and memory management organize exploration under practical constraints.

  • Overview: TREX’s inner loop uses Researcher–Executor collaboration to complete each experiment, while its outer loop represents successive experiments as a search tree.The Researcher formulates plans from objectives and historical records; the Executor carries them out.
  • Iterative exploration: MCTS balances exploiting high-performing solutions with exploring novel fine-tuning schemes when expanding experiment nodes.Each new experimental round creates a new node in the tree.
  • Executor: The Executor implements plans, processes data, trains models, evaluates them, and uses cluster resources through integrated scheduling tools.Its sandbox isolates the experimental environment from cluster data.
  • Researcher: The Researcher moves from high-level improvement strategies to concrete plans covering objectives, data procedures, and training configurations.It can search literature and locate accessible datasets to support planning.
  • Data processing: AIDP supplies high-performance, reliable, deterministic operators for reproducible and parallel LLM training-data pipelines.The library is built on the HuggingFace Datasets ecosystem to reduce computational overhead.
  • Diagnostics: TREX combines task metrics, cross-model comparisons, and bad-case analysis to extract richer feedback for refining subsequent training schemes.The diagnostic mechanism addresses the limited guidance provided by benchmark scores alone.
  • Iterative exploration: MCTS rewards each experiment using the normalized primary evaluation metric and can compare parallel models with slightly varied configurations.The best-performing model in the batch determines the experiment reward.
  • Memory management: Memory context condenses root-to-node trajectories, sibling nodes, and critical tree nodes to preserve useful history while supporting diversity and global insight sharing.This prevents redundant history from exceeding the LLM context window.

4. FT-Bench

FT-Bench fills an evaluation gap for automated LLM fine-tuning by using realistic, controllable tasks and requiring agents to navigate the complete open-environment pipeline.

  • Benchmark gap: Existing benchmarks often evaluate isolated subtasks or traditional machine-learning settings, missing challenges unique to modern LLM training.Examples include instruction formatting, domain-specific evaluation, and full-lifecycle orchestration.
  • Design principles: FT-Bench selects tasks from real-world research or industrial applications to reflect practical LLM fine-tuning proficiency.Realistic task origins are one of the benchmark’s two design principles.
  • Design principles: FT-Bench controls computational and data overhead so multi-round experimental iterations are not bottlenecked by excessively demanding tasks.Controllability is the benchmark’s second stated design principle.
  • Benchmark scope: FT-Bench is dedicated to end-to-end LLM fine-tuning and requires agents to autonomously navigate the complete pipeline in an open environment.This addresses both the isolation and traditional-ML limitations identified in existing benchmarks.

5. Expriment

TREX is evaluated on 10 FT-Bench fine-tuning tasks using controlled experimental settings, normalized relative gains, and iterative exploration. Results show consistent base-model improvement, advantages for Gemini 3 Pro, competitiveness with human-designed fine-tuning, and benefits from MCTS, AIDP, and bad-case analysis.

  • Experimental Setup: Relative performance gains normalize each task’s improvement by the gap between the base model and the superior reference model Qwen3-235B-2507.The task-specific evaluation function measures models before and after autonomous fine-tuning.
  • Main Results: TREX consistently improves the base model across all evaluated tasks, with Gemini 3 Pro outperforming Qwen3-Next-80B-Thinking on most tasks.The results attribute overall efficacy differences to the underlying Researcher LLM’s reasoning capability.
  • Experimental Setup: TREX evaluates Qwen3-1.7B across 10 FT-Bench tasks, comparing Qwen3-Next-80B-Thinking and Gemini 3 Pro as Researcher backends.The Executor uses Claude 4.5 Sonnet, and experiments cap exploration at 20 iterations per task.
  • Main Results: TREX gains 0.498 on TOMG-Bench and 0.205 on OpenFinData, compared with expert-recipe improvements of 0.189 and 0.139 on TOMG-Bench and 0.025 or 0.207 for FEVO on OpenFinData.The comparison uses different fine-tuned models and, for FEVO-R32B, a complex CPT-SFT-RL pipeline.
  • Main Results: Gemini 3 Pro yields greater strategy diversity and execution success than Qwen3-Next-80B-Thinking across experimental trials.TREX uses a coarse-to-fine expansion process that begins with high-level strategies before generating experimental configurations.
  • Ablation Studies: MCTS produces more stable trajectories and more consistent gains than GBFS and SES, while AIDP and bad-case analysis improve subsequent experimentation and final performance.Without AIDP, performance improvements are lower and later training is more interruption-prone because of data-processing failures; analyzing bad cases improves iteration effectiveness.

6. Conclustion

TREX automates the full LLM fine-tuning lifecycle through coordinated planning, execution, and iterative refinement, using tree-based search to explore training strategies under constrained budgets. Across diverse tasks, it consistently improves performance and sometimes matches or surpasses expert-designed pipelines.

  • TREX automates the full lifecycle of LLM fine-tuning through coordinated planning, execution, and iterative refinement.
  • TREX formulates training optimization as a tree-based search problem and uses MCTS to explore strategies under constrained computational budgets.
  • TREX consistently improves model performance across diverse tasks in FT-Bench.
  • In several cases, TREX matches or surpasses expert-designed pipelines.

A.1. Details of the tasks in FT-Bench

FT-Bench includes tasks spanning medical note generation, molecule generation, organic reaction reasoning, cancer-literature classification, and bilingual computer-science question answering, with task-specific metrics and dataset splits.

  • ACI-Bench evaluates AI-assisted medical-note generation from clinical visit dialogues using Rouge-1.
  • TOMG-Bench evaluates text-based open molecule generation with Score = 0.4×Validity+0.6×Accuracy.
  • oMeBench evaluates organic chemical reaction mechanism reasoning with a weighted score combining SMILES Validity, Logical Fidelity Score, oMeS-total, and oMeS-partial.
  • Hoc evaluates cancer-related scientific-literature classification using Macro-F1, while CS-Bench evaluates bilingual computer-science questions using accuracy.
  • For oMeBench, 293 training samples are held out for validation and 2,200 remaining samples are used for training.

A.2. AIDP Package

AIDP equips TREX with high-performance data-processing operators built on HuggingFace Datasets, enabling agents to orchestrate complex data pipelines through function calls in generated scripts.

  • AIDP provides high-performance data-processing operators built on the HuggingFace Datasets ecosystem.
  • TREX accesses AIDP by using exposed Python signatures to orchestrate and implement complex data pipelines through function calls within generated scripts.
  • The toolkit supports experimental robustness, reproducibility, and efficient large-scale data manipulation.

A.3. Additional Experiment Results

Figures 7 and 8 show TREX performance trajectories across FT-Bench tasks, comparing iterative refinement with Gemini 3 pro and Qwen3-next-80b-thinking as the Researcher.

  • Figures 7 and 8 illustrate TREX performance trajectories across all tasks in FT-Bench.
  • Figure 7 compares iterative experimental-refinement trajectories for oMeBench, ACI-Bench, GTA, and HOC.
  • Figure 7 uses Gemini 3 pro as the Researcher in the left panel and Qwen3-next-80b-thinking in the right panel.
  • Figure 8 compares iterative experimental-refinement trajectories for LawBench, EconlogicQA, SST-2, and CS-Bench.
  • Table 5 overviews AIDP primitive operators for data ingestion, formation, analysis, synthesis, and selection.

A.4. Experiment Plan Examples

TREX’s experiment plans explore training scale, batching, learning rate, epochs, LoRA rank, and data composition across iterative molecule-generation trials. The examples also vary QED coverage and the use of natural chemical descriptions to target generalization and instruction naturalness.

  • Experiment trajectory: The molecule-generation trajectory showcases three experimental rounds, including a baseline and later interventions, within TOMG-Bench.The documented examples are Experiment-1, Experiment-4, and Experiment-12.
  • Baseline and hyperparameter search: Grid search varies data volume, batch size, learning rate, training steps, and LoRA rank to establish a baseline and identify reasonable parameter ranges.The stated purpose is to validate the experimental pipeline and determine baseline hyperparameters.
  • Data design: TOMG-Bench supplies 90,000 instruction-tuning samples for molecule editing, optimization, and generation, with random subsets of 10K, 25K, or 50K used in experiments.The dataset is formatted for LLaMA Factory and does not overlap with the evaluation set, according to the cited passage.
  • Baseline and hyperparameter search: The plans compare training epochs of 1 versus 3 and LoRA ranks of 16 versus 32 while holding other settings fixed in the listed groups.The learning-rate groups compare 1e-4 and 5e-4; the plans also specify lora_rank = 8 in several controlled configurations.
  • Data design: A QED-focused experiment constructs 10,000 instruction-tuning samples by randomly selecting 3,333 examples from each of three QED intervals.The intervals are 0.0–0.3, 0.3–0.7, and 0.7–1.0, creating a uniform distribution intended to improve adaptability to diverse QED constraints.
Loading 2604.14116v2…