Context Compression for Auto-regressive Transformers with Sentinel Tokens
This addresses memory and latency issues in long-context LLM inference, offering a practical improvement for efficient deployment.
The paper tackles the quadratic complexity and excessive memory footprint of attention in Transformer-based LLMs during generation by proposing a plug-and-play method to compress intermediate activations, reducing memory and computational costs. Experiments show advantages over sparse attention baselines in fluency, n-gram matching, and semantic similarity, with profiling demonstrating improved system throughput.
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.