DeepSeek開源大模型「記憶」模塊，梁文鋒署名新論文，下一代稀疏模型提前劇透

來源：機器之心

就在十幾個小時前，DeepSeek 發(fā)布了一篇新論文，主題為《Conditional Memory via Scalable Lookup:A New Axis of Sparsity for Large Language Models》，與北京大學合作完成，作者中同樣有梁文鋒署名。

• 論文地址：https://github.com/deepseek-ai/Engram/blob/main/Engram_paper.pdf

簡單總結(jié)一波這項新研究要解決的問題：目前大語言模型主要通過混合專家（MoE）來實現(xiàn)稀疏化，這被稱為「條件計算」。但是，現(xiàn)有的 Transformer 缺少原生的知識查找機制，只能被迫通過計算過程低效地模擬檢索行為。

針對這一現(xiàn)狀，DeepSeek 提出了條件記憶（conditional memory），從而與 MoE 的條件計算互補，并通過引入一個新模塊 Engram 來實現(xiàn)。

目前，模塊「Engram」相關的實現(xiàn)已經(jīng)上傳到了 GitHub。

項目地址：https://github.com/deepseek-ai/Engram

這讓網(wǎng)友們感慨：「DeepSeek is back！」

此外，結(jié)合元旦期間公布的研究《mHC:Manifold-ConstrainedHyper-Connections》，我們可以明確的是 DeepSeek v4 的模樣愈發(fā)清晰，就等上新了！

除了條件計算（MoE），

LLM 還需要一個獨立的條件記憶 Engram

MoE 模型通過條件計算實現(xiàn)了模型容量的擴展，但現(xiàn)有的 Transformer 架構缺乏原生的知識查找原語，只能通過計算過程低效地模擬檢索行為。

為了解決這一問題，DeepSeek 提出了條件記憶（conditional memory）這一與條件計算互補的稀疏化維度，并通過 Engram 模塊加以實現(xiàn)。Engram 在經(jīng)典 -gram 嵌入的基礎上進行了現(xiàn)代化改造，使其能夠以 O (1) 時間復雜度完成知識查找。

通過形式化提出稀疏性分配問題，DeepSeek 還發(fā)現(xiàn)了一條呈 U 型的擴展規(guī)律，用以刻畫神經(jīng)計算（MoE）與靜態(tài)記憶（Engram）之間的最優(yōu)權衡關系。

在這一規(guī)律的指導下，DeepSeek 將 Engram 擴展至 270 億參數(shù)規(guī)模，并在嚴格等參數(shù)量、等 FLOPs 的條件下，其整體性能顯著優(yōu)于純 MoE 基線模型。

尤為值得注意的是，盡管記憶模塊本身主要被用于提升知識檢索能力（如 MMLU 提升 +3.4、CMMLU 提升 +4.0），但 DeepSeek 觀察到其在通用推理能力（如 BBH 提升 +5.0、ARC-Challenge 提升 +3.7）以及代碼與數(shù)學推理任務（HumanEval 提升 +3.0、MATH 提升 +2.4）上帶來了更為顯著的增益。

進一步的分析表明，Engram 能夠?qū)㈧o態(tài)知識的重建負擔從模型的淺層中剝離出來，從而有效加深網(wǎng)絡用于復雜推理的有效深度。此外，通過將局部依賴關系交由查表機制處理，Engram 釋放了注意力機制的容量，使其能夠更專注于全局上下文建模，從而顯著提升了長上下文檢索能力（例如 Multi-Query NIAH 的準確率從 84.2 提升至 97.0）。

最后，Engram 在系統(tǒng)層面同樣展現(xiàn)出基礎設施感知的高效性：其確定性的尋址方式支持在運行時從主機內(nèi)存進行預取，幾乎不會帶來額外的性能開銷。

DeepSeek 認為，條件記憶將成為下一代稀疏大模型中不可或缺的核心建模原語。

