新注意力让大模型上下文内存占用砍半！精度不减还能加速2倍

今天小编(佼昌翰)要和大家分享的是新注意力让大模型上下文内存占用砍半！精度不减还能加速2倍,欢迎阅读~

大模型同样的上下文視窗，只需一半内存就能实现，而且精度无损？

前苹果 ASIC 架构师 Nils Graef，和一名 UC 伯克利在读大學生一起提出了新的注意力机制Slim Attention。

它以标准多头注意力（MHA）为基准，对其中的 value 缓存处理过程进行了调整，实现了更少的内存占用。

具体来说，Slim Attention 既可以让KV 缓存大小减半，也可以在 KV 缓存大小不变的情况下让上下文翻倍，都不会带来精度损失。

此外，在内存带宽受限的场景下，它还可以将模型的推理过程加速1.5-2 倍。

网友评价，Slim Attention 虽然简单，但却是一个很酷的想法。

还有 AI 创业者评论说，这是一项重大突破，可能重塑对模型训练和部署的看法。

K-Cache is All You Need

在标准的 MHA 机制当中，对于输入 X 会通过线性变换，经由三个投影矩阵 W_Q、W_K、W_V 得到 Q ( query ) 、K ( key ) 和 V ( value ) 三个矩阵。

在推理阶段，每个输入 token 计算得到的 K 和 V 向量都需要缓存起来，形成 KV cache 供后续 token 计算时使用。

Slim Attention 的核心思路是，利用 MHA 中 W_K 和 W_V 通常都是方阵的性质，只存储 K 而不直接存储 V，然后实时利用 K 计算出 V。

△原始 MHA（左）与改进版（右）对比

在训练阶段，Slim Attention 与标准 MHA 一样，会对输入 X 计算 Q、K、V 三个矩阵，注意力计算和梯度回传也与标准 MHA 完全一致。

在 W_K 可逆的前提下，Slim Attention 引入一个新的参数矩阵 W_KV：

W_KV = W_K^ ( -1 ) · W_V

据此，可以得到：

V = X · W_V = X · W_K · W_K^ ( -1 ) · W_V = K · W_KV

推理过程则主要分为两个阶段——提示阶段（并行计算）和生成阶段（自回归）。

提示阶段与标准 MHA 一样，将输入的所有 token 并行计算 Q、K 矩阵，但不同的是，这里不直接计算 V，而是将中间结果 K 缓存供后续使用。

生成阶段每个时间步生成一个新 token，首先计算该时间步的 Q 向量 q，然后基于 q 和之前时间步缓存的 K 矩阵，计算注意力得（即 softmax 的输入）。

在 softmax 之前，Slim Attention 通过公式 V = K · W_KV 实时计算 V 矩阵。具体有两种方式 :

直接计算 V，然后将 softmax 结果与 V 相乘（矩阵乘法）得到注意力输出；

先将 softmax 结果与 K 相乘，然后再与 W_KV 相乘，当序列较长时这种方式更高效。

剩余流程（残差连接、前馈层等）与标准 MHA 一致，最后将当前步的 k 向量添加到 K 缓存中，供下一时间步使用。

总之，Slim Attention 是标准 MHA 的精确数学重写，因此与近似方法不同，可确保准确率不会下降。

以此为前提，Slim Attention 实现了 KV 缓存减半或上下文翻倍的效果。

前苹果架构师与 UC 伯克利大學生成果

Slim Attention 的作者是 AI 初创公司 OpenMachine 的创始人兼 CEO Nils Graef，以及 UC 伯克利在读大學生Andrew Wasielewski。

Nils 的主业是机器学习加速器的架构和设计，曾发表两篇 IEEE 期刊论文和 30 多项专利，引用次数超过 900 次。

创立 OpenMachine 前，Nils 在知名推理加速平台 Groq（注意不是马斯克的 Grok）担任芯片架构师。

更早的时候，他先后担任过谷歌 ML 加速器架构 & 设计工程师和苹果 ASIC 架构师。

Andrew Wasielewski 是 UC 伯克利在读大學生，专业是物理和 EECs（电气工程与计算机科学），预计将于明年毕业。

根据论文署名信息显示，Slim Attention 的工作是 Andrew 在 OpenMachine 完成的。

去年 7 月，Nils 和 Andrew 还与其他人合作，发表了一篇名为Flash normalization的论文，提出了一种更快的 RNS 归一化方式。

此外在 Slim Attention 的致谢中还提到，艾伦实验室的 Dirk Groeneveld，以及 SGLang 三作谢志强，对其工作提供了有益讨论；Transformer 作者之一、Character.AI 创始人 Noam Shazeer 给出了积极反馈。

论文地址：

https://arxiv.org/abs/2503.05840

参考链接：

https://x.com/rohanpaul_ai/status/1901092052282339474

关于新注意力让大模型上下文内存占用砍半！精度不减还能加速2倍就分享完了，您有什么想法可以联系小编(佼昌翰)。