UniGist: Towards General and Hardware-aligned Sequence-level Long Context Compression
This addresses a critical memory overhead problem for deploying large language models in general-purpose applications, representing an incremental advancement in compression techniques.
The paper tackles the memory bottleneck of key-value cache in large language models for long-context inputs by introducing UniGist, a sequence-level compression framework that uses special compression tokens to preserve context efficiently, resulting in significant improvements in compression quality across multiple long-context tasks.
Large language models are increasingly capable of handling long-context inputs, but the memory overhead of key-value (KV) cache remains a major bottleneck for general-purpose deployment. While various compression strategies have been explored, sequence-level compression, which drops the full KV caches for certain tokens, is particularly challenging as it can lead to the loss of important contextual information. To address this, we introduce UniGist, a sequence-level long-context compression framework that efficiently preserves context information by replacing raw tokens with special compression tokens (gists) in a fine-grained manner. We adopt a chunk-free training strategy and design an efficient kernel with a gist shift trick, enabling optimized GPU training. Our scheme also supports flexible inference by allowing the actual removal of compressed tokens, resulting in real-time memory savings. Experiments across multiple long-context tasks demonstrate that UniGist significantly improves compression quality, with especially strong performance in detail-recalling tasks and long-range dependency modeling.