用压缩标记替换冗余文本，实现长文本高效记忆压缩
UniGist: Towards General and Hardware-aligned Sequence-level Long Context Compression
- 用特殊压缩标记细粒度替换原始文本，保留关键上下文
- 在长文本任务中显著提升压缩质量，尤其擅长细节与远距离依赖建模
- 支持灵活推理，可实时释放内存，适合资源受限部署
大语言模型处理长上下文的能力日益增强，但键值缓存（KV cache）带来的内存开销仍是通用部署的主要瓶颈。尽管已有多种压缩策略，序列级压缩（即丢弃某些标记的完整KV缓存）尤为困难，可能导致重要上下文丢失。为此，我们提出UniGist，一种序列级长上下文压缩框架，通过细粒度地将原始标记替换为特殊压缩标记（gists），有效保留上下文信息。采用无分块训练策略并设计高效的内核与压缩偏移技巧，实现优化的GPU训练。该方案还支持灵活推理，允许实际移除压缩标记，实现即时内存节省。多组长上下文任务实验表明，UniGist显著提升压缩质量，尤其在细节回忆和长程依赖建模任务中表现优异。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。