ML-Master 2.0: The AI That Builds AI, and Ranked First Doing It

In partnership with

^{SnackOnAI Engineering | Senior AI Systems Researcher | Technical Deep Dive | May 3, 2026}

The prevailing assumption in AI-for-AI research is that better reasoning models produce better ML engineering agents. Give the agent a stronger LLM backbone, get better Kaggle solutions. ML-Master (SJTU, arXiv:2506.16499) challenges this directly. Its central claim: the bottleneck is not reasoning capability. It is the inability of reasoning models to effectively use the experience accumulated during exploration. Long exploration histories overwhelm the LLM context window, causing hallucinations and degraded reasoning quality. ML-Master's answer is not a bigger model. It is a selectively scoped memory mechanism that distills exploration trajectories into structured, bounded insights before passing them to the reasoning process.

ML-Master 2.0 (released December 2025) achieves 56.44% overall medal rate on MLE-Bench, a 92.7% improvement over v1's 29.33%, ranking first on the live leaderboard above Meta's AIRA-dojo, Microsoft's R&D-Agent, and Google's CAIR MLE-STAR-Pro-1.5. The most remarkable result: 152.2% improvement on medium-complexity tasks (20.18% → 50.88%), the tier where the exploration-reasoning integration matters most.

This newsletter dissects ML-Master's architecture as a systems problem: what MCTS-inspired parallel exploration does, how the selectively scoped memory mechanism extracts and bounds insights, what the three-layer Hierarchical Cognitive Caching system in v2 adds, and why medium-complexity tasks were the diagnostic battleground.

Scope: ML-Master 1.0 (arXiv:2506.16499) and 2.0 architectures, the MCTS-based exploration module, selectively scoped memory, HCC layers, MLE-Bench results. Not covered: X-Master (general scientific AI) beyond the thematic comparison.

What It Actually Does

ML-Master (SJTU SAI Agents Lab) is an AI-for-AI (AI4AI) agent: given a machine learning competition task (a Kaggle dataset, a problem statement, and a performance metric), it autonomously designs, implements, trains, and refines ML pipelines to maximize the score. No human writes the feature engineering code. No human selects the model architecture. No human tunes the hyperparameters.

MLE-Bench (OpenAI) is the evaluation environment: 75 real-world Kaggle competitions spanning tabular data, computer vision, and NLP, stratified by complexity (low/medium/high). The "medal rate" metric measures the proportion of tasks where the agent matches or exceeds the Kaggle bronze threshold, which itself requires beating the majority of human competitors. A 56.44% medal rate means ML-Master 2.0 achieves competitive performance (bronze or better) on more than half of 75 real Kaggle competitions, fully autonomously.

MLE-Bench Leaderboard (current, from the official page):

Rank	Agent	LLM	Low %	Medium %	High %	All %	Runtime
1	ML-Master 2.0	DeepSeek-V3.2-Speciale	75.76	50.88	42.22	56.44	24h
2	Leeroo	Gemini-3-Pro-Preview	68.18	44.74	40.00	50.67	24h
4	CAIR MLE-STAR-Pro-1.5	Gemini-2.5-Pro	68.18	34.21	33.33	44.00	24h
11	AIRA-dojo	o3	55.00	21.97	21.67	31.60	24h
13	ML-Master 1.0	deepseek-r1	48.48	20.18	24.44	29.33	12h

The gap on medium tasks (50.88% vs. 44.74% for second place) is the most diagnostic number. Medium-complexity tasks are where the exploration space is large enough that naive search fails, but tractable enough that a well-organized search succeeds. This is where the memory architecture earns its performance.

The Architecture, Unpacked

^{Focus on the bidirectional arrow between the memory mechanism and both the exploration and reasoning modules. Memory is not a one-way log. It actively shapes which parts of the search tree exploration expands next AND constrains what context the reasoning model sees. This bidirectional flow is the core architectural innovation.}

The Code, Annotated

Snippet One: MCTS Node Selection and Parallel Expansion (UCT-based)

# ML-Master's exploration uses MCTS-inspired tree search.
# This snippet reconstructs the core node selection and parallel expansion logic
# from the paper's description and architecture.

