LARK：基于可学习性的轨迹选择方法用于高效推理蒸馏

# LARK：基于可学习性的推理蒸馏轨迹选择
来源：https://arxiv.org/html/2605.30651
Tianrun Yu¹, Kaixiang Zhao¹, Chih\-Chun Chen¹, Amanda Hughes¹ Taylor W\. Killian¹, Fenglong Ma², Weitong Zhang³, Porter Jenkins¹,∗  
¹Brigham Young University  
²The Pennsylvania State University  
³University of North Carolina at Chapel Hill  
\*通讯作者：pjenkins@cs\.byu\.edu

###### 摘要
我们研究推理蒸馏中的轨迹选择问题，即教师生成的推理轨迹被选择性地用作学生模型的监督信号。现有方法依赖启发式标准，如轨迹质量或模型置信度，但它们常常忽略一条轨迹对学生而言是否可学习。在本文中，我们提出 LARK¹¹我们的代码可在https://github.com/Tianrun-Yu/LARK获取。，一种基于可学习性的推理轨迹选择方法。LARK 选择那些学生能够高效学习，同时保留完整训练分布泛化能力的轨迹。LARK 的核心是一个可学习性因子 ρ\\rho，它刻画了学生训练损失下降的速率。为了高效估计该速率并保持泛化性，我们引入了一个可学习性代理和一个 χ2\\chi^\{2\}-正则化选择策略，该策略平衡可学习性与分布覆盖，两者都具有关于其估计误差的强理论保证。实验上，LARK 在多个基础模型和推理任务上持续优于数据选择基线。诊断分析表明，LARK 分数能预测下游训练效用，且 LARK 选择的轨迹能诱导更快的监督微调损失降低。

