PaperlayerView exact revision on arXiv

Source-linked AI summary

Amortizing intractable inference in large language models

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

arXiv:2310.04363v3cs.LGcs.CL

TL;DR

Autoregressive LLMs make many useful tasks difficult because tractable querying is limited to prefix-conditioned sampling, while tasks such as infilling require intractable posterior sampling. The paper fine-tunes LLMs with diversity-seeking GFlowNets to amortize inference, reporting improved fidelity-diversity trade-offs and downstream sample efficiency and generalization versus maximum-likelihood training or reward-maximizing policy optimization.

  • Problem

    Autoregressive LLMs make tractable inference over their knowledge limited to prefix-conditioned sampling, whereas tasks such as infilling and constrained generation require intractable inference.

    Source: §1 · p. 1

  • Method

    The method initializes a GFlowNet policy with a pretrained LLM and trains it using a reward objective evaluable by that same LLM to sample from target distributions.

    Source: §1 · p. 2

  • Results

    GFlowNet fine-tuning improves fidelity-diversity trade-offs, sample efficiency, and generalization versus maximum-likelihood training or reward-maximizing policy optimization; subjectivity classification improves 10.9% absolutely over supervised fine-tuning with 10 labeled examples.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    Amortized inference converts computation into better test-time performance without additional data.

    Source: §6 · p. 9

    Experiments use models up to 6B parameters because of resource constraints, and exploration for problems with more complex latents remains open.

    Source: §6 · p. 9

Abstract

from arXiv · show

Autoregressive large language models (LLMs) compress knowledge from their training data through next-token conditional distributions. This limits tractable querying of this knowledge to start-to-end autoregressive sampling. However, many tasks of interest -- including sequence continuation, infilling, and other forms of constrained generation -- involve sampling from intractable posterior distributions. We address this limitation by using amortized Bayesian inference to sample from these intractable posteriors. Such amortization is algorithmically achieved by fine-tuning LLMs via diversity-seeking reinforcement learning algorithms: generative flow networks (GFlowNets). We empirically demonstrate that this distribution-matching paradigm of LLM fine-tuning can serve as an effective alternative to maximum-likelihood training and reward-maximizing policy optimization. As an important application, we interpret chain-of-thought reasoning as a latent variable modeling problem and demonstrate that our approach enables data-efficient adaptation of LLMs to tasks that require multi-step rationalization and tool use.

1 INTRODUCTION

Autoregressive LLMs make many useful text-generation and reasoning tasks difficult because their tractable sampling is limited to conditioning on a prefix. This work uses GFlowNet fine-tuning for amortized sampling from intractable posteriors, including latent reasoning chains and tool-use sequences, and reports gains over supervised fine-tuning and PPO on selected tasks.

1 INTRODUCTION

  • Autoregressive LLMs support tractable sampling conditioned on a prefix, but infilling, constrained generation, and likely sequence continuation require intractable inference.The posterior over latent reasoning chains is also intractable because the conditional distributions needed to sample one token at a time are not easy to compute.Source: §1 · p. 1
  • Treating chain-of-thought as a latent-variable model makes reasoning posterior inference over token sequences Z conditioned on question-answer context X and Y.The target posterior is pLM(Z|X,Y).Source: §1 · p. 1
  • MCMC proposals are difficult to craft for multimodal language distributions and can make new-input inference slow, while PPO can concentrate on a few modes rather than model distributional diversity.The cited work notes that mode collapse can worsen when the target distribution is misspecified.Source: §1 · p. 2
  • GFlowNet fine-tuning trains a pretrained-LLM policy with a reward evaluable by the same LLM to sample from target distributions, enabling amortized inference.GFlowNets are diversity-seeking reinforcement learning algorithms that train policies to sample objects proportional to a reward such as the LLM joint likelihood.Source: §1 · p. 2
  • 10.9% absolute improvement over supervised fine-tuning on subjectivity classification with 10 labeled examples, and 63% outperformance of supervised fine-tuning and PPO on integer arithmetic with 50 demonstrations.The arithmetic setting also shows notable improvements in out-of-distribution generalization.Source: §1 · p. 2
  • The introduction presents a general algorithm for amortized sampling from intractable LLM posteriors, probabilistic chain-of-thought fine-tuning, and empirical work on continuation, reasoning, arithmetic, tool use, and story infilling.The stated contribution list includes results on integer arithmetic and natural-language reasoning.Source: §1 · p. 2