import math
import concurrent.futures
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class SolutionNode:
    """Represents one state in the ML development search tree."""
    node_id: str
    code: str                          # the current ML pipeline code
    validation_score: float            # score on the competition's validation set
    parent_id: Optional[str] = None
    children: list = field(default_factory=list)
    visit_count: int = 0               # N_node: times this node was selected
    total_reward: float = 0.0          # sum of scores from this node's subtree

def uct_score(node: SolutionNode, parent_visit_count: int, C: float = 1.41) -> float:
    """
    Upper Confidence Bound for Trees score.
    Balances exploitation (high score) with exploration (rarely visited nodes).

    ← THIS is the trick: nodes with high scores AND low visit counts
      get high UCT scores. This drives the agent toward:
      1. Refining solutions that are already working well (exploitation)
      2. Exploring promising but unvisited branches (exploration)
      Without UCT, the agent would greedily exploit the first good solution
      and never discover better solutions in unexplored branches.
    """
    if node.visit_count == 0:
        return float('inf')  # ← unvisited nodes always get selected first

    exploitation = node.total_reward / node.visit_count  # average score
    exploration = C * math.sqrt(math.log(parent_visit_count) / node.visit_count)
    return exploitation + exploration

def select_top_k_nodes(tree: dict, k: int) -> list[SolutionNode]:
    """
    Select top-k nodes by UCT score for parallel expansion.
    ← This is where parallelism enters: instead of picking ONE node
      (standard MCTS), ML-Master picks k nodes, enabling k workers to
      explore different branches of the solution space simultaneously.
    """
    root = tree['root']
    scored_nodes = [
        (uct_score(node, root.visit_count), node)
        for node in tree.values()
        if node.visit_count < max_visits  # don't over-exploit one branch
    ]
    scored_nodes.sort(reverse=True, key=lambda x: x[0])
    # ← return top-k, not top-1: the key parallelism decision
    return [node for _, node in scored_nodes[:k]]

def parallel_expand(selected_nodes: list[SolutionNode],
                    llm_client,
                    competition_spec: dict,
                    memory: 'SelectivelyScopedMemory') -> list[SolutionNode]:
    """
    Expand k nodes in parallel, each generating a new child solution.
    Each worker: (1) calls LLM with bounded memory context, (2) generates code,
    (3) executes it, (4) records the new score.
    """
    def expand_single(node: SolutionNode) -> SolutionNode:
        # ← bounded context: memory is scoped, not raw trajectory dump
        memory_context = memory.get_bounded_context(
            node=node,
            competition_spec=competition_spec,
            max_tokens=4096,  # hard cap prevents context window overflow
        )

        # LLM generates a refined solution based on:
        # - current code + its score
        # - bounded memory (what worked, what failed, cross-trajectory insights)
        new_code = llm_client.generate_refinement(
            current_code=node.code,
            current_score=node.validation_score,
            memory_context=memory_context,
            competition_spec=competition_spec,
        )

        # Execute the new code and measure validation score
        new_score = execute_and_evaluate(new_code, competition_spec)

        return SolutionNode(
            node_id=f"{node.node_id}_child_{hash(new_code)}",
            code=new_code,
            validation_score=new_score,
            parent_id=node.node_id,
        )

    # ← True parallel execution: all k expansions run simultaneously
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(selected_nodes)) as executor:
        futures = [executor.submit(expand_single, node) for node in selected_nodes]
        return [f.result() for f in concurrent.futures.as_completed(futures)]

The top-k selection is the parallelism mechanism. Standard MCTS selects one node. ML-Master selects k nodes and expands them simultaneously, enabling parallel exploration of different solution hypotheses. The bounded memory context passed to each worker (max 4096 tokens) is what prevents context window overflow during parallel calls.

Snippet Two: Selectively Scoped Memory Mechanism

# SelectivelyScopedMemory: the core mechanism that separates ML-Master
# from prior AI4AI agents that pass raw exploration logs to the reasoning model.

from dataclasses import dataclass
from typing import Optional

