KSA通过摘要令牌压缩长序列,降低大模型上下文开销。
Kwai Summary Attention Technical Report

- 用可学习的摘要令牌压缩历史上下文,实现线性内存增长
- 在长序列中保持远距离依赖关系,避免注意力退化
- 适合需要高效长文本理解的推荐与代码智能场景
长上下文能力已成为下一代大语言模型的重要演进方向,尤其在语义理解、代码智能和推荐系统中表现突出。然而,标准softmax注意力随序列长度呈二次时间复杂度,导致极长序列训练与推理成本急剧上升。现有方案主要通过两条路径缓解:一是减少每层的键值缓存(KV cache),如头级压缩的GQA、嵌入维度压缩的MLA,但其仍与序列长度保持1:1线性关系;二是采用友好的KV缓存架构,如局部注意力SWA、线性核GDN,但常在缓存效率与长程建模效果间权衡。我们提出一种未被充分探索的中间路径:在保持KV缓存与序列长度线性关系的同时,通过特定比例k进行语义级压缩。该方法不追求最小化缓存,而是以可接受的内存代价保留完整的、可追溯的远距离依赖。基于此,我们提出Kwai Summary Attention (KSA),通过将历史上下文压缩为可学习的摘要令牌,降低序列建模成本。
原文摘要 · Abstract (English)
Long-context ability, has become one of the most important iteration direction of next-generation Large Language Models, particularly in semantic understanding/reasoning, code agentic intelligence and recommendation system. However, the standard softmax attention exhibits quadratic time complexity with respect to sequence length. As the sequence length increases, this incurs substantial overhead in long-context settings, leading the training and inference costs of extremely long sequences deteriorate rapidly. Existing solutions mitigate this issue through two technique routings: i) Reducing the KV cache per layer, such as from the head-level compression GQA, and the embedding dimension-level compression MLA, but the KV cache remains linearly dependent on the sequence length at a 1:1 ratio. ii) Interleaving with KV Cache friendly architecture, such as local attention SWA, linear kernel GDN, but often involve trade-offs among KV Cache and long-context modeling effectiveness. Besides the two technique routings, we argue that there exists an intermediate path not well explored: {Maintaining a linear relationship between the KV cache and sequence length, but performing semantic-level compression through a specific ratio $k$}. This $O(n/k)$ path does not pursue a ``minimum KV cache'', but rather trades acceptable memory costs for complete, referential, and interpretable retention of long distant dependency. Motivated by this, we propose Kwai Summary Attention (KSA), a novel attention mechanism that reduces sequence modeling cost by compressing historical contexts into learnable summary tokens.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。