Transformer可以改裝成Mamba了：蘋果把推理成本直接打成線性

編輯｜Sia

最近，蘋果又整了個活兒，很工程、也挺關鍵：

把又貴又強的 Transformer，改造成又便宜又差不多強的 Mamba。而且，性能基本沒怎么掉。

聽起來炸不炸？有點像水電煤級別的升級。

為啥要折騰這個？很簡單，Transformer 確實猛，這十年基本就是靠它打天下。但問題也很現實：它越長越貴，而且是平方級變貴。

短文本還好，一旦上到長上下文（比如代碼、agent、多輪推理），那成本就不是有點高，而是直接肉疼。

于是，大家就開始找替代方案。比如線性 attention、RWKV，還有這兩年很火的 Mamba。這些模型的思路都很統一：別再平方爆炸了，改成線性。好處也很直接：更快、更省顯存、推理更絲滑。

但問題來了， 它們不夠強。尤其一旦規模上去，和 Transformer 還是有差距。于是就卡住了一個經典局面：

要性能？用 Transformer（但貴） ；
要便宜？用 Mamba（但弱一點）；

那有沒有可能「既要又要」？能不能不重訓，把 Transformer 的能力，直接搬到 Mamba 上？

直接蒸餾，性能肯定得崩。Transformer 像那種隨時翻資料的學霸， Mamba 更像全靠記憶的選手。你突然讓一個翻書型選手閉卷考試，還不給過渡—— 那基本就是災難現場。

蘋果沒有硬剛，而是搞了一個「兩步走」：

先造一個「中間形態」，讓 Transformer 先變成一個更簡單、更接近 Mamba 的版本。再從這個中間版本，轉成 Mamba。

用來把 Transformer 里原本計算很貴的 Softmax Attention，換成一個更省算力的 Mamba 模塊。中間加了一個過渡形態：先把 Attention 變成一種學出來的線性 Attention（Hedgehog），再結合Mamba，最終得到一個混合模塊，叫HedgeMamba

第一步，把原始的 softmax attention 換成一種線性 attention，同時盡量不丟性能。

問題在于，傳統線性注意力一直有明顯短板，和 softmax attention 的表現差距很大。為了解決這個差距，他們依據Mercer 定理，用了 Hedgehog 方法，用一個小型神經網絡（MLP）去學習一種特征映射，讓線性 attention 盡量模仿原來的行為。

再通過余弦相似度（cosine similarity）蒸餾，讓這個新結構在輸出上對齊 Transformer。這一步做完，相當于得到一個「更便宜但還挺像 Transformer」的中間模型。

第二步，是把這個已經對齊好的線性 attention，進一步嵌入到 Mamba 結構里。

他們做了一件很關鍵的事情：把 attention 里的核心計算方式，對應映射到 Mamba 的內部參數上，讓 Mamba 在初始化的時候，行為就已經接近前一步的模型，而不是從零開始學。這一步本質是在做結構對齊。

不過，還有一個問題需要處理。原始 attention 自帶一種歸一化機制，而線性版本沒有，所以他們額外加了一步歸一化，讓輸出形式更接近原來的 attention，同時又不破壞計算效率。

完成這些之后，才進入真正的訓練階段。

他們對整個模型做微調，用標準的交叉熵訓練，并重新啟用 Mamba 原本的能力，比如卷積和門控（ Gate ）機制。這一步的作用，是讓模型不只是模仿，而是用自己的方式把能力重新學出來。

整套方法的關鍵不在某一個技巧，而在這條路徑本身：先讓兩種模型在「表達方式」上對齊，再做結構轉換，最后通過訓練把能力恢復出來。

也正是因為這樣分步處理，才避免了直接蒸餾時常見的性能崩塌問題。

效果到底咋樣？性能幾乎沒掉，但成本邏輯已經變了。

論文里最關鍵的一張表，把三類1B模型擺在一起對比：Transformer教師模型（Pythia）、傳統蒸餾基線（ Hedgehog ），以及他們的方法（ HedgeMamba ）。在只用10B token（大約是教師訓練數據的 2.7% ）的情況下，結果非常直接——

教師模型的困惑度是 13.86，基線方法掉到 14.89，而HedgeMamba把這個指標拉回到 14.11，已經貼得很近了。

他們用一個大約10B token訓練出來的1B模型做實驗，最后得到的 Mamba 模型能夠保留原始 Pythia-1B Transformer 在下游任務中的性能，其困惑度（perplexity）保持在 14.11，接近老師模型的 13.86。

這件事的含義其實挺重的。

過去大家默認一個前提：只要你把 Transformer 換成另一種架構，性能就會明顯掉一截。

但這篇論文給出的答案是，這個損失可以被大幅追回來。而且，不只是語言建模指標好看，在Arc、PIQA、BoolQ、RACE、LogiQA 這些下游任務上，HedgeMamba基本全面超過基線，同時整體表現已經逼近教師模型。這說明它保留下來的不只是表面的概率分布，而是相當一部分推理能力和語義結構。

更關鍵的是，這種效果不是調出來的，而是有方法論支撐的。他們嘗試過最直接的做法——從 Transformer 一步蒸餾到 Mamba，結果是 PPL 直接炸到 100 以上，幾乎不可用。

換句話說，兩階段蒸餾在這里不是優化，而是繞不過去的結構性條件。

后面的消融和分析，其實是在解釋這條路徑為什么成立。

比如，架構上，真正起作用的不是簡單疊模塊，而是門控機制——也就是讓模型學會該記什么、不該記什么；

訓練策略上，兩階段的數據分配也不是平均最優，而是明顯偏向第二階段，說明中間表示只是過渡，真正的能力是在后半段完成遷移；

再看數據規模，從1B 到 10B token，性能是穩定往上走的，沒有出現不收斂或反復震蕩的情況，這一點很重要，因為它說明這條路線是可以規模化放大的。

架構消融實驗表明，讓 Mamba 好用的關鍵，不是簡單堆結構，而是門控機制。

蒸餾的兩步（S1 和 S2），數據到底該怎么分配才最有效？兩階段蒸餾是必要的，而且最優策略是輕S1 + 重S2。

蒸餾過程中 token 數量（訓練數據量）對效果的影響。似乎只要給足夠多的蒸餾數據，Mamba 可以逼近 Transformer 的性能。

把這些信息合在一起，這篇工作的價值就不在「又做了一個更好的模型」，而在于它提供了一種新的工程可能性。

過去幾年積累的大量 Transformer 模型，并不需要全部推倒重來，而是有機會通過一套流程，被「轉制」為更高效的形態。

如果這件事能穩定復現，那整個開源模型生態、甚至很多公司的自研模型，都有機會被整體降本重構。

https://arxiv.org/abs/2604.14191