PaperlayerView exact revision on arXiv

Source-linked AI summary

Adapting Language Models to Compress Contexts

Alexis Chevalier, Alexander Wettig, Anirudh Ajith, Danqi Chen

arXiv:2305.14788v2cs.CL

TL;DR

AutoCompressors address finite context windows and the cost of processing long documents by adapting pre-trained language models to compress text into reusable summary vectors. They improve long-context language modeling and support efficient in-context learning, retrieval, and passage re-ranking, while retaining less information than full attention in some settings.

  • Problem

    Finite context windows and the high computational cost of processing long documents limit language models' efficiency and scalability.

  • Method

    AutoCompressors recursively produce summary vectors from document segments with an unsupervised objective, using the vectors as soft prompts for subsequent language modeling.

  • Results

    Across long-context modeling, in-context learning, retrieval, and passage re-ranking, summary vectors improve performance or efficiency; retrieval modeling achieves 1.5× the perplexity gains of plain-text passages at equal sequence lengths.

  • Takeaways & Limitations

    Summary vectors can be pre-computed, cached, and reused to reduce attention-window size and inference costs across long-context and retrieval settings.

  • Takeaways & Limitations

    The study evaluates only OPT models up to 2.7B parameters and one 7B Llama model, leaving performance at larger scales unestablished.

Abstract

from arXiv · show

Transformer-based language models (LMs) are powerful and widely-applicable tools, but their usefulness is constrained by a finite context window and the expensive computational cost of processing long text documents. We propose to adapt pre-trained LMs into AutoCompressors. These language models are capable of compressing long contexts into compact summary vectors, which are then accessible to the model as soft prompts. Summary vectors are trained with an unsupervised objective, whereby long documents are processed in segments, and summary vectors from all previous segments are used in language modeling. We fine-tune OPT and Llama-2 models on sequences of up to 30,720 tokens and show that AutoCompressors can utilize long contexts to improve perplexity. We evaluate AutoCompressors on in-context learning by compressing task demonstrations and find that summary vectors are good substitutes for plain-text demonstrations, increasing accuracy while reducing inference costs. Finally, we explore the benefits of pre-computing summary vectors for large corpora by applying summary vectors to retrievalaugmented language modeling and a passage re-ranking task. Overall, AutoCompressors emerge as a simple and inexpensive solution to extend the context window of LMs while speeding up inference over long contexts.

1 Introduction

AutoCompressors adapt pretrained language models to compress long documents into summary vectors that extend context handling and reduce inference costs. The paper develops this approach and evaluates it for long-context modeling, in-context learning, and retrieval applications.

  • 1 Introduction: AutoCompressors compress text into short summary vectors that serve as soft prompts for processing later document segments.The vectors are produced from language-model output states and are one or two orders of magnitude shorter than the original text.
  • 1 Introduction: Summary accumulation concatenates vectors from all segments to represent the entire document, while randomized segmenting supports variable-length contexts.The authors report that these innovations improve long-range information retention and enable reasoning over multiple passages.
  • 1 Introduction: AutoCompressors are fine-tuned from OPT-2.7B and Llama-2-7B on sequences ranging from 6,144 to 30,720 tokens using a single 80GB NVIDIA A100 GPU.The models improve perplexity over long documents and show robustness to domain generalization.
  • 1 Introduction: On 11 classification tasks, summary vectors outperform few-shot in-context learning with comparable in-context tokens on 8 out of 11 tasks.The evaluation includes 7 SuperGLUE tasks and compresses up to 90 demonstrations.
  • 1 Introduction: For retrieval applications, summary vectors achieve 1.5× the perplexity gains of plain-text passages at equal sequence lengths and the best re-ranking performance-throughput trade-off.The approach pre-computes summary vectors for large corpora and applies them to retrieval-augmented language modeling and zero-shot passage re-ranking.

2 Related Work

