SlowGuess/NeurIPS_2023_OCR huggingface dataset:
ChatGPT reason: The news title is about an art exhibit, so Assistant B's response of 'Arts & Culture' is more accurate and specific than Assistant A's 'Culture'.
We also see more cases of runaway generations that do not occur with the positive control model:
Instruction: Provide an example of how a table of contents can be generated automatically in a LATEX document. Make sure your example contains the necessary commands. <G1> Assistant A (pos control): \tableofcontents \chapter{Introduction} \chapter{Theory} \chapter{Experiments} \chapter{Conclusion} Assistant B (LLaMA-7B gist): \tableofcontents \chapter{Introduction} \section{The Basics} \subsection{The Basics} \subsubsection{The Basics} \subsubsection{The Basics} \subsubsection{The Basics} [repeats for hundreds of tokens] \subsubsection{The Basics}
ChatGPT choice: A
ChatGPT reason: Assistant A provides a simple and correct example of generating a table of contents in LaTeX. Assistant B's example is unnecessarily long and does not follow the standard structure of a table of contents.
While it is unclear why only the gist models exhibit this behavior, these issues can likely be mitigated with more careful sampling techniques.
Table 3: Gist efficiency improvements. For different caching strategies (None, Instruction, Gist), we record CUDA wall time and GFLOPs ( $\pm$  std dev). Then we report the absolute/relative improvement of Gist Caching over these alternative strategies, with  $95\%$  confidence intervals in parentheses.
Model	Metric	Caching Strategy	Absolute/Relative Δ
None	Instructioni	Gistii	vs. None	vs. Instruction
LLaMA-7B	CUDA time (ms) ↓	23.4 ± 6.88	22.1 ± 6.58	21.8 ± 6.55	1.60 (1.29, 1.90)	.221 (.140, .302)
GFLOPs ↓	915 ± 936	553 ± 900	552 ± 899	6.8% (5.5, 8.1)	1.0% (.63, 1.4)
362 (337, 387)	.607 (.448, .766)
40% (37, 42)	.11% (.08, .14)
FLAN-T5-XXL	CUDA time (ms) ↓	31.0 ± 5.31	N/A	29.7 ± 5.07	1.30 (1.10, 1.51)	N/A
GFLOPs ↓	716 ± 717	N/A	427 ± 684	4.2% (3.6, 4.9)
289 (268, 310)	N/A
40% (37, 43)
iAverage KV Cache Length = 26	iiAverage KV Cache Length = 1
6 Compute, Memory, and Storage Efficiency
Finally, we return to one of the central motivations of this paper: what kind of efficiency gains does gisting enable? To answer this question, we compare the compute requirements (CUDA wall time and FLOPs) during inference with the single-token gist models using different caching strategies:
1. No caching: just encoding the full prompt  $t$ .
2. Instruction caching: caching the activations of the uncompressed instruction  $t$  (keys and values for all layers) into what is called the KV cache. This is the most common caching behavior for Transformer inference [4, 26] and is supported in libraries like Hugging Face Transformers [41]. However, it is only applicable to decoder-only models, since in models with bidirectional encoders like T5, the instruction representations  $t$  depend on the input  $x$ .
3. Gist caching: Compressing the prompt into the gist prefix  $G(t)$ .
Table 3 displays the results of profiling a single forward pass through the model (i.e. one step of decoding with a single input token) with PyTorch [24] 2.0, averaged across the 252 Human instructions. Gist caching improves significantly over unoptimized models, with  $40\%$  FLOPs savings and  $4 - 7\%$  lower wall time for both models. Note that at these (relatively) small scales, the wall time improvements are smaller than the FLOPs reductions because much of the inference latency is caused by moving tensors from high-bandwidth memory (HBM) to the chip compute cores, i.e. what Pope et al. [26] call the "memory time". Larger sequence lengths and batch sizes will lead to additional speedups, as the overall latency becomes dominated by the actual matrix computations.
For LLaMA-7B, the picture is more nuanced when compared to caching the full instruction. Compute improvements of gist caching are smaller: a negligible decrease in FLOPs  $(0.11\%)$  and a modest  $1\%$  speedup in wall time. This is because the FLOPs required for a Transformer forward pass is dominated by processing of the new input tokens, rather than self-attention with the KV cache. For example, a forward pass through LLaMA-7B with a single input token and a 2000-length KV cache is only  $\sim 10\%$  more expensive than the same forward pass with no KV cache—see Appendix F for more details. Nevertheless, this small decrease in FLOPs leads to a disproportionate decrease in wall time  $(1\%)$ , likely because the self-attention computations are slower relative to their FLOPs contribution.
At large scales and with heavily reused prompts, a  $1\%$  latency speedup can still accumulate into significant cost savings over time. More importantly, however, there are key benefits of gist caching over instruction caching besides latency: compressing 26 tokens into 1 gives more space in the input context window, which is bounded by absolute position embeddings or GPU VRAM. For example, for LLaMA-7B, each token in the KV cache requires 1.05 MB storage. $^3$  While the total contribution of the KV cache relative to the memory needed for LLaMA-7B inference is negligible at the prompt lengths we tested, an increasingly common scenario is developers caching many prompts across a large number of users, where storage costs quickly add up. In these scenarios, gisting allows caching of up to  $26\mathbf{x}$  more prompts than full instruction caching, using the same amount of storage!
34 (fp32 bytes)  $\times 2$  (keys+values)  $\times 32$  (num layers)  $\times 32$  (num attn heads)  $\times 128$  (head dim)  $= 1.05\mathrm{MB}$ .
7 Additional Related Work
Gisting builds upon past work in (parameter-efficient) instruction finetuning and context distillation, as discussed in Section 2. Here we will outline some additional connections to related work:
Adapting LMs without Backprop. As mentioned in Section 2, gisting can be viewed as a way to adapt LMs without gradient descent, by predicting the the prefix ( $\approx$  parameters) of the adapted model. Similarly, HyperTuning [25] predicts the prefix of a model for a task using (input, output) pairs for that task. If HyperTuning is a "few-shot" adaptation method, predicting the prefix from few-shot examples, then gisting can be seen as a "zero-shot" version, predicting the prefix from the language instruction alone. Gisting also has the additional benefit over HyperTuning of being conceptually simpler: instead of training a separate LM to predict the prefix, the LM itself is used as the HyperNetwork [13], with only a tiny change to the attention masks needed for prompt compression.
Compression and memory in transformers. The idea of "compressing" prompts is closely related to previous attempts at storing past representations to improve memory and long-range sequence modeling in Transformers [9, 21, 27, 42, 43]. In particular, the Compressive Transformer [27] compresses transformer activations into a smaller compressed memory using a learned convolutional operator. Gisting can be seen as a variant of the Compressive Transformer with 3 key differences. First, the compression function is not a separately learned function, but the LM's own self-attention mechanism, controlled by an input-dependent gist token. Second, the compression function is learned jointly with instruction finetuning via the standard language modeling loss, not a specialized auxiliary reconstruction loss as in [27]. Finally, our task of interest is not long-range sequence modeling, but caching and reusing instruction following prompts for efficiency reasons.
Sparse attention mechanisms. By restricting attention masks, gisting draws inspiration from efficient/sparse attention methods in Transformers (see [32] for review). For example, some sliding window attention mechanisms [2, 6] may remove the need to keep the entire KV cache around during decoding, but these more general methods are not optimized for caching arbitrary parts of the input sequence of varying length, which prompt compression demands. In light of this, gisting can be viewed as an input-dependent sparse attention mechanism specifically aimed at improving efficiency of the prompting workflow now commonly used in LMs.
In this paper we presented gisting, a framework for prompt compression in LMs, and a simple way of implementing gist models by modifying Transformer attention masks that incurs no additional cost over standard instruction finetuning. Gisting can be seen either as a modified form of instruction finetuning or a method for (meta-)context distillation of an LM. Gist models can compress unseen OOD prompts up to 26x while maintaining output quality, resulting in up to  $40\%$  FLOPs reduction and  $4.2\%$  wall clock speedups over unoptimized models, and enabling new methods for prompt caching in encoder-decoder models. While wall-time improvements for decoder-only LMs are smaller, gisting nevertheless enables caching 1 order of magnitude (26x) more prompts relative to full instructions.
Gisting is a promising method for improving LM efficiency, but carries some limitations. While gisting seems to succeed in capturing the "gist" of instructions (hence the name), achieving such compression necessarily results in some loss of nuance of the original instruction; Secton 5.1 illustrates a concrete failure case we observed. Since the behavior of LMs on edge cases is already not well understood, it is especially important for practitioners to carefully evaluate whether the compute/accuracy tradeoff of gisting is sufficiently safe and robust for their use case, before deployment.
Nevertheless, we believe gisting opens several interesting directions for future work. First, the masking method presented here can be easily integrated into existing instruction finetuning workflows, but another exciting approach is to retrofit an existing, frozen LM by training a smaller model to compress prompts, if finetuning the larger LM is inconvenient. Second, the largest efficiency gains from gisting will result from compressing very long prompts, for example  $k$ -shot prompts for large  $k$  that may even exceed a single context window. Finally, compression performance can likely be improved through "gist pretraining": first learning to compress arbitrary spans of natural language, before then learning prompt compression. Such objectives could be devised by inserting gist tokens into other pretraining objectives, perhaps during language modeling or T5's span corruption objective.
We thank the Stanford Alpaca team, especially Xuechen Li, with assistance with Alpaca data and finetuning, and Gabriel Poesia for the codebase used to collect human evaluations. Additional thanks to Xiuyu Li for help debugging and ensuring reproducibility of the open-source codebase. JM is supported by the Open Philanthropy AI Fellowship and XLL is supported by the Stanford Graduate Fellowship.
[1] A. Askell, Y. Bai, A. Chen, D. Drain, D. Ganguli, T. Henighan, A. Jones, N. Joseph, B. Mann, N. DasSarma, et al. A general language assistant as a laboratory for alignment. arXiv preprint arXiv:2112.00861, 2021.
[2] I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. arXiv:2004.05150, 2020.
[3] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
[4] C. Chen. Transformer inference arithmetic. https://kipp.ly/blog/ transformer-inference-arithmetic/, 2022.
[5] W.-L. Chiang, Z. Li, Z. Lin, Y. Sheng, Z. Wu, H. Zhang, L. Zheng, S. Zhuang, Y. Zhuang, J. E. Gonzalez, I. Stoica, and E. P. Xing. Vicuna: An open-source chatbot impressing GPT-4 with  $90\%$  ChatGPT quality. https://vicuna.lmsys.org/, 2023.
[6] R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
[7] E. Choi, Y. Jo, J. Jang, and M. Seo. Prompt injection: Parameterization of fixed inputs. arXiv preprint arXiv:2206.11349, 2022.
[8] H. W. Chung, L. Hou, S. Longpre, B. Zoph, Y. Tay, W. Fedus, E. Li, X. Wang, M. Dehghani, S. Brahma, et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416, 2022.
[9] Z. Dai, Z. Yang, Y. Yang, J. G. Carbonell, Q. Le, and R. Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978-2988, 2019.
[10] W. J. Dixon and F. J. Massey Jr. Introduction to statistical analysis. McGraw-Hill, 1951.
[11] X. Geng, A. Gudibande, H. Liu, E. Wallace, P. Abbeel, S. Levine, and D. Song. Koala: A dialogue model for academic research. Blog post, April 2023. URL https://bair.berkeley.edu/blog/2023/04/03/koala/.
[12] F. Gilardi, M. Alizadeh, and M. Kubli. ChatGPT outperforms crowd-workers for text-annotation tasks. arXiv preprint arXiv:2303.15056, 2023.
[13] D. Ha, A. Dai, and Q. V. Le. HyperNetworks. In International Conference on Learning Representations, 2017.
[14] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
[15] N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, Q. De Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly. Parameter-efficient transfer learning for NLP. In International Conference on Machine Learning, pages 2790-2799. PMLR, 2019.
[33] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
[34] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[35] J. Wang, Y. Liang, F. Meng, H. Shi, Z. Li, J. Xu, J. Qu, and J. Zhou. Is ChatGPT a good NLG evaluator? A preliminary study. arXiv preprint arXiv:2303.04048, 2023.
[36] Y. Wang, Y. Kordi, S. Mishra, A. Liu, N. A. Smith, D. Khashabi, and H. Hajishirzi. Self-Instruct: Aligning language model with self generated instructions. arXiv preprint arXiv:2212.10560, 2022.
[37] Y. Wang, S. Mishra, P. Alipoormolabashi, Y. Kordi, A. Mirzaei, A. Naik, A. Ashok, A. S. Dhanasekaran, A. Arunkumar, D. Stap, et al. Super-Natural Instructions: Generalization via declarative instructions on  $1600+$  NLP tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 5085-5109, 2022.
[38] J. Wei, M. Bosma, V. Zhao, K. Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q. V. Le. Finetuned language models are zero-shot learners. In International Conference on Learning Representations, 2022.
[39] J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. H. Chi, Q. V. Le, D. Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022.
[40] D. Wingate, M. Shoeybi, and T. Sorensen. Prompt compression and contrastive conditioning for controllability and toxicity reduction in language models. In *Findings of the Association for Computational Linguistics: EMNLP* 2022, pages 5621–5634, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics. URL https://aclanthology.org/2022-findings-emnlp.412.
[41] T. Wolf, L. Debut, V. Sanh, J. Chaumont, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y. Jernite, J. Plu, C. Xu, T. Le Scao, S. Gugger, M. Drame, Q. Lhoest, and A. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online, Oct. 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-demos.6. URL https://aclanthology.org/2020.emnlp-demos.6.
[42] Y. Wu, M. N. Rabe, D. Hutchins, and C. Szegedy. Memorizing transformers. In International Conference on Learning Representations, 2022.
[43] H. Zhang, Y. Gong, Y. Shen, W. Li, J. Lv, N. Duan, and W. Chen. Poolingformer: Long document modeling with pooling attention. In International Conference on Machine Learning, pages 12437-12446. PMLR, 2021.
A Example PyTorch Implementation of Gist Masking
See Listing A.1 for a sample annotated implementation of gist masking. This PyTorch implementation relies on basic NumPy-style tensor operations and can thus be adapted easily to a framework like JAX.
B Data, Training, Evaluation, and Compute Details
Code, data, and model checkpoints are available at https://github.com/jayelm/gisting.
Data. For LLaMA-7B, we used a maximum sequence length of 512 tokens during training and evaluation, except with the Human validation split, where the maximum length was increased to 768 (the Human instructions are longer). Examples longer than this length are truncated from the end. For FLAN-T5-XXL, we set a maximum input length (task  $t +$  input  $x$ ) of 128 and a maximum output length of 256, except again for the Human split, where the maximum input and output lengths were both set to 384. For both models, we set a maximum generation length of 512 tokens. These lengths were chosen such that  $< 1\%$  of examples across the board were truncated during training and evaluation for both models.
Training. Full hyperparameters for training runs are located in Table A.1. These parameters were adapted from previous published work finetuning LLAMA/FLAN-T5. For LLaMA-7B, parameters are identical to those used in training Alpaca Taori et al. [31]. For FLAN-T5-XXL, parameters are identical to those used in training Tk-INSTRUCT [37], except with a 5e-5 learning rate, as used in the Tk-INSTRUCT GitHub repository, rather than the 1e-5 learning rate in the paper.
LLaMA-7B was trained for 3000 steps, while FLAN-T5-XXL was trained for 16000 steps. Since there are about 130k examples in Alpaca+, given the batch sizes in Table A.1 this corresponds to about  $\sim 3$  epochs and  $\sim 2$  epochs of training, respectively. These numbers, again, are identical to Taori et al. [31] and Wang et al. [36]. We note that the training time is relatively flexible; for example, we did not see substantial gains training beyond 1 epoch for FLAN-T5-XXL.
Evaluation. During evaluation and benchmarking, we simply greedily decoded the most likely sequence. We saw limited gains from beam search with beam size  $B = 4$ .
Compute. Experiments were run on a cluster machine with 4xA100-SXM4-80GB NVIDIA GPUs, 480GB RAM, and 16 CPUs, using PyTorch 2.0 [24], Hugging Face Transformers [41], and DeepSpeed [29]. Training runs take about  $\sim 7$  hours to complete for LLaMA-7B and  $\sim 25$  hours for FLAN-T5-XXL. Benchmarking results were obtained on the same machine, but using just 1 of the A100 GPUs.
Table A.1: Hyperparameters for training runs.
LLaMA-7B	FLAN-T5-XXL
num steps	3000	16000
num train epochs	≈ 3	≈ 2
batch size	128	16
learning rate	2e-5	5e-5
warmup ratio	0.03	0
precision	bf16	bf16
optimizer	AdamW	AdamW
Deepspeed
# GPUs (A100 80GB)	4	4
ZeRO stage	3	3
subgroup size	1e9	1e9
max live params	1e9	1e9
max reuse distance	1e9	1e9
<sup>4</sup>https://github.com/yizhongw/Tk-Instruct/blob/1ab6fad/scripts/train_tk_instruct.sh
Listing A.1: Sample PyTorch implementation of gist masking.
import torch   
def reverse_cumsum(x: torch.Tensor) -> torch.Tensor: ""Cumulative sum from right to left. See https://github.com/pytorch/pytorch/issues/33520. return x + torch.sum(x, dim=-1, keepdims=True) - torch.cumsum(x, dim=-1)   
def make_mask_pre_first_gist(input: torch.Tensor, gist_token: int, dtype=torch.int64) -> torch.Tensor: ""Returns a mask where all tokens prior to the first gist token are masked out. Arguments: inputs: a Tensor of input tokens where the last dimension is the sequence length. gist_token: the integer id of the gist token. dtype: the dtype of the mask, default int64. Returns: The requested mask. return ((inputs == gist_token).cumsum(-1) >= 1).typedtype)   
def make_mask_post_last_gist(input: torch.Tensor, gist_token: int, dtype=torch.int64) -> torch.Tensor: ""Returns a mask where all tokens after the last gist token are masked out. Computing the same as mask_pre_first_gist_token, but reverses the sequence before and after the cumsum. Arguments: inputs: a Tensor of input tokens where the last dimension is the sequence length. gist_token: the integer id of the gist token. dtype: the dtype of the mask, default int64. Returns: The requested mask. return (reverse_cumsum(input == gist_token) >= 1).typedtype)   
def make_gist_mask(input: torch.Tensor, gist_token: int, pad_token: int, dtype=torch.int64) -> torch.Tensor: ""Creates a gist mask from supplied inputs and gist/pad tokens. Tokens after the last gist cannot attend to tokens prior to the first gist. Additionally, tokens *before* the last gist cannot attend to tokens *after* the last gist. The gist mask is broadcasted to 4D (with a singleton dim 1) for compatibility with multi-headed attention (where dim 1 is the head dimension). Arguments: inputs: a Tensor of shape (batch_size, seq_len) input tokens. gist_token: the integer id of the gist token. pad_token: the integer id of the pad token. inputs == pad_token are masked out. dtype: the dtype of the mask, default int64. Returns: The requested mask of shape (batch_size, 1, seq_len, seq_len)   
# Attention mask for tokens before the last gist token. pre_gist_mask = make_mask_post_last_gist(inputs, gist_token, dtype=torch(bool)[:, None, None] # Attention mask for tokens after the last gist token. post_gist_mask = make_mask_pre_first_gist(inputs, gist_token, dtype=torch(bool)[:, None, None] # Construct time masks by permuting to time dimension. pre_gist_time_mask = pre_gist_mask.permute((0, 1, 3, 2))   
mask = torch.where(pre_gist_time_mask, pre_gist_mask, post_gist_mask) mask = mask & (inputs != pad_token)[:, None, None] # Mask out pad tokens. return mask.typedtype)
1   
2   
3   
4   
5   
6   
7   
8   
9   
10   
11   
12   
13   
14   
15   
16   
17   
18   
19   
20   
21   
22   
23   
24   
25   
26   
27   
28   
29   
30   
31   
32   
33   
34   
35   
36   
37   
38   
39   
40   
41   
42   
43   
44   
45   
46   
47   
48   
49   
50   
51   
52   
53   
54   
55   
56   
57   
58   
59   
60   
61   
62   
63   
64   
65   
66   
67
C Details of ChatGPT Evaluation
Given a user instruction and two AI assistant responses, your job is to select the response that is more helpful, and give reasons why. Judge responses holistically, paying special attention to whether the response (1) correctly follows the provided instruction and (2) is factually accurate. If both responses are equal, you may call a tie, but please use this response sparingly.
Example 1:  
...  
Instruction: Given the name of a city, tell me the name of the state it is in.  
Input: Los Angeles  
Assistant A: California  
Assistant B: Wyoming  
...
Your output should be:  
{“reason”: "Los Angeles is in California. Only Assistant A is correct.", "choice”: "A"}
Instruction: Give me some examples of fun things to do with the kids on a weekend. Assistant A: For example, what should I do with my kids on a weekend? What if I'm in a city? What if I'm on the beach?
Assistant B: You could try going to the zoo or a museum. If you want to stay at home, you could play board games or bake cookies.
Your output should be:
reason: Assistant A doesn't seem to be following the instruction. Assistant B provides helpful information with examples.
choice: B
Example 3:  
```python
Instruction: Write me a python function that prints "Hello, World". Assistant A: def hello_world(): print("Hello!") Assistant B: "Hello, World"
reason: Assistant B just wrote 'Hello, World'. Assistant A actually wrote a Python function, even if it doesn't exactly print the right thing, so overall Assistant A is better.
choice: A
Instruction: Translate the following sentence from English to French.  
Input: I like cats.
Assistant A: Me gustan los gatos.
Assistant B: 我喜欢猫.
reason: Both assistants got the language wrong.
choice: tie