2 MOTIVATING EXAMPLE: GENERATING RANDOM NUMBERS WITH LLMS

The random-number task tests sampling from a target distribution specified by an unnormalized density. PPO can enforce valid outputs but leaves the distribution skewed, while GFlowNet fine-tuning sharply reduces divergence from the uniform target; supervised fine-tuning could also match it if target samples were available.

2 MOTIVATING EXAMPLE: GENERATING RANDOM NUMBERS WITH LLMS

  • The task instantiates sampling from a target distribution given an unnormalized density, using a prompt for integers uniformly drawn from 0 to 100.Source: §2 · p. 2
  • KL divergence falls from 3.37 for the original LLM to 9.75 · 10^-5 for the GFlowNet-fine-tuned model on support [0, 100].Figure 2 compares empirical distributions from PPO and GFlowNet fine-tuning.Source: Figure 2 · p. 3, §2 · p. 3
  • The pretrained model samples numbers far from uniform, possibly reflecting instruction fine-tuning or training-data frequency biases.The passage gives numbers beginning with “1” as an example of values more frequent in training data.Source: §2 · p. 3
  • PPO produces mostly valid integers, but their distribution remains highly skewed.Rewarding all valid integers equally gives policy-gradient methods an expected gradient of zero.Source: §2 · p. 3
  • Supervised fine-tuning could also match the target here, but requires target-distribution samples, which are generally unavailable for the intractable problems studied.Source: §2 · p. 3

3 FINE-TUNING LLMS TO SAMPLE FROM INTRACTABLE DISTRIBUTIONS

Autoregressive LLMs make left-to-right sampling tractable but leave many useful conditional-generation tasks intractable. The section frames chain-of-thought as posterior inference and develops GFlowNet objectives for amortized sampling.

3.1 PROBLEM: INTRACTABLE INFERENCE IN LARGE LANGUAGE MODELS

  • Autoregressive models support tractable left-to-right sampling, but sampling other conditionals over their sequence distributions is intractable.The section identifies tempered continuation, infilling, and constrained generation as examples.Source: §3.1 · pp. 3–4
  • Infilling samples tokens between fixed prefix X and suffix Y, while reverse generation is the special case with an empty X.The section notes that infilling and reverse generation also feature in methods that optimize instruction sequences or reasoning steps.Source: §3.1 · p. 4

3.2 REASONING THROUGH LATENT VARIABLES

  • Chain-of-thought can be treated as posterior inference over latent token strings Z, sampling likely reasoning chains that lead from X to the desired Y.The latent-variable formulation is useful when modeling pLM(Y|X) is harder than modeling pLM(Z|X) and pLM(Y|XZ).Source: §3.2 · pp. 4–5
  • The latent-object formalism and GFlowNet objectives also apply to trees or sets when a likelihood model p(Y|XZ) is available, but these extensions are not investigated.The authors identify formal reasoning and multi-step inference as potential applications.Source: §3.2 · p. 5
  • Variational EM can combine posterior samples of Z in an E-step with joint-likelihood optimization of the language model in an M-step.The section illustrates this combination in an experiment involving amortized inference and supervised fine-tuning.Source: §3.2 · p. 5

3.3 AMORTIZED INFERENCE WITH GFLOWNET OBJECTIVES

  • GFlowNets learn autoregressive policies whose sequence probabilities are proportional to a reward, allowing posterior sampling by setting the reward to pLM(XZY).The policy samples tokens until a stop symbol is generated.Source: §3.3 · p. 5
  • GFlowNet training uses trajectories from the current policy, a tempered policy, and a replay buffer to explore the combinatorially large sequence space.The objective can be minimized to zero for all trajectories simultaneously given enough model capacity.Source: §3.3 · p. 5
  • For reasoning with unseen inputs, the GFlowNet can condition on X, amortizing sampling over inputs and supporting Bayesian model averaging over sampled Z values.For infilling, where Y is available at test time, the policy is conditioned on both X and Y.Source: §3.3 · p. 6

