Source-linked AI summary
Language Model Cascades
David Dohan, Winnie Xu, Aitor Lewkowycz, Jacob Austin, David Bieber, Raphael Gontijo Lopes, Yuhuai Wu, Henryk Michalewski, Rif A. Saurous, Jascha Sohl-dickstein, Kevin Murphy, Charles Sutton
TL;DR
Language-model reasoning methods lack a single framework for composing probabilistic models with string-valued variables and dynamic structure. The paper formulates them as language model cascades using probabilistic programming, and demonstrates the framework on several reasoning structures, including a Twenty Questions system solving 29% of tasks.
Problem
Existing language-model reasoning approaches use repeated interactions, model composition, and dynamic structures that call for a unified probabilistic framework over complex data such as strings.
Method
The paper represents language-model cascades as probabilistic programs with string-valued random variables, observations, control flow, and composed model structures.
Results
LaMDA 137B solves 29% of Twenty Questions tasks in the reported cascade experiment.
Takeaways & Limitations
The cascade framework places diverse reasoning methods and inference strategies in one formalism and supports multi-step and interactive language-model systems.
Takeaways & Limitations
The paper does not evaluate inference methods beyond rejection sampling, while efficient inference for string-valued variables remains a practical challenge.
Abstract
from arXiv · showhide
Prompted models have demonstrated impressive few-shot learning abilities. Repeated interactions at test-time with a single model, or the composition of multiple models together, further expands capabilities. These compositions are probabilistic models, and may be expressed in the language of graphical models with random variables whose values are complex data types such as strings. Cases with control flow and dynamic structure require techniques from probabilistic programming, which allow implementing disparate model structures and inference strategies in a unified language. We formalize several existing techniques from this perspective, including scratchpads / chain of thought, verifiers, STaR, selection-inference, and tool use. We refer to the resulting programs as language model cascades.
1. Introduction
Language model cascades provide a probabilistic-programming framework for unifying diverse language-model reasoning methods over string-valued variables.
- 1. Introduction: Language model cascades define joint probabilistic models over strings using probabilistic programming languages extended with string-valued variables.The framework conditions these models on string-valued observations to infer unknowns.
- 1. Introduction: The framework captures scratchpads, chain-of-thought prompting, verifiers, selection-inference, bootstrapping, and related language-model reasoning proposals.These approaches use prompting or repeated model interactions to support informal reasoning.
- 1. Introduction: By expressing disparate model structures and inference strategies in one framework, cascades support more complex multi-step reasoning problems.The authors propose that this could enable generic procedures for inference, parameter tuning, and prompt selection based on end-to-end objectives.
2. Related work
Prior work supplies probabilistic-programming foundations, language-model prompting and composition methods, and inference strategies that language model cascades unify over string-valued variables.
- 2. Related work: Probabilistic programming languages support stochastic program structure, sampling, and conditioning, extending graphical models beyond fixed distributions.Their graph size and shape can depend on random outcomes or data-dependent stopping criteria.
- 2. Related work: Language models answer questions through sampling or MAP decoding, with few-shot examples supplied either in prompts or through fine-tuning.The predictive distribution is written as p(A|Q, θ), with prompting or parameter updates incorporating training pairs.
- 2. Related work: Scratchpads and chain of thought introduce an auxiliary thought variable, modeling answers as p(A|T, Q)p(T|Q).The thought is generated before the answer and can be encouraged through prompting or fine-tuning.
- 2. Related work: Beam search approximates latent-thought marginalization with a MAP thought, while Self-Consistency samples multiple thoughts and selects the most common answer.These strategies address the practical impossibility of summing over all possible string-valued thoughts.
- 2. Related work: PromptChainer and Socratic models compose language models with control flow or multimodal interaction, while related work connects language models to Bayesian networks and causal graphical models.PromptChainer nodes may query language models or external systems, and Socratic models target settings lacking a single suitable model.
3. Cascades
Language model cascades represent string-valued language-model interactions as probabilistic programs, unifying reasoning methods, verifiers, selection-inference, tool use, and dynamic control flow. The framework is illustrated through several cascade constructions and a preliminary Twenty Questions experiment achieving 29% task success.
- 3. Cascades: Language model cascades define joint distributions over string-valued variables sampled from language models, supporting reasoning programs with complex control flow and observations.They are implemented as trace-based probabilistic programs embedded in Python, with pretrained language models parameterizing conditional distributions.
- 3.1. Scratchpads and Chain of thought: Chain-of-thought cascades model question, thought, and answer variables, condition answers on few-shot question-thought-answer examples, and support ancestral-sampling inference.The same framework can represent semi-supervised learning by treating thoughts as hidden variables and imputing them for larger question-answer datasets.
- 3.3. Selection-Inference: Selection-inference cascades separate reasoning into selecting relevant facts and inferring new facts from them, with the selection and inference steps repeatable for multistep reasoning.Prompted examples include questions, fact sets, selected subsets, inferences, and answers.
- 3.4. Verifiers: Verifier cascades add validity variables to reasoning traces, allowing sampled solutions to be rejected or ranked by the probability that their reasoning is correct.Verifiers can be binary correctness predictors or produce natural-language critiques, while thought generation can condition on verifier correctness.
- 3.6. Twenty questions: Language model cascades also express interactions with external tools and adaptive conversations, including calculator or browser calls and inference over string-valued actions.In Twenty Questions, forward sampling with early rejection and concept masking yielded 29% success on the evaluated tasks using LaMDA 137B.
4. Discussion
The framework unifies complex string-valued model compositions through probabilistic programming, while highlighting inference efficiency as a practical challenge and multimodal and planning extensions as possibilities.
- Probabilistic programming provides a flexible formalism for composing models into complex probabilistic models over strings and unifying existing algorithms.
- The framework can represent many planning and reinforcement-learning tasks through control as inference.
- Although presented for strings, the ideas are described as applicable to multimodal systems that combine image and text models.
- Efficient inference remains the core technical challenge for language model cascades with string-valued variables.The paper suggests particle-based inference, language-model proposal distributions, and guide networks as directions for future work.
- Probabilistic program induction could search for cascades that solve target tasks instead of assuming a fixed program structure.
A.1. Inference
The inference API reifies unobserved values in a probabilistic model conditioned on observations, with ancestral sampling as the simplest inference strategy.
- Inference reifies specific unobserved values conditioned on observed values in a probabilistic model.
- Ancestral sampling, also called forward sampling, is the simplest inference algorithm described for these programs.
A.2. Code examples
The code examples express language-model reasoning techniques as probabilistic programs with sequential modules, verifiers, dynamic stopping, and selection-inference loops.
- Question answering: The basic question-answering graph directly samples a question and then an answer conditioned on that question.
- Chain of thought: Chain of thought introduces a latent thought between question generation and answer generation.
- Self critique: Self-critique adds a natural-language critique step after generating a question, thought, and answer.
- Verifiers and stopping: A verifier can score reasoning sequences, while a separate model output determines when iterative generation should stop.
- Selection-Inference: Selection-Inference alternates selecting relevant facts and inferring new facts, with a model-controlled stopping condition and final answer.
B.1. Problem definition
The problem-definition example models a bounded two-agent guessing game in which Bob questions Alice until identifying a concept or reaching a stopping condition.
- The program ends when Bob guesses the concept, asks 10 questions unsuccessfully, or fails to answer a question.
- The model correctly identifies 11 of 40 tested concepts, concentrated among the least abstract concepts.
B.2. Examples
The examples illustrate successful and unsuccessful two-agent Twenty Questions conversations, where Bob asks questions and Alice answers while the hidden concept is supplied to Alice. Success occurs when Bob identifies the concept; failure includes malformed or repeated questioning that does not produce the correct answer.
- In the successful example, Bob identifies “love” immediately after Alice reveals the concept in the prompted exchange.
- A second example labels the conversation a success after Alice’s concept is “shoes” and the interaction proceeds through the game prompts.
- The wrong-answer example shows Bob producing malformed behavior, including treating the answer as already known and repeating a question about Monopoly.
B.3. Code for twenty questions
The Twenty Questions program constructs separate prompts for Alice and Bob, iterates through question-answer rounds, and stops or rejects execution based on guessing and response validity. It hides the target concept from Alice’s generated response when necessary and records successful guesses by round number.
- The program builds Alice and Bob prompts, initializes the shared conversation, and iterates for at most max_questions rounds.
- Bob’s response must contain a question mark; otherwise the program rejects the run because no valid question was generated.
- When Bob’s response contains the target concept, the program yields success with the number of rounds used.
- Alice receives the accumulated conversation and target concept, and any generated occurrence of that concept is replaced with “concept” rather than exposed directly.