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

Progress-Aware Distillation for Mitigating Stagnation in Small Language Model Agents. By Yuanpu Cao and five others, from Penn State University and Microsoft CoreAI. Published on the seventeenth of October. Its code repository is public and real: about forty source files.

Agent systems built on large language models reason well and solve complex real-world tasks, but they consume a lot of tokens and add inference latency, which limits deployment. One response has been to transfer agentic capability from large models to small language models, through distillation based on supervised fine-tuning. That works until the interactions get long. As trajectories stretch out, with extended chains of reasoning and action and a lot of feedback from the environment, distilled student agents tend to stagnate. They get trapped in unproductive loops, persisting with strategies that have already failed, and struggle to make meaningful progress. The authors' systematic analysis finds this happens consistently, across diverse small models and task domains.

Progress-Aware Distillation, or ProD, targets stagnation directly. It detects stagnation in the trajectories the student itself generates, and penalizes it. The emphasis on the student's own trajectories is the key choice. Ordinary distillation shows the student the teacher's path. ProD also looks at the paths the student actually takes, which is where the loops occur, and makes getting stuck something the training objective can see and push against.

The penalty is calibrated against the teacher. ProD encodes the stagnation gap between student and teacher into an adaptive margin, and uses that margin to penalize likely stagnation dynamically. The procedure is applied iteratively, so that with each round the student agent moves closer to the teacher's behaviour. The teacher supplies the reference point. Stagnation is judged relative to how the teacher behaves, not against some fixed standard, and the margin is adaptive rather than set once.

The experiments involve eight small language models, ranging from zero point six billion to eight billion parameters, tested on both in-domain and out-of-domain benchmarks. ProD substantially mitigates progress stagnation, and it improves both the success rate of the distilled agents and how efficiently they complete tasks.

The diagnosis is the useful part to carry away. Imitating a teacher's trajectories teaches a student what good steps look like, but says little about what to do once it has wandered into a loop the teacher never entered. ProD measures the student against the teacher on the student's own ground, and the result is a small agent that spends fewer steps going in circles.