4 EMPIRICAL RESULTS

The experiments evaluate GFlowNet fine-tuning on sentence continuation, story infilling, subjectivity classification, and tool-assisted arithmetic. Across these tasks, the reported results favor GFlowNet in likelihood-diversity trade-offs, reference similarity, low-data performance, and generalization.

4.1 SENTENCE CONTINUATION

  • GFlowNet produces higher-likelihood continuations while maintaining more diversity than baselines, including when baselines receive 5× the compute.At lower temperatures it outperforms diverse beam search in likelihood; with five times as many beams, performance remains comparable while diversity is higher.Source: Figure 3 · p. 6, §4.1 · p. 7

4.2 INFILLING STORIES

  • GFlowNet infills are closer to reference infills than baseline outputs, and the authors report that they coherently link story beginnings and endings.The evaluation includes BERTScore, BLEU-4, GLEU-4, and GPT-4 judgments.Source: Table 2 · p. 7, §4.2 · p. 7

4.3 SUBJECTIVITY CLASSIFICATION

  • In the low-data SUBJ regime, GFlowNet fine-tuning consistently outperforms supervised fine-tuning, and an additional supervised step sometimes improves performance.The task uses tens of labeled examples; Table 3 reports test accuracy for an instruction-tuned GPT-J 6B.Source: Table 3 · p. 7, §4.3 · pp. 7–8

4.4 SOLVING ARITHMETIC PROBLEMS STEP BY STEP

  • GFlowNet significantly outperforms on in- and out-of-distribution arithmetic, while PPO performs poorly with miscalibrated rewards and can produce invalid tool calls.The authors attribute the robustness to matching the full reward distribution rather than collapsing to a single reward mode.Source: Table 4 · p. 8, §4.4 · pp. 8–9

5 FURTHER RELATED WORK

Prior work approaches intractable text inference with MCMC, search, sequential Monte Carlo, and masked-language-model sampling. Related work also connects GFlowNets to posterior inference and compares approaches to latent chain-of-thought sampling.

Sampling from intractable marginals

  • For infilling and constrained generation, prior approaches include MCMC [Miao et al., 2019; Zhang et al., 2020a], variable-neighborhood ascent [Malkin et al., 2021], and sequential Monte Carlo [Lew et al., 2023].Masked language models have also been used for Gibbs-sampling variants [Wang & Cho, 2019; Goyal et al., 2022; Yamakoshi et al., 2022] and small-token-set marginals [Torroba Hennigen & Kim, 2023].Source: §5 · p. 9

GFlowNets

  • GFlowNets were proposed for sampling discrete objects from unnormalized rewards, and later work interprets them as variational inference for Bayesian posteriors over structured objects.The cited applications include biological sequence generation [Jain et al., 2022] and posterior sampling over structured objects [Deleu et al., 2022; 2023; van Krieken et al., 2023; Hu et al., 2023].Source: §5 · p. 9

Chain-of-thought reasoning in LLMs

  • Prior chain-of-thought methods use prompts or demonstrations [Kojima et al., 2022], sample latent chains with MCMC [Xu et al., 2023; Zhou et al., 2022], or fine-tune on successful chains [Zelikman et al., 2022].The passage says fine-tuning on successful chains maximizes reward without guaranteeing diversity; Wang et al. (2023b) report improved prediction from aggregating conclusions across chains.Source: §5 · p. 9

6 CONCLUSION

The paper concludes that GFlowNet fine-tuning improves fidelity-diversity trade-offs, sample efficiency, and generalization on downstream tasks relative to maximum-likelihood training or reward-maximizing policy optimization. It identifies open directions and boundaries, including limited model scale and unresolved exploration for complex latent variables.