Engram 架構如下，其設計目標是在結(jié)構上將靜態(tài)模式存儲與動態(tài)計算過程從 Transformer 主干網(wǎng)絡中分離出來，從而對其進行增強。該模塊對序列中每一個位置依次執(zhí)行兩個功能階段：檢索與融合。

在運行過程中，DeepSeek 首先對當前位置的后綴 N-gram 進行提取與壓縮，并通過哈希機制以確定性的方式檢索對應的靜態(tài)嵌入向量。隨后，這些被檢索到的嵌入會在當前隱藏狀態(tài)的調(diào)制下進行動態(tài)調(diào)整，并進一步通過一個輕量級卷積操作加以精煉。最后，Engram 與多分支架構進行集成。

基于哈希 -gram 的稀疏檢索

這一階段的目標是將局部上下文映射到靜態(tài)記憶條目，這一過程主要包括分詞器壓縮以及通過確定性哈希機制來檢索對應的嵌入表示。

分詞器壓縮：為了最大化記憶單元的語義密度，DeepSeek 引入了一層詞表投影（vocabulary projection）。為此，他們預先設計了一個映射函數(shù)
，其將原始 token ID 映射為基于文本規(guī)范化等價關系（例如使用 NFKC 規(guī)范化、統(tǒng)一大小寫等）得到的規(guī)范化標識符（canonical identifiers）。在實際應用中，對于一個規(guī)模為 128k 的分詞器，該過程能夠?qū)⒂行г~表規(guī)模縮減約 23%（詳見附錄 C）。

多頭哈希：直接對所有可能的 N-gram 組合空間進行參數(shù)化在計算和存儲上都是不可行的。借鑒 Tito Svenstrup 等（2017）的工作，DeepSeek 采用了一種基于哈希的近似方法。為了降低哈希沖突的影響，對于每一種 N-gram 階數(shù) n，引入 K 個相互獨立的哈希頭。

上下文感知門控

前一階段通過哈希 -gram 從條件記憶中檢索得到的嵌入向量，本質(zhì)上提供的是一種與具體語境無關的靜態(tài)先驗信息。然而，正因為其靜態(tài)屬性，這些嵌入缺乏對當前上下文的自適應能力，并且在實際應用中可能受到哈希沖突或詞項多義性帶來的噪聲干擾。

為此，DeepSeek 在檢索之后引入了一種上下文感知的門控機制，其設計靈感來源于注意力機制。

系統(tǒng)效率：計算與存儲的解耦

在帶有記憶機制的模型中，規(guī)模擴展往往受到 GPU 高帶寬顯存（HBM）容量有限的制約。然而，Engram 所采用的確定性檢索機制天然支持將參數(shù)存儲與計算資源進行解耦。不同于 MoE 依賴運行時隱藏狀態(tài)進行動態(tài)路由，Engram 的檢索索引完全由輸入 token 序列決定。這種可預測性使得針對訓練與推理階段的專門優(yōu)化策略成為可能，如圖 2 所示。

在訓練階段，為容納大規(guī)模嵌入表，DeepSeek 采用標準的模型并行方案，將嵌入表分片分布在多張 GPU 上。在前向傳播過程中，通過 All-to-All 通信原語收集被激活的嵌入行；在反向傳播階段，則將對應梯度分發(fā)回各個分片，從而使總可用記憶容量能夠隨加速器數(shù)量線性擴展。

在推理階段，這種確定性特性進一步支持一種預取–重疊（prefetch-and-overlap）策略。由于在前向計算開始之前即可確定所需訪問的記憶索引，系統(tǒng)能夠通過 PCIe 從容量充足的主機內(nèi)存中異步地預取嵌入向量。為有效掩蔽通信帶來的延遲，Engram 模塊被放置在主干網(wǎng)絡中的特定層級，利用其前序 Transformer 層的計算作為緩沖，從而避免 GPU 計算停頓。

