One-sentence Summary

Researchers at Mila – Quebec AI Institute, Université de Montréal, and University of Oxford use generative flow networks (GFlowNets) to fine-tune LLMs for amortized Bayesian inference, enabling tractable sampling from intractable posteriors as a data-efficient distribution-matching alternative to maximum-likelihood and reward-maximizing training, with applications in chain-of-thought reasoning and tool use.

Key Contributions

• A Bayesian inference formulation casts querying autoregressive LLMs for sequence continuation, infilling, and constrained generation as sampling from intractable posterior distributions over latent sequences or outputs.

• GFlowNet fine-tuning trains LLMs with diversity-seeking reinforcement learning and flow-matching objectives to sample from target posterior distributions rather than only maximizing likelihood or selecting high-reward outputs.

• Experiments show that GFlowNet fine-tuning yields a better fidelity-diversity trade-off in text generation, improves sample efficiency and downstream generalization over maximum-likelihood training or reward-maximizing policy optimization, and enables data-efficient adaptation to chain-of-thought reasoning and tool use by treating reasoning chains as latent variables.

Introduction

Autoregressive LLMs encode broad world knowledge, yet many useful applications require inference over latent continuations, such as infilling, constrained generation, length or lexical constraints, and finding likely sequence continuations. These tasks can be framed as sampling from intractable posterior distributions, including latent reasoning chains in chain of thought reasoning, where the target posterior is known only up to an unnormalized likelihood. Prior approaches face limitations: prompting does not accurately sample the posterior, MCMC can be slow and sensitive to proposal distributions, and reward maximizing reinforcement learning tends to collapse onto high reward modes instead of covering the distribution. The authors leverage GFlowNet fine tuning to train a language model as an amortized sampler for these LLM posteriors, showing improved fidelity diversity tradeoffs, sample efficiency, out of distribution generalization, and empirical gains on reasoning, infilling, arithmetic, and tool use.

Method

The authors recast several challenging language-modeling tasks as sampling from intractable conditional distributions under an autoregressive large language model. A standard autoregressive model factorizes a sequence as

which makes prefix-conditioned generation tractable. However, many useful tasks require sampling from global conditionals such as a low-temperature continuation, an infilled sequence given both preceding and succeeding context, a reverse generation conditioned on an answer, or a constrained generation problem where

These distributions require normalization over all possible latent token sequences, so directly sampling from them is generally intractable. Instead of relying on tokenwise tempering, beam search, or handcrafted prompt engineering, the method learns an amortized sampler that can generate samples from these target distributions.

The method is built around a latent-variable view of reasoning. For a given input context $X$X and desired output $Y$Y, a latent sequence $Z$Z represents the intermediate reasoning, infill, explanation, or rationale. The target posterior is

Here, $p_{\mathrm{LM}}(XZY)$pLM​(XZY) is the unnormalized joint likelihood assigned by the pretrained language model to the concatenated sequence $XZY$XZY. Sampling from this posterior corresponds to finding plausible latent chains that make the desired output $Y$Y more likely given the input $X$X.

The overall framework is depicted as follows:

The sampler is implemented as a generative flow network. The GFlowNet policy is parametrized as an autoregressive language model that produces $Z$Z token by token. Starting from the empty string, it samples

until it emits a terminal stop symbol. The marginal probability of a completed latent sequence is therefore the product of token-level policy probabilities and the terminal probability:

The training goal is to fit the policy so that this marginal likelihood is proportional to a reward:

For posterior sampling, the reward is set to the language-model likelihood of the full concatenated sequence:

or to an equivalent task-specific score when constraints, penalties, or external evaluations are involved.

To train this amortized sampler, the authors use a modified subtrajectory balance objective that accounts for the fact that a sequence can be terminated at different intermediate states. For a completed sequence $Z=z_{1:n}\top$Z=z1:n​⊤, the objective is

This objective enforces consistency between rewards, continuation probabilities, and termination probabilities across all subtrajectories. At convergence, when the objective can be minimized to zero, the policy samples terminal states in proportion to their rewards.

The training pipeline uses a mixture of trajectories to improve exploration in the combinatorially large token-sequence space. Each batch is composed from samples generated by the current GFlowNet policy, samples generated by a tempered version of that policy, and past trajectories stored in a replay buffer. This mixture is important because posterior sampling requires discovering high-likelihood latent sequences among an exponentially large set of possibilities.

