PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuanzhe Li, Pengxin Wang, Yuxin Ren, Jianing Deng, Jingtong Hu, Song Wang, Jingdi Chen, Huanrui Yang

arXiv:2610.09074v1cs.LG

TL;DR

Long-horizon agents face accumulating inference costs, and existing pruning methods can underuse their compression potential while recovery approaches have distinct limitations. TAP combines trajectory-anchored on-policy recovery with iterative, recovery-informed channel pruning; at 60% FFN removal it retains high success on both benchmarks while reducing GPU time per successful task.

  • Problem

    Long-horizon agents require repeated model calls, while existing structural pruning methods can fail to realize agentic models’ compression potential and recovery approaches face training–inference mismatch or compounding errors.

    Source: §1 · p. 1

  • Method

    TAP alternates structural pruning and trajectory-anchored on-policy recovery, then re-scores channels using gradients of the recovery objective on the recovered student.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    At 60% FFN removal, TAP retains 99.2% and 88.0% of dense 7B success on ALFWorld and WebShop, while reducing GPU time per successful task by approximately 22% and 17%, respectively.

    Source: §1 · p. 2

  • Takeaways & Limitations

    TAP demonstrates effective structural compression of long-horizon agents at substantial FFN removal while lowering cost per successful task.

    Source: §1 · p. 2, §5 · p. 9

    TAP relies on informative teacher trajectories, and anchored recovery does not cover every state the deployed student may visit; the study evaluates FFN channel pruning of Qwen2.5 agents in two environments.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-horizon agents incur accumulating inference costs, while existing structural pruning methods underuse their compression potential and can fail at aggressive ratios. TAP addresses pruning and recovery together, retaining strong task success after substantial FFN removal while reducing cost per successful task.

1 INTRODUCTION

  • Existing structural pruning methods developed mainly for static tasks leave agentic compression potential unrealized; a 7B WebShop agent at 80% FFN removal matches the success rate previous methods attain at 40% removal.The paper attributes the gap to decisions about what to prune and how to recover.Source: §1 · p. 1
  • One-shot importance estimates miss how pruning changes the student, while fixed-prefix distillation misses student-generated continuations and full-trajectory OPD lets errors compound across turns.Token-likelihood criteria also lower out-of-distribution success by 5.5 points relative to distribution matching on agent contexts.Source: §1 · p. 2
  • TAP anchors training contexts to teacher trajectories but lets the student generate complete reasoning–action responses under frozen-teacher token-level supervision.This prevents student actions from altering later training contexts while retaining on-policy supervision within each turn.Source: §1 · p. 2
  • At 60% FFN removal, TAP retains 99.2% and 88.0% of dense 7B success on ALFWorld and WebShop, while reducing GPU time per successful task by approximately 22% and 17%, respectively.At 80% removal, it retains 80.7% and 79.9% of dense success.Source: §1 · p. 2, Figure 1 · p. 2
  • TAP couples trajectory-anchored OPD with iterative pruning and recovery-informed channel recalibration, so the recovery objective guides channel selection at each stage.Experiments cover ALFWorld and WebShop at two model scales, with ablations of iterative pruning, recalibration, and anchoring.Source: §1 · pp. 2–3

2 RELATED WORK

Prior structural pruning and distillation work largely targets static language tasks or non-agentic settings. TAP instead prunes dense RL-trained agents over long horizons, using its recovery divergence to score channels on agent contexts and its recovered student.

2 RELATED WORK

  • Structural-pruning research largely evaluates language models on static tasks, using importance scores, pruning-error compensation, and fine-tuning or distillation for recovery (Ma et al., 2023; Zhang et al., 2024; Muralidharan et al., 2024).Gradual or multi-step pruning with retraining can improve high-compression results (Zhu & Gupta, 2018; Zhang et al., 2024; Pan et al., 2025).Source: §2 · p. 3
  • TAP differs by anchoring every training context to teacher trajectories while supervising student responses, requiring no environment interaction and coupling recovery with iterative channel selection.It scores channels using gradients of the teacher–student divergence minimized during recovery, then re-scores on the recovered student before subsequent cuts.Source: §2 · p. 3
  • On-policy distillation supervises student-generated text, while multi-turn methods address rollout errors through rollout-depth curricula, teacher intervention, or turn-wise budget rebalancing (Gu et al., 2024; Agarwal et al., 2024; Wang et al., 2026a; Li et al., 2026a; Zhou et al., 2026).Other cited work addresses short-to-long rollouts on non-agentic tasks or collects student rollouts in the environment (Zhang et al., 2026).Source: §2 · p. 3