Prior work adapts language models through soft prompts, context distillation, and efficient long-range architectures. AutoCompressors differ by learning to generate reusable summaries for arbitrary contexts while retaining pretrained initialization.

  • 2 Related Work: Soft prompt tuning adapts pretrained Transformers by optimizing newly initialized embeddings while freezing the model’s existing parameters.Prompt compression methods learn a soft prompt for each new context and align predictions conditioned on that prompt with predictions conditioned on the original context.
  • 2 Related Work: AutoCompressors learn to predict their own soft prompts from context, avoiding optimization separately for every new context.This contrasts with prior prompt compression that provides no knowledge transfer between similar contexts.
  • 2 Related Work: Context distillation compresses information such as instructions, while AutoCompressors target any context information, including long documents.A concurrent approach compresses instructions into short key-value attention prefixes.
  • 2 Related Work: Efficient long-range architectures use approaches including sparse or approximate attention, recurrence, conditional computation, and retrieval of previous tokens.The cited approaches include restricting attention, approximating attention, introducing recurrent elements, and retrieving prior tokens at the output layer.
  • 2 Related Work: Many existing architectures require expensive training from scratch or substantially deviate from pretrained initialization, whereas AutoCompressors extend pretrained models to longer sequences.The paper notes that many language models also lack the inductive bias to extrapolate to longer sequences.

3 Method

AutoCompressors adapt pre-trained language models to compress long documents into summary vectors that future segments use for prediction. The method combines accumulated summaries, variable-length training segments, and an unsupervised document-level objective to retain information across long contexts.

  • Summary vectors: AutoCompressors append special summary tokens whose output vectors become soft prompts for subsequent text segments.The summary vectors are produced from preceding context and passed to the next segment as a compact prompt.
  • Summary accumulation: Summary accumulation concatenates summaries from all preceding segments, creating a direct information pathway whose length grows linearly with document length.For segment S_i, the accumulated summary σ<i concatenates σ1 through σi−1 and has length (i−1)κ.
  • Positional embeddings: For OPT-style absolute positional embeddings, summary tokens and vectors receive no positional embeddings, preserving pre-trained context positions and allowing arbitrary compression steps during training.Their order remains encoded through separate summary-token embeddings.
  • Training objective: The model is trained by minimizing cross-entropy over every document token while conditioning predictions on the summaries of preceding segments.This objective encourages summary vectors to store information that future segments can use for token prediction.
  • Randomized segmenting: Randomly varying segment lengths during training improves performance when evaluation uses fixed-length segments.Each sampled segment must still fit within the model’s context window.
  • Efficient training: Backpropagation through time, gradient checkpointing, cached summaries, and gradient stopping after two compression steps reduce the computational graph during training.The stopping rule assumes adjacent-token prediction is sufficient for learning useful compression.

4 Language Modeling Evaluation

AutoCompressors use accumulated summary vectors to retain information from long contexts while reducing the tokens and memory needed for processing. Across OPT and Llama-2 evaluations, they improve long-context perplexity, though longer-context Llama-2 comparisons still favor the base and full-attention models.

  • 8K-token sequences: AutoCompressors consistently outperform RMT on 6,144-token contexts and can benefit from shorter sequences than seen during training.The comparison includes OPT-2.7B models evaluated on long-context language modeling.
  • Long-context training: AutoCompressors trained on up to 30,720-token sequences use summary accumulation to condition later segments on compressed earlier context.OPT models use 20 compression steps for 30,720-token training sequences, while Llama-2 uses 6,144-token sequences with 50 summary tokens.
  • 8K-token sequences: 150 summary vectors process 6,144 context tokens with performance similar to or better than extended full attention using up to 2,048 additional plain-text tokens.AutoCompressors use 50 × 3 = 150 summary vectors, whereas extended full attention cannot extend beyond 4,096 tokens in the reported setting.
  • 30K-token sequences: Both OPT AutoCompressor models use the entire 28K-token context to reduce perplexity, whereas RMT does not benefit from increasing context from 14K to 28K.Stopping gradients reduces CUDA memory and enables fine-tuning the OPT-2.7B AutoCompressor on one 80GB NVIDIA A100 GPU.
  • Scaling to Llama-2: For Llama-2, compressing 4,096 tokens into 100 summary vectors matches extended full attention with 512 plain-text tokens, while 6,144 tokens and 150 vectors improve perplexity slightly further.Summary vectors also preserve perplexity when short contexts are compressed.
  • Scaling to Llama-2: Llama-2 and extended full attention outperform the AutoCompressor when longer contexts are provided, motivating further work on preserving all context information.This limitation is reported for the longer-context Llama-2 evaluation.

5 Compressing Demonstrations for In-Context Learning

