技术摘要：LARK (基于可学习性的轨迹选择)

问题陈述

本文探讨了**推理蒸馏（reasoning distillation）**面临的挑战，即通过微调较小的学生大语言模型（LLM），使其模仿由较大的教师模型生成的推理轨迹。现有的数据选择方法通常在问题层面进行操作，或依赖于轨迹质量、正确性或模型置信度等启发式指标，但它们往往忽略了一个关键因素：可学习性（learnability）。

具体而言，对于同一个问题，可能存在多个候选推理轨迹（由不同的教师模型、采样运行或风格生成）。即使这些轨迹最终都导向相同的正确答案，它们提供的训练信号也截然不同。一条看起来“高质量”或“自然”的轨迹，并不一定是学生模型学习效率最高的轨迹。核心问题在于：我们能否设计一种有原则的轨迹选择准则，以确保蒸馏出的推理轨迹对学生模型是可学习的？

方法论：LARK

作者提出了 LARK（基于可学习性的锚点时间排序法，Learnability-grounded Anchor-time Ranking），该方法基于优化理论中的可学习性定义来选择轨迹，而非使用静态质量指标。

1. 将可学习性作为优化目标

LARK 将轨迹选择建模为一个策略优化问题。它定义了一个可学习性因子 $\rho(\theta; q)$，用于表征学生模型训练损失下降的速率。该因子定义为梯度平方范数与损失的比值：
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
其中 $q$ 是针对 $K$ 个候选轨迹的权重向量。在 Polyak–Łojasiewicz (PL) 条件下，较高的 $\rho$ 意味着更快的收敛速度。

2. 锚点时间近似

直接在整个微调轨迹上优化时间积分的可学习性目标是难以实现的。LARK 引入了锚点相对条件（anchor-relative condition），该条件指出，对于处于“懒惰训练（lazy training）”阶段的过度参数化模型，其可学习性速率与其在初始锚点参数 $\theta_{ref}$ 处的数值成正比。这使得可以通过最大化 $\rho(\theta_{ref}; q)$ 来近似该目标。

3. 局部线性化与前向传播代理

直接最大化 $\rho(\theta_{ref}; q)$ 存在问题，因为它是一个拟凸函数，会坍缩为退化的“one-hot”解（即仅选择单个轨迹）。为了避免这种情况，LARK 采取了以下措施：

• 局部线性化： 使用一阶泰勒展开，在均匀分布 $p$（即不进行选择的标准 SFT）附近对 $\rho$ 进行展开。
• 前向传播代理 ($\hat{g}_k$)： 为了避免为每个候选轨迹计算精确梯度（这需要反向传播），LARK 推导出了一个仅需前向传播即可计算的代理值 $\hat{g}_k$。该代理值通过学生的 Token 级预测残差和交叉熵损失来估计可学习率的局部梯度。
  • 具体而言，$\hat{g}_k$ 倾向于那些学生表现出“有结构的自信错误”的轨迹：即损失 $\ell_k$ 较高（表明预测较差），但残差质量集中在少数几个“近失（near-miss）”Token 上（从而提供清晰的修正信号）。

4. $\chi^2$-正则化选择策略

为了防止选择器通过偏离原始数据分布（这会损害泛化能力）来“欺骗”可学习性准则，LARK 引入了一个 $\chi^2$-正则化项。由此产生了一个闭式、受预算参数化的选择规则：

• 轨迹根据其得分 $\hat{g}_k$ 进行排序。
• 选择前 $B$ 个轨迹。
• 不同于硬性的前 $B$ 个截断，LARK 分配的软权重与第 $(B+1)$ 个轨迹以上的得分裕量成正比。该权重规则旨在最大化可学习性提升的下界，同时控制与均匀先验之间的 $\chi^2$ 散度。

核心贡献

1. 公式化： 作者将推理轨迹选择表述为一个基于可学习性的策略优化问题，引入了锚点时间可学习率 $\rho$ 作为一个有原则的目标。
2. 算法： 他们推导出了一个闭式的、受预算参数化的选择规则，该规则使用前向传播代理 $\hat{g}_k$ 并结合 $\chi^2$-正则化。这避免了昂贵的逐轨迹反向传播，同时在理论上保证了估计误差和泛化能力的边界。
3. 实验验证： 论文证明了 LARK 在多种基座模型（Qwen-2.5-7B/1.5B, Llama-3.2-3B）和推理基准测试（AIME, AMC, GPQA, MATH-500）中，一致优于现有的基线方法（包括基于启发式、质量或学生侧评分的方法）。

实验结果

• 性能： LARK 在所有测试的模型预算设置下均取得了最高的平均 Acc@5。例如，在 Qwen-2.5-7B 且预算为 $B=3$ 的情况下，LARK 比最强的基线 RSR 提升了 7.07 个百分点。
• 效率： LARK 是一种仅需前向传播的方法，其计算成本与成本最低的学生侧基线方法相当，而精确的基于梯度的法需要显著更多的计算量。
• 诊断分析：
  • 损失衰减： LARK 选择的轨迹在训练期间诱导了最快的 SFT 损失下降。
  • 预测有效性： 通过按教师模型聚合后，LAK 的得分 $\hat{g}_k$ 与下游准确率表现出强正相关性（Pearson $r=0.69$）。
  • Token 级行为： 被选中的轨迹表现出较低的正确 Token 概率，但具有更高的“错误 Token 集中度”，这证实了 LARK 选择的是那些学生虽然犯错、但错误具有结构性且可纠正的样本。

重要性与主张

本文声称，可学习性是高效推理蒸馏中一个独特且有效的原则，它独立于轨迹质量或难度。通过将选择过程植根于学生模型的优化动力学，LARK 识别出了能提供最高效学习信号的轨迹。

作者强调，其方法：

• 保持了泛化能力： 通过 $\chi^2$-正则化，防止选择策略过拟合到一组狭窄的“容易”或“被操纵”的轨迹。
• 具有理论依据： 其选择规则是基于优化理论（PL 条件、梯度流）的第一原理推导而来的，并具有明确的误差界限。
• 具有实用性： 它不需要为选择规则本身进行超参数调优（正则化强度由预算 $B$ 决定），并且避免了反向传播带来的计算开销。

这项工作表明，未来的扩展可以将这种基于可学习性的视角应用于噪声更大的候选池和更广泛的推理领域，但目前的范围严格限定在数学问题求解的推理蒸馏中。