給大模型加上第三種記憶格式，把寶貴的參數(shù)從死記硬背知識中解放出來！

中科院院士鄂維南領(lǐng)銜，上海算法創(chuàng)新研究院等團隊推出Memory3，比在參數(shù)中存儲知識以及RAG成本都更低，同時保持比RAG更高的解碼速度。

圖片

在實驗中，僅有2.4B參數(shù)的Memory3模型不僅打敗了許多7B-13B的模型，在專業(yè)領(lǐng)域任務(wù)如醫(yī)學上的表現(xiàn)也超過了傳統(tǒng)的RAG方法，同時推理速度更快，“幻覺”問題也更少。

圖片

目前相關(guān)論文已上傳到arXiv，并引起學術(shù)界關(guān)注。

圖片

知識按使用頻率分類

這一方法受人腦記憶原理啟發(fā)，獨立于存儲在模型參數(shù)中的隱性知識和推理時的短期工作工作記憶，給大模型添加了顯式記憶。

具體來說，人類的記憶大致可以分為三部分:

• 顯式記憶：可以主動回憶的長期記憶，比如讀過的文章。獲取顯式記憶很容易，但提取時需要一定的回憶過程。
• 隱式記憶：無意識使用的長期記憶，比如騎自行車的技能。獲取隱式記憶需要大量重復練習，但使用時毫不費力。
• 外部信息：存在大腦之外的信息，如考試時的備考資料。獲取和使用都很輕松，但遇到新問題時作用有限。

可以看出，三種記憶形式在獲取和使用的效率上形成了鮮明的互補。人腦會根據(jù)知識的使用頻率，巧妙地在它們之間分配存儲位置，從而最小化整體開銷。

反觀大模型，目前主要依賴在參數(shù)中以隱式記憶的形式來存儲知識，這導致兩個問題：

• 知識分配效率低：無論一個知識使用得多頻繁，都一視同仁塞進參數(shù)里，導致大量冷知識占用了寶貴的參數(shù)空間。
• 知識提取效率低：每次使用知識，都得動用大量參數(shù)參與計算。

目前在訓練階段，團隊將大模型比作顯式記憶能力受損的患者，靠學習如何系鞋帶一樣的大量重復練習才能背下一點知識，消耗大量的數(shù)據(jù)和能量。

在推理階段，大模型又好像一個人每寫一個單詞時都要回憶起畢生所學的一切，就很不合理。

基于以上思路，團隊按照知識的預(yù)期使用頻率（橫軸）計算了讀寫成本（縱軸），陰影區(qū)域表示給定記憶格式的最小成本區(qū)域。

結(jié)果發(fā)現(xiàn)，把常用知識塞進模型參數(shù)里成本最低，但容量有限；不常用的知識直接檢索效率最高，但每次讀取都要重新編碼，成本高；而顯式記憶則是個平衡點，對于使用次數(shù)中等的大部分知識最劃算。

圖片

記憶電路理論

團隊進一步在論文中提記憶電路理論，在大模型語境下重新定義知識和記憶，以確定哪些知識更適合存儲為顯式記憶，以及什么樣的模型架構(gòu)適合讀寫顯式記憶。

圖片

通過分析一些已知的大模型內(nèi)部機制，如事實問答、搜索復制粘貼等，團隊認為大模型中的每條知識都可以表示為一個輸入-輸出關(guān)系，加上實現(xiàn)這個關(guān)系的內(nèi)部電路（circuit）。

電路指計算圖中的一個子圖，由一些注意力頭和MLP神經(jīng)元組成，這些電路的輸入輸出具有一定的語義關(guān)聯(lián)。大模型的知識可進一步分為兩類:

• 具體知識（specific knowledge）：電路的輸入和輸出都具有明確的語義，如常識、常見短語等。
• 抽象知識（abstract knowledge）：電路的輸出語義可變，如搜索、復制、粘貼，需要通過輸入推理出輸出。

接下來，作者引入可分離知識（separable knowledge）的概念：如果一個知識可以僅通過文本實現(xiàn)而不必內(nèi)置到模型參數(shù)里，那它就是可分離的。