@dataclass
class TrajectoryInsight:
    """A distilled insight from one exploration trajectory."""
    approach_description: str     # what was tried (human-readable)
    validation_score: float       # score achieved
    key_changes: list[str]        # specific code changes that moved the score
    failure_reason: Optional[str] # if failed, root cause analysis
    is_redundant: bool = False    # if True, a similar trajectory is already in memory

class SelectivelyScopedMemory:
    """
    Distills exploration trajectories into bounded, structured insights.

    The problem with raw trajectory logs:
    - A 12-hour MCTS run generates thousands of code executions
    - Raw logs can be millions of tokens: context window overflow
    - LLM performance degrades with overly long contexts (hallucination risk)

    The solution: selective capture + bounding
    ← THIS is the trick: memory is NOT a log. It is a curated summary.
    """

    def __init__(self, max_insights: int = 20, max_tokens_per_context: int = 4096):
        self.insights: list[TrajectoryInsight] = []
        self.max_insights = max_insights            # hard cap on stored insights
        self.max_tokens_per_context = max_tokens_per_context

    def capture(self, trajectory: dict) -> None:
        """
        Evaluate a trajectory and selectively add to memory.
        Filters out low-information and redundant trajectories.
        """
        insight = TrajectoryInsight(
            approach_description=trajectory['description'],
            validation_score=trajectory['score'],
            key_changes=self._extract_key_changes(trajectory),
            failure_reason=self._analyze_failure(trajectory) if trajectory['score'] < 0.1 else None,
        )

        # ← Selection filter: only capture if it adds information
        if self._is_high_information(insight) and not self._is_redundant(insight):
            self.insights.append(insight)
            # ← Evict lowest-value insight if over capacity
            if len(self.insights) > self.max_insights:
                self._evict_lowest_value()

    def get_bounded_context(self, node, competition_spec, max_tokens: int) -> str:
        """
        Produce a bounded context for the reasoning model.
        Prioritizes high-scoring insights and failure analyses.
        ← Token budget enforced: reasoning model NEVER sees raw trajectory dump
        """
        # Sort by information value: high scores first, then failure analyses
        ranked_insights = sorted(
            self.insights,
            key=lambda x: (x.validation_score, x.failure_reason is not None),
            reverse=True,
        )

        context_parts = []
        token_budget = max_tokens

        for insight in ranked_insights:
            insight_text = self._format_insight(insight)
            insight_tokens = len(insight_text.split())  # rough token estimate

            if insight_tokens > token_budget:
                break  # ← hard stop: never exceed budget

            context_parts.append(insight_text)
            token_budget -= insight_tokens

        return "\n\n".join(context_parts)

    def _is_high_information(self, insight: TrajectoryInsight) -> bool:
        """
        Filter: keep insights that changed the score meaningfully
        OR provide a clear failure root cause for future avoidance.
        """
        # ← Threshold-based: minor score fluctuations are noise, not signal
        is_meaningful_improvement = insight.validation_score > self._current_best * 1.01
        is_informative_failure = insight.failure_reason is not None and len(insight.failure_reason) > 50
        return is_meaningful_improvement or is_informative_failure

    def _is_redundant(self, insight: TrajectoryInsight) -> bool:
        """Check if a similar insight is already in memory."""
        for existing in self.insights:
            if self._approaches_are_similar(insight, existing):
                insight.is_redundant = True
                return True
        return False

The get_bounded_context hard stop is the most important implementation detail. Without a token budget, a reasoning model receiving an unbounded memory dump will hallucinate. The bound is not a soft guideline; it is a hard constraint enforced before every reasoning call.

It In Action: End-to-End Worked Example

Scenario: ML-Master 2.0 running on a medium-complexity tabular Kaggle competition (customer churn prediction, a representative medium-complexity task).

Input: Competition specification, 24-hour time limit, validation metric: AUROC.

Phase 1: Initial exploration (hours 0-3)

Tree state: root node only
Strategy: explore diverse baseline approaches

Worker 1: LightGBM with default features
  → validation AUROC: 0.841
  → memory captures: "LightGBM baseline AUROC 0.841, default features"

Worker 2: XGBoost with interaction features
  → validation AUROC: 0.849
  → memory captures: "XGBoost + interaction features +0.008 over LightGBM"

