Training a latent variable instead of searching for it

AlphaGo Zero learned Go by playing against itself, with a tree search choosing each move and a neural network trained to predict the moves the search chose. At the end of training, the same network could be played two ways. Choosing whichever move it rated most probable, with no search at all, it rated 3,055 Elo, on a rating scale where a 400-point lead means winning about ten games in eleven. With five seconds of search per move guided by that network, it rated 5,185, a difference of 2,130 Elo.

The network had been trained on the search’s own choices for forty days. On its own it still played 89 Elo below AlphaGo Fan, a 2015 program that used search and beat the European champion Fan Hui. Silver et al.

A similar comparison in Atari games comes out very differently. MuZero’s network was trained the same way. Played with one simulation per move (a simulation is one pass down the search tree, adding one position to it), which the paper equates with choosing moves from the network alone, it kept about 95% of the mean score it reached with fifty, across 57 games. MuZero’s authors suggest that its learned model of the game is less accurate in Atari than in Go, and neither paper tests why the domains differ. Schrittwieser et al.

The move sequence in these games is a latent variable: an unobserved quantity that, together with the input, decides the outcome. The same structure appears throughout machine learning. The code an autoencoder assigns to an image, the program behind a set of input-output examples, the 3D scene behind a few photographs, and the chain of thought between a question and its answer are all latent variables. Searching for one means finding it separately for each input, at answer time, by optimizing a score: gradient steps on one image’s code, tree search over moves, sampling many chains and keeping one that a checker accepts. Training means fitting a network once so that it outputs the latent, or a distribution over it, in a single pass for any input. The second is called amortized inference, because the cost of the optimization is paid once and spread across every input the network will see.

Can the latent be trained instead of searched for? It can, and the two routes pay for the same optimization in different places. The evidence below supports six statements.

• How close a trained predictor comes to search depends on the domain. In Atari it keeps about 95% of its search’s mean score; in Go it stands 2,000 Elo short even on positions like those it trained on. Whatever the gap is on familiar inputs, it grew with distance from training wherever that was measured.
• Judged by the objective the search optimizes, a predictor can never beat a search run to its optimum on one input, though it can beat a short search. Judged on something that objective cannot see, it can beat a search when the input alone leaves the latent undetermined, because training supplies what the single input lacks.
• A little search started from the prediction recovers part of the shortfall: much of it for 3D scenes and little in Go. In the one autoencoder study, it recovered much only in a model trained with that search, because otherwise the decoder learned to ignore the latent.
• Where the latent cannot be trained by gradient through the objective and nothing can simulate it, search makes the training targets. In the two games measured, how much search made the targets, and of what kind, mattered more than how much ran at answer time. A network can close much of the gap to a fixed search, and stays behind one that improves alongside it.
• Training a reasoning model against an answer check is this arrangement for the chain of thought. It compresses part of the search into one sample, more on its training problems than on new ones. Its reach, the share of problems solved in at least one of many tries, has risen a few points past the untrained model’s when the objective was built for many-sample search. Under the ordinary objective it has grown only with help from outside the model’s own samples, or on problems from the AIME, a US mathematics competition, when chains as well as answers are judged.
• Training compute and search compute trade at a measurable rate in one board game. For language models no fixed rate has been measured, and in the clearest case the trade stopped at the part of the search that picks an answer.

Two ways to pay for one optimization

The cleanest measurement of what training gives up comes from variational autoencoders, where both routes optimize the same quantity. A variational autoencoder has a decoder that turns a code into an image and an encoder that guesses, from an image, a distribution over the codes that could have produced it. The training objective is a lower bound on how probable the model makes each image, and the bound is tight only when the encoder’s guess matches the true distribution over codes. The shortfall, measured in nats (units of log-probability, where one nat is a factor of about 2.7 in probability), has two parts. One is the amortization gap: how much better the guess for a particular image would be if it were optimized for that image alone. The other is what remains even then, because the family of distributions the encoder can output cannot match the truth.

Cremer, Li and Duvenaud measured both. On 1,000 training images of handwritten digits, a standard encoder’s bound fell 2.77 nats short of the model’s estimated log-likelihood. Optimizing each image’s distribution separately, until it stopped improving, recovered 1.34 nats of that; the other 1.43 nats no distribution of the encoder’s form could recover. So on the digits the amortization gap was 48% of the shortfall. On images of clothing it was 58%, and on 100 small colour photographs reduced to one bit per colour channel, 72%. Cremer, Li and Duvenaud

