Context Compression for Auto-regressive Transformers with Sentinel Tokens (2023)

Abstract: The quadratic complexity of the attention module makes it gradually become the bulk of compute in Transformer-based LLMs during generation. Moreover, the excessive key-value cache that arises when dealing with long inputs also brings severe issues on memory footprint and inference latency. In this work, we propose a plug-and-play approach that is able to incrementally compress the intermediate activation of a specified span of tokens into compact ones, thereby reducing both memory and computational cost when processing subsequent context. Experiments on both in-domain language modeling and zero-shot open-ended document generation demonstrate the advantage of our approach over sparse attention baselines in terms of fluency, n-gram matching, and semantic similarity. At last, we comprehensively profile the benefit of context compression on improving the system throughout. Code is available at https://github.com/DRSY/KV_Compression .

## introduction the transformer architecture @xcite has become the underpinning component of modern large language models (llms) @xcite @xcite in recent years. however, the quadratic computational complexity and memory footprint of the attention mechanism have largely limited applying transformers to increasingly long contexts with constrained computing resources. to mitigate such issues, prior works @xcite @xcite have explored an assortment of efficient transformer variants, mostly by replacing the original quadratic attention operation with various forms of linearized approximation. though promising, large-scale pre-training and specialized * the corresponding author. cuda kernels are typically required for these models to achieve performance comparable to off-theshelf llms and fulfill real efficiency gains. in this work, we aim to improve the efficiency of existing transformer-based llms without any architectural changes. specifically, we focus on the key-value cache, which accounts for the majority of memory footprint and data movement (i/o) cost when dealing with increasingly long input using llms. we propose a plug-and-play approach to incrementally compress the key-value memory of a contiguous span of tokens into compact ones. specifically, we introduce a pair of special sentinel @xmath0. during training, we modify the causal attention mask such that future tokens @xmath1. by continually training llms with the next token prediction objective, the model learns to extract and condense task-relevant information of the bounded span into the ending sentinel token. the reduced context length alleviates both memory and computational costs when processing subsequent tokens, thereby improving system throughput with larger batch sizes and faster decoding speed during inference. we conduct experiments on the wikitext-2 language modeling benchmark and show that our approach is generalizable to llms with various sizes (from 1.3b to 3b) and position encoding schemes such as absolute position embedding @xcite and rotary position encoding @xcite . compared to sparse attention baselines, our approach is able to effectively compress historical key-value cache with significantly reduced degradation in perplexity. moreover, we demonstrate that our approach outperforms sparse attention in zero-shot open-ended document generation across different compression ratios as evaluated by perplexity, rouge @xcite , and bertscore @xcite . finally, we empirically demonstrate that context compression is able to confer considerable improvement in system throughput for text generation. ## background and related work in this section, we present necessary background knowledge as well as drawing distinctions of our approach to existing literatures. ## approach in this section, we present a plug-and-play approach to compress the full-length contextual representations into shorter ones while preserving information necessary to accomplish the end task. ## experiments we mainly evaluate our approach on language modeling task and then explore its zero-shot generalization ability on open-ended document generation. finally, we quantitatively measure the effect of context compression on system throughput. ## conclusion in this work, we propose a simple yet effective approach that enables llms to summarize the keyvalue memory of specified span of tokens. experiments on language modeling and open-ended generation demonstrate that our approach substantially outperforms sparse attention baselines in terms of information compression and generation quality.