早于DeepSeek Engram！用「查表」重置Transformer記憶 | ICLR

新智元報道

編輯：LRST

【新智元導讀】ICLR論文STEM架構率先提出「查表式記憶」架構，早于DeepSeek Engram三個月。它將Transformer的FFN從動態計算改為靜態查表，用token索引的embedding表直接讀取記憶，徹底解耦記憶容量與計算開銷。

近年來，隨著大模型規模與知識密度的持續爆發，研究人員開始重新審視一個底層問題：模型的參數究竟該如何組織，才能最高效地承擔「記憶」的功能？

在傳統的Transformer架構中，前饋神經網絡（FFN）的知識通常隱式地埋藏在up-projection等密集矩陣內 。這種通過輸入進行動態激活的矩陣乘法，雖然保證了表達能力，卻在參數的可尋址性、后期可編輯性以及系統計算效率上存在著天然的局限 。

為了突破這一瓶頸，學術界和工業界逐漸轉向更離散、更結構化的參數組織路徑。

近期DeepSeek推出的engram機制成功引爆了業內對「查表式記憶（lookup-based memory）」的關注 。但令人矚目的是，早于engram問世約三個月前，一篇入選 ICLR 的論文就已經對該方向進行了極其系統的探索 。

項目主頁: https://infini-ai-lab.github.io/STEM/

與傳統的混合專家模型（MoE）在現有稀疏路由上做修補不同，STEM（Scaling Transformers with Embedding Modules）選擇直接對 FFN 結構「動刀」：它摒棄了動態運行時的路由機制，將 up-projection 替換為按token索引的層級 embedding 表，以一種純靜態的方式重構了 Transformer 的記憶訪問路徑 。

從「算地址」到「查地址」

如果用「鍵值對記憶（key-value memory）」的視角來審視標準 Transformer，像 SwiGLU 這樣的 FFN 結構，本質上是通過一次 up-projection 將輸入映射到高維空間，從而生成一個能被 gate 調制的「地址向量」 。這一過程極其依賴輸入相關的密集矩陣乘法，不僅計算昂貴，而且參數高度耦合 。

STEM 團隊提出了一個靈魂拷問：如果 FFN 的核心作用只是「按token訪問記憶」，我們真的需要每次都動態計算這些地址向量嗎？

基于此，STEM給出了一種極致簡單直接的解法：

• 徹底移除up-projection，不再動態計算地址向量 。

• 為模型的每一層單獨維護一個按token索引的embedding表。

• 在前向傳播時，直接根據token id 「查表」，提取對應的靜態向量 作為原先的 。

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

這一看似輕量的模塊替換，實現了一個極其本質的架構跨越：模型的「記憶容量」終于與「單token的計算量」實現了徹底解耦。

連鎖效應

四大維度的全面躍升

雖然僅僅替換了FFN的一個子模塊，STEM 卻在實驗中展現出了驚人的全方位優勢 ：

1. 即插即用的「知識編輯」

這是STEM最硬核的特性之一 。因為每一層的embedding都與特定token id強綁定，研究人員甚至不需要重新訓練，只需替換特定token的STEM向量，就能直接修改模型輸出的事實 。

例如，僅通過互換「Spain」與「Germany」的向量，模型在回答首都問題時就會發生相應的改變 。這為未來的模型內部機制理解與知識編輯打開了全新大門 。

2. 訓練極度穩定（告別動態路由的煩惱）

與依賴運行時路由的MoE不同，STEM是一種靜態稀疏架構 。由于每個token在每一層訪問的 embedding 都是恒定確定的，它完美避開了MoE訓練中令人頭疼的負載傾斜（load skew）和損失突刺（loss spike）問題，且不需要任何all-to-all通信 。

3. 更寬廣的「記憶空間」

從幾何空間分布來看，STEM 的 embedding 表展現出了更大的角度散布（large angular spread） 。這意味著不同token 的向量更趨近于正交，大幅減少了參數間的相互干擾（cross-talk） 。在同等算力下，模型能塞下更多「可尋址的記憶槽位」 。

4. 計算與I/O雙重減負

砍掉up-projection后，每一層都能省下龐大的矩陣乘法開銷（約級別） 。更妙的是，龐大的embedding表完全可以離載（offload）到 CPU 內存中，配合異步預取（prefetch）和緩存策略高效運行 。

實驗與落地

長上下文表現亮眼

團隊在350M和1B規模的模型上對 STEM 進行了嚴密的消融實驗 。數據表明，STEM 相比于 dense 架構基線，整體平均性能提升了約3–4%，在部分知識密集型任務上，提升幅度甚至飆升至9–10% 。特別是在大海撈針（Needle-in-a-Haystack）和LongBench等長文本評測中，上下文越長，STEM的優勢就越顯著。

對于工程落地，論文也給出了避坑指南：

• 替換講究位置：核心在于替換up-projection，如果盲目替換gate-projection，反而會破壞模型的上下文調制能力 。

• 優化存儲與顯存：embedding表可放在CPU，但在訓練時需注意將梯度寫回對應的優化器狀態 。在追求極致性價比時，還可以采用「部分層替換」或混合變體策略來平衡顯存壓力 。

結語

STEM架構向我們清晰地傳達了一個信號：在無腦堆疊算力和參數量之外，通過巧妙重構參數的「組織方式」，我們依然能夠榨取巨大的性能紅利。 在當前基座大模型越發龐大復雜的語境下，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 等強化學習系統研發，并關注基礎模型的推理與決策能力。

參考資料：

https://infini-ai-lab.github.io/STEM/