Worker 3: Neural network (tabular transformer)
  → validation AUROC: 0.823
  → memory captures: "Tabular transformer underperformed, likely small dataset"

Worker 4: Feature engineering (target encoding)
  → validation AUROC: 0.857
  → memory captures: "Target encoding improves by +0.016 over default features"

Memory at hour 3: 4 insights, bounded to 800 tokens
Best score: 0.857 (XGBoost + target encoding)

Phase 2: Reasoning-guided refinement (hours 3-10)

Reasoning model receives bounded memory context (800 tokens):
  "Target encoding improved by +0.016 over defaults.
   XGBoost outperforms LightGBM by +0.008 on this task.
   Tabular transformer underperformed: dataset likely too small for deep models.
   Hypothesis: ensemble of XGBoost + feature engineering variants could gain further."

Reasoning model output:
  1. Ensemble XGBoost + LightGBM (memory shows both competitive)
  2. Add polynomial interaction features on top of target encoding
  3. Try Bayesian hyperparameter optimization for XGBoost

MCTS expansion (top-3 UCT nodes selected):
Worker 1: XGBoost + target encoding + polynomial features → 0.863
Worker 2: Ensemble XGBoost + LightGBM → 0.868
Worker 3: XGBoost Bayesian tuning → 0.871

Memory at hour 10: 7 insights, bounded to 1,400 tokens
Best score: 0.871

Phase 3: HCC L2 trigger and plan revision (hour 10-18)

Score stagnation detected (0.871 → 0.872 over 3 hours, marginal gain)
L2 cache triggered: tactical plan revision

L2 memory synthesizes: "Core XGBoost pipeline saturated. Gains are coming from
  feature engineering, not model architecture. New direction: advanced feature
  extraction via RAPIDS cuDF and feature selection."

New exploration launched:
Worker 1: SHAP-based feature selection → removes 40 low-importance features → 0.875
Worker 2: RAPIDS cuDF acceleration + new temporal features → 0.878
Worker 3: Post-processing calibration (sigmoid) → 0.879

Best score at hour 18: 0.879

Phase 4: Final submission (hours 18-24)

Ensemble best 3 solutions: AUROC 0.883
Test submission: 0.881 (slight variance from validation)
Kaggle bronze threshold for this competition: ~0.875
Result: BRONZE MEDAL → counts toward medal rate

Total tree nodes explored: 847
Memory distillations performed: 31
Peak memory context size: 3,200 tokens (well within 4,096 budget)
LLM API calls: 312 (exploration) + 28 (reasoning) = 340 total

Why This Design Works, and What It Trades Away

The MCTS-based exploration is the correct search algorithm for ML engineering for the same reason it works in game playing: the solution space is too large for exhaustive search, reward (validation score) is delayed until code execution completes, and promising directions can be identified only after observing score improvements across multiple iterations. The UCT formula's exploration bonus prevents the common failure mode of greedy agents that exploit the first good solution and miss significantly better alternatives.

The selectively scoped memory solves a problem that prior AI4AI agents did not explicitly address. Research on long-context LLM degradation shows consistent quality drops as context length increases beyond the model's effective attention range. An agent running for 12-24 hours generates exploration logs orders of magnitude longer than any LLM's effective context window. ML-Master's memory mechanism is not a nice-to-have. It is a necessary component for maintaining reasoning quality across a multi-hour autonomous run.

The HCC three-layer hierarchy addresses the temporal dynamics of ML engineering tasks. Short-term memory (L1) handles immediate execution feedback. Mid-term memory (L2) handles strategic plan revision when the current approach plateaus. Long-term memory (L3) handles cross-task heuristics accumulated over many competition runs. Each layer operates at a different time scale, and the ablation confirms all three are necessary: removing L3 drops MLE-Bench-Lite from 72.7% to 54.5%.

What ML-Master trades away:

Interpretability. The agent's decision process is a combination of MCTS tree navigation, LLM code generation, and memory distillation. No component of this process produces a human-readable explanation of why a specific approach was chosen. The code works; the reasoning behind it is opaque.

Determinism. Parallel MCTS with LLM code generation is non-deterministic. Different runs of the same competition may explore different branches and produce different final solutions. The benchmark reports variance (±1.24% for low complexity, ±2.86% for medium) that reflects this.

