Demonstrate-Search-Predict: Composing retrieval and language models for knowledge-intensive NLP — part 2/4

Cleaned-text code points [15564, 31537). Bibliography omitted here; source.md includes it.

This simple case study illustrates the power of composition in the DSP abstraction. Because the pipeline is a well-defined program in which transformations communicate by passing text attached to Example s, a simple map-and-filter strategy can leverage the LM and RM to bootstrap annotations for a full pipeline from end-task labels. This is an extensible strategy, but even in its simplest form it generalizes the approaches explored recently by Zelikman et al. (2022) , Wei et al. (2022) , Zhang et al. (2022) , and Huang et al. (2022) in which an LM self-generates chain-of-thought rationales for an individual prompt.

By bootstrapping pipelines, Demonstrate makes it easy to explore complex strategies in Search and Predict without writing examples for every transformation. This includes strategies that are challenging to explore without custom annotations in traditional retrieval-augmented NLP. For instance, Khattab et al. (2021a) introduces a pipeline for multi-hop reasoning that is trained with weak supervision, extending work by Lee et al. (2019) and Khattab et al. (2021b) . In it, the target 3 or 4 passages that need to retrieved must be labeled but the system discovers the best order of “hops” automatically.

In contrast, DSP allows us to build complex pipelines without labels for intermediate steps, because we can compose programs out of small transformations. If LM and RM can accurately process such transformations “zero-shot” (i.e., without demonstrations) on at least one or two examples, these examples can be discovered with end-task labels and used as demonstrations.

To draw on our earlier analogy with DNN frameworks like PyTorch, Demonstrate aims to replace the function of backpropagation in extensible ways by simulating the behavior of the program (corresponding to a “forward” pass) and programmatically learning from errors. In doing this with frozen models and with only end-task labels, Demonstrate introduces a high degree of modularity. In particular, without hand-labeling intermediate transformations, developers may swap the training domain, update the training examples, or modify the program’s strategy, and use annotate to automatically populate all of the intermediate fields for demonstrations.

#### Selecting Demonstrations

It is not always possible to fit all of the training examples in the context window of the LM . DSP provides three primitives for selecting a subset of training examples, namely, sample , knn , and crossval .

⬇

As a baseline choice, k k demonstrations can be randomly sampled from train using the sample primitive, an approach used by Brown et al. (2020) and much subsequent work. We can also leverage the RM ’s representations and select from the training set the k k nearest neighbors to the input text, a strategy explored by Liu et al. (2021) . Another strategy is to apply cross-validation to select among a number of sampled sets of demonstrations ( Perez et al., 2021 ) . For example, given | train | = 100 |\texttt{train}|=100 training examples, crossval would select n n subsets of k = 5 k=5 examples each, and return the set with which a transformation evaluate performs best on the remaining 95 95 examples.

#### Compositions & Extensions

By manipulating demonstrations and higher-order transformations, these simple selection and bootstrapping primitives can be combined to conduct larger novel strategies. If the training set is very large (e.g., | train | = 100,000 |\texttt{train}|=100,000 ), we can conduct knn to find the nearest k = 16 k=16 examples and only annotate these, arriving at a system that learns incrementally in real-time. If the training set is moderately large (e.g., | train | = 1000 |\texttt{train}|=1000 ), we can conduct crossval and cache the performance of all prompts it evaluates on each training example. At test time, we can use knn to find k = 50 k=50 similar examples to the test input and select the prompt that performs best on these k k examples, producing an adaptive system that is informed by the quality of its pipeline on different types of examples.

### 2.4 Search

The Search stage gathers passages to support transformations conducted by the LM . We assume a large knowledge corpus—e.g., a snippet of Web, Wikipedia, or arXiv—that is divided into text passages . Providing passages to the LM facilitates factual responses, enables updating the knowledge store without retraining, and presents a transparency contract: when in doubt, users can check whether the system has faithfully used a reliable source in making a prediction.

In the simplest scenarios, Search can directly query the RM , requesting the top- k k passages (from a pre-defined index) that match an input question. This baseline instantiation of Search simulates retrieval in most open-domain question answering systems, which implement a “retrieve-then-read” pipeline, like Lee et al. (2019) , Khattab et al. (2021b) , Lazaridou et al. (2022) , and many others.

⬇

#### Search Strategies

In many scenarios, the complexity of the task demands more sophisticated Search strategies that empower the RM to find relevant passages. Our running example (Figure 2 ) is one such scenario, in which we suspect examples are likely to require multi-hop reasoning in particular. Other settings, for instance, pose conversational challenges, in which the information need expressed by a user can only be resolved by taking into account previous turns in the conversation, or demand more extensive planning ( Zhong et al., 2022 ) .

