新智元報(bào)道

編輯：LRST

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

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

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

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

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

項(xiàng)目主頁: https://infini-ai-lab.github.io/STEM/

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

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

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

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

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

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

• 為模型的每一層單獨(dú)維護(hù)一個(gè)按token索引的embedding表。

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

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

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

連鎖效應(yīng)

四大維度的全面躍升

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

1. 即插即用的「知識(shí)編輯」

這是STEM最硬核的特性之一 。因?yàn)槊恳粚拥膃mbedding都與特定token id強(qiáng)綁定，研究人員甚至不需要重新訓(xùn)練，只需替換特定token的STEM向量，就能直接修改模型輸出的事實(shí) 。

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

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

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

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

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

4. 計(jì)算與I/O雙重減負(fù)

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

實(shí)驗(yàn)與落地

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

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

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

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

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

結(jié)語

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

作者介紹

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

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

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

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

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

參考資料：

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