ICLR 2026｜早于DeepSeek Engram，STEM已重構Transformer記憶

近年來，隨著大語言模型規模與知識密度不斷提升，研究者開始重新思考一個更本質的問題：模型中的參數應如何被組織，才能更高效地充當「記憶」。

在標準 Transformer 的 FFN 中，知識主要隱式存儲在 up-projection 等密集矩陣里，并通過輸入相關的矩陣乘法被動態激活，這種方式在表達力上有效，但在參數可尋址性、可編輯性與系統效率上存在天然限制。

圍繞這一問題，學術界與工業界逐漸將目光投向更離散、更結構化的參數組織方式。以 DeepSeek 提出的 engram 為代表，近期多項工作重新喚起了對「查表式記憶（lookup-based memory）」的關注。而有意思的是，在 engram 發布之前約三個月，這篇 ICLR paper 也在這一方向進行系統探索 。

• 論文鏈接：https://arxiv.org/pdf/2601.10639
• 項目主頁：https://infini-ai-lab.github.io/STEM/

與 MoE 相關工作不同，STEM 并非在現有稀疏路由框架上做局部改進，而是直接從 FFN 結構入手，將 up-projection 替換為按 token 索引的層級 embedding 表，以一種靜態、無運行時路由的方式重構 Transformer 中「記憶」的訪問路徑。

從「算地址」到「查地址」：STEM 的核心直覺

如果從 key-value memory view 來看標準 Transformer，SwiGLU 等 FFN 結構通常通過一次 up-projection，將輸入表示映射到高維空間，生成可被 gate 調制的「地址向量」。

這一過程本質上是輸入相關的密集矩陣乘法，既昂貴，又在參數層面高度耦合。

STEM 的關鍵洞察在于：如果 FFN 的作用更接近于「按 token 訪問記憶」，那么這些地址向量是否真的需要由輸入動態生成？

基于這一問題，STEM 選擇了一條極為直接的路徑：不再通過 up-projection 計算地址向量，而是為每一層維護一個按 token 索引的 embedding 表，在前向傳播時直接根據 token id 查表獲取對應向量。

gate 與 down-projection 則被完整保留，用于對查表得到的向量進行上下文調制與壓縮。

這一替換，使 FFN 的結構發生了一個本質變化：「記憶容量」第一次與「每 token 計算量」被明確解耦。

結構變化雖小，卻帶來一系列連鎖效應

盡管 STEM 只替換了 FFN 中的一個子模塊，但實驗結果表明，這一設計在多個維度上產生了顯著影響。

1. 可編輯性：直接修改「知識向量」

STEM 的一個尤為引人注目的特性，是其顯式的 token–參數對應關系。

由于每一層的 embedding 都直接與 token id 綁定，研究者可以在不進行再訓練的情況下，直接替換某個 token 的 STEM 向量，從而改變模型在相關事實上的輸出行為。比如圖中例子展示僅通過替換「Spain」與「Germany」的向量，就可以使模型對首都問題的回答發生對應變化。

這一能力為模型的知識編輯、和對模型的理解提供了一種新思路。

2. 訓練穩定性：靜態稀疏優于動態路由

與 MoE 等依賴運行時路由的方案不同，STEM 屬于一種靜態稀疏結構：每個 token 在每一層所訪問的 embedding 是確定的，不涉及負載均衡或 all-to-all 通信。這種設計避免了 MoE 中常見的 load skew 與 loss spike 問題。

3. 更高效的「記憶空間」：embedding 幾何結構的改變

從表示空間的角度看，用 token-indexed embedding 替代 up-projection，會顯著改變 FFN 輸入向量的幾何分布。STEM 中的 embedding 表呈現出更大的角度散布（large angular spread），不同 token 的向量更加接近正交，從而有效減少了參數之間的相互干擾（cross-talk）。

這一性質意味著，在相同甚至更低的計算成本下，模型能夠維持更多「可尋址的記憶槽位」，對知識存儲和檢索尤為有利。

4. 計算與系統效率：FLOPs 與 IO 的雙重下降

從計算復雜度上看，移除 up-projection 后，每一層可節省約 d?d_ff 級別的矩陣乘法開銷。更重要的是，大規模的 embedding 表可以被離載到 CPU，并通過異步 prefetch 與緩存策略進行訪問。

實驗結果：規模、任務與長上下文的系統驗證

在 350M 與 1B 兩個模型規模上，STEM 被系統性地與 dense baseline 進行對比。

結果顯示：

• 整體平均性能提升約3–4%，部分知識類任務提升幅度可達9–10%；
• 在 Needle-in-a-Haystack 與 LongBench 等長上下文評測中，STEM 的優勢還會隨著上下文長度的增加而擴大。

工程視角：如何真正落地

論文也對 STEM 的系統實現給出了較為完整的討論：

• 替換位置：僅替換 up-projection 是關鍵，盲目替換 gate-projection 會破壞上下文調制能力；
• 存儲策略：embedding 表可 CPU 離載，訓練時需要將梯度回寫到對應的 optimizer state；
• 折衷方案：在參數占比與顯存壓力之間，可采用部分層替換或混合變體進行平衡。

結語

STEM 展示了一種值得關注的趨勢：通過改變參數的「組織方式」，而非一味增加規模或計算，模型同樣可以獲得顯著能力提升。

在當前大模型架構逐漸走向復雜化的背景下，這種簡潔、穩定且工程友好的設計思路，或許正是下一階段基座模型演進中不可忽視的一條路徑。

作者介紹

本文第一作者 Ranajoy Sadhukhan 為卡內基梅隆大學（CMU）InfiniAI Lab 博士生，師從陳貝迪教授。該工作完成于其在 Meta AI 實習期間，實習導師包括劉澤春、曹晟（Rick Cao）與田淵棟等研究人員。

InfiniAI Lab 由陳貝迪教授創立，致力于模型、系統與硬件協同設計，研究高效且可擴展的 AI 算法與系統，重點方向包括長上下文多模態建模、突破傳統 scaling laws 的新一代模型架構，以及基礎模型的理解與推理能力增強，同時推動算法與系統層面的效率優化，以促進 AI 技術的普及化。

劉澤春為 Meta AI 研究科學家，研究方向涵蓋基座模型訓練，大模型壓縮、稀疏化與端側部署優化，專注于模型高效推理與系統協同設計。

曹晟（Rick Cao）為 Meta AI 研究員，主要研究大模型系統優化與高效推理架構設計，關注大規模模型在真實系統環境中的部署與加速問題。

田淵棟為 Meta AI 資深研究科學家，長期從事強化學習與大模型研究，曾參與 AlphaZero 等強化學習系統研發，并關注基礎模型的推理與決策能力。