LARK: Learnability-Grounded Trajectory Selection for Efficient Reasoning Distillation
This work provides an incremental improvement in data selection for reasoning distillation, benefiting researchers and practitioners working on efficient student model training.
This paper addresses the problem of selecting teacher-generated reasoning trajectories for student model supervision in reasoning distillation. The authors propose LARK, a method that selects trajectories based on their learnability by the student, characterized by a learnability factor $\rho$, and show it consistently outperforms existing data selection baselines across multiple models and reasoning tasks.
We study trajectory selection for reasoning distillation, where teacher-generated reasoning trajectories are selectively used as supervision for a student model. Existing methods rely on heuristics such as trajectory quality or model confidence, but they often overlook whether a trajectory is learnable by the student. In this paper, we present LARK, a learnability-grounded method for reasoning trajectory selection. LARK selects trajectories that the student can learn efficiently while preserving the generalization of the full training distribution. At the core of LARK is a learnability factor $ρ$, which characterizes the rate at which the student's training loss decreases. To estimate this rate efficiently and maintain generalization, we introduce a learnability proxy and a $χ^2$-regularized selection policy that balances learnability and distributional coverage, both with strong theoretical guarantees on their estimation error. Empirically, LARK consistently outperforms data selection baselines across multiple base models and reasoning tasks. Diagnostic analyses show that the LARK score predicts downstream training utility and that LARK-selected trajectories induce faster supervised fine-tuning loss reduction. Our code is available at https://github.com/Tianrun-Yu/LARK.