可模仿知識（imitable knowledge）是可分離知識的一個特例，可以直接用描述這條知識自身的文本去“教會”另一個不具備這條知識的大模型，無需通過參數(shù)來編碼。

一個核心結(jié)論是，具體知識都是可模仿的，因此也是可分離的，都可轉(zhuǎn)化為顯式記憶。論文從理論上給出了（非形式化）證明。

圖片

團隊進一步把具體知識按使用次數(shù)分成“無關(guān)緊要”、專業(yè)知識和常見短語三個等級，不同等級按照讀寫成本分別適合三種不同的記憶格式。

圖片

擁有顯式記憶的大模型Memory3

那么如何實現(xiàn)顯式記憶呢？

以注意力層的key-value向量作為顯式記憶的載體，在推理之前，Memory3模型將所有引用文本轉(zhuǎn)換為顯式記憶，并將它們保存在硬盤或非易失性內(nèi)存設(shè)備上。

在推理時，模型會查詢與當前上下文最相關(guān)的一些顯式記憶，將它們并入注意力機制中，與上下文的key-value向量一起計算注意力分數(shù)，生成下一個token。

圖片

然而，海量文本轉(zhuǎn)化成的顯式記憶不僅需要更多的磁盤空間，而且在推理過程中還會占用GPU內(nèi)存，從而損害LLM生成的吞吐量。

為此，Memory3采取了多維度壓縮優(yōu)化策略：

• layer維度：只有前半部分的注意力層（記憶層）產(chǎn)生和存取顯式記憶，后半部分仍然是普通的注意力層。
• head維度：每層只有少部分head（如1/5）負責處理顯式記憶的key-value，其他head保持原樣。
• token維度：對于每個head，只選取參考文本中最相關(guān)的少量token（如8個），提取其key-value作為顯式記憶。

最后再進一步用向量量化（vector quantization）壓縮每個key和value向量到更短的表示。

多級壓縮的組合，使得顯式記憶的規(guī)模從45.9TB壓縮到4.02TB，壓縮到一個GPU集群通常配備的存儲容量之內(nèi)。

另外，團隊在顯式記憶的讀寫上還有一些值得注意的細節(jié)設(shè)計：

• 推理時為了避免不同文本片段重復檢索顯式記憶，Memory3每隔64個token做一次檢索，中間共享檢索結(jié)果。
• 頻繁調(diào)用顯式記憶會產(chǎn)生IO開銷。為此，Memory3在內(nèi)存中維護了一個固定大小的緩存，存儲最近訪問過的顯式記憶。
• 對于參考文本，模型使用不同的輸入符號（“<s>Reference:”）將其與普通文本區(qū)分開，避免干擾文本理解。
• 對于顯式記憶中的不同文本片段，模型為其分配了同一區(qū)間的位置編碼，保留局部上下文。這種”平行”位置編碼避免了長文本中間部分被忽略的問題。

最終訓練出來的Memory3模型，在HuggingFace排行榜上的評測結(jié)果如下，顯式記憶將平均分數(shù)提高了2.51%。

相比之下Llama2-7B和13B之間的分數(shù)差異為4.91%，而13B模型的非嵌入?yún)?shù)數(shù)量接近7B模型的兩倍。

因此，可以說顯式記憶可以將“有效模型大小”提高了2.51/4.91≈51.1%。如果用Qwen-1.8B和4B來做參考，計算結(jié)果相似，“有效模型大小”提高49.4%。

圖片

在幻覺評估上，Memory3避免了將文本壓縮到模型參數(shù)中可能會導致的信息丟失，表現(xiàn)的比大部分模型要好。

圖片

論文中還詳細報告了從數(shù)據(jù)到訓練、微調(diào)和對齊過程的具體設(shè)置，感興趣的可以查看原文。

論文地址：https://arxiv.org/abs/2407.01178

參考鏈接：[1]https://x.com/rohanpaul_ai/status/1809782336021537094