這也要求一種硬件 — 算法協(xié)同設計（hardware–algorithm co-design）：一方面，將 Engram 放置得更深可以拉長用于隱藏通信延遲的計算窗口；另一方面，從建模效果來看，較早地介入以卸載局部模式的重建更為有利。因此，Engram 的最優(yōu)插入位置必須同時滿足建模性能與系統(tǒng)時延兩方面的約束。

此外，自然語言中的 -gram 天然遵循 Zipfian 分布，即少量高頻模式貢獻了絕大多數(shù)的記憶訪問。這一統(tǒng)計特性啟發(fā)研究者可以構建一種多級緩存層次結(jié)構（Multi-Level Cache Hierarchy）：將高頻訪問的嵌入緩存于更快的存儲介質(zhì)中（如 GPU HBM 或主機 DRAM），而將大量低頻的長尾模式存放在容量更大但速度較慢的存儲介質(zhì)中（如 NVMe SSD）。這種分層設計使 Engram 能夠擴展到極大規(guī)模的記憶容量，同時對有效訪問延遲的影響保持在最低水平。

U 型擴展規(guī)律與稀疏性分配

作為「條件記憶」的一種具體實現(xiàn)，Engram 在結(jié)構上與 MoE 專家提供的「條件計算」形成了互補。本節(jié)旨在探究這種二元特性（Duality）的擴展屬性，以及如何最優(yōu)地分配稀疏容量。

具體而言，本項研究由兩個核心問題驅(qū)動：

• 有限約束下的分配：在總參數(shù)量和訓練計算量固定（即等參數(shù)、等 FLOPs）的情況下，應該如何在 MoE 專家與 Engram 嵌入之間劃分稀疏容量？

• 無限記憶范式：考慮到 Engram 具有不隨規(guī)模增長（Non-scaling）的查找開銷，如果放寬記憶預算或進行激進擴展，Engram 自身會表現(xiàn)出怎樣的擴展行為？

首先來看MoE 與 Engram 之間的最優(yōu)分配比例。在計算匹配公式時，DeepSeek 使用以下三個參數(shù)度量來分析這個權衡：

• P_tot：總的可訓練參數(shù)，不包括詞匯嵌入和語言模型頭。

• P_act：每個 token 激活的參數(shù)。這一量度決定了訓練成本（FLOPs）。

• ：不激活的參數(shù)，表示可用于擴大模型大小而不增加計算成本的「自由」參數(shù)預算（例如未選擇的專家或未檢索的嵌入）。

DeepSeek 在每個 FLOPs 預算內(nèi)保持 P_tot 和 P_act 固定，這樣模型具有相同數(shù)量的參數(shù)和相同的每 token FLOPs。對于 MoE，P_act 由選定的 top-k 專家決定，而未選擇的專家的參數(shù)貢獻給 P_sparse。對于 Engram，每個 token 只檢索固定數(shù)量的槽（slots），因此增加嵌入槽的數(shù)量會增加 P_tot，但不會增加每 token 的 FLOPs。

其次是「在無限內(nèi)存模式下的 Engram」。在固定參數(shù)預算下優(yōu)化分配之外，DeepSeek 探索了互補的設置：激進的內(nèi)存擴展。這個研究的動機來自于 Engram 獨特的能力，能夠?qū)⒋鎯εc計算解耦。

DeepSeek 使用一個固定的 MoE 主干，具有 P_tot ≈ 3B 和 P_act = 568M，并訓練了 100B 個 token 以確保收斂。在此基礎上附加了一個 Engram 表，并調(diào)整了槽的數(shù)量 M 從 2.58 × 10? 到 1.0 × 10?（增加最多約 13 億參數(shù)）。

下圖 3（左）揭示了驗證損失與分配比例 之間一致的 U 形關系。值得注意的是，即使 MoE 分配減少到僅 ≈ 40%（即 5.7B 模型為 46 個專家，9.9B 模型為 43 個專家），Engram 模型仍然達到了與純 MoE 基準（ = 100%）相當?shù)男阅堋?/p>

