Language Model Cascades

Introduction

Language models (LMs) have demonstrated impressive few-shot learning abilities (Brown et al., 2020; Chowdhery et al., 2022). This has led to a number of proposals to use LMs as the basis of informal reasoning, including scratchpads (Nye et al., 2021), chain of thought prompting (Wei et al., 2022; Wang et al., 2022), learned verifiers (Cobbe et al., 2021), selection-inference (Creswell et al., 2022), and bootstrapping (Zelikman et al., 2022). They have also been applied in formal mathematics settings to guide theorem provers (Polu & Sutskever, 2020). These methods involve prompting to encourage step-by-step reasoning, repeated interactions with a single LM, or multiple LMs linked together, with the models being fine-tuned or prompted in different ways.

In this position paper, we argue that a useful unifying framework for understanding and extending this disparate body of work is in terms of probabilistic programming languages (PPL) extended to work with strings, instead of more atomic data types like integers and floats. That is, we use a PPL to define a joint probability model on string-valued random variables, parameterized using LMs, and then condition this model on string-valued observations in order to compute a posterior over string-valued unknowns, which we can then infer. We call such a probabilistic program a language model cascade. We show that this framework captures many recent approaches, and also allows us to tackle more complex multi-step reasoning problems. By implementing many disparate model structures and inference strategies in a single framework, we hope that language model cascades will enable the development of generic procedures to perform inference, tune parameters, and choose prompts based on end-to-end objectives.An implementation is available at model-cascades.github.io

Related work

There is a rich prior literature on probabilistic programming languages (PPLs), which extend probabilistic graphical models to support more complex joint distributions whose size and “shape” can itself be stochastic (e.g., a graph unrolled for a random number of iterations, until a data-dependent stopping criterion is met). PPLs extend traditional programming languages with the ability to sample from distributions and observe values of variables based on data (i.e. condition the model). The semantics of sample and observe vary depending on the inference algorithm. For more details, see van de Meent et al. (2018).

Recently there has been an explosion of interest in large language models, such as GPT-3 (Brown et al., 2020) and PaLM (Chowdhery et al., 2022). These can be used for tasks such as “zero-shot” question-answering. In this setting, we provide the question $Q$ as a prompt to the LM, and then sample answers from the model, which we denote by $p(A|Q,\theta)$, where $\theta$ are the pre-trained model parameters. Alternatively, we can compute the MAP answer, $\hat{A}=\operatorname*{\arg\!\max}_{A}p(A|Q,\theta)$.

To ensure the model “does the right thing”, we can provide a small training set of question-answer pairs, $D=\{(Q^{m},A^{m}):m=1:M\}$ pairs. This can be provided as extra context to the model, provided in the text prompt, followed by sampling from $p(A|Q,D,\theta)$. We refer to this as “few-shot prompting”. We can also fine-tune the model parameters on $D$ to get $\theta^{\prime}$, and then sample from $p(A|Q,\theta^{\prime})$.

We can improve performance by introducing an additional auxiliary “thought” variable, and then extend the model to have the form $p(A,T|Q)=p(A|T,Q)p(T|Q)$, where each conditional is computed using an LM which includes its conditioning variables as a part of its input. Work on scratchpads (Nye et al., 2021) and chain of thought (Wei et al., 2022) illustrate this, and finetune or prompt the LM to produce this auxiliary thought before answering.

We typically condition this on a small set $D_{S}$ of $(A^{m},T^{m},Q^{m})$ triples, and optionally a larger set $D_{L}$ of $(A^{m},Q^{m})$ pairs. We then compute a distribution over answers to a test question using

where $\hat{p}(\cdot)=p(\cdot|D_{L},D_{S},\theta)$ is the prior predictive distribution. (Scratchpad creates its prior predictive by fine-tuning, while Chain of Thought adds $D_{S}$ to the LM prompt.)

In practice, we cannot sum over all possible strings $T$ in Equation 1. The most common approach is to compute the MAP estimate $\hat{T}=\operatorname*{\arg\!\max}\hat{p}(T)$ using beam search, and then to approximate the sum over $T$ with this single value. More recently, Self Consistency (Wang et al., 2022) proposed to sample multiple values for $T$ using forward sampling of $(A,T)$ given $Q$, and then taking the answer $A$ that is most common in this setThis bucketing is practical because most standard benchmarks have answers that are just a couple words..