6 CONCLUSION

  • GFlowNet fine-tuning improves the fidelity-diversity trade-off for text generation and sample efficiency and generalization on downstream tasks compared with maximum-likelihood training or reward-maximizing policy optimization.The authors describe the method as amortized inference that converts computation into better test-time performance without additional data.Source: §6 · p. 9
  • Future work includes cross-task transfer, a shared universal reasoner, stronger base LLMs, epistemic-uncertainty estimates, and richer latent-variable generation.The authors also suggest applications to probabilistic programming, structured chains of thought, program synthesis, and planning with world models.Source: §6 · p. 9

Limitations

  • Experiments use models up to 6B parameters, exploration for more complex latents remains open, and the method improves inference rather than the LM’s knowledge.The authors therefore state that hallucination and miscalibration tied to knowledge representation are not addressed.Source: §6 · p. 9

A.1 GLOSSARY OF RL FOR LLMS

The glossary frames the paper’s reinforcement-learning terms for LLMs, while the learning objective adapts subtrajectory balance to autoregressive generation and learns only a forward sampling policy.

A.1 GLOSSARY OF RL FOR LLMS

  • The glossary frames RL as cumulative-reward agent learning, policies as state-to-action mappings, and GFlowNet rewards as unnormalized posterior targets.In this setting, policies govern LLM text generation, while rewards guide sampling toward the target distribution.Source: §A.1 · p. 17
  • Distribution matching means approximating a target probability distribution, while policy-gradient methods update policy parameters to increase expected reward.The glossary names PPO as an example of a policy-gradient method.Source: §A.1 · p. 17

A.2 LEARNING OBJECTIVE

  • Subtrajectory balance trains GFlowNets over partial trajectories using a forward policy, a backward policy, and a state-flow function.The objective is defined for a partial trajectory from state s_m to s_n.Source: §A.2 · p. 17
  • Autoregressive left-to-right generation forms a tree with one path and one parent per state, making backward transitions trivial.The backward policy is 1 in this setting.Source: §A.2 · p. 17
  • Replacing final state flow with terminal reward removes the separate flow-function parameters, leaving the forward sampling policy q_GFN as the learned object.The objective sums over partial trajectories with equal weight (λ = 1).Source: §A.2 · p. 17

B SENTENCE CONTINUATION

Sentence continuation targets a distribution over complete next sentences, which differs from token-by-token sampling. The section describes a GFlowNet approach that favors high-probability sentences and reports a balance between likelihood and diversity.

B SENTENCE CONTINUATION

  • Sentences are used as the generation unit because they form more self-contained units of meaning than individual words.The section relates sentences to whole thoughts and reasoning steps.Source: B SENTENCE CONTINUATION · p. 17
  • As temperature approaches zero, tokenwise greedy decoding need not produce the highest-probability sentence, unlike the desired sentence-level distribution.The desired low-temperature distribution over sentences concentrates on highly probable sentences and approaches the argmax.Source: B SENTENCE CONTINUATION · pp. 17–18
  • The GFlowNet policy starts from the pretrained LM weights, extends the prompt autoregressively, and terminates at a period or a forced maximum-length boundary.The prompt consists of the text preceding the next sentence to be sampled.Source: B SENTENCE CONTINUATION · p. 18
  • The policy is trained to sample sentences in proportion to their tempered probability under the original LM, with temperatures below 1 favoring higher-probability sentences.At T = 1, the target is proportional to the original LM distribution; the preferential sampling comes from setting 0 < T < 1.Source: B SENTENCE CONTINUATION · p. 18
  • The experiment used 1,000 OpenWebText prompts and a pretrained 1.5B-parameter GPT2-XL, with the GFlowNet fine-tuned through an 80M-parameter LoRA space.Fifty prompts were used for validation.Source: B SENTENCE CONTINUATION · p. 18
  • Training with reward temperature 0.875 achieves a good balance between log-likelihood and diversity in empirical next-sentence samples.The examples compare GFlowNet samples with autoregressive LM samples and diverse beam search.Source: B SENTENCE CONTINUATION · p. 18

C INFILLING STORIES

For story infilling, the section evaluates generated middle passages by story coherence and reports higher average ratings for GFlowNet infills than for baselines. It also identifies a boundary on the assumed pretrained-model knowledge for this task.

