[Paper Review] Amortizing intractable inference in large language models

J. Edward Hu, Moksh Jain+5 more|arXiv (Cornell University)|Oct 6, 2023

Topic Modeling4 citations

TL;DR

This paper proposes amortized Bayesian inference via fine-tuning large language models (LLMs) with generative flow networks (GFlowNets) to enable efficient, diverse sampling from intractable posterior distributions—such as those required for chain-of-thought reasoning and constrained generation. The method outperforms supervised fine-tuning and PPO in data efficiency and sample diversity, achieving 95.2% accuracy on in-distribution integer arithmetic tasks with 50 seed rationales.

ABSTRACT

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.

Motivation & Objective

• To address the limitation of autoregressive LLMs in tractably querying knowledge beyond autoregressive sampling.
• To enable efficient, diverse sampling from intractable posteriors in tasks like sequence continuation, infilling, and constrained generation.
• To improve data efficiency and sample diversity in few-shot reasoning and tool-use adaptation compared to standard fine-tuning and reinforcement learning.
• To frame chain-of-thought reasoning as a latent variable modeling problem amenable to amortized Bayesian inference.

Proposed method

• Fine-tune a pretrained LLM using GFlowNets, a diversity-seeking reinforcement learning algorithm that trains policies to sample objects with probability proportional to a reward function.
• Use the LLM's own likelihood as a reward function to evaluate sequences during training, enabling self-supervised reward modeling.
• Initialize the GFlowNet policy with a pretrained LLM and continue training via a reward objective that matches the joint probability of input, rationale, and output.
• Employ a buffer of seed rationales to improve exploration and initial policy quality, especially in low-data regimes.
• Apply temperature scheduling and curriculum learning to stabilize training and improve convergence.
• Use LoRA for efficient fine-tuning with hyperparameters tuned for sample diversity and reward maximization.

Figure 1: Left: Three problems of reasoning in language – sentence infilling, chain-of-thought reasoning, and problem-solving with external tool use – can all be seen as instances of the latent variable model at the top left, where an input ( $X$ ) generates the output ( $Y$ ) via a latent variable

Experimental results

Research questions

• RQ1 Can amortized inference via GFlowNet fine-tuning enable effective and diverse sampling from intractable posteriors in LLMs?
• RQ2 How does GFlowNet fine-tuning compare to supervised fine-tuning and PPO in data efficiency and sample diversity for reasoning tasks?
• RQ3 Can GFlowNet fine-tuning effectively learn to generate valid, correct rationales for arithmetic reasoning with limited supervision?
• RQ4 Does the use of a seed buffer of rationales significantly improve performance in low-data few-shot reasoning scenarios?
• RQ5 Can GFlowNet fine-tuning reduce mode collapse and over-optimization issues common in PPO-based policy learning?

Key findings

• GFlowNet fine-tuning achieved 95.2% test accuracy on in-distribution integer arithmetic tasks when seeded with 50 rationales, significantly outperforming zero-shot prompting and baseline fine-tuning.
• With 50 seed rationales, the GFlowNet model achieved 75.4% accuracy on out-of-distribution (OOD) arithmetic problems, demonstrating strong generalization beyond training distribution.
• The method produced valid, step-by-step rationales that correctly evaluated expressions, unlike PPO, which often generated repetitive or invalid sequences.
• Performance improved monotonically with more seed rationales, with 0 seed examples yielding only 22.6% accuracy, highlighting the importance of initial supervision for exploration.
• PPO models produced high-reward but invalid rationales (e.g., repeating expressions), while GFlowNets generated correct, diverse reasoning chains, indicating reduced mode collapse.
• The ablation study confirmed that seed rationales are critical for performance, as performance dropped sharply without them, suggesting that exploration is challenging without prior knowledge.

(a) Base model 50.5% of samples are valid numbers.