通過構(gòu)建稀疏性分配問題，團(tuán)隊(duì)發(fā)現(xiàn)了一個(gè) U 形擴(kuò)展規(guī)律，在該規(guī)律的指導(dǎo)下將 Engram 擴(kuò)展到 270 億個(gè)參數(shù)，其性能優(yōu)于嚴(yán)格等參數(shù)、等浮點(diǎn)運(yùn)算次數(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 釋放了注意力資源用于全局上下文，從而顯著提升了長(zhǎng)上下文檢索能力。

Engram 模塊架構(gòu)

Engram 是一個(gè)條件記憶模塊，旨在通過在結(jié)構(gòu)上將靜態(tài)模式存儲(chǔ)與動(dòng)態(tài)計(jì)算分離來增強(qiáng) Transformer 主干網(wǎng)絡(luò)。

團(tuán)隊(duì)實(shí)現(xiàn)了一個(gè)基于規(guī)范化文本等價(jià)性的詞匯投影層標(biāo)識(shí)符。對(duì)于 128k 分詞器，該過程可使有效詞匯量減少 23%。

為了增強(qiáng)表達(dá)能力，他們采用了一種受注意力機(jī)制啟發(fā)的上下文感知門控機(jī)制。針對(duì)擴(kuò)展感受野并增強(qiáng)模型的非線性，團(tuán)隊(duì)引入了一個(gè)短的深度因果卷積。

Engram 系統(tǒng)分為兩個(gè)階段：

• 訓(xùn)練階段：海量嵌入表被分片到可用的 GPU 上。采用全對(duì)全通信原語(yǔ)來跨設(shè)備檢索活動(dòng)嵌入行；

• 推理階段：Engram 表被卸載到主機(jī)內(nèi)存。利用確定性檢索邏輯，主機(jī)異步預(yù)取并傳輸嵌入，與設(shè)備端先前 Transformer 模塊的計(jì)算重疊通信。

團(tuán)隊(duì)采用了先進(jìn)的多分支架構(gòu)作為默認(rèn)骨干網(wǎng)絡(luò)，該架構(gòu)具有卓越的建模能力。

自然語(yǔ)言的 N-gram 本質(zhì)上遵循齊普夫分布，這種統(tǒng)計(jì)特性促使團(tuán)隊(duì)采用多級(jí)緩存層次結(jié)構(gòu)。這種分層結(jié)構(gòu)使得 Engram 能夠擴(kuò)展到海量?jī)?nèi)存容量，同時(shí)將有效延遲的影響降至最低。

大規(guī)模預(yù)訓(xùn)練

基于 Engram 架構(gòu)和經(jīng)驗(yàn)推導(dǎo)的分配規(guī)則，研究人員將 Engram 擴(kuò)展到數(shù)十億參數(shù)規(guī)模。他們訓(xùn)練了四個(gè)模型：

（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 對(duì)比 50k）即可達(dá)到基線模型 LongPPL 的性能，同時(shí)在 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ù)和等浮點(diǎn)運(yùn)算次數(shù)的 MoE-27B 基線模型。這些提升并非僅限于知識(shí)密集型任務(wù)，而包含了一般推理領(lǐng)域以及代碼和數(shù)學(xué)推理領(lǐng)域。

在長(zhǎng)上下文的情景訓(xùn)練中，研究團(tuán)隊(duì)發(fā)現(xiàn)長(zhǎng)上下文性能并非完全由架構(gòu)先驗(yàn)決定。即使在上下文擴(kuò)展階段控制相同的模型架構(gòu)和固定的計(jì)算預(yù)算，長(zhǎng)上下文性能仍然隨著預(yù)訓(xùn)練的進(jìn)行而單調(diào)提升。這表明長(zhǎng)上下文性能與基礎(chǔ)模型的一般建模能力密切相關(guān)。因此，嚴(yán)格的架構(gòu)比較必須通過調(diào)整基礎(chǔ)模型的損失來控制這一混淆變量，而不僅僅是調(diào)整訓(xùn)練步驟。

Engram 的內(nèi)部機(jī)制分析

有效深度

研究人員使用 LogitLens分析各層預(yù)測(cè)結(jié)果的演變。

通過顯式地訪問外部知識(shí)，Engram 減少了所需的計(jì)算步驟，從而在網(wǎng)絡(luò)層級(jí)結(jié)構(gòu)的早期階段就獲得了高置信度的有效預(yù)測(cè)。

通過繞過顯式查找進(jìn)行早期特征合成，Engram 在功能上等同于增加模型的有效深度。

模塊設(shè)置

他們?cè)谑芸丨h(huán)境下對(duì) Engram 進(jìn)行消融，以研究每個(gè)關(guān)鍵模塊設(shè)計(jì)的有效性。

結(jié)果表明，隨著插入點(diǎn)深入，性能逐漸下降。只需一輪注意力機(jī)制即可提供有意義的上下文相關(guān)的 hfor 門控，同時(shí)還能及時(shí)替換骨干網(wǎng)底層的局部聚合。

參數(shù)敏感性

研究人員通過在推理過程中完全抑制稀疏嵌入輸出來評(píng)估模型。

基準(zhǔn)測(cè)試結(jié)果急劇下降，僅保留了原有性能的  29% 至 44%，這證實(shí)了記憶痕跡模塊是參數(shù)知識(shí)的主要存儲(chǔ)庫(kù)。

相反，閱讀理解任務(wù)表現(xiàn)出驚人的恢復(fù)能力，保留了 81%至 93% 的性能，這表明基于上下文的任務(wù)主要依賴于主干網(wǎng)絡(luò)的注意力機(jī)制，而非記憶痕跡。

團(tuán)隊(duì)還發(fā)現(xiàn)，卸載 100B 參數(shù)的嵌入表對(duì)吞吐量的影響可以忽略不計(jì)，在 8B 主干網(wǎng)上的峰值僅為 2.8%。

Engram 將基礎(chǔ)設(shè)施感知效率作為首要設(shè)計(jì)原則。其確定性尋址實(shí)現(xiàn)了存儲(chǔ)和計(jì)算的解耦，從而能夠?qū)⒑Ａ繀��?shù)表卸載到宿主內(nèi)存，而推理開銷卻可以忽略不計(jì)。相關(guān)人員表示，條件內(nèi)存函數(shù)將成為下一代稀疏模型不可或缺的建模原語(yǔ)。

參考資料：

https://github.com/deepseek-ai/Engram/blob/main/Engram_paper.pdf