TLDR: AdaSwitch is a novel knowledge distillation method that dynamically combines on-policy and off-policy generation at the token level for small language models (SLMs). It allows SLMs to explore their own predictions and then selectively integrates teacher guidance based on real-time quality assessment, addressing the training-inference mismatch and low-quality output issues of existing methods. This approach consistently improves accuracy and maintains computational efficiency across various tasks and model pairs.
Small language models (SLMs) are becoming increasingly vital for applications where speed and computational resources are limited, such as in search engines and recommendation systems. However, achieving high performance with these smaller models can be a significant challenge. Knowledge Distillation (KD) is a powerful technique that helps transfer the advanced capabilities of large teacher models to smaller student models. Yet, current KD methods often face a dilemma: off-policy distillation offers high-quality supervision but can lead to a mismatch between training and real-world inference, while on-policy approaches maintain consistency but rely on the student’s potentially lower-quality outputs.
To overcome these limitations, researchers have introduced AdaSwitch, a new approach that intelligently blends on-policy and off-policy generation at the token level. AdaSwitch empowers the student model to first explore its own predictions and then, based on a real-time assessment of prediction quality, selectively incorporates guidance from the more capable teacher model. This dynamic switching mechanism ensures both consistency in training and high-quality supervision.
How AdaSwitch Works
AdaSwitch operates through a two-stage sequence generation process. Initially, the student model enters an ‘exploration stage,’ where it autonomously generates an initial sequence. This stage is crucial for maintaining consistency between how the model is trained and how it will perform during inference. As the student generates tokens, AdaSwitch continuously measures the ‘divergence’ – essentially, how much the student’s prediction for the next token differs from the teacher’s prediction.
To make this assessment adaptive, AdaSwitch uses a ‘sliding window’ to calculate a moving average of recent divergences. This average helps set a dynamic threshold. If the student’s current prediction divergence exceeds this adaptive threshold, AdaSwitch transitions into a ‘guidance stage.’ At this point, the teacher model takes over and generates the remaining sequence of tokens. A key design choice in AdaSwitch is that once the switch to the teacher occurs, all subsequent tokens for that sequence are generated by the teacher, preventing frequent, potentially disruptive, alternations between the two models. This ‘one-time switch’ mechanism helps reduce sequence distortion and prevents the student from overfitting to the teacher, striking a balance between consistency and quality.
Experimental Validation and Performance
The effectiveness of AdaSwitch was rigorously tested across three datasets: DialogSum for dialogue summarization, and GSM and GSM-Plus for arithmetic reasoning. The experiments utilized two prominent large language model families: Qwen 2.5 and Llama 3.1, with larger versions serving as teachers and smaller versions as students.
The results were compelling. AdaSwitch consistently improved performance across most scenarios, demonstrating its robustness and versatility. For instance, on the GSM arithmetic reasoning task, AdaSwitch significantly outperformed the second-best method by 7.2% for Llama models and 11.8% for Qwen models. This highlights AdaSwitch’s ability to effectively integrate the strengths of both on-policy and off-policy learning at a fine-grained, token level.
Beyond accuracy, AdaSwitch also proved to be computationally efficient. While it does incur a modest overhead compared to pure on-policy methods (averaging about 1.3 times the runtime), it is notably more efficient than other mixed methods like SKD, achieving a 10% reduction in time consumption. This acceptable overhead is particularly important as the knowledge distillation process typically occurs offline, meaning it doesn’t add latency to the final student model during real-world inference.
Also Read:
- STEPER: Empowering Smaller Language Models with Advanced Step-by-Step Reasoning
- LightReasoner: How Smaller AI Models Can Enhance Larger Ones
Deeper Insights into the Distillation Process
Further analysis revealed interesting dynamics within AdaSwitch. During distillation, the divergence between the student and teacher models significantly decreased, indicating that the student was progressively learning and aligning with the teacher. The ‘switch rate’ – the proportion of sequences that triggered a switch from student exploration to teacher guidance – also provided insights. More challenging tasks, like GSM, showed consistently high switch rates (above 95%), suggesting a greater reliance on teacher guidance. Easier tasks, like SUMM, exhibited lower rates, allowing for more student exploration. This demonstrates AdaSwitch’s adaptive nature, balancing exploration and guidance based on task difficulty.
Crucially, AdaSwitch also addresses the challenge of low-quality outputs in the early stages of on-policy KD. Experiments showed that AdaSwitch achieved the most rapid improvement rate in test performance during these initial steps, quickly surpassing other methods. This indicates that its token-level mixing effectively ensures high-quality generated sequences and reduces the training-inference discrepancy from the outset.
The research paper, AdaSwitch: Adaptive Switching Generation for Knowledge Distillation, provides a detailed account of this innovative framework.
In conclusion, AdaSwitch presents a practical and effective method for distilling small language models. By dynamically combining the best aspects of on-policy and off-policy learning at the token level, it significantly enhances accuracy and robustness while maintaining acceptable computational overhead, paving the way for more powerful and efficient SLMs in various applications.
