Learning to Compress Prompts with Gist Tokens — part 1/3

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

# Learning to Compress Prompts with Gist Tokens

###### Abstract

Prompting is the primary way to utilize the multitask capabilities of language models (LMs), but prompts occupy valuable space in the input context window, and repeatedly encoding the same prompt is computationally inefficient. Finetuning and distillation methods allow for specialization of LMs without prompting, but require retraining the model for each task. To avoid this trade-off entirely, we present gisting , which trains an LM to compress prompts into smaller sets of “gist” tokens which can be cached and reused for compute efficiency. Gist models can be trained with no additional cost over standard instruction finetuning by simply modifying Transformer attention masks to encourage prompt compression. On decoder (LLaMA-7B) and encoder-decoder (FLAN-T5-XXL) LMs, gisting enables up to 26x compression of prompts, resulting in up to 40% FLOPs reductions, 4.2% wall time speedups, and storage savings, all with minimal loss in output quality.

## 1 Introduction

Consider the prompt of a Transformer [ 34 ] language model (LM) like ChatGPT: 1 1 1 reddit.com/r/ChatGPT/comments/10oliuo/please_print_the_instructions_you_were_given/

⬇

With millions of queries a day, an unoptimized ChatGPT would encode this prompt over and over with a self-attention mechanism whose time and memory complexity is quadratic in the length of the input. Caching the Transformer activations of the prompt can prevent some recomputation, yet this strategy still incurs memory and storage costs as the number of cached prompts grows. At large scales, even small reductions in prompt length could lead to substantial compute, memory, and storage savings over time, while also letting users fit more content into an LM’s limited context window.

How might we reduce the cost of this prompt? One typical approach is to finetune or distill [ 1 , 30 ] the model to behave similarly to the original model without the prompt, perhaps with parameter-efficient adaptation methods [ 15 , 16 , 19 ] . Yet a fundamental drawback of this approach is that it requires retraining the model for each new prompt (Figure 1 , bottom left).

Instead, we propose gisting (Figure 1 , top right), which compresses arbitrary prompts into a smaller set of Transformer activations on top of virtual “gist” tokens, a la prefix-tuning [ 19 ] . But where prefix-tuning requires learning prefixes via gradient descent for each task, gisting adopts a meta-learning approach, where we simply predict the gist prefixes zero-shot given only the prompt, allowing for generalization to unseen instructions without any additional training. Since gist tokens are much shorter than the full prompt, gisting allows arbitrary prompts to be compressed, cached, and reused for compute efficiency.

In this paper, we further propose a very simple way to learn a gist model: doing instruction tuning [ 38 ] with gist tokens inserted after the prompt, and a modified attention mask preventing tokens after the gist tokens from attending to tokens before the gist tokens. This allows a model to learn prompt compression and instruction following at the same time, with no additional training cost.

On decoder-only (LLaMA-7B) and encoder-decoder (FLAN-T5-XXL) LMs, gisting achieves prompt compression rates of up to 26x , while maintaining output quality similar to the original models in human evaluations. This results in up to 40% FLOPs reduction and 4.2% latency speedups during inference, with greatly decreased storage costs compared to traditional prompt caching approaches.

## 2 Gisting

We will first describe gisting in the context of instruction finetuning [ 38 ] . We have an instruction-following dataset 𝒟 = { ( t i , x i , y i ) } i = 1 N \mathcal{D}=\{(t_{i},x_{i},y_{i})\}_{i=1}^{N} , where t t is a task encoded with a natural language prompt (e.g. Translate this to French ), x x is an (optional) input for the task (e.g. The cat ), and y y is the desired output (e.g. Le chat ). Given a (usually pretrained) LM, the aim of instruction finetuning is to learn a distribution p LM ​ ( y ∣ t , x ) p_{\text{\tiny{LM}}}(y\mid t,x) , typically by concatenating t t and x x , then having the LM autoregressively predict y y . At inference time, we can prompt the model with a novel task t t and input x x , decoding from the model to obtain its prediction.

