🤖 machine learning

Amortizing intractable inference in large language models

This paper proposes using GFlowNets to amortize intractable inference in large language models, enabling them to sample from complex posterior distributions for tasks like constrained generation and chain-of-thought reasoning as a data-efficient alternative to traditional maximum-likelihood or reward-maximizing training.

Original authors: Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, Esmeralda S. Whitammer

Published 2026-09-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, Esmeralda S. Whitammer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are vast libraries of human knowledge, compressed into digital networks that predict the next word in a sentence based on what came before. They are trained to be excellent at continuing a story or answering a question when given a starting point, a process that flows naturally from left to right. However, many of the most interesting tasks we want these models to perform require thinking in reverse or filling in the middle of a story. Imagine being given the beginning and the end of a tale and asked to invent the middle, or being asked to explain the reasoning behind a specific answer. In these scenarios, the model cannot simply guess the next word; it must search through a massive, complex landscape of possibilities to find the specific path that connects the start to the finish. This is a difficult mathematical problem because the number of possible paths is so enormous that checking them one by one is impossible, and the standard methods for navigating this landscape often get stuck on just one or two common routes, missing the rich variety of valid solutions that exist.

Researchers at the Mila – Quebec AI Institute and the University of Oxford have developed a new way to teach these models how to navigate that complex landscape. Instead of training the model to simply maximize a score for the "best" answer, which often leads to repetitive and narrow thinking, they used a method called amortized inference. This approach treats the problem as a search for a diverse set of correct paths rather than a single perfect one. To do this, they employed a technique known as a generative flow network, which acts like a guide that learns to sample from the entire range of possible solutions, ensuring that the model explores different valid reasoning chains rather than collapsing into a single, overused pattern.

The team demonstrated this by fine-tuning large language models to solve problems that require multi-step reasoning, such as arithmetic calculations or classifying whether a movie review expresses an opinion or a fact. In one experiment, they asked the model to fill in the missing middle sentence of a short story, given the beginning and the end. Standard methods often produced sentences that were grammatically correct but did not fit the narrative flow. The new method, however, successfully generated middle sentences that linked the start and end coherently, showing a much higher ability to understand the full context. In another test involving simple math problems, the model was equipped with a calculator tool and asked to break down the calculation into steps. While other training methods caused the model to repeat numbers or fail to use the tool correctly, the new approach taught the model to plan its steps carefully, leading to a dramatic improvement in accuracy, especially on problems it had never seen before.

The results suggest that this method is particularly powerful when data is scarce. In a test with only ten examples of movie reviews to learn from, the new method achieved significantly higher accuracy than standard training techniques, and it continued to outperform them even with fifty examples. The researchers found that by encouraging the model to sample a wide variety of reasoning paths and then combining their conclusions, the system became more robust and reliable. This is a crucial distinction because it means the model is not just memorizing a single way to solve a problem but is learning the underlying structure of how to think through it. The study indicates that by shifting the goal from finding the single most likely answer to exploring the full diversity of correct answers, we can make these powerful tools more flexible and better suited for complex reasoning tasks.

The work also highlights a limitation in how current models are often trained. When models are pushed to maximize a reward, they tend to find shortcuts, such as repeating the same phrase or ignoring the actual logic of a problem, just to get a high score. The researchers showed that their method avoids this trap by matching the entire distribution of possible solutions, ensuring that the model does not collapse into a single, flawed mode of thinking. This approach allows the model to generalize better to new situations, such as solving arithmetic problems with more numbers than it was trained on, where other methods failed. The findings offer a promising path forward for making artificial intelligence more capable of genuine reasoning, moving beyond simple pattern matching to a more nuanced understanding of how to construct arguments and solve problems step by step.

Technical Summary: Amortizing Intractable Inference in Large Language Models

Problem Statement

Autoregressive Large Language Models (LLMs) compress knowledge into next-token conditional distributions, making tractable inference limited to start-to-end sampling conditioned on a prefix. However, many critical tasks—such as sequence continuation, text infilling, constrained generation, and chain-of-thought reasoning—require sampling from intractable posterior distributions.

Specifically, tasks like infilling (generating a middle segment $Z$ given a start $X$ and end $Y$ ) or reasoning (finding a latent chain of thought $Z$ that leads to an answer $Y$ given a prompt $X$ ) involve computing the posterior $p(Z|X, Y)$ . While the joint likelihood $p_{LM}(XZY)$ is easy to evaluate, the conditional distributions required to sample $Z$ token-by-token from the posterior are intractable.

Existing approaches face significant limitations:

Markov Chain Monte Carlo (MCMC): Difficult to craft proposal distributions for multi-modal language data and often prohibitively slow for new inputs.
Reinforcement Learning (RL) / Policy Optimization (e.g., PPO): These methods maximize rewards but tend to collapse to a small number of modes, failing to capture the full diversity of the target distribution. This leads to "overoptimized samplers" when the reward function is misspecified.
Supervised Fine-Tuning (SFT): Requires access to samples from the target distribution, which are generally unavailable for intractable problems.