PromptChainer (Wu et al., 2022) proposes a visual interface for composing language models together, specifying control flow and prompting strategies for each node in a chain. Nodes may query language models or external systems. Socratic models (Zeng et al., 2022) extends model chaining to the multimodal setting and demonstrates zero-shot abilities on tasks for which no single model exists.

The Eliciting Latent Knowledge proposal (ELK, 2022) suggests making latent variables explicit, modelled using a Bayesian network, to improve interpretability and safety for advanced AI systems.

Ortega et al. (2021) explains a formalism for LM finetuning with causal graphical models in order to extend the predictive capabilities of AI agents towards more adaptive behaviour. They focus on analysing an auto-regressive action (random variable) prediction scheme in the interactive setting of RL where a model is simultaneously a generator and predictor of data.

Cascades

In this section, we show how to create cascades of LMs to tackle various language-based reasoning problems. A cascade is a probabilistic program that includes string-valued random variables, sampled from an LM. For example, Figure 2 is a simple cascade for question answering. Each of the yield expressions return a string distributed according to the language model S.The first argument to S defines a unique name for the random variable, and the remaining arguments conditions the LM on a string prefix. A variable may be marked as observed within the program, S(’varname’, obs=’observed value’), or at inference time This program defines a joint distribution over the variables question, thought, and answer. Programs with complex control flow and observations are included in Appendix A.

We implement cascades as a trace-based probabilistic programming language embedded in Python via effect handlers, inspired by Bingham et al. (2018); Phan et al. (2019), and via coroutines, inspired by Kochurov et al. (2019). A pretrained LM is used to parameterize all conditional distributions. A cascade supports arbitrary control flow and recursion. While the current presentation is in terms of few-shot prompting of causal language models, we emphasize that the ideas are immediately applicable to finetuned models, masked LM setting, and other complex data types including images.

As our first example, we show how to represent a chain of thought (Nye et al., 2021; Wei et al., 2022) as shown in Figure 1 and subsequent graphical model figures; refer to the corresponding probabilistic programs in Appendix A. We condition the $A$ node not just on the test question $Q$, but also on previous $(Q^{m},T^{m},A^{m})$ triples, which constitute the few-shot prompting part of the model. This is denoted by the shaded nodes inside the plate. Inference can be implemented by ancestral sampling.

2 Semi-supervised learning

In Section 3.1, we provided a manually created set $(Q^{m},T^{m},A^{m})$ triples, where the “thoughts” or “rationalizations” were provided. A more scalable approach is to define a small set $D_{S}$ of such “supervised” triples, but then to provide a larger set $D_{L}$ of $(Q^{m},A^{m})$ pairs, which are easier to gather. We can augment the pairs in $D_{L}$ by adding the hidden $T^{m}$ variable to get a semi-supervised setup, shown in Figure 3.

The Self-Taught Reasoner (STaR) (Zelikman et al., 2022) proposes a procedure for fine-tuning LMs in the chain-of-thought type setting. We can interpret their method as a stochastic EM-like procedure in the cascade of Figure 3. In particular, they first fine-tune on the “fully observed” dataset $D_{S}=\{(Q^{m},T^{m},A^{m})\}$. Then they impute the unknown $T_{i}$ values in the “partially observed” dataset $D_{L}=\{(Q^{m},T^{m}=?,A^{m})\}$ during the “E” step by doing rejection sampling on $p(T,A|Q^{m})$ until finding a thought which leads to the known correct answer. If sampling $(T,A)$ given the question fails to find the correct answer, they sample thoughts from $p(T|Q^{m},A^{m})$. This uses a recognition network to approximately sample from the posterior distribution over thoughts given the known correct answer. They call this approach “rationale generation with rationalization”. They then update the parameters in the “M” step based on these imputed thoughts. By interpreting the rationale generation at this higher level of abstraction, we open up the possibility of applying this tuning method to other types of cascades.

3 Selection-Inference

Selection Inference (Creswell et al., 2022) is a recent example of multiple interacting LM modules. It proposes splitting reasoning into: the selection module which selects a subset of facts given a question, and the inference module which infers new facts given this subset.

It may be represented by the model in Figure 4. Here $S$ is the selection of a subset of “facts” from a pre-specified set of facts, and $I$ is an inference driven by that fact. The $S$ and $I$ nodes can be iterated to do multistep reasoning. The model is “trained” by giving it examples, $D=\{(Q^{m},\{F^{mj}\},S^{m},I^{m},A^{m}):m=1:M\}$, as part of the prompt.