AutoCompressors compress in-context demonstrations into summary vectors and evaluate them across 11 classification and multiple-choice tasks. Summary vectors outperform comparable plain-text demonstrations on most tasks while reducing inference cost, although fine-tuning can hurt some Llama-2 performance relative to initialization.

  • Evaluation: Summary vectors compress 1–3 demonstration segments into 50, 100, or 150 vectors for evaluating in-context learning across 11 datasets.The evaluation compares compressed demonstrations with plain-text prompts containing 150 or 750 tokens’ worth of demonstrations.
  • Results: Summary vectors outperform 150-token plain-text demonstrations on 8/11 tasks and outperform 750-token demonstrations on eight named tasks.The latter tasks are AG News, SST-2, BoolQ, WiC, WSC, CB, COPA, and MultiRC.
  • Results: Summary vectors increase accuracy while reducing inference cost, making them a strong alternative to plain-text demonstrations.This conclusion summarizes the reported in-context learning comparison.
  • Comparison with RMT: The OPT-2.7B AutoCompressor achieves higher accuracy than the RMT baseline on 8 out of 11 tasks, while RMT does not benefit from multiple compression steps.The authors attribute the AutoCompressor result to summary accumulation.
  • Limitations: Fine-tuned Llama-2 AutoCompressors have substantially worse zero-shot accuracy on some tasks and slightly worse ICL performance than the Llama-2 initialization.The authors suspect domain mismatch between fine-tuning data and the Llama-2 pre-training corpus.

6 Compressing Retrieval Corpora for Efficient Inference

AutoCompressors pre-compute compact summary vectors for retrieved passages, enabling multiple documents to be fused for retrieval-augmented inference with reduced computational and storage costs. In language modeling and passage re-ranking, fused summaries provide favorable performance–throughput trade-offs.

  • Retrieval-augmented language modeling: Pre-computed summary vectors can be stored cheaply and fused at retrieval time, allowing multiple documents to be accessed in a single forward pass.This shifts work from inference to corpus preprocessing and reduces the attention-window size.
  • Retrieval-augmented language modeling: Fused Summaries outperforms Fused Passages and REPLUG when retrieving 50-token passages.The comparison uses perplexity gains over a no-retrieval baseline for the OPT-2.7B AutoCompressor.
  • Retrieval-augmented language modeling: Fused Summaries top-10 outperforms REPLUG top-2 with 512-token passages while also gaining a 1.7× throughput increase.Throughput is measured without batching on a single NVIDIA A100 GPU.
  • Retrieval-augmented language modeling: For 10 retrieved documents, Fused Summaries remains inexpensive and stores compressed 10B-token datasets in 5TB per domain using half precision.The authors characterize this as a trade-off between storage costs and throughput.
  • Passage re-ranking: Fused Summaries benefits from fusing summary vectors from up to 10 documents, although REPLUG top-10 still outperforms it.The model was trained to accumulate only 3 sets of summary vectors, and higher-quality summaries remain future work.
  • Passage re-ranking: In passage re-ranking on NQ, pre-computed summary vectors substantially speed up inference despite the passages being only 100 words long.The evaluation measures Recall@20 and unbatched throughput on a single NVIDIA A100 80GB GPU.

7 Conclusion

The paper adapts pre-trained language models into AutoCompressors that recursively compress contexts into summary vectors. Experiments show that these vectors retain contextual information, encode demonstrations, support retrieval, and can be cached for efficiency, while larger-scale models and better summary quality remain open directions.

  • Conclusion: AutoCompressors recursively compress contexts into summary vectors that retain important contextual information.The models are adapted from pre-trained language models through a training strategy introduced in the paper.
  • Conclusion: Summary vectors can encode in-context demonstrations and be used in retrieval settings.These capabilities extend the uses of compressed context beyond long-document language modeling.
  • Conclusion: Pre-computed summary vectors can be cached and re-used, reducing the size of the attention window and offering practical efficiency gains.The conclusion frames caching as a practical benefit of the approach.
  • Conclusion: Future work includes scaling AutoCompressors to bigger models and improving summary-vector quality to further close the gap with full attention over long-range contexts.The authors identify both model scale and compression quality as remaining challenges.

Limitations

The paper identifies limits in model scale, information retention, and computational scaling for AutoCompressors. Summary accumulation remains quadratic in the number of segments, although at a lower rate than full attention.

  • Model scale: Experiments apply AutoCompressors only to OPT models up to 2.7B parameters and one Llama model of 7B parameters.The authors state that performance for larger models remains to be established.
  • Information retention: Summary vectors ignore some useful information accessible through full attention, and adding more summary vectors does not always improve performance.The authors suspect that the training signal may be limited by strong predictions from plaintext tokens in the current segment.
  • Computational complexity: Summary accumulation still has quadratic complexity as the number of segments increases, despite growing more slowly than full attention.The authors suggest more efficient ways to combine many summary vectors as future work.