3 METHOD

TAP defines complete-task performance as its pruning objective and organizes compression as staged pruning, recovery, and re-scoring. Its recovery uses student responses in teacher-trajectory contexts, while its channel scores come from the recovery objective on the adapted student.

3.1 PRELIMINARIES

  • TAP progressively compresses the student through prune–recover–rescore stages, using teacher trajectories as fixed anchor contexts and the same recovery objective to guide subsequent channel selection.Algorithm 1 divides a supervised-response-token budget across stages, with a KD warm-up followed by anchored OPD.Source: §3.1 · p. 4, Figure 2 · p. 4

3.2 TRAJECTORY-ANCHORED ON-POLICY DISTILLATION

  • Anchored OPD samples contexts from teacher trajectories, has the student generate responses, and supervises student-generated prefixes with the frozen teacher.The generated action is not executed; the next training context is sampled from the anchor pool, so recovery requires no additional environment interaction.Source: §3.2 · p. 5

3.3 RECOVERY-GUIDED ITERATIVE PRUNING

  • TAP scores each channel by approximating its effect on the anchored-OPD recovery objective with a first-order Taylor criterion, then removes the lowest-scoring channels within each layer.It re-scores on the recovered student before subsequent cuts; at the first cut, it perturbs the student with small Gaussian noise because KL gradients vanish when student and teacher are identical.Source: §3.3 · pp. 5–6

4 EXPERIMENTS

The experiments evaluate TAP’s task-success retention and deployment cost against pruning baselines and recovery variants on long-horizon agents. They examine task performance, channel recalibration, recovery design, and efficiency across two benchmarks and model scales.

4.1 EXPERIMENTAL SETUP

  • Experiments use ALFWorld and WebShop with RL-trained 1.5B and 7B agents, evaluating in-distribution and out-of-distribution ALFWorld splits and WebShop’s 500-task test set.The 1.5B models are pruned at 40% and 60%, and the 7B models at 60% and 80%.Source: §4.1 · p. 6

4.2 MAIN RESULTS: TASK SUCCESS UNDER AGGRESSIVE PRUNING

  • At 60% FFN removal, TAP retains 99.2% of dense ALFWorld success; at 80%, it retains 79.9% of dense WebShop success, with its lead over baselines growing as removal increases.At 60% removal, TAP retains 69.38% of dense 1.5B WebShop success while no baseline exceeds 9.07%; at 80% removal, it exceeds CFSP + KD by 10.28 and 8.12 retention points on ALFWorld and WebShop.Source: §4.2 · pp. 6–7, Table 1 · p. 7

4.3 WHAT TO PRUNE: RECOVERY-GUIDED CHANNEL RECALIBRATION

  • Recalibration improves five of six metrics over fixed-ranking iteration and all six over one-shot pruning; uniform scheduling trails cubic scheduling on all six metrics.Replacing reverse-KL scoring with token NLL lowers ALFWorld Out SR from 89.80% to 84.33% at 60% removal.Source: §4.3 · pp. 7–8, Table 2 · p. 8

4.4 HOW TO RECOVER: TRAJECTORY-ANCHORED DISTILLATION

  • Anchored OPD improves WebShop SR over logits KD and SeqKD and gives the highest ALFWorld Out SR at both removal ratios; full-trajectory OPD collapses on ALFWorld at 80%.Full-trajectory OPD reaches 49.76% and 41.54% In and Out SR, versus 76.19% and 72.14% for anchored OPD; the paper attributes the collapse to cross-turn drift on longer episodes.Source: Table 3 · p. 8, §4.4 · pp. 8–9