On images the encoder had not been trained on, the gap was two to about seven and a half times larger. The paper’s figure shows this for a model of the clothing images trained on 50,000 images, with the gaps measured on 100 of its training images and 100 held-out ones, in one run. The gap was about 3.5 nats on training images and about 8 on held-out ones. With a wider encoder it was about 2 on training images and about 15 on held-out ones, and the held-out gap kept growing as training went on. A larger encoder narrows the gap on the inputs it was trained on and can widen it on new ones. Cremer, Li and Duvenaud

Margossian and Blei proved when the gap can close at all. An encoder that sees one input can match per-input optimization exactly when each input has its own latent with no links between latents, as in a variational autoencoder. If a latent depends on a neighbouring observation, the encoder must see that observation too. If the latents depend on each other, as in a hidden Markov model, the encoder may need the whole sequence: for some such models, no encoder that sees less than all of it can close the gap, however large the encoder is. Margossian and Blei

On the same objective, a trained encoder can never beat per-input optimization run to its optimum. The optimization is free to give every input its own best answer, while the encoder must produce its answer from the input through one shared function. A short search can still lose to the encoder. In a synthetic test by Kim and colleagues, 20 gradient steps from a random start reached a bound of 22.33 nats on the negative log-likelihood against the encoder’s 21.77, where lower is better. Kim et al.

The prior a trained predictor supplies

Margossian and Blei’s bound is about one objective on one input. A trained predictor can still beat a search when the input’s own objective does not determine the answer, because training across many inputs builds in a prior that a single input cannot supply.

Reconstructing a 3D scene from photographs shows it with a clean comparison. NeRF fits a separate network to each scene by gradient descent until it renders the given photographs, about 14 hours per scene. PixelNeRF instead trains one network on 88 scenes to predict a new scene from its photographs in a single forward pass. Both were scored on how well they rendered views of each scene that neither was given, which the search’s objective cannot see. The score was peak signal-to-noise ratio, on which higher is better. Given the same input views of the same 15 held-out scenes, pixelNeRF scored 15.55 dB from one view and 19.33 dB from three, where the per-scene search scored 8.00 and 9.85. With nine views the search scored higher, 22.14 against 21.10, by less than the spread between scenes. The per-scene search here had no prior of its own. A regularized one would probably narrow that margin; that comparison was not run. Yu et al.

The prior has a cost off the training distribution. MVSNeRF is another network trained on the same 88 scenes to predict a new scene in one pass. In its tests the predictor saw 3 views of each scene and the per-scene search 16, so part of each gap below is information the predictor did not have. It trailed per-scene optimization by 0.38 dB on held-out scenes of the same kind. On photographs of real scenes taken from one side, with the cameras all pointing roughly the same way, it trailed by 4.04 dB, and on synthetic objects seen from all sides by 7.01 dB (this digest’s arithmetic on the paper’s table). Chen et al.

Learned optimizers make the same trade for a network’s weights: a trained network proposes each update, in place of a tuned optimizer searching for good settings. VeLO was meta-trained on tasks of up to 20,000 training steps and then fine-tuned on tasks of up to 200,000. It trained faster than learning-rate-tuned Adam on all 83 of its benchmark tasks. It fell behind tuned baselines on models beyond about 500 million parameters, larger than nearly all of its meta-training tasks, and as runs approached and passed 200,000 steps, the edge of its fine-tuning. Metz et al.

Averaging is a second cost of the prior. A predictor with one output averages when the right latent differs between inputs that look alike. On a family of sine waves, a network pretrained on all of them and then given five points of a new one had a mean squared error of 2.41 after one adjustment step. A network trained to be a good starting point for that step had 0.67. In the authors’ words, “by pre-training on all tasks, the model would learn to output the average output for a particular input value.” Finn, Abbeel and Levine

A little search from the prediction

One fix for both costs is to start a search from the predictor’s output and run a few steps. How much that recovers varies widely.