In the retrieval-augmented NLP literature, multi-hop search ( Xiong et al., 2020 ; Khattab et al., 2021a ) and conversational search ( Del Tredici et al., 2021 ; Raposo et al., 2022 ) pipelines have received much attention. These systems are typically fine-tuned with many hand-labeled query “rewrites” ( Anantha et al., 2020 ) , “decompositions” ( Geva et al., 2021 ; Min et al., 2019 ) , or target hops ( Yang et al., 2018 ; Jiang et al., 2020 ) . Supported with automatic annotations from Demonstrate , the Search stage allows us to simulate many such strategies and many others in terms of passing queries, passages, and demonstrations between the RM and LM . More importantly, Search facilitates our vision of advanced strategies in which the LM and RM cooperate to incrementally plan a research path for which the RM gathers information and the LM identifies next steps.

#### Case Study

Let us build on our running multi-hop example as a case study. We can define multihop_search_v2 (Line 4 in our core program), a slightly more advanced version of the Search transformation from Figure 2 . This transformation simulates the iterative retrieval component of fine-tuned retrieval-augmented systems like IRRR ( Qi et al., 2020 ) , which reads a retrieved passage in every hop and generates a search query (or a termination condition to stop hopping), and Baleen ( Khattab et al., 2021a ) , which summarizes the information from many passages in each hop for inclusion in subsequent hops.

⬇

In multihop_search_v2 , Line 7 calls the generate primitive, which invokes the LM to produce a query for each retrieval hop. The LM is conditioned on a prompt that is prepared using the hop_template template. (We discuss prompt templates and the generate primitive in § 2.5 .) Here, this template may be designed to generate a prompt that has the following format (e.g., for the second hop).

⬇

As shown, the LM is instructed to read the context retrieved in earlier hops and a complex question. It is then prompted to write: (i) a summary of the supplied context and (ii) a search query that gathers information for answering that question. The generated text will be extracted and assigned to the summary and query variables in ( multihop_search_v2 ; Line 7). On Line 10, we terminate the hops if the query is “N/A”. Otherwise, Line 12 retrieves k = 5 k=5 passages using the query and Line 13 assigns the context for the subsequent hop (or for Predict ), setting that to include the summary of all previous hops as well as the passages retrieved in the final hop so far.

#### Comparison with self-ask

It may be instructive to contrast this multi-hop DSP program with the recent “self-ask” ( Press et al., 2022 ) prompting technique, which we compare against in § 3 . Self-ask can be thought of as a simple instantiation of DSP ’s Search stage. In it, the LM asks one or more “follow-up questions”, which are intercepted and sent to a search engine. The search engine’s answers are concatenated into the prompt and are used to answer the question. This is essentially a simplified simulation of IRRR ( Qi et al., 2020 ) .

As a general framework, DSP can express ideas like self-ask and many other, more sophisticated pipelines as we discuss in the present section. More importantly, DSP offers a number of intrinsic advantages that lead to large empirical gains: 80%–290% over self-ask. For instance, DSP programs are deeply modular, which among other things means that DSP programs will annotate and construct their own demonstrations. Thus, they can be developed without labeling any of the intermediate transformations (e.g., the queries generated). In addition, the LM prompts constructed by DSP get automatically updated to align with the training data and retrieval corpus provided. In contrast, approaches like self-ask rely on a hand-written prompt with hard-coded examples.

Moreover, DSP assigns the control flow to an explicit program and facilitates design patterns that invoke the LM (or RM ) to conduct small transformations. This allows us to build steps that are dedicated to generating one or more retrieval queries, summarizing multiple passages per hop, and answering questions. These steps are individually simpler than the self-ask prompt, yet our multi-hop DSP program deliberately composes them to build richer pipelines that are thus more reliable. In contrast, self-ask delegates the control flow to the LM completions, maintaining state within the prompt itself and intercepting follow-up questions to conduct search. We find that this paradigm leads to a “self-distraction” problem (§ 3.5 ) that DSP programs avoid.

#### Fusing Retrieval Results

