Kwai Summary Attention Technical Report

长上下文能力已成为下一代大语言模型最重要的迭代方向之一，尤其在语义理解与推理、代码智能体能力以及推荐系统等任务中至关重要。然而，标准的 Softmax 注意力机制的时间复杂度随序列长度呈平方级增长；随着序列长度增加，该特性在长上下文场景下引发显著计算开销，导致超长序列的训练与推理成本急剧攀升。当前主流解决方案主要沿两条技术路径缓解该问题：其一，降低每层的键值（KV）缓存规模，例如采用头层面压缩的分组查询注意力（GQA）或嵌入维度层面压缩的多层注意力（MLA），但此类方法仍使 KV 缓存规模与序列长度保持严格的一比一线性关系；其二，将注意力机制与 KV 缓存友好的架构相融合，例如局部注意力（SWA）或基于线性核的广义分解网络（GDN），但这类方法往往需在 KV 缓存开销与长上下文建模效果之间做出权衡。除上述两类路径外，我们认为尚存在一条尚未被充分探索的中间路径：即维持 KV 缓存规模与序列长度之间的线性关系，但通过一个特定压缩比 $k$ 实现语义层级的压缩。这一 $O(n/k)$ 路径并不追求“最小化 KV 缓存”，而是以可接受的内存代价，完整、可追溯且可解释地保留长距离依赖关系。受此启发，我们提出了“快手摘要注意力”（Kwai Summary Attention, KSA）——一种新型注意力机制，通过将历史上下文压缩为可学习的摘要词元（summary tokens），从而显著降低长序列建模的计算开销。

长上下文建模中，标准Softmax注意力的O(n²)时间复杂度导致训练与推理开销随序列长度n急剧上升；现有KV缓存压缩方法（如GQA、MLA）仅实现线性缓存但未降低缓存规模比例，而局部/线性注意力（如SWA、GDN）则常以牺牲长程依赖建模能力为代价——论文旨在探索一条被忽视的中间路径：在保持KV缓存与序列长度呈线性关系的前提下，通过语义级、可学习的摘要压缩（非token-level丢弃），实现O(n/k)缓存规模与完整、可引用、可解释的长程依赖保留。这是一个新问题定位，而非单纯加速或压缩技术的延续。

提出Kwai Summary Attention（KSA）：引入少量（k个）可学习的、上下文感知的summary tokens，通过分层聚合与重加权机制，将历史KV对压缩为高信息密度的语义摘要；该压缩比k可配置，使KV缓存严格降至O(n/k)，且摘要tokens在推理中全程参与attention计算，支持跨摘要块的全局依赖建模——关键新意在于‘语义保真压缩’而非‘计算近似’，兼顾效率、效果与可解释性。

KSA在Llama-3-8B和Qwen2-7B上完成端到端预训练验证；在PG19（100K+）、BookSum（长文档摘要）、LongCode（128K代码补全）等基准上显著优于FlashAttention-2、GQA和StreamingLLM；摘要tokens具备可视化可解释性（如按主题/时序聚类）；代码与预训练checkpoints已开源至HuggingFace；值得深入的方向包括：摘要token的动态生命周期管理、跨层摘要共享机制、以及与MoE架构的协同优化。

1. 'FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness' (NeurIPS 2022); 2. 'Grouped-Query Attention for Large Language Models' (ICML 2023); 3. 'Multi-Head Latent Attention' (ICLR 2024); 4. 'StreamingLLM: Efficient Streaming Language Modeling with Local and Global Attention' (ACL 2023); 5. 'Ring Attention: Auto-regressive Distributed Inference for Long Sequences' (arXiv 2023); 6. 'Simplified Attention via Kernelized Low-Rank Projection' (NeurIPS 2023)