本文提出结构化代理蒸馏框架,通过分割大型语言模型代理轨迹为推理和行动片段并施加分段特定监督,在压缩模型时显著提升任务成功率、推理效率和一致性,优于token级基线。
Large Language Model, Reinforcement Learning, Reasoning, Agent, Efficiency
Jun Liu, Zhenglun Kong, Peiyan Dong, Changdi Yang, Tianqi Li, Hao Tang, Geng Yuan, Wei Niu, Wenbin Zhang, Pu Zhao, Xue Lin, Dong Huang, Yanzhi Wang
Carnegie Mellon University, Northeastern University, Harvard University, MIT, Peking University, University of Georgia, Florida International University
Generated by grok-3
Background Problem
大型语言模型(LLMs)作为决策代理在ReAct等框架中通过交错推理和行动展现出强大能力,但其高推理成本和大模型体积限制了实际部署。传统token级蒸馏方法忽略了代理行为中推理和行动的结构化特性,导致学生模型在压缩过程中丢失关键的推理逻辑和行动一致性。本文提出了一种结构化代理蒸馏(Structured Agent Distillation, SAD)框架,旨在通过分割轨迹并施加分段特定监督,解决如何在压缩模型时保留推理和行动一致性的问题。
Method
本文提出的结构化代理蒸馏(SAD)框架核心在于将教师模型生成的轨迹分割为推理([REASON])和行动([ACT])两个部分,并对每个部分施加特定的监督损失以训练学生模型。具体步骤如下:
- 轨迹分割:通过基于规则的解析,将教师代理的轨迹分解为推理和行动两个不相交的片段,使用提示模板确保跨任务一致性。
- 分段感知监督:构建两个二进制token掩码(推理掩码和行动掩码),确保每个token仅属于一个功能片段,并分别计算推理和行动部分的损失。
- 损失函数设计:使用KL散度分别定义推理一致性损失(LCoT)和行动一致性损失(LAct),最终总损失为两者的加权和(权重相等)。
- 课程学习机制:基于轨迹复杂性评分(结合推理和行动长度及教师策略熵)进行课程采样,从简单到复杂逐步训练学生模型。
批判性思考:虽然方法在理论上通过语义解耦来提升学生模型对教师行为的模仿能力,但规则 기반的分割方法可能在复杂任务中缺乏鲁棒性,尤其是在非文本或非ReAct风格的代理任务中。此外,损失函数的加权方式(简单相等)可能未充分优化推理和行动之间的平衡,缺乏对权重选择的深入探讨。
Experiment
实验在三个代理环境(ALFWorld、WebShop、HotPotQA-ReAct)上进行,评估了从GPT-2-1.5B、OPT-13B和LLaMA-13B教师模型蒸馏到不同规模学生模型(120M至7B)的效果。实验设置包括:
- 数据集与任务:涵盖了具身指令跟随、基于Web的工具规划和多跳问答任务,任务选择合理,覆盖了代理行为的多样性。
- 评估指标:包括任务成功率、推理长度(效率)、推理一致性(CoT匹配率)和延迟(每集步骤数),指标设计全面,关注性能与效率。
- 基线对比:与token级蒸馏基线(MiniLLM)对比,显示出一致的改进,例如在120M模型上任务成功率提升约4.3%。
- 结果分析:SAD在所有规模模型和数据集上均优于基线,尤其在小规模模型(120M、340M)上改进显著;推理长度更短,延迟更低,CoT一致性更高。
- 消融研究:验证了推理监督、行动监督和分段分割的重要性,去除任一组件均导致性能下降。
批判性思考:虽然结果显示改进,但与基线的差距(例如任务成功率提升4-5%)是否足够显著以证明方法的实用性值得商榷。实验未充分探讨不同任务类型或模型架构对结果的影响,缺乏对分割规则鲁棒性的验证。此外,延迟指标以步骤数而非实际时间计算,可能无法完全反映真实部署场景中的效率。
Further Thoughts
本文提出的结构化代理蒸馏框架在理论上为代理行为的压缩提供了一个新视角,但其实际应用价值仍需进一步验证。特别是,规则 기반的分割方法可能在面对多模态代理或非结构化任务时遇到挑战,未来可以探索基于学习的分段方法以提高鲁棒性。此外,课程学习机制的潜力未被充分挖掘,结合自适应复杂性评分或与强化学习(RL)方法结合可能进一步提升训练稳定性。另一个有趣的方向是将此框架与其他压缩技术(如量化或剪枝)结合,探索在边缘设备上的部署可能性,尤其是在机器人或实时决策场景中。最后,本文的方法可能与多代理协作框架(如LangChain或AutoGen)结合,研究如何在多代理系统中实现结构化知识转移,这可能对构建高效的分布式AI系统具有重要意义。