AssisTRAG通過集成一個智能信息助手來提升LLMs處理復雜推理任務的能力。該框架由兩個主要組件構成：一個凍結的主語言模型和一個可訓練的助手語言模型。

AssisTRAG與之前的RAG對比

1. 組件

• 主語言模型（Main LLM）：負責根據提供的信息生成答案。這個模型是固定的，不進行訓練。
• 助手語言模型（Assistant LLM）：負責信息管理，包括記憶管理和知識管理。這個模型是可訓練的。

可以看出，相比之前的RAG，該框架主要創新點是引入一個Assistant LLM，下面詳細看下助手語言模型功能。

2. AssisTRAG

AssisTRAG主要分為兩個類別來處理復雜任務：記憶管理和知識管理。記憶管理涉及整合和分析來自內部內存的內容，而知識管理則側重于利用外部知識。這兩個主要功能由AssisTRAG的四個核心能力支持：

• 工具使用（Tool Usage）：從內部記憶和外部知識庫中檢索相關信息。
• 動作執行（Action Execution）：進行推理、分析信息需求并提取知識。
• 記憶構建（Memory Building）：記錄從歷史交互中提取的關鍵知識和推理模式。
• 計劃指定（Plan Specification）：確定在答案生成過程中每個步驟的必要性。

2.1 記憶管理

記憶管理模塊通過存儲和檢索歷史交互來提升主語言模型的性能。有效的記憶管理包括兩個主要過程：捕獲新見解和檢索先前存儲的信息。這兩個過程主要有以下三個能力：

1. Action I: 筆記記錄

筆記記錄功能  負責記錄每個歷史交互中的關鍵信息和推理模式。具體來說，給定主語言模型的歷史交互（包括問題 、參考答案  和生成的答案 ），助手的任務是將答案背后的關鍵推理過程記入記憶槽  中：

這樣，所有先前問題的記憶槽累積起來形成助手的記憶 ，用于后續的記憶檢索。

作用：通過記錄和存儲歷史交互中的關鍵推理過程，助手能夠在面對類似問題時快速檢索和應用這些推理模式，從而提高回答的質量和效率。

2. Tool I: 記憶檢索器

記憶檢索器利用助手的記憶 ，根據當前問題  檢索歷史上相關的記憶。具體來說，記憶檢索器會檢索與當前問題相關的歷史記憶，并將其表示為：

作用：通過從記憶中檢索相關歷史記憶，助手可以為當前問題提供額外的上下文和信息，幫助主語言模型更好地理解和回答問題。

3. Plan I: 評估檢索記憶的有用性

如果問題是全新的，檢索到的記憶可能不僅無助于回答，還可能負面影響主語言模型的響應。因此，助手需要評估檢索到的記憶是否對當前問題有幫助。具體來說，助手使用一個提示來評估檢索到的記憶是否有助于回答當前問題。只有當答案是肯定的時候，檢索到的記憶才會被提供給主語言模型。

實現：通過一個提示，助手評估檢索到的記憶是否對當前問題有益。如果評估結果是肯定的，這些記憶槽將被提供給主語言模型。

作用：這一計劃確保了只有相關的記憶會被提供給主語言模型，避免了無關或負面的記憶干擾回答的質量。

2.2 知識管理

知識管理涉及以下四個能力：

1. 動作 II：問題分解

• 功能：將當前問題分解為多個子查詢，以便從不同方面檢索知識。
• 實現：給定一個問題 ，動作  將其分解為一系列子查詢 。例如，問題“音樂家和諷刺作家Allie Goertz寫了一首關于《辛普森一家》角色Milhouse的歌曲，Matt Groening是根據誰的名字命名的？”可以被分解為：

• 子查詢1：Allie Goertz寫了一首關于《辛普森一家》角色Milhouse的歌曲。
• 子查詢2：Matt Groening是根據誰的名字命名了Milhouse？

2. 工具 II：知識檢索器

• 功能：利用子查詢  從外部知識庫中檢索相關文檔。
• 實現：使用子查詢 ，知識檢索器從外部知識庫  中檢索相關文檔，表示為 。例如，子查詢1和子查詢2將被用來從維基百科等外部數據庫中檢索與Allie Goertz和Milhouse相關的文檔。

3. 動作 III：知識提取

• 功能：從檢索到的文檔中提取關鍵知識。
• 實現：給定問題  和檢索到的文檔 ，動作  負責從中提取相關知識 。例如，從檢索到的文檔中提取出：

• Allie Goertz 是一位美國音樂家和諷刺作家。
• Milhouse 是《辛普森一家》中的一個虛構角色，由Matt Groening創造，并以意大利獨裁者Benito Mussolini命名。

4. 計劃 II：評估提取知識的相關性

• 功能：確保提供給主語言模型的信息準確且相關。
• 實現：助手通過提示評估提取的知識是否與當前問題相關。只有當提取的知識被認為對回答問題有幫助時，才會將其提供給主語言模型。

2.3 AssisTRAG 訓練

AssisTRAG 的訓練過程分為課程助理學習和強化偏好優化兩個階段。課程學習階段通過逐步增加任務復雜性來提升助手在 RAG 過程中的技能。強化偏好優化階段則通過根據主 LLM 的反饋調整助手的輸出，確保助手提供的幫助更加貼合主 LLM 的需求。

2.4 AssisTRAG 推理

AssisTRAG 通過三個步驟啟動其推理過程：

1. 信息檢索與整合

在這個初始階段，AssisTRAG 首先激活 Action II 來理解主 LLM 的信息需求。然后，它使用 Tool I 和 Tool II 分別從內部記憶和外部知識庫中檢索相關信息。接著，它調用 Action III 從檢索到的文檔中提取關鍵知識。

2. 決策

在這一階段，AssisTRAG 決定是否將檢索到的記憶和提取的知識提供給主 LLM。它激活 Plan I 和 Plan II 來評估這些檢索到的記憶和知識對當前問題的相關性和有用性。如果助手認為這些信息有幫助，它們將被提供給主 LLM 以輔助回答生成。

3. 回答生成與記憶更新

在最后階段，AssisTRAG 提示主 LLM 根據問題、其內部知識和助手提供的信息生成答案。之后，AssisTRAG 激活 Action I，利用其筆記功能，捕捉交互中的關鍵推理步驟并將其納入其記憶中。這確保了助手的知識庫保持最新。

實驗效果

參考文獻

Boosting the Potential of Large Language Models with an Intelligent Information Assistant，https://arxiv.org/pdf/2411.06805v1

本文轉載自公眾號大模型自然語言處理  作者：余俊暉

原文鏈接：??https://mp.weixin.qq.com/s/fLO1i9uW8R4OQbF9nYGP9Q??