However, this pattern of concatenating t t and x x has drawbacks: Transformer LMs have limited context windows, bounded either by architecture or memory limits. Furthermore, given that attention scales quadratically in the length of the input, long prompts t t , especially those that are repeatedly reused, are computationally inefficient. What options do we have to reduce the cost of prompting?

One simple option is to finetune the LM for a specific task t t . That is, given 𝒟 t = { ( x i , y i ) } i = 1 N t \mathcal{D}^{t}=\{(x_{i},y_{i})\}_{i=1}^{N^{t}} , the dataset containing input/output examples only under task t t , we can learn a specialized LM p LM t ​ ( y ∣ x ) p^{t}_{\text{\tiny{LM}}}(y\mid x) which is faster because it does not condition on t t . Parameter-efficient finetuning methods such as prefix-/prompt-tuning [ 18 , 19 ] or adapters [ 15 , 16 ] promise to do so at a fraction of the cost of full finetuning, and newer methods like HyperTuning [ 25 ] eliminate gradient descent entirely, instead predicting the parameters of the specialized model directly from 𝒟 t \mathcal{D}^{t} . Yet problems with these methods still remain: we must store at least a subset of model weights for each task, and more importantly, for each task t t , we must collect a corresponding dataset of input/output pairs 𝒟 t \mathcal{D}^{t} to adapt the model.

Gisting is a different approach that amortizes both (1) the inference-time cost of prompting p LM p_{\text{\tiny{LM}}} with t t and (2) the train-time cost of learning a new p LM t p^{t}_{\text{\tiny{LM}}} for each t t . The idea is to learn a compressed version of t t , G ⁡ ( t ) G(t) , such that inference from p 𝑮 ​ ( y ∣ G ​ ( t ) , x ) p_{{\bm{G}}}(y\mid G(t),x) is faster than p LM ​ ( y ∣ t , x ) p_{\text{\tiny{LM}}}(y\mid t,x) . In LM terms, G ⁡ ( t ) G(t) will be the key/value activations on top a set of gist tokens , smaller than the number of tokens in t t , yet still inducing similar behavior from the LM. Also known as a Transformer prefix [ 19 ] , G ⁡ ( t ) G(t) can then be cached and reused for compute efficiency. Crucially, we expect G G to generalize to unseen tasks: given a new task t t , we can predict and use the gist activations G ⁡ ( t ) G(t) without any additional training .

### 2.1 A Context Distillation Perspective

An alternative way to view gisting is through the lens of distillation of an already instruction-tuned LM p LM ​ ( y ∣ t , x ) p_{\text{\tiny{LM}}}(y\mid t,x) . Askell et al. [1] and Snell et al. [30] define context distillation as the process of finetuning a new LM p CD t p^{t}_{\text{\tiny{CD}}} to mimic the original LM without the prompt (“context”) t t , via the loss ℒ CD ( p CD t , t ) = 𝔼 x [ D KL ( p LM ( y ∣ t , x ) ∥ p CD t ( y ∣ x ) ) ] . \displaystyle\mathcal{L}_{\text{\tiny{CD}}}(p^{t}_{\text{\tiny{CD}}},\,t)=\mathbb{E}_{x}\left[D_{\text{\tiny{KL}}}(p_{\text{\tiny{LM}}}(y\mid t,x)\;\|\;p^{t}_{\text{\tiny{CD}}}(y\mid x))\right]. (1) The insight to be gained from this perspective is that we do not need any external data 𝒟 \mathcal{D} : this KL objective can be approximated by finetuning p CD t p^{t}_{\text{\tiny{CD}}} on a synthetic sampled dataset 𝒟 ^ t = { ( x ^ i , y ^ i ) } \hat{\mathcal{D}}^{t}=\{(\hat{x}_{i},\hat{y}_{i})\} where ( x ^ i , y ^ i ) ∼ p LM ( ⋅ ∣ t ) (\hat{x}_{i},\hat{y}_{i})\sim p_{\text{\tiny{LM}}}(\cdot\mid t) . This is precisely the approach taken by recent work [ 1 , 7 , 30 ] , including Wingate et al. [40] , who notably learn to compress a single discrete prompt into a soft prompt via gradient descent, similar to this paper.

