17岁高中生实现Ilya预言！Kimi团队“旋转”注意力机制，训练效率提升25%，马斯克都来围观

17岁高中生实现Ilya预言：Kimi团队“旋转”注意力机制，训练效率提升25%

一位17岁的高中生，以共同第一作者的身份，在Kimi团队将OpenAI首席科学家Ilya Sutskever提出的一个设想变成了现实。

Ilya曾有一个著名的预言：将按时间顺序处理数据的LSTM网络“旋转90度”——即把时间轴替换为模型的深度轴——就演变成了当今主流的残差网络（ResNet）。

Kimi团队从中获得启发：既然时间维度上的LSTM能对应深度维度上的残差连接，那么后来在序列模型中取代了LSTM的“注意力机制”，是否也可以进行同样的“旋转”操作？

他们新提出的“注意力残差”（Attention Residuals）技术，正是将注意力机制“旋转了90度”，从序列的时间维度移植到了网络的深度维度。

应用这一新机制后，模型在计算当前层时，能够智能地“回头看”，根据需求自由地提取前面任何一层的有用信息，而非像传统残差连接那样无差别地累加。

这项研究成果引起了广泛关注，连埃隆·马斯克也前来围观，并评论称“令人印象深刻”。

除了马斯克，AI领域的大神安德烈·卡帕西（Andrej Karpathy）也对此进行了思考，直言我们对《Attention is All You Need》这篇Transformer开山之作的理解仍然不够深入。

将这种新机制置于Kimi自家的Kimi Linear 48B大模型（激活参数3B）上进行验证，结果显示训练效率提升了25%，而推理延迟的增加被控制在2%以内。

残差连接的“记忆负担”

要理解这项创新，首先需要回顾标准残差连接的工作原理。

传统做法是：第N层的输出 = 第N层的计算结果 + 第N-1层的输出。信息以此方式逐层累加，使得每一层理论上都包含了前面所有层的信息。

但问题随之而来。在当前大模型主流的PreNorm（前置归一化）范式下，残差连接中所有历史层的贡献都是以相等权重进行累加的。

这就像一个“记忆力过好”的人，把所有的经历都以相同的强度存入大脑。导致早期输入的信息在后续层中被逐步稀释，难以被有效检索。大量层甚至可以被剪枝而只带来微小损失，这被称为“PreNorm稀释问题”。

更棘手的是，隐藏状态的范数会随着网络深度不断增长。研究人员发现，在深层网络中，这种无约束的增长会导致训练不稳定。

月之暗面团队转换了思路：既然问题根源在于“无差别累加”，那么就让网络自己学会“选择性回忆”。

用注意力实现“选择性回忆”

团队观察到一个关键的对偶性：网络的深度维度与序列的时间维度，在本质上具有相似的结构。

在Transformer处理序列时，利用注意力机制让当前位置能够“选择性关注”序列中之前的位置。那么，在深度维度上，为何不能让当前层“选择性关注”之前的层呢？

“注意力残差”机制便应运而生：
* 将当前层的一个可学习伪查询向量作为query。
* 将所有前面层的输出作为key和value。
* 利用注意力机制对前面层的信息进行加权聚合。

这样，网络就能学会判断哪些历史层的信息对当前计算最重要，并赋予其更高权重；而不相关的层，其权重则会自然降低。

然而，这引入了一个新挑战：计算复杂度爆炸。对于一个L层的网络，如果每一层都需要对前面所有层做完整的注意力计算，复杂度将达到O(L²)，在实践中无法承受。

分块注意力残差：化繁为简

论文中提出的解决方案是“分块注意力残差”。

其核心思想是将连续的若干层打包成一个块，并对块内部的信息进行压缩，只保留一个“摘要向量”。

具体操作如下：
* 将L层网络划分为B个块，每个块包含若干层。
* 在每个块结束时，将该块内所有层的信息压缩成一个单一的块级表征向量。
* 后续层进行注意力计算时，只需关注这些块级表征以及当前块内的实时层输出，而非全部L个独立层。

通过这种方法，注意力计算的复杂度从O(L²)降低到了O(L·B)。在实践中，B可以设置得很小（论文中为8-16），从而大幅降低了计算开销。

此外，团队还实施了一系列工程优化，如缓存式流水线通信、序列分片预填充、KV缓存粒度优化等，以进一步提升效率。

在Kimi Linear上的验证：1.25倍效率提升

理论需要实践验证。团队在其自研的Kimi Linear架构上进行了大规模测试。这是一个采用线性注意力的大模型，总参数量480亿，激活参数量30亿。

实验表明，在相同的计算预算下，采用注意力残差的模型能获得更好的下游任务性能。反之，要达到相同的性能水平，所需的训练计算量减少了约20%，相当于获得了1.25倍的训练效率优势。

在具体任务上，包括数学推理（MATH、GSM8K）、代码生成（HumanEval、MBPP）在内的多项评测均表现持平或略有提升，多语言理解的一致性也有所改善。

更重要的是，注意力残差是一个“即插即用”的模块，无需修改网络其他部分，可以直接替换原有的残差连接。

论文还提供了一个有趣的视角。团队将这项工作称为“时间-深度对偶性”的应用。在他们看来，深度神经网络中的“层”与循环神经网络中的“时间步”，本质上都是对信息的迭代处理。Transformer的成功在于用注意力机制取代了RNN中固定的循环结构。那么，在深度维度上，是否也应该用注意力机制来取代固定的残差连接呢？

17岁高中生成为共同一作

更引人注目的是，这篇令马斯克、卡帕西等人都为之瞩目的论文，其共同第一作者之一是一名年仅17岁的高中生——陈广宇。

另外两位共同一作，分别是Kimi团队的关键人物、旋转位置编码（RoPE）的提出者苏剑林，以及Kimi Linear的第一作者张宇。

尽管注意力残差是团队协作的成果，但一名高中生能与两位资深研究者并列一作，已足够令人惊叹。

a16z创始人马克·安德森、Thinking Machines联合创始人等业界知名人士都在社交媒体上关注了他的账号。

大约一年前才刚刚开始了解大模型的陈广宇，是从北京的一场黑客松起步，最终走向了硅谷。之后回国时，他选择加入了月之暗面。

也正是顺着这条线，他开始一路向更底层钻研，从阅读论文、查看实现，到研究 Triton kernel、理解注意力机制为何能被如此重写与加速。

在月之暗面，这条路仿佛绕了一圈又回到原点——他最初被底层技术所吸引，最终从事的也正是最底层、最核心的工作。

陈广宇的经历，与其说是一个“少年天才一路开挂”的故事，不如说是一种不同的成长路径：先被时代最前沿的技术所震撼，再一步步将兴趣打磨成能力，最终将这份能力带到了真正的大模型研发现场。

相关链接
– 论文地址：https://github.com/MoonshotAI/Attention-Residuals/
– 参考链接：[1] https://mp.weixin.qq.com/s/gRR99pEDWb5qsk2a2hwe2w
– 参考链接：[2] https://nathanchen.me/public/About%20me.html

关注“鲸栖”小程序，掌握最新AI资讯