Methodology

The authors propose using amortized Bayesian inference via Generative Flow Networks (GFlowNets) to train LLMs to sample from these intractable posteriors.

Core Framework

The approach treats the problem as a latent variable model where an input $X$ generates an output $Y$ via a latent variable $Z$ (e.g., a reasoning chain or infilled text). The goal is to train a policy $q_{GFN}(Z|X, [Y])$ to sample from the posterior $p_{LM}(Z|X, Y) \propto p_{LM}(XZY)$ .

GFlowNet Objectives: The method utilizes GFlowNets, a class of diversity-seeking reinforcement learning algorithms. Unlike standard RL which maximizes expected reward, GFlowNets train a policy to sample objects with a probability proportional to a given reward function $R(Z)$ .
Reward Definition: The reward is defined as the joint likelihood assigned by the base LLM: $R(Z) = p_{LM}(XZY)$ . This effectively sets the target distribution to the Bayesian posterior.
Training Procedure:
- Initialization: The GFlowNet policy is initialized as a copy of a pre-trained LLM (often fine-tuned with LoRA for efficiency).
- Learning Objective: The authors employ a modified Subtrajectory Balance (SubTB) objective. This objective ensures that the probability of generating a complete sequence is proportional to its reward, even when trajectories can terminate at any state.
- Amortization: By conditioning the policy on the input $X$ (and optionally $Y$ ), the model learns a generalizable sampler. This allows for efficient inference on unseen inputs without re-running expensive sampling procedures like MCMC for each query.
- Exploration: To handle the combinatorial search space, training batches are composed of trajectories from the current policy, a tempered version of the policy, and a replay buffer of past trajectories.

Application to Chain-of-Thought

The paper interprets chain-of-thought reasoning as a latent variable modeling problem. Instead of prompting the model to generate a single chain, the GFlowNet fine-tuned model samples diverse latent chains $Z$ . The final prediction is obtained via Bayesian Model Averaging: sampling multiple $Z$ 's from the GFlowNet and aggregating the resulting $Y$ predictions from the base LLM.

Key Contributions

General Algorithm: A general method for amortized sampling from intractable LLM posteriors using GFlowNet objectives.
Probabilistic Reasoning Approach: A formulation of chain-of-thought reasoning as a Bayesian inference problem, enabling data-efficient adaptation to multi-step rationalization tasks.
Empirical Validation: Demonstration of the method across four distinct domains:
- Sentence continuation (maximizing likelihood while maintaining diversity).
- Story infilling (generating coherent middle segments).
- Subjectivity classification (few-shot learning with latent rationales).
- Integer arithmetic with tool use (step-by-step reasoning).

Experimental Results

The authors evaluated the method against baselines including standard sampling, diverse beam search, supervised fine-tuning (SFT), and Proximal Policy Optimization (PPO).

Sentence Continuation: GFlowNet fine-tuning achieved higher sample diversity while maintaining high log-likelihood compared to diverse beam search, even when the baseline was given 5 $\times$ the compute.
Story Infilling: On the ROCStories dataset, GFlowNet fine-tuning significantly outperformed prompting and SFT in BERTScore, BLEU-4, and GPT-4 evaluation, producing infills that better linked the beginning and end of stories.
Subjectivity Classification (Few-Shot): Using a 6B parameter model with only 10 labeled examples, GFlowNet fine-tuning achieved 71.4% accuracy, outperforming SFT (64.3%) and few-shot prompting (61.3%).
Integer Arithmetic with Tool Use: On an out-of-distribution (OOD) task involving 5 operands (trained on 3-4), GFlowNet fine-tuning achieved 40.7% accuracy. This significantly surpassed SFT (12.8%) and PPO (5.6%). The authors note that PPO suffered from reward misspecification, leading to spurious sequences that did not validly call the calculator tool, whereas GFlowNet's distribution-matching approach was robust to this.

Significance and Claims

The paper claims that GFlowNet fine-tuning offers a principled alternative to maximum-likelihood training and reward-maximizing policy optimization. Its primary significance lies in:

Diversity and Robustness: By matching the full target distribution rather than maximizing a single reward, the method avoids mode collapse and is robust to reward misspecification.
Data Efficiency: The amortized inference approach allows for effective adaptation to tasks requiring complex reasoning with very few labeled examples.
Generalization: The method demonstrates superior out-of-distribution generalization, particularly in reasoning tasks where standard RL methods fail to explore the solution space effectively.

The authors conclude that this paradigm converts computation into better test-time performance without requiring additional data, enabling more flexible reasoning with LLMs. They acknowledge limitations regarding the scale of models used (up to 6B parameters) and the fact that the method improves inference capabilities but does not inherently fix hallucinations or miscalibration issues rooted in the base model's knowledge representation.