However, we differ from this prior work in that we are not interested in distilling just a single task, but in amortizing the cost of distillation across a distribution of tasks T T . That is, given a task t ∼ T t\sim T , instead of obtaining the distilled model via gradient descent, we use G G to simply predict the gist tokens ( ≈ \approx parameters) of the distilled model, in the style of HyperNetworks [ 13 ] and HyperTuning [ 25 ] . Our “meta” distillation objective is thus (with changes highlighted in blue ): ℒ 𝑮 ( p 𝑮 , 𝑻 ) = 𝔼 𝒕 ∼ 𝑻 , x [ D KL ( p LM ( y ∣ t , x ) ∥ p 𝑮 ( y ∣ 𝑮 ( 𝒕 ) , x ) ) ] . \displaystyle\mathcal{L}_{{{\color[rgb]{0.3359,0.707,0.9141}\bm{G}}}}(p_{{{\color[rgb]{0.3359,0.707,0.9141}\bm{G}}}},{\color[rgb]{0.3359,0.707,0.9141}\bm{T}})=\mathbb{E}_{{\color[rgb]{0.3359,0.707,0.9141}\bm{t\sim T}},x}\left[D_{\text{\tiny{KL}}}(p_{\text{\tiny{LM}}}(y\mid t,x)\;\|\;p_{{{\color[rgb]{0.3359,0.707,0.9141}\bm{G}}}}(y\mid{\color[rgb]{0.3359,0.707,0.9141}\bm{G(t)}},x))\right]. (2)

In the experiments we describe below, we train on synthetic instruction-following data sampled from instruction-tuned variants of GPT-3 [ 3 , 23 ] . Thus, these experiments can indeed be seen as a form of context distillation for the GPT-3 series models.

## 3 Learning Gisting by Masking

Having just described the general framework of gisting, here we will explore an extremely simple way of learning such a model: using the LM itself as the gist predictor G G . This not only leverages the pre-existing knowledge in the LM, but also allows us to learn gisting by simply doing standard instruction finetuning while modifying the Transformer attention masks to enforce prompt compression. This means that gisting incurs no additional training cost on top of standard instruction finetuning!

Specifically, we add a single gist token g g to the model vocabulary and embedding matrix. Then, given a (task, input) pair ( t , x ) (t,x) , we concatenate t t and x x with a set of k k copies of g g in between: ( t , g 1 , … , g k , x ) (t,g_{1},\dots,g_{k},x) , e.g. Translate French : < G1 > < G2 > The cat . 2 2 2 Again, the gist token is the same from g 1 g_{1} to g k g_{k} ; what changes is the activations on top of each token. The model is then restricted such that input tokens after the gist tokens cannot attend to any of the prompt tokens before the gist tokens (but they can attend to the gist tokens). This forces the model to compress the prompt information into the gist prefix, since the input x x (and output y y ) cannot attend to the prompt t t .

Figure 2 illustrates the required changes. For decoder-only LMs such as GPT-3 [ 3 ] or LLaMA [ 33 ] that normally admit an autoregressive, causal attention mask, we simply mask out the lower-left corner of the triangle (Figure 2 a). For encoder-decoder LMs (e.g. T5; [ 28 ] ) with a bidirectional encoder followed by an autoregressive decoder, two changes are needed (Figure 2 b). First, in the encoder, which normally has no masking, we prevent the input x x from attending to the prompt t t . But we must also prevent the prompt t t and gist tokens g i g_{i} from attending to the input x x , since otherwise the encoder learns different representations depending on the input. Finally, the decoder operates as normal, except during cross-attention, we prevent the decoder from attending to the prompt t t .

