技术深度解析
大模型推理的核心瓶颈在于KV缓存。在自回归生成过程中，每个Transformer层会存储之前token的键（K）和值（V）张量，用于计算当前token的注意力分数。对于一个拥有L层、H个注意力头、上下文长度为N token的模型，KV缓存大小约为2 * L * H * N * d_k（其中d_k是每个头的维度）。以Llama 3.1 405B为例，该模型使用128层和64个注意力头，仅处理32K token时缓存就会膨胀到数百GB——远超GPU内存容量。
KV缓存共享通过允许多个注意力头共享相同的缓存键值对来解决这一问题。其洞察在于：许多注意力头学习到了冗余或互补的模式。通过将注意力头分组到共享的KV池中——通常通过学习的路由机制或简单的平均化实现——内存使用量可降低与共享比例相等的倍数。早期实验表明，4倍共享比例可将KV缓存大小减少75%，同时在标准基准测试上的准确率下降不到0.5%。
多头压缩（MHC）则更进一步。MHC不是共享，而是通过一个学习的线性投影或小型Transformer模块，跨注意力头压缩KV缓存。可以将其视为一个瓶颈层，从所有头中提炼出最重要的信息，形成紧凑的表示。压缩后的缓存在注意力计算过程中被实时解压。某主要研究实验室近期的一篇论文证明，MHC在长上下文任务上可实现8倍压缩，而困惑度仅下降1-2%。GitHub仓库`mhc-attention`（目前拥有2300颗星）提供了基于PyTorch的参考实现，支持从头训练和微调现有模型。
压缩注意力机制是降低标准注意力二次复杂度的架构变革。滑动窗口注意力（用于Mistral 7B和Mixtral 8x7B）将每个token的注意力限制在固定大小的前序token窗口内，使复杂度降为O(N * W)，其中W为窗口大小。稀疏注意力（如BigBird、Longformer）使用预定义的稀疏模式——全局token、滑动窗口和随机连接——实现O(N log N)或O(N)复杂度。近期关于线性注意力的研究（如Mamba、RWKV）则完全用循环或状态空间模型替代了softmax注意力，实现了真正的O(N)复杂度，但通常以降低某些任务的表达能力为代价。
| 方法 | 内存缩减 | 复杂度缩放 | 困惑度下降（vs. 全注意力） | 示例模型 |
|---|---|---|---|---|
| KV缓存共享（4倍） | 75% | O(N²)（与全注意力相同） | <0.5% | 定制Llama 3.1 8B |
| 多头压缩（8倍） | 87.5% | O(N²) | 1-2% | MHC-Llama 7B |
| 滑动窗口（W=4096） | 50%（8K上下文时） | O(N * W) | 2-3%（长距离任务） | Mistral 7B |
| 稀疏注意力（BigBird） | 60-80% | O(N log N) | 1-3% | Longformer, BigBird |
| 线性注意力（Mamba） | 90%以上 | O(N) | 3-5%（检索任务） | Mamba 2.8B |
数据要点： 没有任何单一方法占据主导地位。KV共享和MHC能最好地保持全注意力质量，但仍面临二次计算成本。滑动窗口和稀疏注意力提供更好的扩展性，但在需要长距离依赖的任务上表现下降。线性注意力提供最佳扩展性，但在召回密集型任务上表现不佳。最优解决方案很可能结合多种技术——例如，使用MHC实现内存效率，使用滑动窗口实现计算效率。
关键玩家与案例研究
Mistral AI一直是实用压缩注意力领域的先驱。其Mistral 7B模型使用窗口大小为4096 token的滑动窗口注意力，能够在消费级GPU上实现高效推理。该公司的Mixtral 8x7B混合专家模型在此基础上扩展了稀疏MoE层，以极低的成本实现了GPT-3.5级别的性能。Mistral的方法非常务实：他们牺牲部分长距离能力以换取显著的推理速度提升，这一权衡已被证明在商业上非常成功。
Anthropic则选择了不同的路径。其Claude 3.5 Sonnet模型据称使用了多头压缩的变体，但具体细节仍属专有。内部基准测试表明，Claude能在超过20万token的上下文中保持连贯性——这远超滑动窗口单独所能达到的效果。Anthropic押注的是：长上下文保真度对于法律文档审查和代码库分析等企业应用至关重要，即使这需要更复杂的压缩技术。
Google DeepMind通过其Ring Attention和Blockwise Parallel Transformer技术贡献了基础性研究，这些技术将KV缓存分布到多个设备上，从而实现近乎无限的上下文长度。其Gemini 1.5 Pro模型结合了环形注意力和稀疏门控机制，展示了1000万token的上下文窗口。虽然尚未广泛部署，但这项工作展示了将KV缓存分布到数千个加速器上、同时保持训练和推理效率的潜力。