Progress-Aware Distillation for Mitigating Stagnation in Small Language Model Agents

AI summary

Overview

Research area: Machine learning — agent distillation, small language models (SLMs), and efficient LLM agent systems.

Technical level: Advanced (research-level, though the core intuition is explained accessibly).

Scope: This paper identifies a failure mode called progress stagnation in SLM agents distilled from large teachers, and proposes Progress-Aware Distillation (ProD), a training objective that detects and penalizes stagnant student behavior using an adaptive margin derived from teacher–student stagnation gaps.

What This Paper Is About

Agents built on large models solve multi-step tasks well but are expensive in tokens and latency, so researchers distill their behavior into small models via supervised fine-tuning on teacher trajectories. The paper shows that as these distilled SLMs run longer Reason–Act–Observe loops, they get trapped in unproductive cycles, repeating failed searches or code instead of using feedback to change strategy. The goal is to train SLM agents that keep making forward progress and stay closer to teacher behavior.

Key Contributions

1. Identifies and characterizes progress stagnation as a pervasive failure mode of SFT-based agent distillation, documented across model families (Qwen3, Llama3, Gemma3), task domains, and fine-tuning paradigms, and worsening as model scale decreases and trajectories lengthen.

2. Defines three turn-level stagnation scores — lexical (ROUGE-L against all prior turns), semantic (cosine similarity of all-mpnet-base-v2 embeddings), and LLM-as-judge — and shows they agree closely (Pearson r > 0.92), so they capture the same underlying failure.

3. Proposes Progress-Aware Distillation (ProD), which encodes the student–teacher stagnation gap plus answer correctness into an instance-dependent adaptive margin inside a SimPO-style preference objective, jointly optimized with the original SFT loss.

4. Demonstrates gains on 8 SLMs from 0.6B to 8B parameters across in-domain and out-of-domain benchmarks, with both higher accuracy and lower interaction cost, and shows iterative ProD further closes the gap.

Main Findings

• Stagnation is severe and systematic: SFT-distilled students show substantially higher stagnation than the GPT-4.1 teacher, with lexical stagnation scores reaching close to 0.9 — meaning student actions closely resemble earlier failed attempts. The disparity grows with shorter models and longer trajectories.

• The three metrics converge: Pairwise Pearson correlations on Qwen SFT agents range from 0.929 to 0.979 (all r > 0.92), averaged across HotpotQA and MATH.

• ProD improves accuracy over base and SFT models: Averaged over 8 students, ProD improves in-domain accuracy by 29.2% over base models and 4.7% over SFT, and out-of-domain accuracy by 18.5% and 4.9% respectively. Relative improvement over SFT is approximately 10% in-domain and 17% out-of-domain.

• ProD nearly matches the teacher on some settings: On HotpotQA, Qwen3-4B with ProD reaches 70.8% versus the GPT-4.1 teacher's 71.0%. On MATH the same student reaches 85.2% versus SFT's 81.9% and the teacher's 90.4%.

• Gains hold across model families and scales: Examples from Table 3 include Llama-3.1-8B-Instruct in-domain average 61.8 (SFT) to 67.2 (ProD), Llama-3.2-1B-Instruct 35.9 to 41.6, Gemma-3-4b-it 56.4 to 61.0, and Gemma-3-1b-it 28.6 to 33.2.

• Inference is cheaper: Averaged over 8 students in ID and OOD settings, average turns fall from 4.144 to 3.663 (11.6% reduction) and average tokens from 16206 to 12847 (20.7% reduction).

• Repetition-control baselines do not solve it: On HotpotQA with Qwen3-4B, SFT scores 66.6, repetition penalty at strengths 1.1/1.2/1.3 scores 66.2/65.6/61.6, Concise SFT scores 66.4, and ProD scores 70.8. For Llama3.2-3B the corresponding ProD result is 67.0 versus 61.6 for SFT.

• Both components matter: Removing the SFT loss and adaptive margin (effectively SimPO on student-generated trajectories) is unstable and underperforms SFT on Qwen3-4B and Gemma3-4B-it. Keeping SFT but using a fixed margin consistently beats SFT but remains below full ProD.

• A balanced margin works best: Across Qwen3-4B, Llama3-8B, and Gemma3-4B, the best average accuracy is consistently at α = 0.5 (e.g., Qwen3-4B average 78.0 versus 74.2 for SFT), while α = 1 (stagnation gap only) still substantially beats SFT.

