NeurIPS 2024 | 解鎖大模型知識記憶編輯的新路徑，浙大用「WISE」對抗幻覺

机器之心發表於2024-10-23
大模型
圖片

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com


本篇工作已被 NeurIPS（2024 Conference on Neural Information Processing Systems）會議接收，文章第一作者為浙江大學軟體學院碩士生王鵬，師從張寧豫副教授。

在當前人工智慧的迅猛發展中，大模型的知識記憶能力成為了提升智慧系統理解和推理能力的關鍵。然而，與人類記憶相比，機器記憶缺乏靈活性和可控性，難以在動態環境中實現有效的知識更新與編輯。人類的大腦擁有高度適應性的記憶機制，能夠根據外部環境變化及時進行資訊的篩選、修正與增強。這種能力不僅使得我們能夠精準地獲取資訊，還可以根據任務需求高效地呼叫相關知識。

相比之下，現有的大模型主要依賴固定的引數和資料來儲存知識，一旦訓練完成，修改和更新特定知識的代價極大，常常因知識謬誤導致模型輸出不準確或引發「幻覺」現象。因此，如何對大模型的知識記憶進行精確控制和編輯，成為當前研究的前沿熱點。

本文借鑑認知科學和人類記憶的機制，探討了大模型終身知識編輯問題，提出了一種基於雙重記憶機制的大模型知識編輯方法 WISE, 旨在持續更新大語言模型的世界知識和糾正其幻覺性輸出。此工作結合引數化長期記憶和工作記憶，在保持語言模型通用能力的同時可成功對模型進行數千次連續編輯。
圖片

  • 論文標題：WISE: Rethinking the Knowledge Memory for Lifelong Model Editing of Large Language Models

  • 論文連結: https://arxiv.org/abs/2405.14768
  • 程式碼連結: https://github.com/zjunlp/EasyEdit

1. 背景與挑戰

隨著大模型（LLMs）的廣泛應用，持續更新其世界知識和糾正幻覺性輸出成為一個關鍵問題。過去的方法在長期模型知識編輯中往往無法同時實現可靠性、泛化性和區域性性，這被稱為「不可實現三角」（如下圖）。
圖片
圖 1 可靠性、泛化性和區域性性之間的度量三角

2. 理論基礎

2.1 終生模型知識編輯定義

終生模型知識編輯問題專注於對 LLMs 進行連續的、大量的編輯操作，目的是使模型的輸出能夠與人類預期保持一致，同時保留模型先前的知識與能力 (如圖 2 所示)。具體來說，就是透過一系列時間序列上的編輯操作，逐步改進模型對特定查詢的處理能力，這些編輯操作由一個不斷變化的編輯資料集圖片來驅動 [3,16,17]。
圖片
圖 2 終生模型編輯任務示意圖

終生模型編輯的目標是實現以下三個關鍵特性：

  • 可靠性（Reliability）：模型能夠在順序編輯後記住當前和之前的編輯操作。

  • 泛化能力（Generalization）：模型不僅僅記住查詢-目標對，而是能夠理解並泛化到不同型別的查詢。

  • 區域性性（Locality）：模型編輯操作不會影響與編輯知識無關的預訓練知識。

其過程可以描述為：給定一個已經在圖片上預訓練的模型圖片, 當模型需要糾正錯誤或注入新知識時會使用一個隨時間變化的編輯資料集圖片來進行編輯操作。在第 T 個編輯步驟中，模型知識編輯器（Model Editor, ME）接收第 T 個編輯樣例和 T-1 步的的模型圖片，併產生修正後的 LLM 模型圖片。遵循以下等式：
圖片圖片是當前編輯步驟的輸入，圖片是期望輸出；同時保留對過去編輯步驟中輸入圖片的記憶，並維持對不相關資料圖片的處理能力。

