通過構(gòu)建稀疏性分配問題，團(tuán)隊發(fā)現(xiàn)了一個 U 形擴(kuò)展規(guī)律，在該規(guī)律的指導(dǎo)下將 Engram 擴(kuò)展到 270 億個參數(shù)，其性能優(yōu)于嚴(yán)格等參數(shù)、等浮點運算次數(shù)的 MoE 基線模型。
值得注意的是，記憶模塊在通用推理和代碼/數(shù)學(xué)領(lǐng)域的提升更為顯著。機(jī)制分析表明，Engram 減輕了骨干網(wǎng)絡(luò)早期層靜態(tài)重構(gòu)的負(fù)擔(dān)，有效加深了網(wǎng)絡(luò)以進(jìn)行復(fù)雜推理。
通過將局部依賴關(guān)系委托給查找操作，Engram 釋放了注意力資源用于全局上下文，從而顯著提升了長上下文檢索能力。
Engram 模塊架構(gòu)
Engram 是一個條件記憶模塊，旨在通過在結(jié)構(gòu)上將靜態(tài)模式存儲與動態(tài)計算分離來增強(qiáng) Transformer 主干網(wǎng)絡(luò)。
團(tuán)隊實現(xiàn)了一個基于規(guī)范化文本等價性的詞匯投影層標(biāo)識符。對于 128k 分詞器，該過程可使有效詞匯量減少 23%。
為了增強(qiáng)表達(dá)能力，他們采用了一種受注意力機(jī)制啟發(fā)的上下文感知門控機(jī)制。針對擴(kuò)展感受野并增強(qiáng)模型的非線性，團(tuán)隊引入了一個短的深度因果卷積。
Engram 系統(tǒng)分為兩個階段：
團(tuán)隊采用了先進(jìn)的多分支架構(gòu)作為默認(rèn)骨干網(wǎng)絡(luò)，該架構(gòu)具有卓越的建模能力。
自然語言的 N-gram 本質(zhì)上遵循齊普夫分布，這種統(tǒng)計特性促使團(tuán)隊采用多級緩存層次結(jié)構(gòu)。這種分層結(jié)構(gòu)使得 Engram 能夠擴(kuò)展到海量內(nèi)存容量，同時將有效延遲的影響降至最低。
大規(guī)模預(yù)訓(xùn)練
基于 Engram 架構(gòu)和經(jīng)驗推導(dǎo)的分配規(guī)則，研究人員將 Engram 擴(kuò)展到數(shù)十億參數(shù)規(guī)模。他們訓(xùn)練了四個模型：
（1）Dense-4B（總參數(shù) 41 億）
（2）MoE-27B（總參數(shù)量為 267 億）（3）Engram-27B（總參數(shù)量為 267 億）（4）Engram-40B（總參數(shù)量為 395 億）
結(jié)果顯示，Engram-27B 僅使用 82% 的預(yù)訓(xùn)練 FLOPs（41k 對比 50k）即可達(dá)到基線模型 LongPPL 的性能，同時在 RULER 任務(wù)上取得了顯著更高的準(zhǔn)確率。
在相同的預(yù)訓(xùn)練損失 (46k) 和相同的預(yù)訓(xùn)練 FLOPs (50k) 設(shè)置下，Engram-27B 在所有指標(biāo)上均顯著優(yōu)于基線模型。
更重要的是，Engram-27B 的性能始終優(yōu)于基于等參數(shù)和等浮點運算次數(shù)的 MoE-27B 基線模型。這些提升并非僅限于知識密集型任務(wù)，而包含了一般推理領(lǐng)域以及代碼和數(shù)學(xué)推理領(lǐng)域。
在長上下文的情景訓(xùn)練中，研究團(tuán)隊發(fā)現(xiàn)長上下文性能并非完全由架構(gòu)先驗決定。即使在上下文擴(kuò)展階段控制相同的模型架構(gòu)和固定的計算預(yù)算，長上下文性能仍然隨著預(yù)訓(xùn)練的進(jìn)行而單調(diào)提升。這表明長上下文性能與基礎(chǔ)模型的一般建模能力密切相關(guān)。因此，嚴(yán)格的架構(gòu)比較必須通過調(diào)整基礎(chǔ)模型的損失來控制這一混淆變量，而不僅僅是調(diào)整訓(xùn)練步驟。
Engram 的內(nèi)部機(jī)制分析
有效深度
研究人員使用 LogitLens分析各層預(yù)測結(jié)果的演變。
通過顯式地訪問外部知識，Engram 減少了所需的計算步驟，從而在網(wǎng)絡(luò)層級結(jié)構(gòu)的早期階段就獲得了高置信度的有效預(yù)測。
通過繞過顯式查找進(jìn)行早期特征合成，Engram 在功能上等同于增加模型的有效深度。
模塊設(shè)置
他們在受控環(huán)境下對 Engram 進(jìn)行消融，以研究每個關(guān)鍵模塊設(shè)計的有效性。
結(jié)果表明，隨著插入點深入，性能逐漸下降。只需一輪注意力機(jī)制即可提供有意義的上下文相關(guān)的 hfor 門控，同時還能及時替換骨干網(wǎng)底層的局部聚合。
參數(shù)敏感性
研究人員通過在推理過程中完全抑制稀疏嵌入輸出來評估模型。
基準(zhǔn)測試結(jié)果急劇下降，僅保留了原有性能的 29% 至 44%，這證實了記憶痕跡模塊是參數(shù)知識的主要存儲庫。
相反，閱讀理解任務(wù)表現(xiàn)出驚人的恢復(fù)能力，保留了 81%至 93% 的性能，這表明基于上下文的任務(wù)主要依賴于主干網(wǎng)絡(luò)的注意力機(jī)制，而非記憶痕跡。
團(tuán)隊還發(fā)現(xiàn)，卸載 100B 參數(shù)的嵌入表對吞吐量的影響可以忽略不計，在 8B 主干網(wǎng)上的峰值僅為 2.8%。
Engram 將基礎(chǔ)設(shè)施感知效率作為首要設(shè)計原則。其確定性尋址實現(xiàn)了存儲和計算的解耦，從而能夠?qū)⒑Ａ繀?shù)表卸載到宿主內(nèi)存，而推理開銷卻可以忽略不計。相關(guān)人員表示，條件內(nèi)存函數(shù)將成為下一代稀疏模型不可或缺的建模原語。
參考資料：
https://github.com/deepseek-ai/Engram/blob/main/Engram_paper.pdf
·
分享