In autoencoders, it depended on whether the latent was worth searching for. Kim and colleagues trained a text autoencoder whose encoder output was refined by 20 gradient steps during training. Ten refinement steps reached the plateau, a gain of 2.1 nats per 80-token example. The same refinement applied to a standard autoencoder gained about 0.26 (read off the paper’s figure; the paper reports perplexity, which this digest converted to nats). The standard autoencoder’s decoder had learned to ignore the latent, so searching for a better latent, from the encoder or from scratch, found little to gain. Training through the search is what made the latent worth searching for. Kim et al.

For scenes, refinement helped a great deal even though the predictors were not trained for it. MVSNeRF’s authors gave NeRF 10.2 hours per scene. Fifteen minutes of optimization started from MVSNeRF’s prediction, 1/41 of that time and with the same 16 views, closed 87% of the gap on the one-sided photographs and 49% on objects seen from all sides (this digest’s arithmetic). Chen et al.

A larger predictor, Long-LRM, was compared with a 13-minute per-scene optimization by Gaussian splatting, a faster method than NeRF. From 32 views it matched that optimization on peak signal-to-noise ratio in about a second. After 10 refinement steps taking 37 seconds, it passed the optimization on that measure and on structural similarity, while staying behind on a perceptual error measure. Chen Ziwen et al.

In Go a small search recovers little; most of the gain needs hundreds of simulations. On the paper’s own Elo scale, MuZero’s Go network rated 2,396 with no search. Ten simulations per move made it worse, 2,217. A hundred brought it to 2,860, 37% of the way to the 3,660 it reached with 800. Schrittwieser et al.

Gumbel MuZero is a variant that samples which moves to search at the root and trains the network’s move probabilities toward a target built from the search’s value estimates. A network for the smaller 9×9 board, trained with 200 simulations per move, recovered 12% to 16% of its gap to 800 simulations when given anywhere from 2 to 32, and about 84% when given 200 (read off the paper’s figure, on that paper’s own Elo scale). Danihelka et al.

Meta-learning shows when a small search from a learned start matters: when the prediction alone does not pin down the task. MAML trains a network’s starting weights so that a few gradient steps on a new task’s examples adapt it, a small per-task search from a learned start. On standard few-shot image benchmarks the search barely matters. Freezing every convolutional layer during adaptation moved accuracy only from 46.9% to 46.3% on one benchmark. Dropping the gradient steps entirely, and comparing the new examples by similarity instead, gave 48.0% (three seeds). The search matters when the task is revealed only through feedback. On a simulated robot that must discover which direction to run, the learned start alone averaged a return of −50.69. One policy-gradient step, which makes actions more probable in proportion to the return they earned, raised it to 293.19. Raghu et al., Finn, Abbeel and Levine

Search to make the targets

A predictor has to be trained on something. When the latent can be trained by gradient through the objective, as in an autoencoder or a scene model, or the pairs of latent and input can be simulated from a known generative model, the predictor trains directly and no search is needed. When neither holds, usually because the latent is discrete, training needs either a noisy gradient estimate or targets found by search. In the comparisons below the searched targets did better, and in the two games where it was measured, this is the role in which search mattered most.

How much search goes into the training targets, and of what kind, can matter more than how much runs at answer time. A MuZero network trained in Ms. Pac-Man with 5 simulations per move scored 29,147 when it searched with 50 at answer time. A network trained with 50 simulations scored 243,401 under the same answer-time search, and 229,489, 94% of that, choosing moves from its own policy with no search. So in Ms. Pac-Man search mattered far more for making the targets than for answering. Schrittwieser et al.

In 9×9 Go, MuZero trained with 2 or 16 simulations per move reached only about 150 and 330 Elo when evaluated with 800 simulations. Gumbel MuZero, whose search provably improves the policy when its value estimates are exact, reached about 2,050 when trained with 2 (read off the paper’s figure, two seeds). Danihelka et al.

Search targets also beat the other training signals tried. On 9×9 Hex, expert iteration trains a network on tree-search moves and then uses the improved network to guide the next search. At an equal number of network evaluations it improved its network by about 535 Elo, against about 90 for policy-gradient training on game outcomes (read off the paper’s figure, five runs). AlphaGo Zero’s self-play player beat a player of the same architecture trained on human games within 24 hours, although the human-trained network predicted professional moves better. Anthony, Tian and Barber, Silver et al.