4 Verifiers

Although adding explicit “thought” variables to a model has been found to improve performance, models still arrive at incorrect answers, or the correct answer for an erroneous reason. An intuitive way to improve model performance is to train it to judge whether an answer and thought are likely to be “valid”. Cobbe et al. (2021) propose using a separate model as a verifier to filter solutions to reasoning tasks.

We can create a “labeled” training set of the form $D=\{(Q^{m},T^{m},A^{m},V^{m}\}$, where we add a “verification” label $V^{m}\in\{0,1\}$, representing whether the thought $T^{m}$ is a valid form of reasoning for deriving $A^{m}$ from $Q^{m}$, and $A^{m}$ is the correct answer. This can be particularly helpful in settings where there may be more than one way of deriving the answer. The verifiers may be used to reject incorrect examples in ancestral sampling, and the thought generator may itself be conditioned on the verifiers being correct by finetuning or prompting, reminiscent of RL as inference (Levine, 2018) and goal-conditioned policies such as decision-transformer (Chen et al., 2021).

We can extend this to $N$-step reasoning as follows (where we drop conditioning on $D$ for brevity):

We can represent this as shown in Figure 5.

To see why such a verification model can be useful, consider (for simplicity) the case where $N=1$. Suppose we have trained the model to generate valid thoughts and answers by giving it suitable training examples, and then we generate $K$ samples $(T^{k},V^{k},A^{k})\sim p(T,V,A|Q,D)$. We can then rank the samples for validity by computing $r^{k}=p(V^{k}=1|A^{k},Q,D)$, and then picking the $A^{k}$ with largest score $r^{k}$.

Cobbe et al. (2021) train the verifier to predict a binary correctness label. Scheurer et al. (2022) incorporates natural language feedback, and finds that learning is significantly more sample efficient. Preliminary evidence suggests that LMs are capable of critiquing their own chain of reasoning in language, in which case the verifier produces natural language and $p(V_{1:N}=1|Q,A,T_{1:N})$ becomes the likelihood of the verifier taking on a particular string value, such as $p(V_{1:N}=\text{"The reasoning and solution are correct."}|...)$. Saunders et al. (2022) study model generated critiques in the context of summarization.

5 Tool-use

The applications discussed so far involve iterating a language model, within some control flow, without external feedback. There are many tasks of interest in which a model is interacting with external systems. Cobbe et al. (2021) has an LM use a calculator to solve math tasks, while Nakano et al. (2021) put an LM in a loop with a web browser to answer questions. Using PPLs to represent these probabilistic models allows easily representing these cases, by writing the call to the external tool, such as the calculator, directly into the program. Then techniques from simulation based inference, for example, can be applied to do inference in such situations (Cranmer et al., 2020).

6 Twenty questions

In this section, we discuss experimental results using cascades to solve the “Twenty Questions” task from BigBench (BIG-bench collaboration, 2021). This task involves a conversation between two agents, Alice and Bob. Both agents are presented with the rules of the game, and Alice is additionally presented with a concept (e.g. ‘apple’) to describe. Bob has to guess the concept by asking a series of questions $B_{t}$ of the form “Is it X?”, to which Alice answers $A_{t}\in\{\text{`Yes.'},\text{`No.'}\}$. We repeat this process until Bob guesses correctly, or we hit the limit of $T$ rounds. This can be thought of as a pair of interacting Markov chains, which exchange strings, until some final end state is reached, as illustrated in Figure 6.

The goal is to infer what questions Bob should ask to guess the concept as quickly as possible. This can be cast as a reinforcement learning problem with string-valued actions, or equivalently as an inference problem where we condition on the goal state that $A_{T}=\text{`yes'}$ for the soonest possible $T$ (c.f., planning as inference (Levine, 2018)).

In our current preliminary experiments, we use a forward sampling approach (aka ancestral sampling), in which we sample 50 conversations per concept with temperature $1.0$. We consider a trial successful if the target concept appears in $B_{t}$. (i.e., Bob guesses the right answer). We reject a sampling chain early if it is “malformed” (e.g., Bob generates a response that is not a question).

Bob’s turn starts with ‘Is the concept’ which we complete with the LM. Then we let Alice generate an answer; we post-process Alice’s response by replacing all mentions of the true concept with the generic word “concept”, to prevent information leakage. Using the LaMDA 137B large LM (Thoppilan et al., 2022), we find that the model is able to solve $29\%$ of the tasks. See Appendix B for more details.

Discussion

We have shown how probabilistic programming provides a flexible formalism for composing models together to define complex probabilistic models over strings, placing many existing algorithms in a unified framework. While this suggests the possibility of applying a variety of existing inference and train-time techniques to the resulting models, the present work does not evaluate methods beyond rejection sampling.

We can also cast many planning and RL tasks in our framework, by using the perspective of control as inference. While we restrict presentation to the string setting, the ideas presented here are applicable to multimodal settings as well, allowing us to combine image and text models into a larger system.

A challenge applying cascades in practice is the difficulty of probabilistic inference in models with string-valued variables. Previous work in particle based inference for probabilistic programs provides some hope in this direction (Tolpin et al., 2016).

The core technical challenge is efficient inference, as is usually the case with PPLs. A key insight, which we intend to explore in future work, is that we can emulate posterior inference by training the LM to “fill in the blanks”, corresponding to the unknown variables. A similar idea is explored in foundation posteriors (Wu & Goodman, 2022), applied to Stan probabilistic programs, demonstrating that LMs are applicable to numerical data types as well. In other words, we can use LMs as proposal distributions, or guide networks. We also intend to explore fine-tuning methods, going beyond the few-shot prompting approach described here.

Recent advances in program synthesis suggest the possibility of probabilistic program induction (Lake et al., 2015; Goodman et al., 2014) to search for cascades which solve a target task, rather than assuming a fixed probabilistic program structure.

Acknowledgements

We thank Alex Gray, Andreas Stuhlmüller, Ben Poole, Du Phan, Ellen Jiang, Maarten Bosma, Matt Hoffman, Michael Terry, Sharad Vikram, Sherry Tongshuang Wu,and Tuan Anh Le for helpful discussions.

References

Appendix A Implementation

Given a program representing a probabilistic model, inference reifies specific unobserved values conditioned on observed values. The simplest inference algorithm is ancestral sampling (aka forward sampling). The basic inference API is:

A.2 Code examples

In each example below, S is a string distribution. It consists of turning the input values into a prompt, together with any examples provided as few-shot examples to the ‘infer’ method, and sampling until some stopping criterion.

The basic question answering graph directly generates the answer given the question:

Chain of thought introduces a latent thought before producing an answer:

Self critique introduces a step in which the model critiques its own reasoning in natural language:

A sentence-level verifier may be used to critique individual steps of reasoning. Furthermore, when to halt generation may itself be a random variable:

Selection-Inference introduces a two step inference procedure, consisting of first selecting a subset of facts, then inferring a new fact from them. Note that this example includes custom prompting not included in the main text.

Appendix B More details on Twenty Questions

In this task there are two agents: Alice and Bob. Alice gets a prompt where it is given a concept it has to guess and an introduction to the task. Bob gets a prompt where it is instructed on the task. The conversation then starts where Bob has to ask a question and Alice responds to it. If Alice’s response includes the key concept, we change it to the word ‘concept‘ (alternatively, one might reject the trace). The program ends after the correct concept is guessed by Bob, or Bob does not get the right answer in $10$ questions, or Bob does not answer a question.

The 40 concepts that we test the model on are: [’apple’, ’television’, ’dinosaur’, ’airplane’, ’house’, ’tree’, ’coat’, ’shoes’, ’car’, ’train’, ’shower’, ’frisbee’, ’cow’, ’cosmic crisp apple’, ’giganotosaurus’, ’siberian huskey’, ’glass micropipette’, ’jog’, ’catch’, ’defenestrate’, ’eat’, ’apologize’, ’operate’, ’pretend’, ’anger’, ’love’, ’hate’, ’contentment’, ’jealousy’, ’surprise’, ’disgust’, ’hopefulness’, ’global poverty’, ’phase transition’, ’positive sum game’, ’beauty’, ’representative democracy’, ’tall’, ’mauve’, ’perspicacious’]. The model gets right the 11 of the least abstract concepts: [’apple’, ’dinosaur’, ’airplane’, ’house’, ’tree’, ’shoes’, ’car’, ’cow’, ’eat’, ’love’, ’beauty’].

B.2 Examples

B.3 Code for twenty questions

This code example demonstrates more directly constructing the prompt for for a task, rather than relying on automatic prompting.