Structural pruning cuts large language agents costs with minimal performance loss
TAP: Efficient Long-Horizon Agent Pruning via Trajectory-Anchored Recovery
Machine Learning
Summary
Long tasks handled by language-based agents require many model runs, making them slow and expensive. The authors found that existing pruning methods hurt agent tasks because they don't properly account for how models adapt after pruning or how to recover. They propose TAP, a new way to prune these agents by carefully selecting parts to remove and using a special recovery technique that relies on example trajectories from a large teacher model. This approach cuts the size of the model significantly while keeping most of its task success and reducing computing time.
What this means in practice
- •For ai engineering teams: Reduce computational costs when deploying long-horizon language agents by pruning unnecessary model components with minimal loss in performance.
- •For cloud service providers: Offer more efficient agentic AI services by integrating pruning techniques that lower GPU time for long sequential tasks.
Authors
Yuanzhe Li, Pengxin Wang, Yuxin Ren, Jianing Deng, Jingtong Hu, Song Wang, Jingdi Chen, Huanrui Yang
Abstract
Emerging long-horizon agentic tasks require repeated model calls, worsening the inference cost of already-costly language models. While narrow agentic tasks suggest potential for aggressive model pruning without performance drop, empirical results show existing methods proposed for question answering tasks severely degrade task performance when applied to agentic models. We trace this failure to two decisions: what to prune and how to recover. For pruning, one-shot importance estimates fail to track how the pruned model adapts. For recovery, offline distillation covers only teacher prefixes, while full-trajectory on-policy distillation causes student errors to compound across turns. In this work, we propose Trajectory-Anchored Pruning (TAP), the first structural pruning framework for reinforcement learning (RL)-trained agents. TAP couples structural pruning with efficient on-policy recovery, anchoring interactions to teacher trajectories while allowing the student to generate each reasoning-action response. A frozen dense teacher supervises the student's response prefixes, addressing within-response training-inference mismatch while preventing student-induced deviations from propagating across training turns. Instead of one-shot pruning, TAP re-scores channels using gradients of the recovery objective on the recovered student, connecting iterative channel selection to the evolving policy. With 60% of FFN channels removed, TAP retains 99.2% and 88.0% of the dense 7B agents' task success rates on ALFWorld and WebShop, respectively, while reducing GPU time per successful task by approximately 22% and 17%. These results demonstrate effective structural compression of long-horizon agents under a limited recovery budget.