Reasoning as Gradient: Scaling MLE Agents Beyond Tree Search

Yifei Zhang; Xu Yang; Xiao Yang; Bowen Xian; Qizheng Li; Shikai Fang; Jingyuan Li; Jian Wang

cs.LGMar 2, 2026

Reasoning as Gradient: Scaling MLE Agents Beyond Tree Search

Authors: Yifei Zhang, Xu Yang, Xiao Yang, Bowen Xian, Qizheng Li, Shikai Fang, Jingyuan Li, Jian Wang,

Organizations: Nanjing University · Microsoft Research Asia · Peking University · Zhejiang University ·

Abstract

LLM-based agents for machine learning engineering (MLE) predominantly rely on tree search, a form of gradient-free optimization that uses scalar validation scores to rank candidates. As LLM reasoning capabilities improve, exhaustive enumeration becomes increasingly inefficient compared to directed updates, analogous to how accurate gradients enable efficient descent over random search. We introduce Gome, an MLE agent that operationalizes gradient-based optimization. Gome maps structured diagnostic reasoning to gradient computation, success memory to momentum, and multi-trace execution to distributed optimization. Under a closed-world protocol that isolates architectural effects from external knowledge, Gome achieves a state-of-the-art 35.1% any-medal rate on MLE-Bench with a restricted 12-hour budget on a single V100 GPU. Scaling experiments across 10 models reveal a critical crossover: with weaker models, tree search retains advantages by compensating for unreliable reasoning through exhaustive exploration; as reasoning capability strengthens, gradient-based optimization progressively outperforms, with the gap widening at frontier-tier models. Given the rapid advancement of reasoning-oriented LLMs, this positions gradient-based optimization as an increasingly favorable paradigm. We release our codebase and GPT-5 traces at https://github.com/microsoft/RD-Agent.

Figures & tables

Figure 1: Two paradigms for utilizing execution feedback. (a) Score-centric: feedback is compressed into scalar rewards that guide node/branch selection (which may itself be directed, e.g., via UCT/PUCT). (b) Gradient-based: structured reasoning over execution feedback determines how to update the current solution; stronger reasoning yields more accurate signals.

Method	Structure	Feedback Role	Ext. Know.	Code
AIDE	Tree	Ranking	✗	✓
ML-Master	Tree	Ranking	✗	✓
AIRA	Graph	Ranking	✗	✓
MLE-STAR	Chain	Ranking	✓	✓
InternAgent-MLE	Graph	Ranking	✓	✗
KompeteAI	Tree	Ranking	✓	✗

Table 1: Comparison of LLM-based MLE agents. Feedback Role indicates how execution feedback drives the primary optimization loop: Ranking denotes score-centric selection among candidates; Update denotes diagnostic reasoning over structured feedback to generate state-conditioned code modifications.

Concept	Gome Component	Functional Role
Gradient $\nabla L$	Structured Reasoning	Signal: decides how to update
Momentum	Success Memory	Acceleration via proven patterns
Distributed SGD	Multi-trace Optimization	Parallelism with knowledge sharing

Table 2: Analogy between Gome and gradient-based optimization.

Figure 2: Overview of the Gome framework. Multiple traces optimize in parallel, synchronizing through a global shared success memory $\mathcal{M}$ . Each trace iteratively executes solutions, validates improvements, updates shared memory, and reasons over local-global feedback to generate the next hypothesis. Right panels illustrate the gradient-based optimization analogy and key module details.

			Medal rate in different complexity (%)				Other evaluation dimensions (%)
Agent	Time	GPU	Low	Medium	High	All	Valid	Median+	Gold
MLAB
gpt-4o-24-08	24h	A10 †	4.2±1.5	0.0±0.0	0.0±0.0	1.3±0.5	44.3±2.6	1.9±0.7	0.8±0.5
OpenHands
gpt-4o-24-08	24h	A10 †	11.5±3.4	2.2±1.3	1.9±1.9	5.1±1.3	52.0±3.3	7.1±1.7	2.7±1.1
AIDE

