Kimi新架構(gòu)讓馬斯克嘆服！17歲高中生作者一戰(zhàn)成名

• 克雷西 發(fā)自 凹非寺
  量子位 | 公眾號(hào) QbitAI

17歲高中生，以一作身份，在Kimi團(tuán)隊(duì)把Ilya提出的設(shè)想，變成了現(xiàn)實(shí)。

Ilya之前有個(gè)預(yù)言，把按時(shí)間先后順序處理數(shù)據(jù)的LSTM網(wǎng)絡(luò)“旋轉(zhuǎn)90度”，也就是把時(shí)間軸換成模型深度軸，就變成了現(xiàn)在的殘差網(wǎng)絡(luò)。

Kimi團(tuán)隊(duì)認(rèn)為，既然時(shí)間上的LSTM能對(duì)應(yīng)深度上的殘差，那后來淘汰了LSTM的“注意力機(jī)制”自然也可以照做。

他們新搞出的Attention Residuals技術(shù)，就相當(dāng)于把注意力機(jī)制也“旋轉(zhuǎn)了90度”。

用了這套新方法后，模型在計(jì)算當(dāng)前層時(shí)可以聰明地“回頭看”，根據(jù)需要自由決定去提取前面哪一層的信息。

這篇論文讓馬斯克也來圍觀，表示令人印象深刻。

除了馬斯克，這篇論文也引發(fā)了大神Karpathy的思考，直言我們對(duì)Attention is All You Need這篇Transformer開山之作的理解還是不夠。

這種新機(jī)制放到Kimi自家的Kimi Linear 48B大模型（3B激活參數(shù)）上驗(yàn)證，訓(xùn)練效率提升25%，推理延遲增加不到2%。

殘差連接的“記憶負(fù)擔(dān)”

先回顧一下殘差連接的工作原理。

傳統(tǒng)做法是：第N層的輸出 = 第N層的計(jì)算結(jié)果 + 第N-1層的輸出。這樣一路累加下去，每一層都能“記住”前面所有層的信息。

問題來了，在大模型PreNorm主流范式下，殘差連接中所有層的貢獻(xiàn)都是等權(quán)累加。

就像一個(gè)“記憶力太好的人”，把所有經(jīng)歷都以相同權(quán)重存進(jìn)大腦。貢獻(xiàn)被逐步稀釋，早期信息難以檢索，且大量層可被剪枝而損失微小，稱之為“PreNorm dilution problem”。

更麻煩的是，隱藏狀態(tài)的范數(shù)會(huì)隨著深度不斷增長(zhǎng)。研究人員發(fā)現(xiàn)，在深層網(wǎng)絡(luò)中，這種unbounded growth會(huì)導(dǎo)致訓(xùn)練不穩(wěn)定。

月之暗面團(tuán)隊(duì)換了個(gè)思路：既然問題出在“無差別累加”，那就讓網(wǎng)絡(luò)自己決定該回憶什么。

用注意力“選擇性回憶”

團(tuán)隊(duì)觀察到一個(gè)有趣的對(duì)偶性：網(wǎng)絡(luò)的深度維度和序列的時(shí)間維度，本質(zhì)上是同構(gòu)的。

在Transformer處理序列時(shí)，用注意力機(jī)制讓當(dāng)前位置“選擇性關(guān)注”之前的位置。那么在深度維度上，為什么不能讓當(dāng)前層“選擇性關(guān)注”之前的層？

Attention Residuals就這么來的：

• 當(dāng)前層的可學(xué)習(xí)偽查詢向量作為query（learnable pseudo-query）
• 所有前層的輸出作為key和value
• 用注意力機(jī)制加權(quán)聚合

這樣一來，網(wǎng)絡(luò)可以學(xué)會(huì)哪些層的信息對(duì)當(dāng)前計(jì)算最重要，就多關(guān)注一點(diǎn)；不相關(guān)的層，權(quán)重自然降低。

但這帶來一個(gè)新問題：計(jì)算量爆炸。

