LARK: Learnability-Grounded Trajectory Selection for Efficient Reasoning Distillation

2026-05-28

AI Summary

This work addresses a critical limitation in existing reasoning distillation methods, which overlook the actual learnability of teacher-generated trajectories by the student model. To remedy this, the authors propose a learnability-aware trajectory selection mechanism that introduces a learnability factor ρ and an efficient proxy metric, combined with a χ² regularization strategy. This approach selects trajectories that are not only representative of the output distribution but also more amenable to efficient learning by the student. Evaluated across multiple foundation models and reasoning tasks, the method significantly outperforms current data selection baselines, yielding faster training loss reduction. Furthermore, the proposed LARK score effectively predicts downstream training utility, offering both theoretical error bounds and strong empirical performance.

Abstract

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.

Problem

• reasoning distillation
• trajectory selection
• learnability
• student model
• supervision

Innovation

• learnability
• reasoning distillation
• trajectory selection
• data efficiency
• generalization