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

Published
Source
arXiv
Paper number
1223
Field
Research Agents
arXiv ID
2604.14116

Key points

  • Problem: Existing research agents iteratively optimize targets that fit in short text, such as hyperparameters or code patches, whereas LLM fine-tuning is open-ended and requires jointly choosing the training data distribution, algorithms, and hyperparameters. Training data is too large to place in an agent's context, and each training and evaluation run costs substantial time and compute, so evolutionary frameworks that generate and verify proposals in large batches become inefficient or infeasible.
  • Input and output: This is a system paper that proposes TREX, a multi-agent system automating the whole LLM fine-tuning cycle, together with the FT-Bench benchmark. TREX receives a task description with the objective, evaluation data and metrics, optional training data, and experimental settings, gathers and builds data, fine-tunes Qwen3-1.7B, and returns the model and score of its best round. FT-Bench consists of 10 tasks drawn from real scenarios, including clinical note generation, molecule generation, chemical mechanism reasoning, finance, law, and tool use.
  • Method: The Researcher agent searches arXiv papers and Hugging Face datasets, sets a high-level improvement strategy, and writes an experimental plan that usually holds 3 to 5 parallel training configurations, while an Executor agent built on OpenHands runs data processing, training, and evaluation on a GPU cluster. Rounds of experiments are recorded as a tree, and MCTS (Monte Carlo Tree Search) picks the node with the highest UCT score to refine next, with the normalized primary task metric as the reward. Each experiment gets a bad-case analysis and attribution, which is condensed with the root path, sibling nodes, and critical nodes into a memory context for the next design.
  • Implementation: The Executor wraps cluster task-scheduling APIs as tools on top of OpenHands, and the authors built AIDP, a HuggingFace Datasets based data-processing library of loader, scorer, generator, and filter operators that the agent composes into data pipelines. In experiments the Researcher backend is Qwen3-Next-80B-Thinking or Gemini 3 Pro and the Executor uses Claude 4.5 Sonnet, with at most 20 experiment rounds per task and at most 50,000 training samples per experiment. Performance is compared as a normalized relative gain with Qwen3-235B-2507 as the reference model.
  • Results: TREX raises the score of the base model Qwen3-1.7B on all 10 tasks, and with Gemini 3 Pro as the Researcher, ACI-Bench Rouge-1 rises from 0.205 to 0.502 and HoC Macro-F1 from 0.462 to 0.897. CS-Bench, where the training set must be built from scratch, moves only from 0.532 to 0.581, a 15% normalized gain. On TOMG-Bench, TREX's gain of 0.498 exceeds the 0.189 gain of OpenMolIns-Large, which the TOMG-Bench authors fine-tuned on Llama3.1-8B.
  • Limits: Every task uses the single base model Qwen3-1.7B, within at most 20 rounds per task and 50,000 training samples per experiment. Ablations of the tree-search strategy, AIDP, and bad-case analysis cover only oMeBench and GTA with Gemini 3 Pro and are reported as score-curve figures. The comparison with human experts covers TOMG-Bench and OpenFinData, and the human baselines are gains measured on different base models such as Llama3.1-8B and Qwen2.5-32B.

Paper links

External research summaries. These are not HDATF publications or measured product results.

Read original (opens in a new tab)