如果一個(gè)100層的網(wǎng)絡(luò)，每一層都要對(duì)前面99層做full attention residual，復(fù)雜度是O(L2)，根本跑不動(dòng)。

Block AttnRes：分塊壓縮

論文中的解決方案是Block AttnRes。

核心思想是把連續(xù)的若干層打包成一個(gè)block，對(duì)block內(nèi)部的輸出做壓縮，只保留一個(gè)“摘要向量”。

具體操作如下：

• 把L層網(wǎng)絡(luò)分成B個(gè)block，每個(gè)block包含若干層
• 每個(gè)block結(jié)束時(shí)，把block內(nèi)的信息壓縮成單個(gè)向量
• 后續(xù)層做attention時(shí)，只需要關(guān)注塊間表征+塊內(nèi)實(shí)時(shí)層輸出，而非全部L個(gè)層

這樣一來，attention的復(fù)雜度從O(L2)降到了O(L·B)，在實(shí)踐中B可以設(shè)得很小（論文用的是8-16）。

此外，團(tuán)隊(duì)還做了數(shù)個(gè)工程優(yōu)化：緩存式流水線通信、序列分片預(yù)填充、KV 緩存粒度優(yōu)化等等。

Kimi Linear驗(yàn)證：1.25倍效率提升

理論說得通，但真正讓人信服的是大規(guī)模驗(yàn)證。

團(tuán)隊(duì)在自家的Kimi Linear架構(gòu)上做了測(cè)試。這是一個(gè)采用線性注意力的大模型，總參數(shù)48B，激活參數(shù)3B（MoE架構(gòu)）。

同等計(jì)算預(yù)算下，Attention Residuals能獲得更好的下游性能；反過來說，達(dá)到相同性能需要的訓(xùn)練計(jì)算量減少了約20%，相當(dāng)于獲得了1.25倍的效率優(yōu)勢(shì)。

在具體任務(wù)上，數(shù)學(xué)推理（MATH、GSM8K）、代碼生成（HumanEval、MBPP）均持平或略優(yōu)，多語言理解的一致性也有所改善。

更重要的是，Attention Residuals是一個(gè)drop-in replacement，不需要修改網(wǎng)絡(luò)其他部分，直接替換殘差連接即可。

論文里還講到一個(gè)有意思的視角。

團(tuán)隊(duì)把這項(xiàng)工作稱為“時(shí)間-深度對(duì)偶性”（time-depth duality）的應(yīng)用。

在他們看來，深度神經(jīng)網(wǎng)絡(luò)的“層”和循環(huán)神經(jīng)網(wǎng)絡(luò)的“時(shí)間步”，本質(zhì)上是都是對(duì)信息的迭代處理。

Transformer之所以成功，是因?yàn)橛胊ttention替代了RNN中固定的recurrence。

那么在深度維度上，是不是也該用attention替代固定的residual？

17歲高中生入列共同一作

更有意思的是，這篇讓馬斯克、Karpathy等人都為之一震的論文，共同一作之一是一名年僅17歲的高中生——陳廣宇（Nathan）。

另外兩名共同一作，分別是Kimi的關(guān)鍵人物之一、RoPE（旋轉(zhuǎn)位置編碼）的提出者蘇神（蘇劍林），以及Kimi Linear的第一作者張宇。

誠(chéng)然Attention Residuals是團(tuán)隊(duì)協(xié)作取得的成果，但一名高中生出現(xiàn)在這樣的團(tuán)隊(duì)之中，還與兩位大神共列一作，已經(jīng)足夠震撼。

a16z創(chuàng)始人Marc Andreessen、Thinking Machines的聯(lián)創(chuàng)等人都關(guān)注了他的X賬號(hào)。

一年前才剛剛開始了解大模型的陳廣宇，是從北京的一場(chǎng)黑客松開始，一路走向硅谷的。

后來回國(guó)時(shí)，他選擇加入了Kimi。

經(jīng)手過月之暗面投資的奇績(jī)創(chuàng)壇（原YC中國(guó)）創(chuàng)始成員董科含，也曾在其個(gè)人公眾號(hào)上刊載過陳廣宇的一份自傳。