Task scope. ML-Master is specialized for structured ML competition tasks: feature engineering, model selection, hyperparameter optimization. It does not generalize to arbitrary scientific tasks (see X-Master for that direction) or to software engineering beyond ML pipelines.

Technical Moats

The UCT scoring with bounded parallel expansion is operationally complex. Running 4-8 parallel workers each executing ML training jobs (potentially 10-30 minutes each), managing the tree state, and coordinating memory updates across workers without race conditions requires careful distributed systems engineering. The Docker deployment (sjtuagents/ml-master:latest) with --gpus all --shm-size=64g signals the infrastructure requirements. Getting this right at 24-hour run durations with GPU training loops is non-trivial.

The memory distillation quality determines everything. If the _is_high_information threshold is too aggressive, useful insights get filtered and the reasoning model is context-starved. Too lenient, and the context window fills with noise. The correct calibration is task-dependent and competition-dependent. The fact that ML-Master achieves 92.7% improvement over v1 while using the same LLM backbone (roughly equivalent reasoning capacity) is evidence that the memory calibration is the primary lever.

MLE-Bench requires genuine ML expertise in code generation. The competitions involve feature engineering decisions (target encoding, polynomial features, SHAP selection), model selection (XGBoost vs. LightGBM vs. neural), hyperparameter tuning (Bayesian optimization), and ensemble construction. An agent that generates superficially correct code but makes poor ML engineering decisions will score below the bronze threshold. The fact that ML-Master 2.0 achieves 75.76% on low-complexity tasks means it generates genuinely competitive ML code, not just syntactically valid code.

Insights

Insight One: ML-Master's 92.7% improvement from v1 to v2 using a stronger LLM backbone is not evidence that better LLMs are the answer. It is evidence that the memory architecture scales with model capability better than the previous architecture did.

The community narrative around AI4AI agents is LLM-centric: better foundation models produce better agents. ML-Master's results are more nuanced. v1 used DeepSeek-R1 and achieved 29.33%. v2 uses DeepSeek-V3.2-Speciale and achieves 56.44%. But the architectural change from v1 to v2, specifically the HCC three-layer memory system, also contributes to the improvement. The ablation showing 72.7% vs. 54.5% without L3 suggests the architecture accounts for a substantial fraction of the gain. Attributing the entire improvement to the stronger LLM is incorrect. The memory architecture enables the model to use its capabilities more effectively.

Insight Two: The 152.2% improvement on medium-complexity tasks, not the overall result, is the most diagnostic number in the paper, and the community has largely missed why.

Low-complexity tasks (75.76%) are competitive with naive approaches because the solution space is small enough that even basic exploration finds good solutions. High-complexity tasks (42.22%) involve problem structures where 24 hours may simply be insufficient regardless of architecture. Medium-complexity tasks are the regime where the search architecture matters: complex enough that naive approaches fail, tractable enough that intelligent search succeeds. The 152.2% improvement at medium complexity (20.18% → 50.88%) is the clearest evidence that MCTS-based exploration plus selectively scoped memory is specifically solving the search efficiency problem, not just benefiting from a stronger LLM.

Takeaway

ML-Master 1.0 achieved 29.3% medal rate with a 12-hour time constraint, while prior baselines with 24-hour limits scored below that. The memory architecture, not the extra 12 hours, was responsible for the v1 performance advantage.

This is the implicit claim in the v1 paper and the explicit comparison in the leaderboard: ML-Master 1.0 at 12 hours outperforms AIRA-dojo (o3, 24 hours) at 31.60% only marginally, but the 12-hour performance is achieved at half the compute budget. At 24 hours, AIRA-dojo's performance with o3 should theoretically benefit from more exploration time. The fact that ML-Master 2.0 at 24 hours dominates every competitor at every complexity tier suggests the selectively scoped memory mechanism compresses exploration efficiency beyond what additional wall-clock time alone provides. Time is not the bottleneck. Memory quality is.

TL;DR For Engineers

