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

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

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

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

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

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

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

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

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

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

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

這一替換，使 FFN 的結(jié)構(gòu)發(fā)生了一個本質(zhì)變化：「記憶容量」第一次與「每 token 計算量」被明確解耦。

結(jié)構(gòu)變化雖小，卻帶來一系列連鎖效應

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

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

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

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

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

2. 訓練穩(wěn)定性：靜態(tài)稀疏優(yōu)于動態(tài)路由

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

3. 更高效的「記憶空間」：embedding 幾何結(jié)構(gòu)的改變

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

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

4. 計算與系統(tǒng)效率：FLOPs 與 IO 的雙重下降

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

實驗結(jié)果：規(guī)模、任務與長上下文的系統(tǒng)驗證

在 350M 與 1B 兩個模型規(guī)模上，STEM 被系統(tǒng)性地與 dense baseline 進行對比。

結(jié)果顯示：

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

工程視角：如何真正落地

論文也對 STEM 的系統(tǒng)實現(xiàn)給出了較為完整的討論：

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

結(jié)語

STEM 展示了一種值得關注的趨勢：通過改變參數(shù)的「組織方式」，而非一味增加規(guī)?；蛴嬎?，模型同樣可以獲得顯著能力提升。

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

作者介紹

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

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

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

曹晟（Rick Cao）為 Meta AI 研究員，主要研究大模型系統(tǒng)優(yōu)化與高效推理架構(gòu)設計，關注大規(guī)模模型在真實系統(tǒng)環(huán)境中的部署與加速問題。

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