Amortizing intractable inference in large language models

Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, Nikolay Malkin

Research output: Contribution to conferencePosterpeer-review

Abstract / Description of output

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.
Original languageEnglish
Pages1-31
Number of pages31
Publication statusPublished - 8 May 2024
EventThe Twelfth International Conference on Learning Representations - Vienna, Austria
Duration: 7 May 202411 May 2024
https://iclr.cc/

Conference

ConferenceThe Twelfth International Conference on Learning Representations
Abbreviated titleICLR 2024
Country/TerritoryAustria
CityVienna
Period7/05/2411/05/24
Internet address

Fingerprint

Dive into the research topics of 'Amortizing intractable inference in large language models'. Together they form a unique fingerprint.

Cite this