For improved recall and robustness, we can also fuse the retrieval across multiple generated queries. Fusion has a long history in information retrieval ( Fox & Shaw, 1994 ; Xue & Croft, 2013 ; Kurland & Culpepper, 2018 ) and sequentially processing multiple queries was explored recently by Gao et al. (2022) for retroactively attributing text generated by LMs to citations. Inspired by these, we include a fused_retrieval primitive to DSP to offer a versatile mechanism for interacting with frozen retrievers. It accepts an optional fusion function that maps multiple retrieval lists into one. By default, DSP uses a variant of CombSUM ( Fox & Shaw, 1994 ) , assigning each passage the sum of its probabilities across retrieval lists.

To illustrate, the modification below generates n = 10 n=10 queries for the transformation multihop_search_v2 .

⬇

#### Compositions & Extensions

To illustrate a simple composition, we can equip a chatbot with the capacity for conversational multi-hop search by combining a query rewriting step, which produces a query that encompasses all of the relevant conversational context, with the multi-hop transformation, as follows.

⬇

Similar approaches can be used for correcting spelling mistakes or implementing pseudo-relevance feedback ( Cao et al., 2008 ; Wang et al., 2022a ) , in which retrieved passages are used to inform a better search query, though this has not been attempted with pretrained LMs to our knowledge.

### 2.5 Predict

The Predict stage generates the system output using demonstrations (e.g., in x.demos ) and passages (e.g., in x.context ). Predict tackles the challenges of reliably solving the downstream task, which integrates much of the work on in-context learning in general. Within DSP , it also has the more specialized function of systematically aggregating information across a large number of demonstrations, passages, and candidate predictions.

#### Generating Candidates

Generally, Predict has to produce one or more candidate predictions for the end-task. To this end, the basic primitive in Predict is generate , which accepts a Template and (via currying) an Example and queries the LM to produce one or more completions, as explored earlier in § 2.4 . A corresponding primitive that uses the RM in this stage is rank , which accepts a query and one or more passages and returns their relevance scores.

⬇

A Template is an object that can produce prompts, that is, map an Example to a string, and extract fields out of completions. For instance, we can map an example x that has a question and retrieved passages to the following prompt:

⬇

As this illustrates, the LM will be asked to generate a chain-of-thought rationale (CoT; Wei et al. 2022 ; Kojima et al. 2022 ) and an answer, and the generated text will be extracted back into the rationale and answer keys of each completion.

Each invocation to the LM can sample multiple candidate predictions. Selecting a “best” prediction is the subject of much work on decoding ( Wiher et al., 2022 ; Li et al., 2022 ) , but a frozen and general-purpose LM may not support custom modifications to decoding. Within these constraints, we present several high-level strategies for selecting predictions and aggregating information in DSP via the LM and RM .

#### Selecting Predictions

Among multiple candidates, we can simply extract the most popular prediction. When a CoT is used to arrive at the answer, this is the self-consistency method of Wang et al. (2022c) , which seeks to identify predictions at which multiple distinct rationales arrive.

⬇

DSP generalizes this in two ways. First, we can sample multiple “pipelines of transformations” (PoT) within the program, rather than locally with “chains of thought” (CoT) in one transformation. These chains may even invoke different paths in the program, as illustrated below.

⬇

In the snippet above, Line 10 invokes the primitive branch which samples n n different PoTs with a high temperature (e.g., t = 0.7 t=0.7 ) and accumulates their intermediate and final predictions. In this example, our pipeline invokes multihop_search_v2 (§ 2.4 ), which applies a variable number of retrieval hops depending on the questions generated, before doing Predict . That is, PoT_program potentially invokes multiple distinct paths in the program (i.e., with different multi-hop queries and number of hops in each) across branches. It then selects the majority answer overall.

DSP generalizes self-consistency in a second way. When sampling our CoTs or PoTs provides multiple candidates, we can select the top- k k (e.g., top-4) predictions and then compare them directly. For instance, we may prompt the LM to compare these choices as MCQ candidates, a transformation for which Demonstrate can automatically prepare exemplars. This effectively simulates the LM recursion of Levine et al. (2022) , though unlike their approach it does not require a large training set or updating any (prompt-tuning) weights. One such implementation is illustrated in openqa_predict below.

⬇

As an alternative comparison approach, we can invoke the RM via rank to find the prediction that is most grounded in a retrieved contexts (i.e., most similar to the concatenation of the retrieved passages) or, given an RM that can score completions ( Krishna et al., 2022 ) , simply the prediction that has the highest score given the prompt.

#### Aggregating Information

When only a few demonstrations or passages are selected, we can simply concatenate them all into the prompt. For instance, GPT-3.5 text-davinci-002 has a context window of 4097 tokens, which we find to be reasonably large for accommodating several (e.g., 3–5) demonstrations, which individually include their own passages and rationales.