A network can close much of the gap to a fixed search, in one case all of it, and keeps falling short of one that improves as the network does. A network trained to imitate a fixed tree search on Hex played it to parity without searching, winning 87 of 162 games. In the full loop, where each search uses the latest network, the search stayed ahead of the network it trained. In a run that also trained a network to predict who wins from a position, the lead was about 300 to 400 Elo throughout. In a run without one it shrank from about 400 to about 185 (read off the paper’s figure, one run each). A transformer distilling the chess engine Stockfish, which does not change, narrowed the gap from 686 to 452 to 412 Elo as it grew from 9 million to 136 million to 270 million parameters. Anthony, Tian and Barber, Ruoss et al.

Program synthesis shows the same structure with a discrete latent, and there the search at answer time carries much of the accuracy. DreamCoder trains a network to predict, from a task, which program parts its solution is likely to use, on programs its search has found. The network then guides the next search. At a fixed time limit per task, the network was worth 34 to 51 points of tasks solved across five domains, for example about 92% against about 44% on list-processing tasks (read off the paper’s figure, five runs). It works by reordering the search, and no version was ever run without search. In RobustFill, whose network was trained on randomly generated programs rather than searched ones, the single most probable program was right about 55% of the time (read off the paper’s figure). Checking 100 of its candidates against the examples raised that to 89%, and 1,000 candidates to 92%. Ellis et al., Devlin et al.

Discrete latents are also why the training step itself is hard. A discrete choice has no gradient, so training it directly needs a score-function estimator or a relaxation. A score-function estimator raises the probability of each sampled choice in proportion to its score, and so has high variance. A relaxation replaces the choice with a continuous stand-in. At one sample per step, predicting the lower half of handwritten digits from the upper half, the score-function estimator reached a negative log-likelihood of 72.0 nats where the Gumbel-softmax relaxation reached 58.5 (one run each). Jang, Gu and Poole

The other way round the missing gradient is to train on targets. Hard expectation maximization lists the candidate latents that reproduce the answer and trains on the one the model itself rates most probable. Where most candidates reproduced the answer by a wrong route, on a numeric reading-comprehension task and a text-to-SQL task, it beat averaging over them by about 13 points with one model each, though a second model on the reading task gained only 1.2. On four other tasks the gain was 1.6 to 3.5 (no variation across seeds reported). Min et al.

The chain of thought as a latent

A reasoning model’s chain of thought is a latent variable in exactly this sense: the unobserved steps between a question and its answer. Searching for it means sampling many chains and keeping one that a checker accepts. The share of problems solved by at least one of k samples, called pass@k, measures how far that search reaches, its reach for short, when an exact answer check picks the sample. At large k it also counts right answers reached by wrong chains.

Training means reinforcement learning against the checker, which raises the probability of chains that pass. The trained model is the amortized search, and the results above predict three things. One sample should hold part of the search, more on training problems than on new ones. Training should not reach past what the search found. And search at answer time should still add to it.

It compresses the search into a single sample, partly. Yue and colleagues trained Qwen2.5-7B on 2,000 math problems. On held-out problems, the share solved on the first try rose from 10.2% to 28.3%, while the untrained (base) model solved 69.1% given 256 tries and an exact answer check. Training closed 31% of that distance. Zhou reports the same base model’s full curve on the same problems, and on it one trained sample did what about six untrained samples did (this digest’s arithmetic). Neither figure counts the roughly 900,000 chains sampled during training. Yue et al., Zhou

On the problems the model trained on, Yue’s run closed 57% of the distance, against 31% on held-out problems. That is the same pattern as the autoencoder’s larger gap on images it was not trained on. In ProRL, a longer run by Liu and colleagues from a 1.5-billion-parameter model already distilled from a stronger one, one sample closed 22% to 51% of the same distance across six math benchmarks (this digest’s arithmetic on the paper’s figure). So the compression stays partial over long training too. Yue et al., Liu et al.

It does not, by default, extend the search’s reach. In Yue’s run the share solved within 256 tries fell from 69.1% to 63.9%.

