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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

STEM 的一個(gè)尤為引人注目的特性，是其顯式的 token–參數(shù)對(duì)應(yīng)關(guān)系。

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)結(jié)果：規(guī)模、任務(wù)與長(zhǎng)上下文的系統(tǒng)驗(yàn)證

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

結(jié)果顯示：

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

工程視角：如何真正落地

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

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

結(jié)語

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

在當(dāng)前大模型架構(gòu)逐漸走向復(fù)雜化的背景下，這種簡(jiǎn)潔、穩(wěn)定且工程友好的設(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)方向包括長(zhǎng)上下文多模態(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é)家，長(zhǎng)期從事強(qiáng)化學(xué)習(xí)與大模型研究，曾參與 AlphaZero 等強(qiáng)化學(xué)習(xí)系統(tǒng)研發(fā)，并關(guān)注基礎(chǔ)模型的推理與決策能力。