注意，使用終生模型知識編輯技術並非必須一直編輯大模型，如積累大量新資料後可透過全量微調繼續更新大模型的知識。終生模型知識編輯技術適用於小資料持續知識更新和謬誤修正。

2.2 語言模型中的知識記憶

在人類認知中，工作記憶為生物大腦提供了暫時儲存資訊的能力，以便以適應不斷變化的環境的方式執行對話、推理和數學等任務。相似地，過去的文獻 [8, 9, 10] 表明語言模型的記憶可分為長期（情節性的）記憶和工作記憶（短期）：工作記憶可能儲存在神經元的持續啟用（推理時的 Activation）中，長期記憶可能儲存在模型引數（Weight）中。

我們發現更新的知識駐留在記憶中的位置會影響編輯效能，現有方法可以大致分為兩類：編輯長期記憶和編輯工作記憶。長期記憶是透過直接編輯模型引數來更新通用的引數化知識，這種方法會與之前的預訓練知識產生衝突，導致區域性性較差 (例如 FT-EWC [1]、ROME [2])；而工作記憶則是在推理時透過檢索替換神經網路的啟用/表徵，不修改模型引數。儘管工作記憶方法在可靠性和區域性性上表現優異，但其檢索到的表徵難以實現泛化，導致編輯的知識無法有效推廣（例如 GRACE [3]、SERAC [4]）。這些揭示了長期記憶和工作記憶對於終身模型編輯都有缺點。

此外，儘管有一些針對 LLM 架構的特殊記憶設計，如 MemorryLLM [6] 和 Memoria [7]，它們改變了模型架構（大部分 Train from scratch）且不能直接應用於不同的 LLMs。
圖片
圖 3 當前模型編輯方法的比較

這啟發我們提出一個關鍵科學問題：如何設計適配大模型的知識記憶更新機制，以打破終生知識編輯中的不可能三角？

3.WISE 方法介紹
圖片
人類大腦的左右半球在不同任務中的分工給了我們靈感，這啟發我們設計了 WISE，一個具備雙引數記憶機制的框架。WISE 透過主記憶儲存預訓練知識，並引入側記憶來專門儲存編輯後的知識。側記憶可以被視為一種中期記憶，它結合了長時記憶的泛化能力和基於檢索的工作記憶的可靠性與區域性性。我們僅在側記憶中進行編輯，並訓練一個路由器來決定在處理查詢時應使用哪種記憶。

為了實現連續編輯，WISE 還設計了一種知識分片機制，將不同的編輯集合儲存在獨立的、正交的子空間中，最後將這些編輯合併為統一的側記憶。主記憶儲存模型在預訓練階段學到的知識：

1. 側記憶（Wv’）作為一個副本，記錄模型在編輯後的更新資訊。
圖片
2. 知識分片：將側記憶劃分成不同的隨機子空間來儲存編輯資訊。具體來說，對於第 i 個編輯碎片，我們為其生成一個隨機梯度掩碼 Mi。這些掩碼確保了每次編輯都僅在側記憶的特定子空間中進行，從而實現了編輯的區域性化和正交化。
圖片
3. 自適應 Gate：採用基於啟用的門控策略來決定在給定查詢時使用主記憶還是側記憶。門控啟用指示器的計算方式是比較側記憶和主記憶的啟用差異（如下列公式所示）。我們設計了基於邊界的損失函式，確保編輯查詢的啟用指標比無關查詢大，具體目標是：編輯查詢的啟用值應大於無關查詢，且兩者之間的差異超過設定的閾值 γ
圖片
4. 知識合併：透過 Ties-Merge [5] 技術將各個子空間的知識合併為一致的表徵，實現引數的高效利用。

4. 實驗結果

實驗結果表明，直接修改模型權重會覆蓋預訓練的知識，導致新舊知識衝突，破壞區域性性，影響模型對非編輯領域的保留。

