MLExpert: Training an Agent to Run and Improve ML Experiments

Introduction

Generating a training script is only one small part of machine-learning engineering. A capable engineer must decide how to prepare data, choose a model and optimization strategy, run an experiment, inspect the result, and determine what to try next. These decisions form a long chain in which an early choice can affect many later outcomes. MLExpert is aimed at automating this broader process.

The project introduces a 35B mixture-of-experts agent and a post-training framework for autonomous, long-horizon MLE optimization. Instead of depending primarily on a hand-written workflow, the model uses tools to organize experiments and adapts its strategy according to execution feedback.

Key points

• An executable training environment. The Kaggle-MLE Factory constructs machine-learning tasks and provides an interaction harness. The agent must write code, launch experiments, observe results, and continue working rather than produce a single static answer.
• Hierarchical agentic SFT. Hierarchical Agentic Supervised Fine-Tuning transfers general coding and tool-use abilities into the MLE setting. It also uses outcome-guided supervision to select more useful signals from long interaction traces.
• Windowed credit assignment. Windowed Interaction Reinforcement Learning, or WIRL, treats the multiple model decisions between two successive experiment results as one execution window. A shared outcome-based reward is then assigned to that window, allowing intermediate actions to learn from the result they jointly produced.
• Updates before a run ends. Asynchronous collection and policy correction are designed to support policy updates before a complete trajectory has finished. This matters because MLE runs can contain many expensive experiments and may otherwise delay learning until the end.
• Reported benchmark gains. On MLE-Bench Lite, MLExpert improves the base model’s average medal rate from 37.9% to 72.7%, reaching performance comparable to larger frontier models cited by the authors. The paper also reports experiments on NatureBench Lite, agent scaffold optimization, and algorithm optimization.

Why it matters

The central contribution is a way to frame machine-learning engineering as sequential decision making rather than code completion. In a real experiment loop, a failed run is not simply a bad answer: it can provide information about the next search direction. Conversely, a final score often reflects several earlier choices. WIRL addresses this delayed-credit problem by grouping related decisions around observable experiment outcomes.

The approach is also relevant to AI-for-AI research. If an agent can reliably manage parts of model development, it may reduce the amount of manual orchestration required for experimentation. This could make automated model and algorithm search more practical, particularly when workflows involve repeated tool calls and adaptive planning.

The results should still be read within their stated scope. The reported improvements come from benchmark and harness settings, and they do not by themselves establish robust performance in open-ended research. Cost control, recovery from faulty code or misleading feedback, reproducibility, and the reliability of conclusions remain important questions. Transfer across domains also needs broader evidence.

The direction nevertheless shifts attention from whether a model can write machine-learning code to whether it can manage an evolving, expensive experiment program over time.

Source: Hugging Face Daily Papers