## 1 引言
推理蒸馏已发展成为一种有效的范式，用于将链式推理能力从较大的教师大语言模型（LLMs）转移到较小的学生模型中（Hsieh 等人，2023 (https://arxiv.org/html/2605.30651#bib.bib37)；Yuan 等人，2023 (https://arxiv.org/html/2605.30651#bib.bib36)）。在此设定下，教师模型生成推理轨迹，学生模型通过微调来模仿它们。近期研究表明，*一组精心挑选的少量推理示例*能带来显著的性能提升（Ye 等人，2025 (https://arxiv.org/html/2605.30651#bib.bib1)；Muennighoff 等人，2025 (https://arxiv.org/html/2605.30651#bib.bib2)），常常能与使用更大训练集的效果相媲美。这些发现表明，推理蒸馏既需要*生成*足够的推理数据，也需要*识别*出对学生最有用的监督信号。

许多现有的推理蒸馏数据选择方法仍然继承了 LLM 微调中数据选择的经典公式（Xia 等人，2024 (https://arxiv.org/html/2605.30651#bib.bib10)；Xiao and Chen，2025 (https://arxiv.org/html/2605.30651#bib.bib9)），其中选择是在问题或样本级别进行的（Yu 等人，2023 (https://arxiv.org/html/2605.30651#bib.bib38)；Zhang 等人，2025b (https://arxiv.org/html/2605.30651#bib.bib6)；Liu 等人，2024 (https://arxiv.org/html/2605.30651#bib.bib7)）。然而，推理蒸馏提出了一个更细粒度的选择问题。对于单个问题，我们通常有多个候选推理轨迹，由不同的教师模型、采样运行或推理风格生成。即使多条轨迹都能得到相同的正确答案，它们为当前学生模型提供的训练信号也可能大相径庭。一条看似自然、高质量或与学生模型高度对齐的轨迹，并不一定是学生能最高效学习的轨迹。

现有的轨迹选择方法大多依赖启发式标准。一些方法使用外部验证器或 LLM-as-a-judge 分数来评估推理质量（Zheng 等人，2023 (https://arxiv.org/html/2605.30651#bib.bib39)；Yang 等人，2025 (https://arxiv.org/html/2605.30651#bib.bib24)），而 GRAPE、Local Naturalness 和 RSR 则使用学生模型本身对候选轨迹进行评分（Zhang 等人，2025a (https://arxiv.org/html/2605.30651#bib.bib4)；Just 等人，2025 (https://arxiv.org/html/2605.30651#bib.bib8)；Yang 等人，2026 (https://arxiv.org/html/2605.30651#bib.bib3)）。尽管这些方法很有用，但它们并未明确衡量一条轨迹是否被学生模型*可学习*。这一差距引出了核心问题：*我们能否设计一个原则性的轨迹选择标准，确保蒸馏后的推理轨迹对学生模型是可学习的？*

我们通过提出 LARK—Learnability\-grounded Anchor\-time Ranking（基于可学习性的锚点时间排序）来回答这个问题，这是一种基于可学习性的推理轨迹选择方法，如图 1 (https://arxiv.org/html/2605.30651#S1.F1) 所示。LARK 识别出学生模型能最高效学习的轨迹子集。其核心是，LARK 引入了一个原则性的选择标准，从优化角度刻画轨迹的可学习性，同时通过 χ2\\chi^\{2\}-正则化保持泛化性。

我们的贡献如下：
- •我们将推理轨迹选择形式化为一个基于可学习性的策略优化问题。我们引入了锚点时间可学习性速率 ρ\\rho，它刻画了训练后损失的衰减速率，从而将可学习性转化为一个原则性的优化目标。
- •为了高效估计可学习性速率 ρ\\rho，同时防止选择器滥用可学习性标准，我们使用围绕未选择数据分布的一阶泰勒展开，并推导出一个 χ2\\chi^\{2\}-正则化的策略优化问题。理论上，我们证明该策略优化隐式地增加了数据的可学习性，并且在固定预算的轨迹选择下可以闭式求解。
- •实验上，我们展示了 LARK 在多个基础模型和任务上优于现有基线。诊断分析进一步支持了我们的理论主张，并验证了基于可学习性视角进行轨迹选择的合理性。

参见标题
图 1：LARK 流程概览。对于每个问题，多条教师生成的推理轨迹构成一个候选池。LARK 使用学生模型对这些轨迹进行评分，按分数排序，并选择前 BB 条轨迹进行加权 SFT。在图中，aka\_\{k\} 表示排序后的第 kk 条轨迹，gkg\_\{k\} 表示实际的 LARK 分数 g^k\\hat\{g\}\_\{k\}，qkq\_\{k\} 表示其训练权重。排序面板中的索引在排序后重新编号。

##### 符号。向量用小写粗体字母表示，例如 x\\mathbf\{x\}。对于序列响应，y0T\>0 表示总微调时间。对于固定权重向量 q\\mathbf\{q\}，梯度流 φq:\[0,T\]→R|θ|\\phi\_\{\\mathbf\{q\}\}:\[0,T\]\\to\\mathbb\{R\}^\{\|\\boldsymbol\{\\theta\}\|\} 描述了学生模型在加权目标 L\(⋅;q\)\\mathcal\{L\}\(\\cdot;\\,\\mathbf\{q\}\) 上微调时参数的演化，定义为 ODE 的解 φ ̇q\(s\)=−∇θL\(φq\(s\);q\)，φq\(0\)=θref，s∈\[0,T\]。\\displaystyle\\dot\{\\phi\}\_\{\\mathbf\{q\}\}\(s\)=\-\\nabla\_\{\\boldsymbol\{\\theta\}\}\\mathcal\{L\}\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)，\\qquad\\phi\_\{\\mathbf\{q\}\}\(0\)=\{\\boldsymbol\{\\theta\}\}\_\{\\mathrm\{ref\}\}，\\qquad s\\in\[0,T\]。这里 φq\(0\)=θref\\phi\_\{\\mathbf\{q\}\}\(0\)=\\theta\_\{\\mathrm\{ref\}\} 是微调的起始点，φq\(T\)\\phi\_\{\\mathbf\{q\}\}\(T\) 是时间 TT 时的参数，ss 是连续梯度流时间变量。根据链式法则和方程 (2 (https://arxiv.org/html/2605.30651#S3.E2))，我们有 ddsL\(φq\(s\);q\)=−‖∇θL\(φq\(s\);q\)‖2=−ρ\(θ,q\)L\(θ,q\)，φq\(0\)=θref。\\displaystyle\\tfrac\{\\mathrm\{d\}\}\{\\mathrm\{d\}s\}\\mathcal\{L\}\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)=\-\\\|\\nabla\_\{\\boldsymbol\{\\theta\}\}\\mathcal\{L\}\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)\\\|^\{2\}=\-\\rho\(\\boldsymbol\{\\theta\},\\mathbf\{q\}\)\\mathcal\{L\}\(\\boldsymbol\{\\theta\},\\mathbf\{q\}\)，\\quad\\phi\_\{\\mathbf\{q\}\}\(0\)=\\theta\_\{\\mathrm\{ref\}\}。解此常微分方程得到 L\(φq\(T\);q\)=L\(θref;q\)⋅exp⁡\(−∫0Tρ\(φq\(s\);q\)ds\)，\\displaystyle\\textstyle\{\\mathcal\{L\}\(\\phi\_\{\\mathbf\{q\}\}\(T\);\\mathbf\{q\}\)=\\mathcal\{L\}\(\\theta\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)\\cdot\\exp\\left\(\-\\int\_\{0\}^\{T\}\\rho\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)\\mathrm\{d\}s\\right\)，\}\(3\) 这意味着积分 ∫0Tρ\(φq\(s\);q\)ds\\int\_\{0\}^\{T\}\\rho\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)\\mathrm\{d\}s 在控制损失衰减中扮演重要角色。

## 4 提出方法
在本节中，我们介绍 LARK 的方法论。具体来说，第 4.1 节 (https://arxiv.org/html/2605.30651#S4.SS1) 建立了方程 (3 (https://arxiv.org/html/2605.30651#S3.E3)) 中的时间积分可学习性目标可以由在初始学生模型处评估的锚点时间可学习性速率 ρ\(⋅\)\\rho\(\\cdot\) 控制。第 4.2 节 (https://arxiv.org/html/2605.30651#S4.SS2) 然后通过围绕均匀先验进行局部线性化来近似优化 ρ\(θref;q\)\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)，开发了一个前向传递代理 g^k\\hat\{g\}\_\{k\}，该代理无需反向传播即可估计 ρ\\rho 的局部梯度，并通过 χ2\\chi^\{2\}-正则化项来限制残差误差，以保持泛化性并防止对可学习性标准的奖励黑客。第 4.3 节 (https://arxiv.org/html/2605.30651#S4.SS3) 将这些部分组合成一个闭合形式、预算参数化的选择规则，无需调整超参数。

### 4.1 使用锚点相对条件估计可学习性目标
如方程 (3 (https://arxiv.org/html/2605.30651#S3.E3)) 所示，选择对学生模型*可学习*的轨迹可以形式化为优化时间积分可学习性目标 maxq∈ΔK∫0Tρ\(φq\(s\);q\)ds\\max\_\{\\mathbf\{q\}\\in\\Delta^\{K\}\}\\int\_\{0\}^\{T\}\\rho\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)\\mathrm\{d\}s，这有利于在 SFT 期间训练后损失快速衰减的轨迹分布。然而，直接优化这个目标是不可行的，因为它需要跟踪因子 ρ\\rho 沿整个 SFT 优化轨迹。为了解决这个挑战，我们引入以下结构条件，它允许可学习性目标从一个固定的锚点模型近似。

###### 条件 1。如果存在一个绝对常数 κ\>0\\kappa\>0 使得对于每个 q∈ΔK\\mathbf\{q\}\\in\\Delta^\{K\} 和每个 s∈\[0,T\]s\\in\[0,T\]，有 ρ\(φq\(s\);q\)≥κρ\(θref;q\)\\rho\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)\\geq\\kappa\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)，则称该模型满足锚点相对条件。

在条件 1 (https://arxiv.org/html/2605.30651#Thmcondition1) 下，在初始学生参数 θref\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\} 处评估的可学习性速率提供了一个易于处理的完整时间积分目标的代理。这个锚点相对条件与过度参数化的 LLM 在 SFT 过程中经常观察到的懒惰微调行为一致。在这个机制中，模型参数保持接近初始化 θref\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\}，因此可学习性速率 ρ\\rho 保持接近其锚点时间值 ρ\(θref;q\)\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)。我们在附录 A.2 (https://arxiv.org/html/2605.30651#A1.SS2) 中提供了在神经正切核机制下（NTK；Jacot 等人，2018 (https://arxiv.org/html/2605.30651#bib.bib43)）该条件的理论证明，并在附录 A.3 (https://arxiv.org/html/2605.30651#A1.SS3) 中提供了实证验证。以下命题自然成立；其证明见附录 A.1 (https://arxiv.org/html/2605.30651#A1.SS1)。

###### 命题 1。在条件 1 (https://arxiv.org/html/2605.30651#Thmcondition1) 下，对于每个 q∈ΔK\\mathbf\{q\}\\in\\Delta^\{K\}，学生模型的训练后损失满足 L\(φq\(T\);q\)≤L\(θref;q\)⋅exp⁡\(−κT⋅ρ\(θref;q\)\)。\\mathcal\{L\}\(\\phi\_\{\\mathbf\{q\}\}\(T\);\\mathbf\{q\}\)\\leq\\mathcal\{L\}\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)\\cdot\\exp\\bigl\(\-\\kappa T\\cdot\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)\\bigr\)。

命题 1 (https://arxiv.org/html/2605.30651#Thmproposition1) 表明，在锚点相对条件下，SFT 损失的衰减可以由在锚点模型 θref\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\} 处评估的可学习性速率控制。因此，我们可以优化锚点时间替代目标 maxq∈ΔK⁡ρ\(θref;q\)\\max\_\{\\mathbf\{q\}\\in\\Delta^\{K\}\}\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\) 以鼓励学生模型更快的损失衰减，而不是直接优化棘手的轨迹整体目标 ∫0Tρ\(φq\(s\);q\)ds\\int\_\{0\}^\{T\}\\rho\(\\phi\_\{\\mathbf\{q\}\}\(s\);\\mathbf\{q\}\)\\mathrm\{d\}s。

### 4.2 通过局部线性化近似优化可学习性 ρ\(θref;q\)\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)
在实践中，直接优化 ρ\(θref;q\)\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\) 是有问题的。如附录 B.1 (https://arxiv.org/html/2605.30651#A2.SS1) 所示，该目标在单纯形 ΔK\\Delta^\{K\} 上诱导出一个拟凸最大化问题，其最优解在单纯形的一个顶点处达到。因此，精确最大化会导致退化解，将所有概率质量放在单条轨迹上，从而“黑掉”可学习性标准，而不是产生有用的训练分布。这一观察激发我们将可学习性标准 ρ\\rho 视为围绕均匀分布 p=\(1K,⋯,1K\)∈RK\\mathbf\{p\}=\(\\tfrac\{1\}\{K\},\\cdots,\\tfrac\{1\}\{K\}\)\\in\\mathbb\{R\}^\{K\} 的*局部*指示器，该均匀分布对应于没有轨迹选择的标准 SFT。我们不追求全局最大化 ρ\(θref;q\)\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)，而是将注意力限制在 p\\mathbf\{p\} 邻域内的分布 q\\mathbf\{q\}，并围绕 p\\mathbf\{p\} 对目标进行局部线性化。一阶泰勒展开给出

ρ\(θref;q\)−ρ\(θref;p\)\\displaystyle\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)\-\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{p\}\)  
=⟨∇qρ\(θref;q\),q−p⟩\+o\(‖q−p‖2\)\\displaystyle=\\textstyle\{\\langle\\nabla\_\{\\mathbf\{q\}\}\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)，\\mathbf\{q\}\-\\mathbf\{p\}\\rangle\+o\(\\\|\\mathbf\{q\}\-\\mathbf\{p\}\\\|\_\{2\}\)\}  
=∑k=1K∂∂qkρ\(θref;q\)∣q=p⏟gk∗⋅\(qk−pk\)\+o\(‖q−p‖2\)⏟R2\(p,q\)。\\displaystyle=\\textstyle\{\\sum\_\{k=1\}^\{K\}\}\\underbrace\{\\tfrac\{\\partial\}\{\\partial q\_\{k\}\}\\rho\(\\boldsymbol\{\\theta\}\_\{\\mathrm\{ref\}\};\\mathbf\{q\}\)\\mid\_\{\\mathbf\{q\}=\\mathbf\{p\}\}\}\_\{g\_\{k\}^\{\*\}\}\\cdot\(q\_\{k\}\-p\_\{k\}\)\+\\underbrace\{o\(\\\|\\mathbf\{q\}\-\\mathbf\{p\}\\\|\_\{2\}\)\}\_\{R\_\{2\}\(\\mathbf\{p\},\\mathbf\{q\}\)\}。\(4\)

优化方程 (4 (https://arxiv.org/html/2605.30651#S4.E4)) 需要两步。首先，我们需要以计算高效的方式估计一阶导数 gk∗g\_\{k\}^\{\*\}。其次，我们需要仔细限制残差 R2\(p,q\)R\_\{2\}\(\\mathbf\{p\},\\mathbf\{q\}\)。

#### 4.2.1 估计梯度 gk∗g\_\{k\}^\{\*\}
根据 ρ\\rho 的定义，gk∗=∂∂qkρ\(θref;q\