[Paper] 多头低秩注意力

Source: arXiv - 2603.02188v1

概述

本文介绍了 Multi-Head Low-Rank Attention (MLRA)，一种旨在加速大型语言模型（LLMs）长上下文解码的新注意力机制。通过使潜在表示能够在多个 GPU 之间 partitionable，MLRA 减少了通常在逐标记生成时成为瓶颈的内存流量，实现了最高 2.8× 的加速，同时保持（甚至提升）模型质量。

关键贡献

• 可分区的潜在状态： 与多头潜在注意力（MLA）不同，MLRA 的低秩潜在向量可以跨设备拆分，实现高效的张量并行（TP）解码。
• 四路 TP 友好设计： 该架构允许每个 GPU 只加载其 KV 缓存的切片，显著降低芯片外存储带宽使用。
• 领先的性能表现： 实验表明，MLRA 在困惑度和下游基准（如 QA、摘要）上匹配或超越 MLA。
• 解码加速： 在相同硬件配置下，解码速度提升 2.8× 相较于 MLA。
• 开源发布： 代码、预训练权重以及训练/评估脚本已公开，便于可重复性和社区采纳。

方法论

1. 低秩注意力因式分解

   • 传统的自注意力为每个 token 存储完整的键值（KV）矩阵，在每个生成步骤都必须从高带宽内存（HBM）中读取。
   • MLRA 将 KV 缓存因式分解为一个 潜在 低维表示（秩 r）和一组 投影 矩阵，从而降低整体缓存大小。

2. 支持分片的多头设计

   • 现在每个注意力头都有自己的低秩潜在状态。由于潜在维度在各头之间拆分，该状态可以在 TP 设置中 分片 到多个 GPU 上。
   • 在解码过程中，每个设备仅加载其拥有的潜在缓存切片，而投影矩阵则保留在芯片上。

3. 训练流水线

   • 作者在标准的仅解码器 Transformer（类似 LLaMA）中插入 MLRA 模块进行预训练。
   • 他们使用因果语言建模损失与正则化项的组合，鼓励低秩因子捕获大部分注意力信息。

4. 评估

   • 基准测试包括标准语料库上的语言模型困惑度，以及来自 MMLU、GSM‑8K 和摘要套件的下游任务。
   • 解码速度在 4 GPU 节点（NVIDIA A100，80 GB）上使用贪婪搜索和束搜索两种方式进行测量。

结果与发现

• 质量： MLRA 稍微提升了困惑度和下游准确率，相比基线和 MLA，证实低秩分解并未牺牲表达能力。
• 速度： 4 路 TP 实现将 KV 缓存流量降低约 65%，从而相比 MLA 实现 2.8 倍 的解码加速，并相较完整 KV 基线提升约 2 倍。
• 可扩展性： 从 2 到 8 GPU 的实验显示出近线性的吞吐量提升，证明分片设计如预期般有效。

实际影响

• 更快的长上下文应用推理： 聊天机器人、代码助手以及需要在上下文中保留数千个 token 的文档分析工具，现在可以更快地生成响应，显著降低终端用户的延迟。
• 每个 token 的硬件成本更低： 由于每块 GPU 只需获取 KV 缓存的一小部分，相同的推理吞吐量可以在更便宜的 GPU 集群上实现，甚至在单个大内存 GPU 上也能减少芯片间流量。
• 提升 TP 利用率： 现有的张量并行流水线（例如 DeepSpeed、Megatron‑LM）可以在几乎不做改动的情况下采用 MLRA，既获得内存效率和速度提升，又不牺牲模型并行带来的权重分片等优势。
• 更易在边缘类服务器上部署 LLM： 降低的内存带宽需求使得在 HBM 受限的服务器上运行大模型成为可能，为本地或私有云的 LLM 服务打开了大门。

限制与未来工作

• 秩选择敏感性： 低秩维度 (r) 必须仔细调节；过低会影响准确性，过高则削弱内存带宽提升。自动化的秩选择策略尚未探索。
• 仅关注解码器‑Only Transformer： 论文仅评估因果语言模型。将 MLRA 扩展到编码器‑解码器架构（如 T5）或视觉‑语言模型仍是一个未解之题。
• 硬件特定优化： 报告的加速在 A100 GPU 上测得；在其他加速器（如 AMD GPU、TPU）上的性能可能不同，需要专门的 kernel 调优。
• 训练开销： 引入低秩分解会增加额外的投影层，略微提升训练计算量。未来工作可以研究更高效的训练技巧或混合精度方案。

总体而言，多头低秩注意力为需要高吞吐、长上下文 LLM 推理且不牺牲模型质量的开发者提供了一个有吸引力的方案。随着代码和预训练权重已发布，社区可以立即开始实验。

作者

• Songtao Liu
• Hongwu Peng
• Zhiwei Zhang
• Zhengyu Chen
• Yue Guo

论文信息

• arXiv ID: 2603.02188v1
• 分类: cs.LG
• 出版日期: 2026年3月2日
• PDF: 下载 PDF