• A second iteration helps: Two ProD iterations further reduce stagnation and improve accuracy on MATH, particularly for Llama3.2-3B and Llama3.2-1B, where a gap remained after one iteration.

Methodology in Plain English

The researchers first collect successful problem-solving trajectories from GPT-4.1 on the training sets of MATH and HotpotQA — 3,000 trajectories per benchmark source — and keep only trajectories that reach the correct answer. A small model is fine-tuned on these trajectories with LoRA (rank 64, batch size 32, AdamW, linear learning-rate schedule with 3% warmup), with the loss applied only to the model-generated thoughts and actions, not the environment observations.

Next, the distilled student is run on the same questions to produce its own trajectories. Each turn is scored for how much it repeats earlier turns: by word overlap (ROUGE-L against the most similar prior turn), by embedding similarity, or by an LLM judge. The average student score minus the average teacher score gives a stagnation gap.

That gap, combined with whether the student's final answer was wrong, becomes an adaptive margin: students that stagnate more and answer incorrectly get a larger margin and thus a stronger corrective push, while students already progressing smoothly get a gentle update. Training uses a SimPO-style sigmoid objective that raises the relative likelihood of the teacher trajectory over the student trajectory by at least that margin, plus the original SFT loss to prevent both likelihoods from collapsing. The whole procedure can be repeated, with each round generating fresh student trajectories and recomputing penalties.

Why This Matters

Impact on research: The paper reframes agent distillation as a problem of behavioral dynamics over long horizons rather than simple imitation accuracy. It introduces a measurable diagnostic (the teacher–student stagnation gap) and shows that training-time intervention is needed, since decoding-level repetition penalties and shortening teacher demonstrations both fail to address the issue.

Real-world applications:

• Web information-retrieval agents that must issue new, informative search queries instead of repeating failed ones.
• Software engineering agents that iterate on code and error logs without looping on the same broken fix.
• Mathematical and scientific problem-solving assistants that use code execution and tool calls across multiple steps.
• Scientific report generation and other long-horizon tool-use pipelines where token cost and latency matter.

Industry relevance: SLM agents are attractive because they are cheaper and faster than GPT-4.1-class systems, but stagnation erodes that advantage through wasted turns and tokens. ProD's reported 11.6% reduction in turns and 20.7% reduction in tokens, alongside higher success rates, directly targets deployment cost. It also offers a path to running capable agents on smaller, locally hosted models.

Future Directions

• Developing finer-grained progress metrics and more precise methods for localizing exactly which actions are stagnant, so corrections can be targeted rather than trajectory-wide.
• Extending iterative ProD beyond the two iterations studied, and understanding when a single round suffices versus when repeated rounds are needed.
• Testing the adaptive-margin formulation with the semantic and LLM-as-judge stagnation scores as the primary signal rather than the lexical score used in the main results.
• Reducing the training overhead of ProD, which requires generating student trajectories and computing stagnation scores, and analyzing that cost trade-off in more depth.

Target Audience

Machine learning researchers working on agent distillation, small language models, or preference optimization; practitioners deploying tool-using agents under cost and latency constraints; and graduate students interested in failure modes of imitation-based training for multi-turn agents. Readers should be comfortable with supervised fine-tuning, preference-style objectives such as SimPO, and the ReAct/CodeAct interaction format.

Authors’ abstract

While LLM-based agent systems have demonstrated remarkable proficiency in reasoning and solving complex real-world tasks, they also come with substantial token consumption and inference latency, limiting their practical deployment. To enhance the efficiency of agent systems, recent efforts have focused on transferring agentic capabilities from large-scale models to small language models (SLMs) via supervised fine-tuning-based distillation. Nevertheless, as interaction trajectories grow longer, involving extended reasoning-action chains and extensive environmental feedback, distilled student agents tend to experience progress stagnation: becoming trapped in unproductive loops characterized by persistent failed strategies and difficulty making meaningful progress. Our systematic analysis shows that this phenomenon consistently occurs across diverse SLMs and task domains. To address this issue, we propose Progress-Aware Distillation (ProD), which explicitly detects and penalizes stagnation in student-generated trajectories. ProD encodes the stagnation gap between student and teacher into an adaptive margin for dynamically penalizing possible stagnation. By iteratively applying this procedure, ProD enables student agents to progressively approximate teacher behavior. Extensive experiments involving 8 SLMs, ranging from 0.6B to 8B parameters, on both in-domain and out-of-domain benchmarks, demonstrate that ProD substantially mitigates progress stagnation while enhancing both the success rates and task-completion efficiency of distilled SLM agents.