LARK: Learnability-Grounded Trajectory Selection for Efficient Reasoning Distillation
作者: Tianrun Yu, Kaixiang Zhao, Chih-Chun Chen, Amanda Hughes, Taylor W. Killian, Fenglong Ma, Weitong Zhang, Porter Jenkins
分类: cs.LG, cs.AI
发布日期: 2026-05-28
备注: 43 pages, 9 figures, 2 tables
🔗 代码/项目: GITHUB
💡 一句话要点
提出LARK以解决推理蒸馏中的轨迹选择问题
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 推理蒸馏 轨迹选择 可学习性 模型训练 机器学习
📋 核心要点
- 现有的推理蒸馏方法多依赖启发式策略,忽视了轨迹的可学习性,导致选择效率低下。
- LARK通过引入可学习性因子和正则化选择策略,优化了轨迹选择过程,确保学生模型的学习效率。
- 实验结果表明,LARK在多个任务中显著提升了模型的训练效果,且加速了监督微调的损失降低。
📝 摘要(中文)
我们研究了推理蒸馏中的轨迹选择问题,其中教师生成的推理轨迹被选择性地用作学生模型的监督。现有方法依赖于轨迹质量或模型置信度等启发式方法,但往往忽视了轨迹是否可被学生学习。本文提出了LARK,一种基于可学习性的推理轨迹选择方法。LARK选择学生能够高效学习的轨迹,同时保留完整训练分布的泛化能力。LARK的核心是一个可学习性因子$ρ$,它表征学生训练损失下降的速率。为了高效估计这一速率并保持泛化性,我们引入了可学习性代理和$χ^2$正则化选择策略,平衡可学习性和分布覆盖,且在估计误差上具有强理论保证。实验证明,LARK在多个基础模型和推理任务中始终优于数据选择基线。
🔬 方法详解
问题定义:本文旨在解决推理蒸馏中轨迹选择的效率问题。现有方法往往依赖于轨迹质量或模型置信度,未能考虑轨迹的可学习性,导致选择的轨迹可能不适合学生模型的学习。
核心思路:LARK提出了一种基于可学习性的轨迹选择方法,旨在选择学生能够高效学习的轨迹,同时保持训练分布的泛化能力。通过引入可学习性因子$ρ$,LARK能够更好地评估轨迹的学习潜力。
技术框架:LARK的整体架构包括可学习性代理和$χ^2$正则化选择策略。可学习性代理用于估计学生模型的学习速率,而正则化选择策略则在选择轨迹时平衡可学习性和分布覆盖。
关键创新:LARK的主要创新在于引入了可学习性因子$ρ$和正则化选择策略,这与传统方法的启发式选择形成了鲜明对比,确保了选择的轨迹更符合学生模型的学习需求。
关键设计:在设计中,LARK通过可学习性代理来动态评估轨迹的学习潜力,并使用$χ^2$正则化来控制选择过程中的估计误差,确保选择的轨迹既具备学习价值,又能覆盖训练分布。
🖼️ 关键图片
📊 实验亮点
实验结果显示,LARK在多个基线模型和推理任务中均表现优异,显著提升了模型的训练效果。具体而言,LARK选择的轨迹能够加速监督微调过程中的损失降低,且在多个任务中均超过了传统数据选择方法,提升幅度达到了XX%。
🎯 应用场景
该研究在推理蒸馏、模型训练优化等领域具有广泛的应用潜力。通过提高轨迹选择的效率,LARK能够加速学生模型的训练过程,提升模型在实际任务中的表现,尤其是在需要快速适应新数据的场景中。未来,LARK的理念也可扩展到其他机器学习任务中,促进更高效的学习策略的开发。
📄 摘要(原文)
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.