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

新智元報道

編輯：LRST

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

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

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

為了突破這一瓶頸，學(xué)術(shù)界和工業(yè)界逐漸轉(zhuǎn)向更離散、更結(jié)構(gòu)化的參數(shù)組織路徑。

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

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

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

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

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

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

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

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

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

• 在前向傳播時，直接根據(jù)token id 「查表」，提取對應(yīng)的靜態(tài)向量 作為原先的 。

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

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

連鎖效應(yīng)

四大維度的全面躍升

雖然僅僅替換了FFN的一個子模塊，STEM 卻在實驗中展現(xiàn)出了驚人的全方位優(yōu)勢 ：

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

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

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

2. 訓(xùn)練極度穩(wěn)定（告別動態(tài)路由的煩惱）

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

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

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

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

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

實驗與落地

長上下文表現(xiàn)亮眼

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

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

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

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

結(jié)語

STEM架構(gòu)向我們清晰地傳達了一個信號：在無腦堆疊算力和參數(shù)量之外，通過巧妙重構(gòu)參數(shù)的「組織方式」，我們依然能夠榨取巨大的性能紅利。 在當前基座大模型越發(fā)龐大復(fù)雜的語境下，STEM這種簡潔、優(yōu)雅且工程友好的設(shè)計，無疑是下一代模型演進路線上的一座重要燈塔

作者介紹

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

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

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

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

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

參考資料：

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