此外，純 MoE 基準證明是次優(yōu)的：將大約 20%-25% 的稀疏參數(shù)預算重新分配給 Engram 獲得最佳性能。定量分析中，在 10B 范圍內(nèi)（ = 6 × 102?），驗證損失從 1.7248（ = 100%）改善到 1.7109，接近 ≈ 80% 時的最優(yōu)值（Δ = 0.0139）。值得注意的是，這一最優(yōu)點的位置在不同的范圍內(nèi)穩(wěn)定（ ≈ 75%-80%），表明在固定稀疏性下，各個規(guī)模之間有一個穩(wěn)健的分配偏好。這一觀察到的 U 形確認了兩種模塊之間的結(jié)構互補性。

圖 3（右）展示了增加內(nèi)存槽數(shù)量會顯著改善驗證損失，并且這一改進在整個范圍內(nèi)持續(xù)穩(wěn)定。該曲線遵循嚴格的冪律（在對數(shù)空間中線性），這表明 Engram 提供了一個可預測的擴展旋鈕：更大的內(nèi)存在不需要額外計算的情況下繼續(xù)帶來收益。

關鍵一點是，在擴展效率方面：雖然 OverEncoding 通過更大的內(nèi)存表受益，但 Engram 在相同的內(nèi)存預算下釋放了更大的擴展?jié)摿Α?/p>

結(jié)合分配規(guī)律來看，這些結(jié)果驗證了條件記憶作為稀疏容量的獨立、可擴展軸的作用，它補充了 MoE 的條件計算。

實驗結(jié)果

通過提出的 Engram 架構以及經(jīng)驗推導出的分配法則，DeepSeek 將 Engram 擴展至數(shù)十億參數(shù)規(guī)模，以驗證其在真實語言模型預訓練中的有效性。

總共訓練了以下四種模型：

• Dense-4B（總參數(shù)量 41 億），

• MoE-27B（總參數(shù)量 267 億），

• Engram-27B（總參數(shù)量 267 億），

• 以及 Engram-40B（總參數(shù)量 395 億）。

所有模型均采用完全相同的數(shù)據(jù)訓練流程（相同的 token 預算及順序），且在激活參數(shù)量上嚴格匹配。

關于實驗設置，所有模型均在包含 2620 億 token 的語料庫上進行預訓練，并采用了 DeepSeek-v3 的分詞器，其詞表大小為 128k。DeepSeek 在涵蓋語言建模、知識、推理、閱讀理解以及代碼 / 數(shù)學的多樣化基準測試集上對模型進行評估。對于每項基準測試，均遵循標準的提示詞協(xié)議和評估指標。

先來看大規(guī)模預訓練的實驗結(jié)果，如下表 1 所示，稀疏架構展示了比密集模型更優(yōu)的擴展規(guī)律。

在相同的訓練計算預算下，所有三種稀疏變體（MoE-27B，Engram-27B/40B）在所有基準測試中顯著超越了 iso-FLOPs 的 Dense-4B 基準。

更重要的是，Engram-27B 在 iso - 參數(shù)和 iso-FLOPs 的 MoE-27B 基準上持續(xù)取得改進。有趣的是，這些提升并不限于知識密集型任務（例如，MMLU: +3.0，MMLU-Pro: +1.8，CMMLU: +4.0），在這些任務中，內(nèi)存容量直觀上是有益的。此外還觀察到，在一般推理領域（例如，BBH: +5.0，ARC-Challenge: +3.7，DROP: +3.3）以及代碼和數(shù)學推理任務（例如，HumanEval: +3.0，MBPP: +1.6，GSM8K: +2.2，MATH: +2.4）中，改進更加顯著。

擴展到 Engram-40B 進一步減少了預訓練損失，并提高了大多數(shù)基準測試的性能。盡管它尚未在每個任務上嚴格超越 Engram-27B，但這可能是由于訓練不足的結(jié)果。此外，Engram-40B 與基準模型之間的訓練損失差距在訓練結(jié)束時繼續(xù)擴大，表明擴展的內(nèi)存容量尚未在當前的 token 預算內(nèi)完全飽和。

