Aligning reasoning steps to help smaller AI models learn better
From Solver Trajectories to Teaching Trajectories: Cognition-Aligned Reasoning Distillation
Artificial Intelligence
Summary
Large AI models often show how they solve problems step-by-step, but these detailed steps can be hard for smaller AI models to copy and learn from directly. The authors propose a method called CARD that reshapes these problem-solving steps to better fit the smaller model's way of understanding. This makes the smaller models learn more effectively and generalize better to new kinds of problems. CARD works by tweaking existing problem-solving examples using rewards that encourage the smaller model to follow clearer and more understandable steps.
What this means in practice
- •For machine learning engineers: Improve smaller AI model training by transforming teacher model reasoning steps to better match the smaller model's learning style.
- •For ai system architects: Enhance robustness of AI reasoning systems by producing teaching data that helps smaller models generalize beyond training examples.
Authors
Yashuo Luo, Qianren Mao, Ziqi Qin, Junnan Liu, Haitao Li, Qili Zhang, Yuan Gao, Zhijun Chen, Xu Wang
Abstract
Knowledge distillation from large reasoning models (LRMs) often uses raw teacher trajectories as supervision. However, this practice creates a mismatch between the goal of distillation and the nature of teacher trajectories. Advanced LRMs typically improve their reasoning capabilities through reward-driven post-training, such as RLHF. This training optimizes them as strong solvers, rather than as teachers that produce reasoning processes smaller models can easily learn. Since smaller models have different "cognitive capacities" compared to their larger teachers, directly imitating raw teacher trajectories can sometimes be ineffective and often requires a substantial number of high-quality samples. To bridge this mismatch, we propose Cognition-Aligned Reasoning Distillation (CARD), an RL-based framework that transforms solver trajectories into teaching trajectories aligned with the student's cognition. Avoiding expensive resampling, our framework directly transforms existing trajectories. We optimize CARD with two reward signals: (1) a Distribution Alignment Reward that matches the student's predictive distribution, and (2) a Logical Coherence Reward that reduces misleading failed reasoning steps. Comprehensive evaluations across multiple student models and settings demonstrate that CARD outperforms distillation from raw teacher trajectories in in-domain settings while achieving substantial gains in out-of-distribution performance. Code is available at https://github.com/LiteSSLHub/CARD.