论文

TAP:以轨迹锚定恢复支持长程Agent剪枝

TAP: Efficient Long-Horizon Agent Pruning via Trajectory-Anchored Recovery

模型优化智能体系统Agent任务评测剪枝蒸馏

摘要

新兴的长期 Agentic 任务需要重复的模型调用,这使得本已昂贵的语言模型的推理成本进一步恶化。虽然狭窄的 Agentic 任务表明有可能在不降低性能的情况下进行积极的模型修剪,但实证结果表明,为问答任务提出的现有方法在应用于 Agentic 模型时会严重降低任务性能。我们将这一失败归因于两个决定:修剪什么以及如何恢复。对于剪枝,一次性重要性估计无法跟踪剪枝后的模型如何适应。对于恢复,离线蒸馏仅涵盖教师前缀,而全轨迹同策略蒸馏会导致学生错误在轮流中复合。在这项工作中,我们提出了轨迹锚定剪枝(TAP),这是第一个用于强化学习(RL)训练的Agent的结构剪枝框架。 TAP 将结构修剪与有效的策略恢复结合起来,将交互锚定到教师轨迹,同时允许学生生成每个推理-行动响应。冻结的密集教师监督学生的响应前缀,解决响应内训练推理不匹配的问题,同时防止学生引起的偏差在训练轮次中传播。 TAP 不是一次性剪枝,而是使用恢复学生的恢复目标梯度对通道进行重新评分,将迭代通道选择与不断发展的策略联系起来。删除 60% 的 FFN 通道后,TAP 分别保留了 ALFWorld 和 WebShop 上密集 7B 智能体的 99.2% 和 88.0% 的任务成功率,同时将每个成功任务的 GPU 时间减少了约 22% 和 17%。这些结果证明了在有限的恢复预算下长视野Agent的有效结构压缩。

TAP:以轨迹锚定恢复支持长程Agent剪枝:论文原图
图 2:TAP 通过一个目标结合了如何恢复和修剪什么。上图:教师轨迹提供了固定的锚点上下文池。左(如何恢复):在锚上下文中,学生生成推理和动作标记,冻结的教师监督每个生成的标记,并最小化 ℒAOPD\mathcal{L}_{\mathrm{AOPD}} 更新学生。右(要修剪什么):在恢复的学生的最后一个恢复批次得分通道上相同目标的梯度,并删除得分最低的通道。重复修剪-恢复-重新评分循环,直到达到目标去除率。