4.5 DEPLOYMENT EFFICIENCY

  • At 60% removal, end-to-end speedups reach 1.30× on ALFWorld and 1.38× on WebShop, while CPS falls to 0.78× and 0.83× the dense baseline, respectively.At 80% removal, 1.8× more steps per task on ALFWorld outweigh faster inference; WebShop CPS falls to 0.71× the dense baseline.Source: Table 4 · p. 9, §4.5 · p. 9

5 CONCLUSION

TAP couples trajectory-anchored recovery with recovery-informed channel re-scoring for pruning RL-trained long-horizon agents. The conclusion reports retained success at 60% and 80% FFN removal, while noting the method’s trajectory and evaluation scope.

5 CONCLUSION

  • With 60% of FFN channels removed, TAP retains 99.2% and 88.0% of dense 7B-agent success on ALFWorld and WebShop; at 80% removal, it retains 80.7% and 79.9%.TAP supervises student-generated responses within teacher interaction contexts and re-scores channels with gradients of the recovery objective before each cut.Source: §5 · p. 9
  • TAP relies on informative teacher trajectories, and anchored recovery may not cover every state the deployed student visits; evaluation is limited to Qwen2.5 FFN pruning on two environments.The authors identify other architectures, pruning granularities, and agentic domains as directions for future work.Source: §5 · p. 10

A PRUNING WITH CONTINUED RL

  • In the continued-RL setup, each optimizer step removes one percentage point of FFN channels, with channel masks enforced during training.The evaluation uses SPEAR Qwen2.5-1.5B checkpoints with GiGPO continuation, rank-8 LoRA, and a learning rate of 10^-5.Source: A PRUNING WITH CONTINUED RL · p. 15

B SUPPLEMENTARY RESULTS

The supplementary results examine SlimGPT at lower pruning ratios, channel-mask changes under recalibration, and whether falling distillation loss guarantees executable behavior. Recalibration changes retained-channel selections, while lower loss can coexist with malformed generation.

B.1 SLIMGPT AT LOWER PRUNING RATIOS ON SPEAR-1.5B

  • At 40% removal, SlimGPT reduces the SPEAR ALFWorld 1.5B FFN width from 8960 to 5376, with zero In and Out success rates.The 20% removal result is reported in Table 5; the 40% results are reported in Table 1.Source: §B.1 · p. 15

B.2 CHANNEL-MASK CHANGES UNDER RECALIBRATION

  • 14–22% of retained channels are replaced by recalibration despite substantial overlap between its deleted set and one-shot selection.The comparison matches final removal ratios and uses Qwen2.5-7B runs.Source: Table 6 · p. 16

B.3 DECREASING KD LOSS DOES NOT ENSURE EXECUTABLE AGENT BEHAVIOR

  • For CFSP+KD, response-token-weighted forward KL falls from 1.580 to 0.250; for LLM-Pruner+KD, it falls from 2.062 to 0.326.These values compare the first and last 20 updates in the WebShop Qwen2.5-1.5B setting.Source: §B.3 · p. 16
  • Despite decreasing teacher-forced loss, baselines often produce incomplete or empty action blocks, and most executed actions are marked invalid.Improved fitting under teacher prefixes can coexist with malformed free generation.Source: §B.3 · p. 16

B.4 ITERATIVE PRUNING ABLATION PROTOCOL

The ablation protocol compares one-shot and three-iteration pruning while holding initialization, initial scores, anchor contexts, final removal ratio, and recovery-token budget fixed.

B.4 ITERATIVE PRUNING ABLATION PROTOCOL

  • All Table 2 variants share dense initialization, initial response-only reverse-KL Taylor scores, the anchor-context pool, final removal ratio, and total recovery-token budget Q.Source: §B.4 · p. 16
  • Variant A uses 20% of Q for KD and 80% for anchored OPD after one cut; B, C, and U divide Q equally across three iterations.Source: §B.4 · p. 16
  • Each iteration in B, C, and U uses a 20% KD warm-up followed by 80% anchored OPD; U removes cumulative ratios ρ/3, 2ρ/3, and ρ.Source: §B.4 · p. 16

C.1 DENSE CHECKPOINTS AND EVALUATION PROTOCOL

The protocol specifies the evaluated checkpoints and settings, FFN channel-removal mechanics, channel scoring and calibration, and uniform-allocation adaptations for baselines.