WISE 在多個任務（如問答、幻覺修正、分佈外資料）上表現出色，尤其是在 LLaMA、GPT 等架構中，WISE 大幅超越現有編輯方法。透過評估可靠性、泛化性和區域性性三項指標，WISE 在長期編輯中能夠有效解決模型衝突問題，並展示了優異的穩定性和擴充套件性。
圖片
5. 實驗分析

處理長序列持續編輯的潛力
圖片
WISE 在 3K 次編輯下表現出色，尤其是 WISE-Retrieve 透過高效的子空間組織和路由機制，能夠在較少效能下降的情況下應對大量編輯。

路由啟用視覺化
圖片
WISE 透過啟用指標準確區分編輯查詢與非相關查詢，確保編輯的區域性性，併成功將相關查詢路由到側記憶，避免干擾預訓練知識。

在 LLM 中的引入位置
圖片
應在 LLM 的中間到後期層引入側記憶。這些層被認為能夠更好地處理高階語言現象，並且透過殘差連線保持了較低層次的語義資訊，使得編輯操作能夠更有效地影響模型的輸出。

WISE 的額外開銷
圖片
在編輯次數 3K 時，僅增加了 0.64% 的引數量和 4% 的 GPU 視訊記憶體需求，且推理時間開銷較小，具有較高的計算效率。

6. 總結與展望

本文為長期模型知識編輯提供了一種新穎的解決思路，透過側記憶設計和知識分片技術，在不犧牲模型效能的情況下，實現了知識的有效更新。未來的研究可以進一步最佳化路由策略，提升側記憶的檢索效率；探索更好的記憶架構，以應對更加複雜的編輯場景。

當前階段，針對事實和例項記憶等型別的知識編輯，通常採用以下幾種方法：外部記憶更新（如 RAG、Memory 等 [12][13]）、區域性引數更新（如 ROME [2]、AlphaEdit [11]）或全域性引數更新（如微調或對齊）。而對於更抽象的知識型別，如安全性、人格或自我認知等，還可使用執行時干預（Steering [12][13]）或慢思考方法（如藉助 o1 思想進行錯誤修正）。

不斷提升大模型的知識處理能力，進而實現通用人工智慧（AGI），是學術界與工業界的共同目標。大模型知識編輯技術的突破，不僅能夠促進大模型對新知識和新技能的快速、永久習得，還可以實現神經與符號知識之間的高效轉換與處理。此外，當大模型出現致命錯誤或安全隱患時，基於知識編輯技術可以快速定位問題根源，並實現及時的干預和控制。這種技術對確保大模型的可信與安全至關重要。

此外，大模型的知識編輯技術不僅能有效最佳化模型的表現，還能促進對大模型知識機理的深入研究。透過對引數進行干預與分析，研究人員可以進一步解構並理解 「電子大腦」的運作原理。

參考文獻
[1] Overcoming catastrophic forgetting in neural networks.
[2] Locating and Editing Factual Associations in GPT.
[3] Aging with GRACE: Lifelong Model Editing with Discrete Key-Value Adaptors.
[4] Memory-Based Model Editing at Scale.
[5] TIES-Merging: Resolving Interference When Merging Models.
[6] MEMORYLLM: Towards Self-Updatable Large Language Models.
[7] Memoria: Resolving Fateful Forgetting Problem through Human-Inspired Memory Architecture.
[8] FROST: A Distributed Neurocomputational Model of Working Memory Maintenance.
[9] Large Language Models with Controllable Working Memory.
[10] Adaptive semiparametric language models
[11] AlphaEdit: Null-Space Constrained Knowledge Editing for Language Models
[12] Retrieval-augmented generation for large language models: A survey
[13] Towards LifeSpan Cognitive Systems
[14] Word Embeddings Are Steers for Language Models
[15] Steering Llama 2 via Contrastive Activation Addition
[16] Enhance Lifelong Model Editing with Continuous Data-Adapter Association
[17] WilKE: Wise-Layer Knowledge Editor for Lifelong Knowledge Editing

相關文章