Table 3: Percentage of achieving any medals across different ML task complexity levels (left) and other evaluation dimensions (right) on MLE-Bench. Reporting results are mean

\pm

SEM over 3 seeds. Valid, Median+, and Gold indicate the percentage of submissions with valid score, above median score, and gold medal. Best performances are marked in bold ; second best are underlined . † GPU configuration from official MLE-Bench evaluation setup.

Agent	Any-Medal (%)
KompeteAI † (Gemini-2.5-flash,6h)	51.5±1.5
AIRA (o3,24h)	55.0±1.5
InternAgent-MLE † (DeepSeek-R1,12h)	62.1±3.0
FM Agent † (Gemini-2.5-pro,24h)	62.1±1.5
Thesis † (GPT-5-Codex,24h)	65.2±1.5
Leeroo † (Gemini-3-pro,24h)	68.2±2.6

Table 4: Performance comparison on MLE-Bench-Lite. Best performances are marked in bold . † Methods using open-world protocol with external resource access.

Configuration	Medal	Gold	Impr.	IC
Gome (full)	35.1	16.4	41.1	0.92
w/o Structured Reasoning	25.8	13.3	22.6	0.83
w/o Success Memory	28.9	16.9	36.2	0.87
w/o Multi-trace Optimization	32.4	15.1	41.3	0.88

Table 5: Ablation study on Gome components using GPT-5 under 12-hour budget (averaged over 3 runs). Impr.: improvement rate, the percentage of iterations where the proposed solution passes hierarchical validation. IC: Spearman correlation between validation and test improvements.

Panel of Figure 3 — Figure 3: Scaling analysis. (a) As reasoning capability improves, gradient-based optimization increasingly outperforms search baselines. (b) On weaker models, rapid early convergence is followed by a plateau due to noisy feedback; stronger reasoning enables sustained improvement throughout the 12h budget.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Gome	Gome-MCTS	ML-Master	Runs
Efficiency Tier Models
GPT-4o-mini	3.3	5.3	6.0	2
GPT-5-nano	10.7	12.0	12.0	2
GPT-4o	12.0	13.8	13.3	3
Advanced Tier Models
GPT-5-mini	19.3	18.7	16.0	2

Appendix

Table 6: Complete per-model scaling results. All values are any-medal rates (%) on MLE-Bench. Bold indicates the best performance per model. Runs indicates the number of independent trials.

Figure 4: Performance gap between Gome and the best baseline across model capability tiers. Negative values (red) indicate Gome underperforms; positive values (green) indicate Gome outperforms. The crossover occurs at the Advanced tier, with the gap widening progressively in the Frontier tier.

Model	12h	24h	Improvement
o3	32.5	36.4	+3.9 (+12.0%)
GPT-5	35.1	40.4	+5.3 (+15.1%)

Appendix

Table 7: Gome performance under half-time (12h) and full-time (24h) budgets (averaged over 3 runs). All values are any-medal rates (%) on MLE-Bench. See Table 8 for full results.

				Medal rate in different complexity (%)
Agent	w/o Ext. Know.	Time	GPU	Low	Medium	High	All
MLAB
gpt-4o-24-08	✓	24h	A10 ‡	4.2±1.5	0.0±0.0	0.0±0.0	1.3±0.5
OpenHands
gpt-4o-24-08	✓	24h	A10 ‡	11.5±3.4	2.2±1.3	1.9±1.9	5.1±1.3
AIDE

Appendix

Table 8: Percentage of achieving any medals across different ML task complexity levels on MLE-Bench. Reporting results are mean

\pm

SEM over 3 seeds. Best performances are marked in bold ; second best are underlined ; third best are marked with ⋆ . “–” indicates results not reported in the original work. ‡ GPU configuration from official MLE-Bench evaluation setup.

Model	Input	Output	Cost/Task	Medal Rate
DeepSeek-v3.2	$0.28 ($ 0.028)	$0.42	$1.89	28.0%
GPT-5	$1.25	$10.00	$20.74	35.1%
o3	$2.00	$8.00	$26.53	32.5%

