端到端優化所有能力,字節跳動提出強化學習LLM Agent框架AGILE
大語言模型(Large Language Models, LLMs)的強大能力推動了 LLM Agent 的迅速發展。圍繞增強 LLM Agent 的能力,近期相關研究提出了若干關鍵組件或工作流。然而,如何將核心要素集成到一個統一的框架中,能夠進行端到端優化,仍然是一個亟待解決的問題。
來自字節跳動 ByteDance Research 的研究人員提出了基于強化學習(Reinforcement Learning, RL)的 LLM Agent 框架 ——AGILE。該框架下,Agent 能夠擁有記憶、工具使用、規劃、反思、與外界環境交互、主動求助專家等多種能力,并且通過強化學習實現所有能力的端到端訓練。尤其值得注意的是,AGILE 框架允許 Agent 在不自信時主動向人類專家尋求建議。這帶來了兩大優勢:首先,Agent 在處理復雜問題時能夠持續保持高準確率;其次,通過向人類學習,增強了其快速適應新任務的泛化能力。
- 論文地址:https://arxiv.org/abs/2405.14751
- 代碼地址:https://github.com/bytarnish/AGILE
研究者讓 AGILE 框架在復雜問答任務中進行了驗證。在 ProductQA 與 MedMCQA 任務上,經過 RL 訓練后的 13B 模型或者 7B 模型能夠超越提示工程構建的 GPT-4 Agent。
AGILE Agent 框架
如圖 1 (a) 所示,AGILE 框架包含四個核心模塊:LLM、記憶(Memory)、工具(Tools)和執行器(Executor)。LLM 負責預測動作;記憶模塊記錄 Agent 的軌跡;工具模塊提供 Agent 可以調用的外部 API;執行器會監聽 LLM 的輸出,根據 LLM 的輸出調用相應的模塊執行具體的操作,并將執行操作得到的結果添加到 LLM 的上下文中。
圖 1 (b) 展示了 AGILE Agent 在電商問答場景中的一個示例。電商問答是一個復雜的實際應用場景,面臨的挑戰包括需要海量商品的領域知識、靈活運用商品檢索工具、以及快速適應不斷涌現的新商品。如圖 1 (b) 所示,AGILE Agent 會根據用戶的問題檢索記憶,如果無法確定問題的答案,Agent 會向人類專家尋求幫助。在獲得專家的反饋后,Agent 會反思并將新掌握的領域知識存儲在記憶中。在未來,面對新的問題時,Agent 能夠從記憶中檢索到這條知識,并基于這些知識直接給出準確的答案。除此之外,AGILE Agent 也會根據用戶的問題選擇是否調用外部工具(如搜索、數據庫檢索),輔助生成最終的回答。
強化學習定義:LLM Agents 被定義為一個 token-level MDP(Markov Decision Process)。動作空間(Action space)由 LLM 的詞表構成,LLM 生成的每一個 token 是一個動作,LLM 本身則作為 Agent 的策略模型(Policy model)。Agent 的狀態(State)由 LLM 上下文和記憶組成。在每個時刻,LLM 預測動作,執行器根據預定義的邏輯完成狀態轉移,同時環境給予 Agent 相應的獎勵(Reward)。
在 AGILE 框架下,Agent 有兩種策略學習方法。第一種是模仿學習,通過收集人類軌跡數據或更高級別 Agent 的軌跡數據,對 LLM 進行 SFT 訓練。第二種是強化學習,通過定義獎勵函數,利用強化學習算法來訓練 LLM。
此外,LLM Agent 可能會產生跨越數百萬個 tokens 的軌跡,這為直接訓練帶來了挑戰。為了應對這種長程依賴的問題,研究人員提出了一種片段級別的優化算法。
主動尋求幫助:AGILE 框架允許 Agent 主動向外部的人類專家尋求幫助。這種機制有兩個優勢:首先,當 Agent 遇到不確定的情況時,通過求助人類專家,確保其在實際應用中達到高準確率。其次,Agent 能夠通過對人類的反饋反思并積累知識,從而更快適應新環境,提升其泛化能力。然而決定何時尋求幫助是一個復雜決策,它涉及到 Agent 的自我評估、人類反饋對未來的價值以及人類專家的成本。因此,標注何時應該求助是很難的。但在強化學習框架中,可以通過定義相關獎勵,將這種求助能力作為策略模型的一部分,在端到端訓練中得到提升。
實驗結果
ProductQA
ProductQA 是一個商品問答任務。該數據集包含 26 個對應不同商品類別的 QA 任務,每個任務平均包含 3,393 個問題。ProductQA 包括基于事實的問題、推理問題和商品推薦問題,它能夠全面評估 Agent 處理歷史信息和累積知識、利用工具、向人求助、自我評估和反思的能力。此外,訓練和測試集由不同的任務構成,以評估 Agent 在新商品問答上的泛化能力。
在商品問答(ProductQA)任務上,基于 Vicuna-13b 訓練的 AGILE Agent(agile-vic13b-ppo)表現超過了 GPT-4(gpt4-prompt)與提升工程構建的 GPT-4 Agent(agile-gpt4-prompt)。在使用了相當的求助比例(Advice Rate)的情況下,agile-vic13b-ppo 的 acc 相比于 agile-gpt4-prompt 提升了 7.4%,在 Total Score 上提升了 9.2%。
從上圖可以看出,在執行包含上千個問答的任務整個過程中,agile-vic13b-ppo 的 acc 持續穩定地高于 agile-gpt4-prompt。同時尋求人類幫助的頻率(Advice Rate)隨著問答輪數的增加逐漸下降。此外,通過調整人類的咨詢成本(Seeking Advice Cost)和進行強化學習訓練,AGILE Agent 可以有效實現準確率與專家成本的權衡。
消融實驗結果顯示,記憶、反思、咨詢人類建議、工具使用、RL 訓練在實現高性能的 AGILE Agent 中均具有重要作用。
MedMCQA
MedMCQA 是一個多項選擇的問答數據集,其問題來自醫學院入學考試。在 MedMCQA 任務上,基于 Meerkat-7b 訓練的 AGILE Agent(agile-mek7b-ppo)表現優于 GPT-4 Agent。準確率達到了 85.2%,超過了之前的 SOTA 方法 ——GPT 4-MedPrompt 的 79.1%。消融實驗結果進一步驗證了記憶、咨詢人類建議、反思和 RL 訓練的重要性。
更多研究細節,可參考原論文。
總結
AGILE是一種基于強化學習的LLM Agent框架。AGILE Agent具備擁有長期記憶、使用外部工具、向人類專家進行咨詢、反思等能力,并且所有能力可以進行端到端的統一優化。AGILE的一個重要特點是Agent在遇到無法自行解決的問題時可以向人類專家咨詢,這種機制保證了應用場景中對高準確率的要求,同時增強了Agent的學習與泛化能力。實驗表明,經過強化學習訓練的13B或7B模型的Agent,其能力可以超過GPT-4 Agent。
團隊介紹
ByteDance Research 專注于人工智能領域的前沿技術研究,涵蓋了機器翻譯、視頻生成基礎模型、機器人研究、機器學習公平性、量子化學、AI 制藥、分子動力學等多技術研究領域,同時致力于將研究成果落地,為公司現有的產品和業務提供核心技術支持和服務。