Small language models improve by avoiding repetitive stalls during tasks
Progress-Aware Distillation for Mitigating Stagnation in Small Language Model Agents
Machine Learning
Summary
When small language models try to perform tasks requiring many steps, they often get stuck repeating the same mistakes without making progress. The authors studied this problem and found it happens widely across models and tasks. They created a method called Progress-Aware Distillation that helps these smaller models recognize and avoid these stuck points by learning from larger, more capable models. This method helps the small models get better at completing tasks more efficiently and successfully.
What this means in practice
- •For ai engineers: Train smaller AI agents that avoid unproductive loops, improving task-solving efficiency and reliability
- •For chatbot developers: Build chatbots that handle complex interactions without getting stuck repeating ineffective responses
Authors
Yuanpu Cao, Saket Sathe, Hanyu Wang, Ziyi Yin, Fenglong Ma, Jinghui Chen
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.