接下來是長上下文訓練。通過將局部依賴建模卸載至靜態(tài)查找，Engram 架構為處理全局上下文保留了寶貴的注意力容量。DeepSeek 通過進行長文本擴展訓練，對這一結(jié)構性優(yōu)勢進行了實驗驗證。通過采用嚴密的評估協(xié)議，將架構設計帶來的貢獻與基礎模型本身的能力剝離開來，證明了 Engram 在長程檢索和推理任務中帶來了顯著的性能增益。

DeepSeek 首先解耦基礎模型能力與架構設計之間的影響，其次進行受控對照分析，結(jié)果如下表 2 所示，主要得出了以下兩個結(jié)論：

一是超越注意力機制的長文本能力。

雖然注意力機制和位置編碼為上下文處理提供了結(jié)構基礎，但實驗結(jié)果表明，長文本性能并非僅由架構先驗決定。通過觀察 Engram 的演進軌跡（從 41k 步到 50k 步），即使在控制相同模型架構和固定長文本擴展階段計算預算的前提下，長文本性能仍隨預訓練進程單調(diào)提升。這表明長文本性能與基礎模型的通用建模能力存在內(nèi)在耦合。因此，嚴謹?shù)募軜媽Ρ缺仨毻ㄟ^對齊「基礎模型損失（Loss）」而非僅僅對齊「訓練步數(shù)」來控制這一混淆變量。

二是受控設置下的架構優(yōu)越性。

基于上述原則，DeepSeek 將 Engram 與 MoE 基準模型進行了對比測試。在控制基礎能力的前提下，Engram 模塊的效率增益變得十分顯著：

• 等損耗設置（Iso-Loss Setting，41k 步 vs. 基準）：該設置嚴格分離了架構效率的影響。當對比 Engram-27B（46k 步）與完整訓練的 MoE-27B（50k 步），即預訓練損失完全對齊的兩個模型時，Engram 表現(xiàn)出顯著增益。具體而言，它在復雜檢索任務中大幅超越基準模型（例如，多查詢「大海撈針」 NIAH：97.0 vs. 84.2；變量跟蹤 VT：87.2 vs. 77.0）。

• 等計算量設置（Iso-FLOPs Setting，50k 步 vs. 基準）：在標準的等計算預算下，Engram-27B（50k 步）進一步拉大了差距，在所有指標上均實現(xiàn)了頂尖性能。

• 極端設置（約 82% 計算量）：即使是提前停止訓練的 Engram-27B（41k 步），在面對完整訓練的 MoE-27B（50k 步）時依然極具競爭力。它在 LongPPL 指標上與基準持平，并在 RULER 測試中實現(xiàn)超越，這充分證明了 Engram 架構的內(nèi)在優(yōu)越性。

最后，下圖 4 是對表示對齊與收斂速度的分析。(a) 基于 LogitLens 的逐層 KL 散度分析。在模型淺層，KL 散度持續(xù)保持在較低水平，這表明 Engram 加速了預測的收斂。(b-c) 為基于 CKA 計算的相似度熱力圖。高相似度對角線顯著的向上偏移表明，Engram 的淺層在功能上等效于 MoE 模型的深層，從而有效地增加了模型的深度。

更多細節(jié)請參考原論文。

閱讀最新前沿科技趨勢報告，請訪問歐米伽研究所的“未來知識庫”

https://wx.zsxq.com/group/454854145828

未來知識庫是“ 歐米伽 未來研究所”建立的在線知識庫平臺，收藏的資料范圍包括人工智能、腦科學、互聯(lián)網(wǎng)、超級智能，數(shù)智大腦、能源、軍事、經(jīng)濟、人類風險等等領域的前沿進展與未來趨勢。目前擁有超過8000篇重要資料。每周更新不少于100篇世界范圍最新研究資料。 歡迎掃描二維碼或訪問https://wx.zsxq.com/group/454854145828進入。