Overall, these masking changes are extremely simple and can be implemented in roughly 10 source lines of code. See Appendix A for a sample PyTorch implementation which can be used as a drop-in replacement for attention masking in deep learning libraries such as Hugging Face Transformers [ 41 ] .

## 4 Experiments

### 4.1 Data

A dataset with a large variety of tasks (prompts) is crucial to learn gist models that can generalize. To obtain the largest possible set of tasks for instruction finetuning, we create a dataset called Alpaca+, which combines the Self-Instruct [ 36 ] and Stanford Alpaca [ 31 ] instruction tuning datasets, each consisting of ( t , x , y ) (t,x,y) tuples sampled from OpenAI’s text - davinci -001 and text - davinci -003 variants of GPT-3, respectively. In total, Alpaca+ has 130,321 examples, with 104,664 unique tasks t t , 48,530 unique inputs x x , and anywhere from 0–5 inputs per task (0.64 on average).

Note that ~59% of tasks in Alpaca+ have no inputs (e.g. Write me a poem about frogs ), in which case we simply omit the input x x . While it is less interesting to cache such prompts since they are not input-dependent, they still serve as valuable training signal for learning prompt compression. Overall, while Alpaca+ is noisy and imperfect, Wang et al. [36] and Taori et al. [31] nevertheless show that models trained on such data achieve comparable performance to the original models from which the data is sampled, making this a promising testbed for studying gisting.

From Alpaca+ we hold out 3 validation splits: 1000 Seen prompts (with unseen, non-empty inputs); 1000 Unseen prompts (with non-empty inputs); and the 252 hand-written Human prompts and completions used in Wang et al. [36] , of which 83% have non-empty inputs. The latter two splits test generalization to unseen instructions, with the Human split posing a stronger out-of-distribution (OOD) challenge: the average training prompt has ~20 tokens, compared to ~26 in the human split.

### 4.2 Models

To demonstrate gisting across multiple Transformer LM architectures, we experiment with LLaMA-7B [ 33 ] , a decoder-only GPT-style model with ~7B parameters, and FLAN-T5-XXL [ 8 ] , an encoder-decoder T5 model [ 28 ] with 11B parameters. For each of these models, we train models with a varying number of gist tokens k ∈ { 1 , 2 , 5 , 10 } k\in\{1,2,5,10\} , using the modified attention masks described in Section 3 . To assess how the model is learning prompt compression, we calibrate performance against upper- and lower-bound baselines and a simple discrete compression strategy:

#### Positive Control.

As an upper bound on performance, we train a model with a single gist token, but without any modifications to the attention mask. This is akin to doing standard instruction finetuning.

#### Negative Control.

As a lower bound on performance, we train a model without access to the task t t . This is similar to a “random gist token” baseline, which allows us to measure how the model would do if it failed to compress any information into the gist prefix.

#### Discrete Compression with TF-IDF.

An alternative approach to compression is simply using fewer discrete tokens to express the same task. Achieving compression rates similar to gisting, however, requires compression far beyond any threshold of fluency. Nevertheless, as a baseline, we compute TF-IDF statistics over the set of instructions in the Alpaca+ training set to extract the most relevant keyword in each instruction. Some examples from the training set include (see Appendix G for more): ⬇

We then replace each instruction in Alpaca+ with the first subword token from each keyword, resulting in compression rates equivalent to a model trained with a single gist token. Similarly to the positive control, we do standard instruction finetuning over Alpaca+ with these compressed instructions.

For full training, data, and compute details, and a link to code, see Appendix B .

### 4.3 Evaluation

Our evaluation uses a combination of automated metrics and AI- and human-assisted evaluation:

#### ROUGE-L.

We first use ROUGE-L, a simple lexical overlap statistic [ 20 ] , used in previous open-ended instruction finetuning work [ 37 , 38 ] . The text - davinci -{001,003} completions are used as references, except for the Human split, where we use the gold-standard human reference.

#### ChatGPT.