ML-Master 2.0 ranks first on MLE-Bench with 56.44% medal rate (92.7% improvement over v1), using MCTS-based parallel exploration, selectively scoped memory, and three-layer Hierarchical Cognitive Caching. The LLM backbone is DeepSeek-V3.2-Speciale.
The MCTS explores solutions as a tree (nodes = code states, edges = refinements). Top-k UCT selection enables parallel expansion across multiple workers. The UCT formula balances exploitation of high-scoring nodes with exploration of unvisited branches.
Selectively scoped memory distills exploration trajectories into bounded structured insights before passing them to the reasoning model. Hard token cap (4096 tokens per context call) prevents context window overflow and hallucination.
HCC three layers: L1 (within-phase execution feedback), L2 (cross-phase tactical revision triggered by score stagnation), L3 (cross-task strategic heuristics). Ablation: removing L3 drops MLE-Bench-Lite from 72.7% to 54.5%.
Medium-complexity improvement (+152.2%) is the most diagnostic result: this is the regime where search architecture matters, proving the MCTS + memory approach specifically solves search efficiency, not just benefiting from better LLMs.

Memory Is the Architecture. The LLM Is the Tool.

ML-Master's core insight is that AI4AI agents fail not because their reasoning models are too weak, but because their memory systems are too naive. Raw exploration logs overflow LLM context windows and degrade reasoning quality. Selectively scoped memory distills those logs into bounded, structured insights that enable high-quality reasoning to guide further exploration. The MCTS tree structure provides the search framework. The HCC layers provide temporal breadth. But the memory mechanism is what makes both of those work. Without it, the MCTS generates a tree of increasingly confused reasoning calls. With it, each reasoning call gets a clean, bounded summary of what the agent has learned.

The first-place MLE-Bench result is the proof. Not "plausible performance on benchmark tasks." First place, above Meta, Microsoft, and Google systems, on 75 real Kaggle competitions, without any human intervention.

References

ML-Master project page, SJTU SAI Agents Lab
ML-Master GitHub Repository, official implementation
ML-Master: Towards AI-for-AI via Integration of Exploration and Reasoning, arXiv:2506.16499, Liu et al., SJTU, June 2025
Toward Ultra-Long-Horizon Agentic Science: Cognitive Accumulation for Machine Learning Engineering, arXiv:2601.10402, Zhu et al., Jan 2026 (ML-Master 2.0 foundation paper)
MLE-Bench: Evaluating Machine Learning Agents on Machine Learning Engineering, OpenAI
X-Master: Towards General-Purpose Scientific AI Agents, related SJTU work on general scientific AI
AIRA-dojo (Meta), rank 11 baseline on MLE-Bench
R&D-Agent (Microsoft), ML engineering agent benchmark comparison
ML-Master 2.0 MLE-Bench Leaderboard, live standings

Summary

ML-Master 2.0 (SJTU SAI Agents Lab) is an AI-for-AI agent that ranks first on MLE-Bench with 56.44% medal rate (92.7% over v1's 29.33%), using MCTS-based parallel exploration of ML solution trees, a selectively scoped memory mechanism that distills exploration trajectories into bounded structured insights before passing them to a reasoning LLM (DeepSeek-V3.2-Speciale), and Hierarchical Cognitive Caching (HCC) with three memory layers operating at within-phase, cross-phase, and cross-task time horizons. The most diagnostic result is the 152.2% medium-complexity task improvement (20.18% → 50.88%), demonstrating that the architecture specifically solves search efficiency in regimes where naive exploration fails, not merely benefiting from a stronger LLM backbone.

Sponsored Ad

If you enjoy practical AI insights, check out SnackOnAI and support the newsletter by subscribing, sharing, and exploring our sponsored ad — it helps us keep building and delivering value 🚀

Your prompts are leaving out 80% of what you're thinking.

When you type a prompt, you summarize. When you speak one, you explain. Wispr Flow captures your full reasoning — constraints, edge cases, examples, tone — and turns it into clean, structured text you paste into ChatGPT, Claude, or any AI tool. The difference shows up immediately. More context in, fewer follow-ups out.

89% of messages sent with zero edits. Used by teams at OpenAI, Vercel, and Clay. Try Wispr Flow free — works on Mac, Windows, and iPhone.

Start flowing free