C.1 DENSE CHECKPOINTS AND EVALUATION PROTOCOL

  • The evaluation uses Qwen2.5-1.5B-Instruct and Qwen2.5-7B-Instruct checkpoints for both ALFWorld and WebShop.Source: §C.1 · p. 16
  • Table 8 rates use 500 WebShop test goals; action-block and invalid-action rates are per step, while repeated-action and step-limit rates are per episode, and flags can overlap.Source: Table 8 · p. 17
  • ALFWorld episodes are capped at 50 steps and WebShop episodes at 15, with decoding temperature 0.4.Source: §C.1 · p. 17

C.2 FFN CHANNEL REMOVAL AND PRUNING RATIO

  • Removing an intermediate channel deletes its rows from both input projections and its column from the output projection; physically resized tensors yield smaller dense operations.Source: §C.2 · p. 17
  • The reported cumulative channel-removal ratio ρ exceeds total parameter reduction because pruning removes only FFN channels.Source: §C.2 · p. 17

C.3 CHANNEL IMPORTANCE SCORING

  • Before each subsequent cut, TAP scores channels using the recovery-objective gradient on the last anchored-OPD batch, requiring one forward–backward pass without new generation or environment interaction.The loss averages reverse KL over supervised reasoning and action tokens, excluding prompt and padding tokens.Source: §C.3 · p. 17
  • For the initial cut, a separate student copy receives Gaussian weight noise with standard deviation 0.01 times each matrix’s RMS; later scoring uses the recovered student without perturbation.Source: §C.3 · p. 17

C.4 BASELINE ADAPTATIONS FOR UNIFORM LAYER-WISE PRUNING

  • All methods receive the same target FFN-removal ratio in every layer; the baseline adaptations retain their within-layer selection or local-reconstruction procedures while replacing other allocation schemes.Source: §C.4 · p. 17

C.5 RECOVERY OBJECTIVES FOR ABLATION STUDIES

The ablation recovery objectives contrast fixed teacher-pair distillation with full-trajectory on-policy distillation, where student actions determine later contexts.

C.5 RECOVERY OBJECTIVES FOR ABLATION STUDIES

  • Logits KD matches the teacher’s full-vocabulary distribution on fixed teacher-generated response prefixes using forward KL.The objectives exclude prompt, environment-observation, and padding tokens from supervision.Source: §C.5 · p. 18
  • Sequence KD uses the same fixed teacher-generated pairs but supervises the teacher’s response tokens as hard targets.Source: §C.5 · p. 18
  • Full-trajectory OPD samples student trajectories and executes student actions, so subsequent interaction contexts depend on preceding actions and resulting observations.It minimizes reverse KL on student response prefixes.Source: §C.5 · p. 18

C.6 RUN CONFIGURATION AND DATA SPLITS

The appendix specifies task splits and TAP recovery settings, controls the response-token budgets and comparisons, and defines model-inference and end-to-end agent-cost measurements.

C.6 RUN CONFIGURATION AND DATA SPLITS

  • Calibration and recovery use training tasks only, with a fixed teacher pool and shared response caps, action parsing, and context handling.Anchored OPD uses batches of 16 prompts, temperature 0.7, top-p 0.9, and LoRA rank 8; Taylor calibration uses anchored responses at each pruning stage.Source: §C.6 · p. 18

C.7 RECOVERY BUDGETS AND CONTROLLED COMPARISONS

  • All distillation methods share response-token budget Q and recovery settings; iterative pruning splits Q equally across three stages, each with 20% KD warm-up and 80% anchored OPD.Q counts response tokens, including reasoning, action, and naturally terminating EOS tokens, but excludes prompts and padding.Source: §C.7 · pp. 18–19, Table 9 · p. 19

C.8 EFFICIENCY MEASUREMENT PROTOCOL

  • Inference benchmarks use matched 512-token inputs and 128 generated tokens, while agent cost per success aggregates evaluation time over successful tasks, including failed attempts.Model inference averages three synchronized measurements after one warm-up; agent evaluations use one GPU and matched task sets.Source: §C.8 · p. 19

Related papers

Loading 2610.09074v1…