C INFILLING STORIES

  • The paper assumes the base model already contains the task knowledge, but pretrained models were not good at assigning high likelihoods to plausible stories.The authors therefore fine-tuned GPT-2 Large on the stories dataset and used it as the reward model.Source: C INFILLING STORIES · p. 18
  • Training samples provide story context and infill targets, and the evaluation prompt frames each test item as a beginning, middle, and end.The hyperparameters are listed in Table C.2.Source: C INFILLING STORIES · p. 18, Table C.1 · p. 20, Table C.2 · p. 20
  • 4.3 is the average rating for reference infills, treated as an upper bound; GFlowNet infills receive much higher average coherence ratings than baselines.The reference stories may potentially be present in GPT-4 training data.Source: C INFILLING STORIES · p. 21
  • Coherence ratings were averaged over 10 infills per method and 100 test examples; the rubric scores narrative structure and logical order on a 1–5 scale.Grammatical errors or logical inconsistencies lower the coherence score.Source: C INFILLING STORIES · p. 21

D SUBJECTIVITY CLASSIFICATION

The subjectivity-classification appendix details GFlowNet training and replay-buffer choices, alongside supervised fine-tuning settings and ablations.

D SUBJECTIVITY CLASSIFICATION

  • The appendix points to SUBJ training examples in Table D.1 and prompts used for all baselines in Table D.2.GFlowNet training runs for 1,000 steps; supervised fine-tuning settings are also specified.Source: D SUBJECTIVITY CLASSIFICATION · p. 22
  • GFlowNet training uses batches of 512 samples, with 8 queries drawn with replacement and 64 rationales sampled per query for rewards from pLM(ZY|X).Training runs for 1,000 steps with a 200-step warmup and learning rate 0.0005; reward temperature anneals from 1.2 to 1 over 150 steps.Source: D SUBJECTIVITY CLASSIFICATION · p. 22
  • Supervised fine-tuning uses 8 queries per batch, runs for 100 steps with a 20-step warmup, and reports the best result from a learning-rate sweep over 0.001 and 0.0001.The described settings apply both to supervised fine-tuning alone and after GFlowNet fine-tuning.Source: D SUBJECTIVITY CLASSIFICATION · p. 22
  • The replay buffer stores inverse-prompted chains in a reward-indexed priority queue, evicts underperforming chains, and is evaluated through ablations of seeding and multi-chain aggregation.The ablations are reported in Table D.4.Source: D SUBJECTIVITY CLASSIFICATION · p. 22

E INTEGER ARITHMETIC

For integer arithmetic, the appendix examines rationale-buffer seeding, reward calibration, and tool use, and contrasts GFlowNet generations with PPO while documenting remaining reasoning errors.

E INTEGER ARITHMETIC

  • To improve reward calibration, the method prepends three hand-crafted dataset rationales when computing the reward, while the GFlowNet policy is conditioned only on the question.The examples are used only for the reward.Source: E INTEGER ARITHMETIC · p. 22
  • When generating a rationale, the system evaluates the last two terms in Python whenever it generates “=”, and evaluation aggregates responses across multiple samples.For all methods, the authors aggregate responses over multiple samples at a fixed temperature.Source: E INTEGER ARITHMETIC · p. 22
  • PPO produces high-reward sequences that are invalid expressions and do not call the tool, whereas GFlowNets generate correct evaluation rationales.The analysis reports that PPO instead learns to repeat the expression.Source: E INTEGER ARITHMETIC · p. 29, Table E.4 · p. 31
  • GFlowNets still make errors, mainly on longer out-of-distribution problems through missing or miscopied operands or operators.The authors attribute these errors potentially to the reward model assigning equally high rewards to rationales with minor mistakes, and suggest better or larger reward models could reduce them.Source: E INTEGER ARITHMETIC · p. 29, Table E.5 · p. 31
  • Performance generally improves as more dataset rationales seed the replay buffer, while using no seeded rationales performs poorly.The authors hypothesize that exploration for good rationales is challenging without good examples; Table E.3 summarizes the ablation.Source: E INTEGER ARITHMETIC · p. 22, Table E.3 · p. 30
Loading 2310.04363v3…