Appendix

Table 9: API pricing (per million tokens) and average cost per competition for Frontier-tier models. Cost/Task is averaged over 3 independent runs. Parenthetical value indicates cache read price.

Task	Metric	DeepSeek-R1	o3	GPT-5
Image Tasks
aptos2019-blindness-detection	QWK $\uparrow$	0.8938	0.9243	0.9158
inaturalist-2019-fgvc6	Top-1 Error $\downarrow$	0.2951	0.3201	0.2261
plant-pathology-2021-fgvc8	F1 $\uparrow$	0.9100	0.8949	0.9044
Text Tasks
detecting-insults-in-social-commentary	AUC $\uparrow$	0.9384	0.9452	0.9176

Appendix

Table 10: Score comparison on 9 MLE-Bench tasks across three data modalities. Best result for each task is highlighted in bold .

\uparrow

: higher is better;

\downarrow

: lower is better.

Hyperparameter	Description	Default
LLM Configuration
chat_model	Base LLM for reasoning	GPT-5
temperature	LLM decoding temperature	1.0
max_retry	Maximum API retry attempts	12000
retry_wait_seconds	Wait time between retries (s)	5
full_runtime	Total Running timeout (h)	12

Appendix

Table 11: Gome Hyperparameter Configuration.

Model	Hardware	Medal Rate	Source
DeepSeek-R1	A100	29.3	Official
DeepSeek-R1	V100	22.7	Our Rerun

Appendix

Table 12: ML-Master reproduction validation on DeepSeek-R1. Original results are from [ 19 ] on A100 GPUs; ours are averaged over 3 independent runs under closed-world protocol on V100 GPUs. All values are any-medal rates (%) on MLE-Bench with 12-hour budget.

Metric	Value
Total optimization iterations	90
Overfitting cases identified	9
Correctly rejected	6 (66.7%)
Incorrectly accepted	3 (33.3%)

Appendix

Table 13: Overfitting detection performance on Stanford COVID Vaccine.

Panel of Figure 5 — Figure 5: Overfitting detection on Stanford COVID Vaccine. The validation score appeared to improve by 57.6%, but Gome ’s hierarchical validation analyzed the code changes and training dynamics to correctly reject the solution, avoiding a catastrophic 137.5% test regression.

Loop	Component	$\boldsymbol{\Delta h_{\mathrm{val}}}$	$\boldsymbol{\Delta h_{\mathrm{test}}}$	Decision	Correct
23	Model	$-63.8\%$	$+0.3\%$	Reject	✓
26	FeatureEng	$-5.1\%$	$+1.4\%$	Reject	✓
32	Model	$-57.6\%$	$+137.5\%$	Reject	✓
40	Model	$-0.2\%$	$+6.3\%$	Accept	✗
47	Ensemble	$-11.9\%$	$+0.4\%$	Reject	✓
49	Model	$-5.8\%$	$+1.2\%$	Accept	✗

Appendix

Table 14: All overfitting cases in Stanford COVID Vaccine.

\Delta h_{\mathrm{val}}

: validation change (negative is improvement).

\Delta h_{\mathrm{test}}

: test change (positive is degradation).

Aspect	MLE-STAR	Gome
Feedback role	Ranking	Update
Plan generation	Multiple candidates	Single hypothesis
Selection mechanism	$\arg\max$ score	Reasoning gate
Block identification	Ablation study	Structured analysis
Subtle overfitting detection	0% (score-driven)	66.7% (reasoning)

Appendix

Table 15: Key distinctions between MLE-STAR and Gome . Overfitting detection rate is based on the case study in § D.1 , where score-driven methods would accept all 9 overfitting cases.

Loop	Val RMSLE	Key Modification	Outcome
0	0.485	Initial baseline	Accepted
2	0.354	Two-stage hurdle architecture	Accepted
7	0.327	Rolling statistics & lag features	Accepted
28	0.318	Store-local holiday processing	Accepted
37	0.318	Horizon bucketing & calibration	Accepted (best)
30	(0.547)	Naive ensemble strategy	Rejected

