Can Monte-Carlo Tree Search build a better optimizer — and could a better optimizer help large models learn more?
A research exploration: most optimizers take a local step (gradient descent goes down, Newton goes to the nearest critical point). This project asks whether a learned, planning optimizer — a small policy + value network taught by MCTS — can instead aim at a global minimum direction, escaping traps that local methods cannot. The motivating bet is that optimizer quality is a lever on what a model can ultimately learn: if the same architecture and data can be driven to a better optimum, that is "free" capability.
This is one of my side projects and I have very limited time. I've taken it far enough to learn what works and what doesn't (see Findings below), and I'm suspending active development here. The code is left in a working, documented state, with an honest record of the negative results and a roadmap for whoever (maybe future-me) picks it up.
— Richard (Xuefeng Ding)
Write the parameter update as a direction. The methods we know choose it locally:
| method | step | uses | what it finds |
|---|---|---|---|
| SGD / momentum | −g |
gradient | downhill — gets stuck in the nearest basin |
| Newton | −H⁻¹g |
gradient + curvature | the nearest critical point (still local; undefined where the Hessian isn't positive-definite) |
| this project | a learned direction, planned by MCTS | gradient + a learned value + multi-step lookahead | aims at a global optimum |
The two networks split the job the way AlphaGo's do:
- Value network → "where is the global optimum?" Not an estimate of the (known) loss, but a reachability-optimism signal: how good is the best point you can reach from here. This is what lets the search tolerate going uphill over a barrier.
- Policy network → "which way do I go?" A distribution over update primitives; the prior that steers the search, and — crucially — the part that could be deployed on its own.
Newton needs H⁻¹ (impossible at billions of parameters) and only sees local curvature; MCTS
replaces the local curvature inverse with planning over a horizon, and replaces the
expensive Hessian with a learned value that points beyond the local basin.
MCTS is expensive — each tree leaf needs a loss/gradient evaluation, many per step. You can never run the search inside the training loop of a 500B-parameter model. So the search is not the product. The plan is:
- Meta-train (offline): MCTS + the value net teach the small policy on cheap problems.
- Deploy (cheap): run the bare policy as the optimizer — no search, no value net, roughly Adam's per-step cost.
This is the "learned optimizer" (L2O) idea; MCTS's contribution over a one-step meta-gradient is the multi-step planning that finds escapes/schedules a myopic method misses.
| Area | Where | What |
|---|---|---|
| 1D AlphaZero demos | examples/alphazero_v1_discrete/ |
the clearest result: policy+value trained on the fly by MCTS escape a local minimum that gradient descent cannot. Adam-like affine action space. |
| Interactive explainer | docs/how-alphazero-works.html |
full walkthrough: the AdamW idea → our action space → how the nets evolve vs f(x) → does it optimize better → honest limits. Open in a browser. |
| Evolution movie | regenerate via examples/alphazero_v1_discrete/evolution_movie.py |
HTML5 animation of the policy/value reshaping over 500 episodes (gitignored, ~6 MB). |
| Real-model testbed | examples/juno_*.py, src/ |
a JUNO muon-reconstruction transformer with tuned SGD / AdamW / Muon baselines, used to sanity-check against real training. |
| Spectral v2 + original package | src/mcts_optimizer/ |
the earlier MCTS-optimizer package (24 primitive×LR actions, self-play challenger) and an SVD-based spectral action space. Toy functions only. |
The action space in the current 1D demo is a deliberately Adam-like affine update whose coefficients the policy chooses per state:
m' = β·m + f'(x) # momentum buffer (part of the search state)
x' = x·(1 + c) − a·m' + b # a = step size, c = weight decay/amplify, b = escape offset
So MCTS acts as a learned, planning replacement for the learning-rate scheduler and weight
decay, plus an offset b (the escape move when f'≈0) that no scheduler has.
The value of this project is as much in the negative results as the positive ones.
Worked:
- A non-myopic value + MCTS escapes a 1D local minimum where gradient descent (and Newton)
get stuck. Verified: from the local basin, the trained loop reaches the global min
(
f=−1.49vs the true−1.55); plain GD stalls at the local min (f=1.44). - The policy/value can be trained on the fly (self-play): escape rate climbs over episodes vs a frozen-init control.
- Distillation of the non-myopic search into a cheap, search-free controller works at ~GD per-step cost.
Failed / limits (these constrain any viable path):
- The value net is load-bearing, and it must be optimistic. Using the exact
−f(x)as the leaf value, or a residualV = −f + ΔVthat starts at the loss, both fail — they are pessimistic along the barrier, so a cheap search never crosses and the policy collapses onto the gradient step. A chicken-and-egg the optimistic value net breaks. - The result is landscape-specific. The nets cache this well's escape; a mirrored well needs retraining. The general skill lives in the search; the nets only amortize it.
- On-the-fly learning is noisy with the momentum-based affine actions (≈4/5 seeds learn at a low budget).
- Tuned Muon ≈ AdamW at equal wall-time on the real JUNO transformer; greedy best-of-k over spectral actions lost to fixed Muon. No free lunch yet at real scale.
- Feature choice is a toy artifact. In 1D,
(loss, gradient)is ambiguous, so absolute position is needed to disambiguate — but absolute weight value doesn't transfer across models, so a single fixed well is a poor proxy for the real generalization question.
- Train on a distribution of problems, not one fixed well — random barrier widths / locations / depths — and measure transfer to held-out problems. This is the test that actually decides whether anything is learned vs memorized.
- Scale-invariant local features (normalized gradient, momentum, curvature proxies — Adam's own ingredients), not absolute position, so a deployed policy can transfer.
- Deploy the bare policy search-free and measure steps/FLOPs-to-target vs AdamW, counting the per-step network overhead honestly.
- 2D → 100D → the matrix-update space of a real neural-network optimizer, where features must become dimension-agnostic and the action set must avoid combinatorial blow-up.
- Continuous
(a, b, c)action (progressive widening + entropy-regularized policy) instead of the discrete menu.
The open crux, shared with all learned optimizers: generalization across scale and architecture. That is exactly what "help on a 500B model you never meta-trained on" demands, and it is unsolved. This is why the project pauses here rather than claiming a win.
git clone git@github.com:neutralino-ai/mcts-optimizer.git
cd mcts-optimizer
python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev,benchmarks]"# the on-the-fly AlphaZero optimizer (escape-rate + value/policy figures)
python examples/alphazero_v1_discrete/alphazero.py
# explanation figures + the evolution movie / convergence analysis
python examples/alphazero_v1_discrete/affine_figures.py
python examples/alphazero_v1_discrete/evolution_movie.py # writes docs/evolution-movie.html
# rebuild the HTML explainer (embeds the figures)
python docs/build_alphazero_html.py # writes docs/how-alphazero-works.htmlpytest -qexamples/alphazero_v1_discrete/ 1D AlphaZero demos (the main result)
examples/juno_*.py real-model transformer testbed baselines
src/mcts_optimizer/ original MCTS-optimizer package + spectral v2 (toy functions)
docs/ HTML explainer + figure/movie generators
docs/superpowers/specs/ design specs
benchmarks/results/ figures and logs
IDEA*.md brainstorming notes
See LICENSE.