去年二月，北京的一場(chǎng)中學(xué)生黑客松上，陳廣宇展示了一個(gè)關(guān)于“人類第三只機(jī)械輔助手”的創(chuàng)新構(gòu)想——ThirdArm。

也正是這個(gè)項(xiàng)目，讓他結(jié)識(shí)了黑客松評(píng)委董科含，后者也成為了他的創(chuàng)業(yè)導(dǎo)師。

當(dāng)時(shí)，董科含追問他，未來是否會(huì)深耕這項(xiàng)技術(shù)，這促使他開始重新審視自己的職業(yè)方向。

隨后他入選了董科含發(fā)起的只有極少數(shù)人入選的青年計(jì)劃，開始接觸IOI（國(guó)際信息學(xué)奧林匹克）金牌得主及資深科研人員。

此前他曾嘗試經(jīng)營(yíng)Shopify跨境電商店鋪、運(yùn)營(yíng)短視頻賬號(hào)，但經(jīng)過董科含的建議，他決定轉(zhuǎn)向理解時(shí)代的底層技術(shù)。

當(dāng)時(shí)還不知道Transformer是什么的他，在DeepSeek研究員袁境陽的指導(dǎo)下，利用Gemini作為輔助工具，通過研讀經(jīng)典論文、追蹤GitHub開源項(xiàng)目等方式逐步建立認(rèn)知。

有一次他在推特上分享了對(duì)一篇博客的反思后，獲得了作者的回復(fù)，這篇帖子也因此引起了一家硅谷AI初創(chuàng)公司CEO的關(guān)注。

該公司于2024年底成立，2025年初完成了800萬美元種子輪融資，資方背景涉及OpenAI與Anthropic。

在通過一項(xiàng)限時(shí)通宵完成的實(shí)驗(yàn)測(cè)試后，他拿到了對(duì)方的錄用通知。

暑假期間，他前往舊金山開啟了為期七周的實(shí)習(xí)。其中前兩周，他負(fù)責(zé)定義并推進(jìn)一個(gè)涉及144張H100顯卡的探索性項(xiàng)目。

在CEO直接指導(dǎo)下，他的工作延伸至運(yùn)營(yíng)層面，參與了招聘系統(tǒng)搭建、技術(shù)內(nèi)容輸出及融資策略討論，并獲得與早期投資者Vinod Khosla交流的機(jī)會(huì)。

在硅谷期間，他維持著高強(qiáng)度工作節(jié)奏，通過咖啡社交與英偉達(dá)工程師及初創(chuàng)創(chuàng)始人建立聯(lián)系。這次經(jīng)歷讓他將科研視為一種支撐創(chuàng)造的底層能力。

實(shí)習(xí)結(jié)束后，陳廣宇回到國(guó)內(nèi)，并于去年11月加入月之暗面。

把他吸引進(jìn)去的，正是Kimi一直做的Flash Linear Attention這一類高效attention工作。

實(shí)際上，正是GitHub上的FLA項(xiàng)目，吸引了他對(duì)機(jī)器學(xué)習(xí)的興趣并被邀請(qǐng)加入Kimi團(tuán)隊(duì)。

也正是順著這條線，他開始一路往更底層鉆，從讀論文、看實(shí)現(xiàn)，到研究 Triton kernel、理解attention為什么能被這樣重寫、這樣加速。

到了月之暗面，這條路也算是繞了一圈又落回原點(diǎn)——

他最初是被底層技術(shù)吸引，最后做的也正是最底層、最核心的那部分事。

相比于講一個(gè)“少年天才一路開掛”的故事，陳廣宇的經(jīng)歷更像是另一種成長(zhǎng)路徑——

先被時(shí)代最前沿的技術(shù)擊中，再一步步把興趣磨成能力，把能力帶到真正的大模型研發(fā)現(xiàn)場(chǎng)里。

論文地址：
https://github.com/MoonshotAI/Attention-Residuals/
[1]https://mp.weixin.qq.com/s/gRR99pEDWb5qsk2a2hwe2w
[2]https://nathanchen.me/public/About%20me.html