TeacherGRPO improves reasoning in small AI models with teacher alignment
TeacherGRPO: Closing the Capacity Gap in Reasoning Distillation via Teacher Alignment
Computation and Language
Summary
Smaller AI models often struggle to learn from smarter larger models because the bigger models’ complex reasoning is too hard to copy. The authors found that adjusting the larger teacher models to better fit what the smaller student models can handle helps bridge this gap without losing reasoning quality. They created a method called TeacherGRPO, which uses a special training approach to carefully align the teacher’s knowledge with the student’s learning ability. This approach outperforms earlier methods that either filtered data or weakened teacher models, and works well across several reasoning tasks.
What this means in practice
- •For machine learning engineers: Improve performance of small AI models by training larger models to better match their reasoning style before knowledge transfer.
- •For ai product developers: Develop smarter but more efficient AI services that rely on smaller models with reasoning ability closer to large models.$Commercial implications: Enables AI products that deliver advanced reasoning while requiring less computing power, appealing to cost-sensitive customers.
Authors
Zhenyu Lei, Zihan Chen, Yaochen Zhu, Shangbin Feng, Zaiyi Zheng, Ruocheng Guo, Yushun Dong, Jundong Li
Abstract
Reasoning distillation from powerful teacher models to smaller students faces the Gap Curse: as teachers grow more sophisticated, their complex distributions increasingly diverge from what students can approximate, causing performance degradation. Existing mitigation strategies either filter out challenging examples through data selection or introduce weaker intermediate assistant models, inherently compromising supervision coverage or quality. We propose Teacher Alignment, which directly adapts the teacher toward the student's distribution without discarding data or degrading reasoning quality. However, naive alignment through standard knowledge distillation triggers catastrophic collapse of the teacher's reasoning capabilities. To address this, we reformulate teacher alignment as reinforcement learning and introduce TeacherGRPO, built on Group Relative Policy Optimization with two key innovations: (i) Curriculum Selective Alignment applies dual token- and distribution-level curricula to focus rewards on high-signal reasoning gaps while filtering noise from trivial tokens and uncertain tail distributions, and (ii) Importance-Adaptive Length Regularization selectively penalizes verbose redundancy while preserving pedagogically critical reasoning steps. The aligned teacher then distills knowledge to students via standard pipelines. Extensive experiments show TeacherGRPO significantly outperforms baselines across diverse reasoning benchmarks and distillation methods. Our code is available at https://github.com/LzyFischer/TeacherGRPO.