
DeepSeek 節前開始蓄力!
最新論文直接給 Transformer 加上“條件記憶”(Conditional Memory),補上了原生缺乏的知識查找機制。
結論中明寫道:我們將條件記憶視為下一代稀疏模型不可或缺的建模原語。
?
還是梁文鋒署名,并與北京大學王選所趙東巖、張輝帥團隊合作。
?
論文中不僅提出了條件記憶這個全新范式,并給出了具體實現方案 Engram 模塊,實驗中讓 27B 參數碾壓同規模純 MoE 模型,甚至變相提升了大模型的推理能力:
讓原來 Transformer 要用 6 層注意力才能干的簡單任務壓縮到 1-2 層搞定,省出來的資源就可以用于更難的推理任務了。
條件記憶的原理其實也非常“原始”:不靠計算,回歸查表,用上了傳統 N-gram 方法。
給大模型一個巨大的詞表,專門存那些固定的實體名稱和兩三個詞的短語,不管詞表多大,找信息都是 O (1) 速度。
關鍵就在于,如此前大模型時代的玩法,DeepSeek 如何解決傳統 N-gram 模型存儲爆炸和多義性問題,又是讓它和現代 Transformer 結合起來的?
讓注意力干“苦力活”太浪費了
團隊的核心觀察是,語言建模其實包含兩種性質完全不同的任務,一種是需要深度動態計算的組合推理,另一種則是檢索靜態知識。
問題在于,現有的 Transformer 架構缺乏原生的知識查找機制。
當模型需要識別一個實體時,它得消耗好幾層注意力和前饋網絡,逐層拼湊特征,最終才能完成。
論文中引用了一個具體案例:”Diana, Princess of Wales”
模型需要經過 6 層才能完成這個識別過程,前幾層還在糾結”Wales 是英國的一個地區”、”Princess of Wales 是某種頭銜”這些中間狀態,最終才能“想起來”這是指戴安娜王妃。
?
本質上是在用昂貴的運行時計算來重建一個靜態查找表,那些本可以用于更高層推理的網絡深度,被浪費在了識別概念這種“苦力活”上。
回歸查表,回歸 N-gram
Engram 的設計思路相當直接:既然經典的 N-gram 模型就能用 O (1) 的時間復雜度捕獲這些局部依賴,那為什么不把這個能力直接嵌入 Transformer?
具體實現上,團隊在原有的 Transformer 層之間插入 Engram 模塊。每個位置的輸入會觸發一次哈希查找:把當前 token 和前面幾個 token 組成的 N-gram 映射到一個巨大的嵌入表中,直接取出對應的向量。
?
為了處理哈希沖突和多義性問題,團隊引入了上下文感知的門控機制,用當前的隱藏狀態作為 Query,檢索到的記憶作為 Key 和 Value,計算一個 0 到 1 之間的標量門控值。
如果檢索到的內容和當前上下文不匹配,門控值就趨近于零,相當于自動屏蔽噪聲。
下圖中,顏色越深說明 Engram 越判斷當前文本片段是“固定靜態模式”,傾向于調用記憶庫中的對應信息。
顏色越淺代表這段文本越動態靈活,主要靠模型的注意力機制處理。
比如只看到“張”是一個常見姓氏,但是“張仲景”三個字湊一起就是固定歷史人物實體了。
?
接下來還要解決傳統 N-gram 模型的兩個痛點。
●
語義重復,同一個詞的不同形式(比如 Apple、apple、?pple)被當成不同 token,浪費存儲。
●
存儲爆炸,所有可能的 N-gram(比如 2 詞、3 詞組合)數量太多,比如 128k 詞表就要存 128k^3 種組合,直接存儲根本存不下。
DeepSeek 團隊首先壓縮 tokenizer,把語義相同但形式不同的 token 歸為一類,128k 詞表的有效規模直接減少 23%,相同語義的 token 聚在一起,查找更高效。
再用多個哈希函數把 N-gram 映射成 embedding 表的索引,
這既解決了存儲爆炸:不管有多少種 N-gram,都通過哈希函數映射到一個固定大小的 embedding 表里,表的大小是質數。
又減少查找沖突:給每種 N-gram 階數(比如 2-gram、3-gram)配 K 個不同的哈希頭,每個哈希頭對應一個獨立的 embedding 表,把所有 N-gram 階數、所有哈希頭取出來的 embedding 向量拼在一起,形成最終的“記憶向量”e?,供后續模塊使用。
?
?
U 型曲線:MoE 和記憶的最優配比
論文最核心的部分是對”稀疏性分配問題”的系統研究。
團隊設計了一個嚴格的實驗框架:固定總參數量和每 token 的激活參數量(也就是計算量),然后在 MoE 專家和 Engram 記憶之間重新分配”閑置參數”預算。
分配比例 ρ 從 100%(純 MoE)逐步降到 40%,實驗結果畫出了一條清晰的 U 型曲線:
?
純 MoE 反而不是最優解,把大約 20% 到 25% 的稀疏參數預算分給 Engram 記憶時,模型驗證集 loss 達到最低點。
在 100 億參數規模下,最優配置比純 MoE 基線的 loss 降低了 0.0139。
更重要的是,這個最優分配點在不同計算預算下都相當穩定,大約在 ρ=75% 到 80% 之間。
團隊解釋了 U 型曲線兩端的含義:
●
MoE 主導時,模型缺乏靜態模式的專用記憶,被迫通過網絡深度和大量計算來低效重建。
●
Engram 主導時,模型丟失了條件計算能力,在需要動態推理的任務上表現下降。
總之,記憶無法替代計算,計算也無法高效模擬記憶。
27B 規模驗證:推理能力提升超預期
按照 U 型曲線的指導,團隊把 Engram 擴展到更大參數規模進行驗證,并對比純 MoE 模型和純密集模型。
所有模型訓練條件一致,激活參數量都是 38 億,訓練 token 都是 2620 億,差異僅在“稀疏能力分配”。
●
Dense-4B:純密集模型。
●
MoE-27B:純混合專家模型,72 個路由專家 +2 個共享專家,所有稀疏參數都給 MoE。
●
Engram-27B:MoE+Engram 混合模型,55 個路由專家 +2 個共享專家,把 5.7B 稀疏參數分配給 Engram 記憶模塊。
●
Engram-40B:進一步擴展 Engram 模塊,保持專家數量不變,Engram 記憶參數增至 18.5B,總參數 39.5B。
?
結果 MoE-27B 和 Engram-27B 對比,知識密集型任務的提升在預期之內:比如 MMLU 提升 3 分,CMMLU 提升 4.0 分,TriviaQA 提升 1.9 分。
但出乎意料的是,通用推理和代碼數學領域的提升幅度也很大:BBH 大幅提升 5.0 分,ARC-Challenge 提升 3.7 分,DROP 提升 3.3 分,HumanEval 提升 3.0 分,MATH 提升 2.4 分,GSM8K 提升 2.2 分。
?
團隊用 LogitLens 和 CKA 分析揭示了原因。
Engram 讓模型的早期層不再需要做特征組合的“苦力活”,KL 散度曲線顯示 Engram 模型的預測收斂速度明顯更快。更直觀的證據來自 CKA 相似度矩陣,Engram-27B 第 5 層的表征,和 MoE 基線第 12 層的表征最為相似。
這意味著 Engram 實際上“加深”了網絡的有效深度,省下來的層數被用于更復雜的推理任務。
?
Engram-40B 進一步增加記憶參數后,大部分任務性能持續提升,且訓練后期損失仍在下降,說明記憶容量還未飽和,后續可繼續擴大。
另外長上下文場景的提升尤為顯著。
在 RULER 測試集上,Multi-Query NIAH 從 84.2 躍升到 97.0,Variable Tracking 從 77.0 提升到 89.0。
?
論文解釋說,Engram 把局部依賴建模卸載給了查找操作,釋放了注意力容量去關注全局上下文。
百億參數表放 CPU 上,延遲幾乎沒影響
接下來又到了喜聞樂見的軟硬結合工程優化環節。
在訓練階段,詞表規模會高達 100B 參數,單個 GPU 存不下,必須拆分到多個 GPU 上,需要 All-to-All 通信機制,讓所有 GPU 之間互相傳遞需要的記憶片段。
在推理階段把詞表卸載到 CPU 內存,同時又不能讓記憶調用拖慢計算節奏。
?
和 MoE 的動態路由不同,Engram 的查找索引只取決于輸入 token 序列,完全可以提前計算。
這個確定性讓團隊能夠把巨大的嵌入表放到 CPU 內存里,用 PCIe 異步預取,讓通信和前面層的計算重疊。
具體通過把 Engram 模塊插在 Transformer 網絡的特定層,GPU 計算前一層的同時,CPU 預取當前層需要的 Engram 記憶,等 GPU 算完前一層,所需的記憶也已經傳輸到位。
實驗直接把一個 1000 億參數的 Engram 表放到 CPU 內存,在 H800 上跑推理。4B 密集模型的吞吐量從 9031 token/s 降到 8858 token/s,8B Dense 模型從 6315 token/s 降到 6140 token/s,額外開銷都在 3% 以內。
?
自然語言 N-gram 天然遵循 Zipfian 分布,極少數高頻模式占據絕大多數訪問量。這意味著可以設計多級緩存:高頻嵌入放 GPU 顯存,中頻放 CPU 內存,長尾放 NVMe SSD,把有效延遲進一步壓縮。
DeepSeek 團隊在結論中寫道:
Engram 將“硬件感知效率”確立為核心設計原則:其確定性尋址機制支持存儲與計算的解耦,能夠將海量參數表卸載至主機內存,且推理開銷可忽略不計。我們認為,條件記憶將成為下一代稀疏模型中不可或缺的建模基元。
DeepSeek 的下一代稀疏模型,已被曝光將在春節前發布,敬請期待。