One study traces the loss to the objective and the sampling budget rather than to training as such. The training method samples a group of 8 chains for each problem and learns from how the chains in the group differ. A problem on which 1% of the model’s chains are correct produces no correct chain in 92% of its groups, and those groups give no gradient from the reward. The study found that such rarely solved problems were the ones training lost, 654 of them in a 3,000-problem split. Consistent with that, raising the budget to 32 chains in Yue’s runs helped, though it did not remove the loss, so the budget is not the whole cause. Zhou, Yue et al.

In two studies, objectives changed to serve the many-sample search that follows kept or raised the reach. Training Gemma 2B for success within 32 samples raised its pass@32 on MATH, a benchmark of competition problems, from 60.0% to 67.0% (one run). The second study used the same held-out problems as Yue’s run. It kept ordinary training only for problems the untrained model had solved at least once in 8 tries, and on the rest replaced it with a penalty for drifting from the untrained model. That raised pass@256 to 73.0%. The untrained model scored 69.1%, and the study’s own ordinary training 68.3% (means of three seeds; the ordinary run’s first seed matches Yue’s checkpoint at step 150 exactly). Chow et al., Zhou

Search at answer time still adds to the trained model, as it does in games. A majority vote over 40 chains added 6.6 to 9.1 points on grade school math to PaLM 2-M with prompts tuned against the answer check. That tuning learns a sequence of input vectors placed before the question and leaves the weights unchanged. Phan et al.

Search on the untrained model can even stand in for most of the training. A Markov chain Monte Carlo sampler that favours sequences the base model itself rates highly uses no checker and no training. With Qwen2.5-Math-7B on MATH500, a 500-problem subset of MATH, it scored 74.8%, against 49.6% for ordinary sampling and 78.5% after reinforcement learning (single runs), and it spent about 8.8 times the tokens per question. That is the trade in its plainest form: most of what training bought once, the sampler buys again on every question. Karan and Du

Models trained on written-out search traces learn the search itself. They reach about the level of what they imitated, and in one of two studies passed it after training against the outcome. A model trained on traces of heuristic searches for the arithmetic game Countdown solved 51.3% of held-out problems. The same model trained on optimal solutions alone solved 25.7%, and the heuristic searches themselves averaged about 57% (read off the paper’s figure). Training against the outcome brought it level with them: 55.6% after three rounds of self-training on its own correct solutions. Gandhi et al.

In another Countdown study, a model trained on traces of a depth-first search, guided by one of the first study’s heuristics, was then trained against the outcome. It solved nearly 70%, above the 57% the search itself solved. In the same study, a model trained on solution paths alone, without the search, went the other way under the same training. Its first-try success rose from about 17% to 41%, while its reach fell from about 69% to about 48% (read off the paper’s figure; one run, 200 test problems). Qin et al.

A check on the answer alone does not pin the chain down. After 400 steps, on the median training question whose every sampled answer was right, a language-model judge found about 30% of the chains behind those answers wrong. Training had no gradient left to fix them, since every answer already passed.

Part of the untrained model’s reach on problems from the AIME, a US mathematics competition, was right answers from wrong chains. Judged on chains as well as answers, the trained model’s reach there exceeded its base’s (one pair of models, 60 problems, a language-model judge). On two easier benchmarks the base caught up even under that judge. Wen et al.

Under the ordinary objective, the one study here in which reach grew past the base on answer checks alone started from a model already distilled from a stronger model’s chains, and its clearest gain came under checks that make guessing implausible. Under unit tests, prolonged training raised pass@256 on a programming benchmark from 88.1% to 97.9%. In the same study, part of the measured growth on puzzles was the model learning an answer format its base had been ignoring. Liu et al.

The extreme of amortization is to train the written chain away entirely, so that the model computes the answer without writing any steps. That kept accuracy on a fixed algorithm and lost much of it on word problems. GPT-2 Small multiplied 9-digit numbers with 0.99 accuracy and no written steps; on grade school math, Mistral 7B fell from 0.68 with a chain to 0.51 without one. Deng, Choi and Shieber

The exchange rate between training and search

If training and search pay for the same optimization, there should be a rate at which one substitutes for the other. In 9×9 Hex there is. Jones trained agents of AlphaZero, AlphaGo Zero’s general-game successor, across three orders of magnitude of training compute, and played each with search trees of many sizes. Fitting the cheapest combinations that reach each playing strength, each tenfold increase in training compute cut the search compute per move by a factor of about 16 (one run per configuration). The paper gives no uncertainty; refitting its released data for this digest gives a 95% interval of about 10 to 21 times. Jones