A.1 OPT Experiments on 8K Tokens

The experiments fine-tune OPT and Llama-2 AutoCompressors on long sequences from large text corpora and evaluate language modeling across in-domain and out-of-domain documents. The reported results show long-context benefits and near-baseline performance without additional summary context.

  • Training setup: OPT models are trained on documents up to 30,720 tokens and evaluated on 30,720-token documents from Books3 and Gutenberg.Books3 is treated as in-domain and Gutenberg as out-of-domain.
  • Llama-2 experiments: Llama-2-7B AutoCompressors and the extended full-attention baseline are fine-tuned on 6,144-token sequences using LoRA on attention heads.The AutoCompressor also optimizes newly initialized summary-token embeddings.
  • Evaluation: Evaluation covers 8,192-token sequences from ArXiv, Books, C4, GitHub, StackExchange, and CommonCrawl.These domains include both specialized and broad web text.
  • Results: Without summary vectors or additional context, AutoCompressor perplexity increases by less than 1% relative to the OPT-2.7B fine-tuned baseline.It also performs slightly better than RMT and significantly better than the extended full-attention model in this setting.

C AutoCompressor Ablations

The ablation varies the number of summary tokens κ for OPT AutoCompressors and evaluates held-out perplexity across domains. Among the tested settings, κ = 50 performs best overall.

  • Ablation setup: The ablation compares OPT AutoCompressors trained with different numbers of summary tokens.
  • Ablation setup: κ = 20, 50, 70, and 100 are evaluated on 8,192-token documents across all evaluation domains.The models are trained as described in Section 4.1.
  • Results: κ = 50 performs best overall on held-out perplexity across the evaluation domains.

D Token-level AutoCompressor Analysis

Token-level analyses show that summary vectors reduce perplexity across an entire 2,048-token segment and capture long-range dependencies. Their benefits include copying or inferring interpretable information such as names, dates, locations, and domain-associated concepts.

  • Token-level perplexity: Summary vectors reduce perplexity over the entire 2,048-token segment, rather than only helping continue the preceding sequence.
  • Token-level perplexity: Extended attention benefits more from additional context at sequence starts, while the AutoCompressor achieves stronger gains at sequence ends.
  • Token-level perplexity: The AutoCompressor nearly matches extended attention at sequence starts and outperforms it at sequence ends.The comparison tracks perplexity gains over OPT-2.7B across 2,048-token sequences.
  • Analysis procedure: The analysis evaluates examples by compressing the first 300 tokens and measuring perplexity on the following 100 tokens.For each example, the analysis records the top 3–5 most improved token predictions.
  • Interpretability: Tokens benefiting most from summary vectors often include interpretable names, dates, and locations copied through the compressed context.Examples are reported for Wikipedia, FreeLaw, and HackerNews.
  • Interpretability: In out-of-domain examples, summary vectors can support inferred concepts and associated-domain information, including “Tokyo” and “Jim Parsons.”The examples connect “Tokyo” to discussion of Japan and “Jim Parsons” to a co-star discussed in context.

F Fused Retrieval-augmented Language Modeling

Fused summary vectors support retrieval-augmented language modeling while avoiding the rapid context-window limits of concatenated passages. Re-ranking retrieved summaries by their distance to the query summary is computationally cheap but produces higher perplexities in the reported ablation.

  • Re-ranking: Re-ranking consistently produces higher perplexities than using retrieval scores when fusing summary vectors.Table 8 reports PPL gains over the no-retrieval baseline for fused summaries with and without re-ranking.
  • Fused summary vectors: Summary vectors from retrieved passages are concatenated in retrieval-score order and fused in a single forward pass.
  • Fused summary vectors: The fused-summary probability combines conditioned and unconditioned probabilities as p(y | x, D) = p(y | Concat[σD, x]) + p(y | x).
  • Re-ranking: Re-ranking passages by the smallest ℓ2 distance between their summary vectors and σx incurs negligible overhead because σx is computed during the same forward pass as p(y | x).
  • Fused passages: Fusing plain-text passages is quickly limited by the pre-trained model’s context window, especially when retrieving many long passages.
  • Evaluation context: The broader evaluation includes OPT-2.7B, Llama-2-7B, AutoCompressor, and RMT models across 11 in-context learning tasks.
Loading 2305.14788v2…