Appendix

Table 16: Optimization trajectory on the Store Sales competition. All intermediate scores are internal validation RMSLE (15-day holdout). The trajectory shows steady improvement from 0.485 to 0.318 over 37 iterations, with the gate mechanism rejecting degrading modifications (Loops 30, 43). The final submission achieved a leaderboard RMSLE of 0.431. Note that the human leaderboard leader (0.378) employs multi-level blending of other participants’ submissions.

Panel of Figure 6 — Figure 6: Codebase comparison. (a) High-ranking public notebooks often use blending scripts that aggregate prediction CSVs submitted by other participants, a form of meta-optimization that, while effective for leaderboard climbing, involves no model training. Gome constructs the full pipeline from raw relational tables under the closed-world protocol without any external knowledge, including feature engineering, horizon-bucketed model training, and Platt calibration.

Category	Count	Percentage
Successful Updates	31	34.4%
Gradient Hallucination	35	38.9%
Resource Constraints	21	23.3%
Validation Blindspots	3	3.3%
Total Failures	59	65.6%

Appendix

Table 17: Failure mode distribution on stanford-covid-vaccine (90 iterations, GPT-5).

Large Language Model-Guided Optimization

Large Language Model-Guided Optimization uses language models to propose, steer, and refine solutions to hard search and design problems.

See all 605 papers

$\nabla$ -Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space
Mar 5, 2026Peihao Wang, Ruisi Cai, Zhen Wang +4LLM Reasoning Test-Time Optimization
The University of Texas at Austin · UC San Diego · TTIC +1
Revisiting Tree Search for LLMs: Gumbel and Sequential Halving for Budget-Scalable Reasoning
Mar 22, 2026Leonid Ugadiarov, Yuri Kuratov, Aleksandr Panov +1Monte Carlo Tree Search Sequential Decision Making
AXXX, Moscow, Russia · MIRAI, Moscow, Russia
Can Bayesian Optimization Efficiently Find a Strong Single Expert in Neural Thickets?
Aug 11, 2026Nigel Bastian Cendra, Abdelhamid Ezzerg, Fernando Julio Cendra +2Bayesian Optimization LLM Post-Training
University College London · Institut Polytechnique de Paris · University of Oxford

Mar 5, 2026J/K move · Enter open · S save

$\nabla$ -Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space

Peihao Wang, Ruisi Cai, Zhen Wang, Hongyuan Mei+3

The University of Texas at Austin · UC San Diego · TTIC+1

LLM Reasoning Test-Time Optimization Efficient Language Model Reasoning Inference-Time Language Model Alignment

Open paper AskPDF arXiv

Scaling inference-time compute for Large Language Models (LLMs) has unlocked unprecedented reasoning capabilities. However, existing inference-time scaling methods typically rely on inefficient and suboptimal discrete search algorithms or trial-and-error prompting to improve the online policy. In this paper, we propose

\nabla

-Reasoner, an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop to refine the policy on the fly. Our core component, Differentiable Textual Optimization (DTO), leverages gradient signals from both the LLM's likelihood and a reward model to refine textual representations.

\nabla

-Reasoner further incorporates rejection sampling and acceleration design to robustify and speed up decoding. Theoretically, we show that performing inference-time gradient descent in the sample space to maximize reward is dual to aligning an LLM policy via KL-regularized reinforcement learning. Empirically,

\nabla

-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark, while reducing number of model calls by approximately 10-40% compared to strong baselines. Overall, our work introduces a paradigm shift from zeroth-order search to first-order optimization at test time, offering a cost-effective path to amplify LLM reasoning.

Abstract

Large Language Model-Guided Optimization

Explore similar work

∇\nabla∇-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space

Revisiting Tree Search for LLMs: Gumbel and Sequential Halving for Budget-Scalable Reasoning

Can Bayesian Optimization Efficiently Find a Strong Single Expert in Neural Thickets?

$\nabla$ -Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space