Conditioning depends on the task. When the latent sequence is used for generalization, such as chain-of-thought reasoning or classification, the GFlowNet policy is conditioned only on the input $X$X. It is initialized as a copy of the pretrained language model and then fine-tuned with the GFlowNet objective. At inference time, one draws latent sequences $Z$Z from $q_{\mathrm{GFN}}(Z\mid X)$qGFN​(Z∣X) and then samples the output from $p_{\mathrm{LM}}(Y\mid XZ)$pLM​(Y∣XZ), which corresponds to averaging over latent reasoning chains. When the target output $Y$Y is already known and the goal is to sample the latent sequence itself, as in infilling or reverse generation, the GFlowNet policy is conditioned on both $X$X and $Y$Y, often through a prompt that contains both contexts. In this setting, the policy directly samples from the posterior $p_{\mathrm{LM}}(Z\mid X,Y)$pLM​(Z∣X,Y).

The same posterior sampler can also support a variational expectation maximization style training loop. In the expectation step, latent sequences $Z$Z are drawn from the amortized GFlowNet posterior. In the maximization step, the original language model can be updated to increase the likelihood of the constructed joint sequence $XZY$XZY. Thus, the framework couples amortized inference over intractable latent variables with language-model fine-tuning, allowing the model to learn to sample plausible reasoning chains, infills, or other latent structures under the desired target distribution.

Experiment

The experiments evaluate GFlowNet fine-tuning on text generation and reasoning tasks, including sentence continuation, story infilling, low-data subjectivity classification, and step-by-step arithmetic problem solving. Across these settings, the method demonstrates that matching the language model posterior can produce high-quality samples, more coherent generated text, and stronger performance than standard prompting or fine-tuning baselines. In particular, it improves diversity and likelihood in continuation generation, supports latent rationale search for classification, and generalizes more robustly than reward-optimizing approaches in tool-assisted reasoning.

The table frames several language model sampling tasks as inference over a mechanism variable conditioned on a question or context and an effect or answer. It presents a pretrained language model as providing unnormalized scores, while a GFlowNet is trained to approximate and amortize the otherwise intractable conditional posterior. The arithmetic experiments support this view by showing that matching the full target distribution with GFlowNet fine-tuning yields more reliable tool use reasoning than prompting, supervised fine-tuning, and PPO. The posterior over reasoning chains or mechanisms given a prefix and target continuation is intractable, so amortized sampling is needed. Tempered sampling, infilling, reverse generation, and constrained generation can be viewed as sampling from language model conditional or constrained distributions. In arithmetic tool use, GFlowNet fine-tuning generalizes better than prompting, supervised fine-tuning, and PPO by avoiding collapse to spurious high reward rationales.

Across automatic metrics and a GPT-4 judge, GFlowNet fine-tuning produces story infills that are closer to reference infills than prompting or supervised fine-tuning. The improvement suggests that sampling conditioned on both the beginning and ending helps the model create more coherent story middles, including turning points. GFlowNet fine-tuning has the highest BERTScore, BLEU-4, GLEU-4, and GPT-4 evaluation scores among the compared methods. It outperforms prompting and supervised fine-tuning on all reported evaluation metrics. The gains match the task goal of generating a fourth sentence that connects the provided beginning and ending.

On SUBJ, zero-shot prompting performs near chance, while few-shot prompting gives only modest gains. GFlowNet fine-tuning provides large improvements over prompting, especially when only a small number of labeled examples are available. Adding supervised fine-tuning to GFlowNet fine-tuning produces the strongest overall test accuracy. GFlowNet fine-tuning greatly outperforms zero-shot and few-shot prompting. GFlowNet fine-tuning is most advantageous in the low-data setting. Combining GFlowNet fine-tuning with supervised fine-tuning achieves the best reported accuracy.

The results show that zero-shot chain-of-thought prompting is weak on integer arithmetic, while adding in-context examples improves accuracy as the number of examples increases. Supervised fine-tuning gives larger gains on training-length expressions than on longer out-of-distribution expressions, and PPO fine-tuning performs poorly because its reward model cannot reliably distinguish valid reasoning from spurious outputs. GFlowNet fine-tuning is reported to be more robust to reward misspecification and stronger on both in-distribution and out-of-distribution examples. In-context examples consistently improve arithmetic accuracy over zero-shot prompting. Supervised fine-tuning generalizes poorly to longer expressions compared with in-distribution expressions. PPO is hurt by poorly calibrated rewards and can produce high-reward but invalid tool-use sequences. GFlowNet avoids reward-mode collapse and yields stronger in-distribution and out-of-distribution performance.

The experiments evaluate GFlowNet fine-tuning across arithmetic tool use, story infilling, and subjective text classification, comparing it with prompting, supervised fine-tuning, and PPO. The results show that GFlowNet more effectively samples from intractable conditional distributions, producing more coherent infills, improving low-data performance, and handling constrained reasoning more reliably. Overall, the findings indicate that amortized posterior sampling reduces reliance on spurious reward signals and yields stronger generalization, especially when paired with supervised fine-tuning in smaller-data settings.