Sub-goal Distillation: A Method to Improve Small Language Agents

📄 arXiv: 2405.02749v1 📥 PDF

作者: Maryam Hashemzadeh, Elias Stengel-Eskin, Sarath Chandar, Marc-Alexandre Cote

分类: cs.LG

发布日期: 2024-05-04

💡 一句话要点

提出子目标蒸馏方法，提升小语言模型在交互式任务中的性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 语言代理 交互式任务 子目标规划 模仿学习

📋 核心要点

  1. 大型语言模型计算成本高，调用次数受限，难以应用于长时程交互任务。
  2. 提出子目标蒸馏方法，利用大型语言模型指导小模型学习子目标规划和执行。
  3. 在ScienceWorld环境中，该方法超越了传统模仿学习方法16.7%。

📝 摘要（中文）

大型语言模型（LLM）在交互式任务中展现出巨大潜力，但其高昂的计算成本和有限的调用次数限制了其在实际应用中的效用，尤其是在长时程交互任务（如决策）或涉及持续进行任务的场景中。为了解决这些限制，我们提出了一种方法，将具有数十亿参数的LLM的性能转移到更小的语言模型（7.7亿参数）。我们的方法包括构建一个分层代理，该代理包含一个规划模块（通过知识蒸馏从LLM学习以生成子目标）和一个执行模块（学习使用基本动作完成这些子目标）。具体来说，我们利用LLM来注释一个oracle路径，生成一系列子目标以完成目标。随后，我们利用这些带注释的数据来微调规划和执行模块。重要的是，在推理过程中，这两个模块都不依赖于对LLM的实时访问，从而显著降低了与LLM交互相关的总体成本，使其成为一个固定成本。在具有挑战性的多任务交互式文本环境ScienceWorld中，我们的方法超越了仅基于基本动作的标准模仿学习16.7%（绝对值）。我们的分析突出了我们的方法相对于其他基于LLM的方法的效率。我们的代码和用于蒸馏的带注释的数据可以在GitHub上找到。

🔬 方法详解

问题定义：现有大型语言模型（LLM）作为智能体在交互式任务中表现出色，但其巨大的计算资源需求和API调用限制使其难以在实际应用中部署，尤其是在需要长期规划和持续交互的复杂任务中。传统的模仿学习方法虽然可以训练较小的模型，但往往难以学习到有效的策略，导致性能不佳。因此，如何将LLM的强大能力迁移到更小、更高效的模型上，是亟待解决的问题。

核心思路：本论文的核心思路是利用知识蒸馏，将大型语言模型（LLM）的规划能力转移到小型语言模型上。具体而言，通过让LLM生成完成任务的子目标序列，并将这些子目标作为监督信号，训练小型语言模型的规划模块。同时，训练一个执行模块，使其能够根据规划模块生成的子目标，执行相应的基本动作。这种分层结构使得小型模型能够有效地学习到复杂的策略，从而在交互式任务中取得良好的性能。

技术框架：该方法构建了一个分层代理，包含两个主要模块：规划模块和执行模块。首先，利用LLM对oracle路径进行标注，生成一系列子目标。然后，使用这些带标注的数据来训练规划模块，使其能够预测给定状态下的下一个子目标。接着，训练执行模块，使其能够根据当前状态和规划模块生成的子目标，选择合适的动作。在推理阶段，规划模块和小模型执行模块独立运行，无需实时访问LLM。

关键创新：该方法最重要的创新点在于提出了子目标蒸馏的概念，通过将LLM的规划能力分解为一系列子目标，并将其作为监督信号，有效地指导小型语言模型的学习。与传统的模仿学习方法相比，该方法能够更好地利用LLM的知识，从而提高小型模型的性能。此外，该方法的分层结构也使得模型更易于训练和理解。

关键设计：论文中关键的设计包括：1）使用LLM生成高质量的子目标序列；2）设计合适的损失函数来训练规划模块和执行模块，例如可以使用交叉熵损失函数来训练规划模块预测正确的子目标，并使用行为克隆损失函数来训练执行模块模仿oracle动作；3）选择合适的网络结构来实现规划模块和执行模块，例如可以使用Transformer模型来实现规划模块，并使用简单的线性层来实现执行模块。具体的参数设置和网络结构细节需要在实际应用中进行调整。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，在ScienceWorld环境中，该方法超越了仅基于基本动作的标准模仿学习16.7%（绝对值）。这表明子目标蒸馏能够有效地将大型语言模型的规划能力迁移到小型模型，从而显著提高其在复杂交互式任务中的性能。该方法在降低计算成本的同时，保持了较高的性能水平。

🎯 应用场景

该研究成果可应用于各种需要长期规划和交互的智能体任务，例如游戏AI、机器人导航、任务导向对话系统等。通过将大型语言模型的知识迁移到小型模型，可以降低计算成本，提高部署效率，从而推动智能体技术在实际场景中的应用。此外，该方法还可以扩展到其他领域，例如教育和医疗，用于个性化学习和辅助诊断。

📄 摘要（原文）

While Large Language Models (LLMs) have demonstrated significant promise as agents in interactive tasks, their substantial computational requirements and restricted number of calls constrain their practical utility, especially in long-horizon interactive tasks such as decision-making or in scenarios involving continuous ongoing tasks. To address these constraints, we propose a method for transferring the performance of an LLM with billions of parameters to a much smaller language model (770M parameters). Our approach involves constructing a hierarchical agent comprising a planning module, which learns through Knowledge Distillation from an LLM to generate sub-goals, and an execution module, which learns to accomplish these sub-goals using elementary actions. In detail, we leverage an LLM to annotate an oracle path with a sequence of sub-goals towards completing a goal. Subsequently, we utilize this annotated data to fine-tune both the planning and execution modules. Importantly, neither module relies on real-time access to an LLM during inference, significantly reducing the overall cost associated with LLM interactions to a fixed cost. In ScienceWorld, a challenging and multi-task interactive text environment, our method surpasses standard imitation learning based solely on elementary actions by 16.7% (absolute). Our analysis highlights the efficiency of our approach compared to other LLM-based methods. Our code and annotated data for distillation can be found on GitHub.