回顧 AGI 的爆發，從最初的 pre-training (model/data) scaling，到 post-training (SFT/RLHF) scaling，再到 reasoning (RL) scaling，找到正確的 scaling 維度始終是問題的本質。2017 年發布的 Transformer 架構沿用至今，離不開 Transformer 強大的 “無損記憶” 能力，當然也需要付出巨大的 KV 緩存代價。換句話說，Transformer 架構具有強大的 memory scaling 能力。

DeepSeek NSA 通過三種方式壓縮 “KV” 實現 sparse attention，但這只是一種可以工作但不優雅的折中方案。因為它在壓縮 Transfromer 的記憶能力，以換取效率。

另一方面，大概從 2023 年火到今天的線性序列建模方法（包括 linear attention 類，Mamba 系列，RWKV 系列）則是另一個極端，只維護一份固定大小 dxd 的 RNN memory state，然后加 gate，改更新規則，但這種方式始終面臨較低的性能上限，所以才會有各種混合架構的同樣可以工作但不優雅的折中方案。

我們認為，未來的模型架構一定具有兩點特性：強大的 memory scaling 能力 + 關于序列長度的低復雜度。后者可以通過高效注意力機制實現，比如：linear 或者 sparse attention，是實現長序列建模的必備性質。而前者仍然是一個有待探索的重要課題，我們把給出的方案稱為 “sparse memory”。

這促使我們設計了 MoM: Mixture-of-Memories，它讓我們從目前主流線性序列建模方法改 gate 和 RNN 更新規則的套路中跳脫出來，稀疏且無限制地擴大 memory 大小。MoM 通過 router 分發 token（靈感來自 MoE）維護多個 KV memory，實現 memory 維度 scaling。每個 memory 又可以進行 RNN-style 計算，所以整體具有關于序列長度線性的訓練復雜度，推理又是常數級復雜度。此外，我們又設計了 shared memory 和 local memory 合作分別處理全局和局部信息。實驗表現相當驚艷，尤其是在目前 linear 類方法效果不好的 recall-instensive 任務上表現格外好，甚至在 1.3B 模型上已經和 Transformer 架構旗鼓相當。

• 論文地址：https://arxiv.org/abs/2502.13685
• 代碼地址：https://github.com/OpenSparseLLMs/MoM
• 未來還會集成在：https://github.com/OpenSparseLLMs/Linear-MoE
• 模型權重開源在：https://huggingface.co/linear-moe-hub

方法細節

Linear Recurrent Memory

對于這部分內容，熟悉線性序列建模的小伙伴可以跳過了。

輸入 經過 query key value proj 得到：

最簡潔的 recurrent 形式線性序列建模方法（對標最基礎的 linear attention）按照下面公式做 RNN 更新：

這里，我們不得不提一下，各種各樣的 Gate 形式（前面的）和更新規則（  右邊的）就是在魔改上面的一行公式，各種具體形式如下表：

（各種方法本身有不同的符號，像 Mamba, HGRN 就不用 q k v，這里為了統一對比全部對標到 linear attention 形式。其中Titans的形式，把 memory update rule 看作 optimzier update 的話，最核心的還是 SGD 形式，暫時忽略momentum/weight decay ，只一個公式表達的話寫成這種梯度更新的形式是合理的。）

其實這些方法又可以進一步細分為不同類別（很多地方都粗略的統一稱為 linear RNN 或者 RNN），這里論文暫時沒提：

• Linear Attention, Lightning Attention, RetNet, GLA, DeltaNet, Gated DeltaNet 屬于 linear attention 類
• Mamba2 屬于 SSM 類，HGRN2 屬于 linear RNN 類
• TTT, Titans 屬于 Test-Time Training 類

Mixture-of-Memories

MoM 思路非常簡單，和 MoE 一樣按照 token 分發，通過 router 為每個 token 選擇 topk 的 memories 并計算各自權重：

所有激活的 topk memories 按照各自權重加權求和得到一份混合記憶：

然后就又回到了 linear 類方法一貫的輸出計算：

另外，這里我們額外引入了 shared memory 的概念，即每個 token 都會經過這個永遠激活的 memory，有利于模型獲取全局信息。相對而言，其他稀疏激活的 memory 更擅長獲取局部信息。消融實驗表明，shared memory 的存在對模型效果有明確的積極作用。

硬件高效實現

MoM的硬件高效Triton算子可以很方便地實現，其輸出的計算可以簡單寫作：

也就是說 MoM 中每個 memory 的計算過程可以復用現有的單個算子，再把所有 memory 的輸出加權求和起來。和直接在算子內先求和再算輸出是數學等價的。

實驗結果

in-context recall-instensive tasks

一直以來，線性序列建模方法因為自身非常有限的 memory 大小，在這類 in-context recall-intensive 任務上表現不好。同時 Transformer 模型得益于其強大的無損記憶能力，非常擅長這類任務。所以已經出現了各種層間 hybrid 的模型，來提升 linear 類模型在這類任務上的效果。

我們首先重點測試了這類任務（結果見下表），使用 Gated DeltaNet 作為 MoM 的 memory 計算形式（在 Memory 更新過程中，每個 memory 都使用 Gated DeltaNet 的 gate 和更新規則），總共 4 個 local sparse memory，激活 2 個，還有一個 shared memory。其中標 ? 的模型來自開源項目（https://huggingface.co/fla-hub），沒標 ?的是我們從頭預訓練的模型。

結果還算相當不錯，在沒有數據污染或任何套路的情況下，結果顯示 MoM 就是單純地效果好。這也和預期一致，翻倍擴展 memory 大小，效果好過其他 linear 類方法。有一些意外的是，在 1.3B 的結果里，MoM 基本可以和 Transformer 相媲美。

其他評測效果

其他評測結果效果也不錯：

推理效率

推理效率是線性序列建模方法的重點，結果顯示 MoM 在常數級復雜度推理速度和顯存占用方面，表現出強大的優勢。

消融實驗

Loss 曲線