At any one strength the trade spans at most one and a half orders of magnitude of training compute. At one end the search is down to a single node, so more training has no search left to remove. At the other, each network’s strength stops rising as its search tree grows, so more search cannot stand in for more training.

The two budgets are different kinds of cost. Training is paid once and search on every answer, so whether trading one for the other saves anything depends on how many answers there will be. In Hex, at one strength level, moving from 10¹⁴ to 10¹⁵ multiply-adds of training saves about 2.3 million multiply-adds per move. It pays for itself after about 400 million moves (this digest’s arithmetic on Jones’s fit, whose range starts slightly above 10¹⁴). Jones’s training compute already counts the self-play search that made the training targets, and leaves out only the gradient updates.

For language models the trade has not been measured as a fixed rate. The closest comparisons trade model size, which stands for training compute, against samples at answer time. Snell and colleagues compared a small model given extra answer-time compute, revising its answers with a verifier choosing among them, against a model 14 times larger decoded once, at matched compute. The small model won clearly on the easiest two-fifths of questions at the larger two of three budgets. On the hardest fifth it lost at every budget, by 2 to 5 of 100 questions (read off the paper’s figure). Neither the verifier’s training nor the revision training was counted. Snell et al.

In the clearest language-model case, what stopped the trade was the part of the search that picks an answer, not the part that proposes them. One Llama-3-70B sample costs about as much as nine Llama-3-8B samples. On 128 MATH problems, with a checker that recognises a correct answer, nine 8B samples solved 59.5% against 47.9% for one 70B sample (this digest’s arithmetic on the released samples). With majority voting or a reward model choosing instead, no number of 8B samples got above 41.4%, although at 10,000 samples at least one sample was right on 98.4% of the problems. Brown et al.

What is settled and what is open

The six statements, restated with their scopes:

• A trained predictor can stand in for a search for a latent variable, and how much it gives up depends on the domain: little in Atari, a great deal in Go. In every domain where it was measured (autoencoders, scenes, learned optimizers and reasoning models), the shortfall grew with distance from training.
• Judged by the search’s own objective, it cannot beat a search run to its optimum. Judged on what that objective cannot see, it beat a per-scene search that had no prior of its own, where the input underdetermined the latent and training supplied the prior.
• A search started from the prediction recovers part of the gap: much of it for scenes, little in Go, and in autoencoders only in a model trained with the search.
• Where the latent cannot be trained by gradient and nothing simulates it, search makes the targets. In the two games measured, how much search made the targets, and of what kind, mattered more than how much ran at answer time. A network closes much of the gap to a fixed search and stays behind one that improves alongside it.
• For reasoning models, training against an answer check compresses part of the untrained model’s many-sample search into one sample, more on training problems than on new ones. It keeps, loses or raises reach depending on the objective and the sampling budget; objectives built for many-sample search raised it a few points. Under the ordinary objective and answer checks alone, it has extended reach only with help from outside the model’s own samples. Judged on chains as well as answers, the trained model’s reach on AIME problems exceeded its base’s.
• Training compute and search compute trade at a measurable rate in one board game. For language models no fixed rate has been measured, and in the clearest case the trade stopped at the part of the search that picks an answer.

Open, with the cheapest experiment for each:

• What is the exchange rate between training and answer-time search for a language model with an exact checker? Repeat Jones’s design on one task with unit tests: several training budgets, each checkpoint scored at several sample counts, and the curves of equal accuracy fitted with uncertainty.
• Does ordinary training against an answer check extend reach to problems the untrained model never solves, starting from an undistilled base? The one unit-test result here started from a distilled model in a single run. Score an undistilled base and its trained model at large k on programming tasks with hidden tests, with chains judged as well as answers, over several seeds.
• How does a trained network compare with search at matched total compute, training included, outside board games? In Hex, Jones’s ledger counts the self-play search but not the gradient updates. No program-synthesis study here counts both, and the language-model comparison that counts pretraining leaves out the verifier’s training. Repeating the Hex accounting for DreamCoder on one domain, with the network’s training and the per-task search both logged, would settle it for programs.

Most of the comparisons here are single runs, several numbers were read off figures, and some findings come from one domain. The sizes should be read as examples.