LARK: Learnability-Grounded Trajectory Selection for Efficient Reasoning Distillation¶

会议: NeurIPS 2026
arXiv: 2605.30651
代码: https://github.com/Tianrun-Yu/LARK
领域: 模型压缩
关键词: 推理蒸馏, 轨迹选择, 可学习性, 梯度流优化, 卡方正则化

一句话总结¶

针对推理蒸馏中仅凭表面质量或学生似然度筛选教师思维链轨迹的启发式缺陷，LARK 从优化理论出发，提出基于初态梯度流损失衰减率的前向无反向代理评分 \(\hat{g}_k\) 与带闭式解的 \(\chi^2\) 正则化预算加权规则，实现高效且泛化保真的推理监督选择。

研究背景与动机¶

在大语言模型推理蒸馏领域，通过将大模型教师生成的思维链（Chain-of-Thought, CoT）推理轨迹迁移给小参数量学生模型进行监督微调（SFT），是赋予小模型复杂逻辑与多步推导能力的成熟范式。然而，在同一道题目下，不同教师模型或多次采样往往会产生大量终解正确但推导风格、长度与表述各异的推理轨迹；甚至多个候选轨迹均达到相同答案，它们提供给学生模型的训练梯度与优化指导信号却截然不同。以往的指令数据筛选研究大多停留在样本或题目层面的粗粒度过滤，无法有效应对同一题目内多条候选轨迹的细粒度甄别问题。

现有的轨迹级选择方法通常依赖启发式准则：一类采用外部验证器或 LLM-as-a-Judge 对文本质量与规范性打分，另一类则利用学生模型自身的反馈（如全句平均对数似然 GRAPE、局部上下文自然度 Local Naturalness，或基于排序与意外度的对齐指标 RSR）。这些指标虽然能筛掉部分低质或生硬文本，却完全忽视了关键的底层动力学问题——该轨迹对当前学生模型而言是否“具有可学习性”（Learnable）。表面最通顺或学生预测置信度最高的轨迹，往往是学生早已掌握的冗余信息，而过于晦涩复杂的轨迹又超出了学生的学习区。

这一瓶颈的核心矛盾在于：从优化理论角度衡量一条轨迹能否带来最快的经验风险收敛，严格依赖于参数空间中的梯度流损失衰减率，但这在全参数空间中需要对每条候选轨迹执行昂贵的反向传播计算，且全局无约束优化极易收敛至单点脉冲分布导致分布崩溃。核心 idea：本文将轨迹选择形式化为锚点时刻梯度流损失衰减率最大化问题，利用最后一层 Softmax 的结构性质推导出免反向传播的前向代理评分 \(\hat{g}_k\)，并通过 \(\chi^2\) 散度正则化约束推导出包含选定预算 \(B\) 的封闭解析解加权规则，使得学生在固定算力预算下精准聚焦于结构化纠错轨迹。

方法详解¶

整体框架¶

LARK（Learnability-grounded Anchor-time Ranking）由候选评分、排序截断与闭式权重分配三个阶段组成。对于每个推理问题，教师池生成多条最终答案正确的候选推理轨迹。首先，基准参考学生模型 \(\boldsymbol{\theta}_{\text{ref}}\) 仅需对每条候选轨迹执行单次前向推理，计算出序列平均交叉熵损失 \(\ell_k\) 以及基于预测残差的 Brier 分数，进而合成免反向传播的局部可学习性梯度代理值 \(\hat{g}_k\)；接着，系统根据 \(\hat{g}_k\) 对所有候选进行降序排列；最后，在预定的轨迹预算 \(B\) 下，利用卡方正则化策略优化在单纯形空间中直接解析计算出软截断训练权重 \(\hat{q}_k\)，输入学生模型进行加权 SFT 训练。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["题目及多教师候选轨迹池<br/>K条正确推导轨迹"] --> B["前向可学习性估计<br/>单次前向评估残差与损失"]
    B --> C["局部梯度代理评分<br/>相对均匀分布的一阶导数"]
    C --> D["闭式预算截断加权<br/>卡方正则化Top-B边际加权"]
    D --> E["加权SFT优化更新<br/>加速学生模型推理蒸馏"]

关键设计¶

1. 前向可学习性估计：利用输出层残差避免昂贵的反向传播

从多任务 Polyak-Łojasiewicz（PL）优化视角出发，轨迹分布 \(\mathbf{q} \in \Delta^K\) 的损失衰减速率由梯度范数与损失比值决定：\(\rho(\boldsymbol{\theta}; \mathbf{q}) \triangleq \|\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}; \mathbf{q})\|^2 / \mathcal{L}(\boldsymbol{\theta}; \mathbf{q})\)。若要在训练前直接为每条轨迹计算真实的单轨迹衰减率 \(\rho_k^* = \|\mathbf{g}_k\|^2 / \ell_k\)，必须在整个网络参数上执行反向传播计算梯度范数 \(\|\mathbf{g}_k\|^2\)，在大规模候选池中计算开销不可承受。

针对这一痛点，LARK 深入剖析了自回归语言模型最后一层线性映射与 Softmax 交叉熵损失的代数特性。设第 \(t\) 步学生的输出概率分布为 \(\boldsymbol{\pi}_t^k\)，教师真实标注的独热向量为 \(\boldsymbol{\delta}(y_t^k)\)，则输出层针对 Logits 的导数正是预测残差 \(\boldsymbol{\pi}_t^k - \boldsymbol{\delta}(y_t^k)\)。在过参数化 Transformer 满足局部隐藏状态范数稳定与雅可比有界条件下，全网梯度的范数平方被输出层梯度紧紧夹逼。因此，LARK 提出了仅依赖单次前向计算的代理指标：

\[\hat{\rho}_k \triangleq \frac{\sum_{t=1}^{|a_k|} \|\boldsymbol{\pi}_t^k - \boldsymbol{\delta}(y_t^k)\|^2}{\sum_{t=1}^{|a_k|} -\log \pi_{\boldsymbol{\theta}}(y_t^k \mid \mathbf{x}, \mathbf{y}_{<t}^k)}\]

理论证明存在常数使得 \(\hat{\rho}_k \in [C_1 \rho_k^*, C_2 \rho_k^*]\)，从而在完全不启动梯度反向传播的前提下，以极高的保真度刻画了单个轨迹对全网参数的动态更新势能。

2. 局部梯度代理评分：在未选择均匀分布邻域做一阶展开

直接在概率单纯形 \(\Delta^K\) 上全局最大化 \(\rho(\boldsymbol{\theta}_{\text{ref}}; \mathbf{q})\) 是一个拟凸优化（Quasi-convex）问题，其理论极值必然落在单纯形的极点上，导致模型将全部权重压在单一轨迹上，诱发严重的作弊式崩溃。为保留全分布的多样性与泛化覆盖度，LARK 将基准点设定在未做筛选的离散均匀先验 \(\mathbf{p} = (\frac{1}{K}, \dots, \frac{1}{K})\)，对 \(\rho\) 进行局部一阶泰勒展开。

通过假定高维参数空间中不同轨迹梯度的渐近正交性（\(\mathbf{g}_k^\top \mathbf{g}_j \approx 0, \forall k \neq j\)），精确方向导数可被解耦并用前向代理量替代，得到每条轨迹的局部边际贡献得分：

\[\hat{g}_k = \frac{\ell_k}{\sum_{i} \ell_i} \left( 2\hat{\rho}_k - \frac{\sum_i \hat{\rho}_i \ell_i}{\sum_i \ell_i} \right)\]

该得分具有直观的物理内涵：它偏好那些学生模型“自信地犯结构化错误”的轨迹。此时轨迹的总交叉熵损失 \(\ell_k\) 并不趋近于零（说明学生尚未完全学会），而残差能量 \(\|\boldsymbol{\pi}_t^k - \boldsymbol{\delta}(y_t^k)\|^2\) 却高度集中于少数极其接近真实 token 的近邻候选上，能够向网络注入高信噪比的局部校正引导。

3. 闭式预算截断加权：卡方正则化与无参 Top-B 软分配

在推导出一阶代理得分后，泰勒展开的高阶残差项必须得到严格控制，以防优化解远离有效展开邻域。理论证明，由代理替换引入的一阶误差和展开余项均受制于所选分布 \(\mathbf{q}\) 与先验均匀分布 \(\mathbf{p}\) 之间的卡方散度 \(\chi^2(\mathbf{q} \parallel \mathbf{p})\)。在实际蒸馏落地中，工程上通常需要将候选轨迹数严格限制在固定预算 \(B \le K\) 之内（即支撑集大小为 \(B\)）。

LARK 将该任务形式化为包含 \(\chi^2\) 散度惩罚的受限凸二次规划。对于排序后的得分 \(\hat{g}_1 \ge \dots \ge \hat{g}_K\)，系统给出了无需超参数搜索的解析最优解：

\[\hat{q}_i = \begin{cases} \dfrac{\hat{g}_i - \hat{g}_{B+1}}{\sum_{j=1}^B (\hat{g}_j - \hat{g}_{B+1})}, & \text{if } i \le B \\ 0, & \text{if } i > B \end{cases}\]

伴生的正则化强度 \(\tau(B) = \frac{1}{K}\sum_{j=1}^B (\hat{g}_j - \hat{g}_{B+1})\) 完全由得分与预算自适应内生确定。与粗暴分配等权 \(1/B\) 的硬截断相比，这种软截断机制依据得分与截断阈值 \(\hat{g}_{B+1}\) 之间的差额进行比例加权，有效缓解了边界处的排序噪声扰动。

实验关键数据¶

主实验¶

实验基于 NuminaMath 中 5,000 道数学难题构建候选池，涵盖 11 个开源与闭源顶尖教师模型（每个教师 3 次采样生成 33 条全正确轨迹）。评估采用 Acc@5（贪婪解码与多采样覆盖率评估），在 3 种不同系列与尺寸的学生模型（Qwen-2.5-7B, Llama-3.2-3B, Qwen-2.5-1.5B）上对比了单轨迹（\(B=1\)）与多轨迹（\(B=3\)）设置。

学生模型	评估配置 (\(B=1\) / \(B=3\))	AIME	AMC	GPQA	MATH-500	平均得分
Qwen-2.5-7B	BaseModel (零样本基线)	18.89 / 18.89	47.79 / 47.79	39.56 / 39.56	56.47 / 56.47	40.68 / 40.68
	Random (随机抽取)	21.11 / 24.44	51.00 / 61.45	56.90 / 59.76	64.67 / 69.67	48.42 / 53.83
	Token Length (选最长轨迹)	13.33 / 20.00	50.60 / 57.03	41.08 / 43.77	60.93 / 65.73	41.49 / 46.63
	Rule-based Quality (LIMO启发式)	21.11 / 28.89	56.63 / 67.07	48.15 / 53.54	65.27 / 70.73	47.79 / 55.06
	LLM-judged Quality (Qwen3-32B打分)	22.22 / 26.67	56.22 / 62.25	50.34 / 53.03	65.80 / 69.60	48.65 / 52.89
	GRAPE (全局负交叉熵)	16.67 / 25.56	55.02 / 63.05	49.33 / 55.05	64.67 / 71.00	46.42 / 53.66
	Local Naturalness (局部掩码似然)	23.33 / 27.78	57.83 / 64.26	51.18 / 56.23	65.80 / 70.87	49.54 / 54.78
	RSR (排序意外度比率基线)	24.44 / 32.22	59.44 / 66.27	52.19 / 56.73	67.33 / 71.47	50.85 / 56.67
	LARK (本文方法)	30.00 / 36.67	65.86 / 76.31	65.99 / 67.34	70.20 / 74.67	58.01 / 63.74
Llama-3.2-3B	BaseModel (零样本基线)	2.22 / 2.22	11.24 / 11.24	11.45 / 11.45	18.13 / 18.13	10.76 / 10.76
	RSR (最强对比基线)	6.67 / 5.56	15.66 / 17.27	42.09 / 43.77	25.73 / 27.87	22.54 / 23.62
	LARK (本文方法)	8.89 / 8.89	19.28 / 22.09	48.32 / 54.55	28.67 / 31.00	26.29 / 29.13
Qwen-2.5-1.5B	BaseModel (零样本基线)	0.00 / 0.00	24.50 / 24.50	30.64 / 30.64	38.53 / 38.53	23.42 / 23.42
	RSR (最强对比基线)	4.44 / 8.89	32.13 / 36.14	47.14 / 46.30	45.93 / 47.73	32.41 / 34.77
	LARK (本文方法)	7.78 / 10.00	40.56 / 42.57	63.13 / 66.67	49.33 / 53.33	41.09 / 43.14

消融实验¶

在 Qwen-2.5-7B 模型与 \(B=3\) 预算配置下，针对评分函数形态与权重分配机制的独立消融对比：

实验组别	变体设计	核心机制与计算形式	平均推理精度 (Acc@5)	表现分析与相对完整版落差点
评分函数消融 (固定 \(\chi^2\text{-}B\) 加权)	倒数似然 (\(1/\ell_k\))	仅考虑全局预测似然度高低	54.21%	效果最差，证明单纯拟合学生既有习惯无法带来显著性能增益 (-9.53%)
	纯残差项 (\(\text{Brier}\))	仅计算残差范数 \(\sum_t \\|\boldsymbol{\pi}_t^k - \boldsymbol{\delta}(y_t^k)\\|^2\)	57.65%	缺乏损失归一化导致难以区分绝对难度与相对增益 (-6.09%)
	轨迹代理率 (\(\hat{\rho}_k\))	单轨迹前向可学习率估计量	60.12%	考虑残差与损失比值，已显著超越传统基线 (-3.62%)
	完整 \(\hat{g}_k\) 得分	考虑单纯形均匀先验下边际梯度贡献	63.74%	全面捕捉各轨迹相对于未选分布的边际优化加速度
权重分配消融 (固定 \(\hat{g}_k\) 评分)	均匀硬截断 (\(\text{Uniform } 1/B\))	选出前 \(B\) 个后直接赋予等权重 \(1/3\)	59.88%	忽略得分差额，受边缘候选排序噪声影响显著 (-3.86%)
	线性比例加权 (\(\text{Proportional}\))	按分数值绝对大小直接归一化	61.05%	未设截断阈值边际差，权重区分度被截距稀释 (-2.69%)
	归一化 Softmax	引入指数平滑温度分布	61.42%	温度超参敏感且容易过度放大单峰分布 (-2.32%)
	闭式 \(\chi^2\text{-}B\) 加权	基于超阈值差额 \((\hat{g}_i - \hat{g}_{B+1})\) 闭式解	63.74%	理论推导最优分布约束，鲁棒吸收截断边界噪声

关键发现¶

可学习性显著超越表面质量与学生自洽性：以 Qwen-2.5-7B 为例，LARK 在 \(B=1\) 与 \(B=3\) 设置下分别以 58.01% 和 63.74% 的总均分刷新了 SOTA，相比于此前最强学生侧自评分基线 RSR 取得了超过 7.0 个百分点的全面领先，在难度极高的竞赛数学 AMC 上更是大幅领先近 10 个百分点。
Token 级别机制印证：诊断统计显示，相较于随机候选池，LARK 选拔出的轨迹在学生模型初态下的正确 Token 平均概率显著更低（0.282 vs. 0.449），但在错误 Token 上的平方浓度 \(\sum_{v \neq y_t} p(v)^2\) 却激增了 26%（0.107 vs. 0.085）。这实证了 LARK 确实精确定位到了“学生充满自信地在临门一脚犯错”的轨迹，赋予其最高信息增益的纠偏信号。
前向打分极其轻量：LARK 仅需单次前向传递和 \(O(|\mathcal{V}|)\) 的轻量残差统计，其计算复杂度与最基础的 GRAPE 一致，完全免去了 RSR 的 Top-K 排序或 Local Naturalness 的多次掩码前向计算开销，具备极强的工业规模化落地可行性。

亮点与洞察¶

将轨迹选择与梯度流损失衰减率建立闭环联系：巧妙利用 PL 条件将多轨迹加权转化为优化问题，跳脱出传统“主观好坏”或“单纯困惑度”的经验主义框框，为推理蒸馏的数据工程提供了坚实的数学底座。
Softmax 输出层雅可比技巧消除反向传播：巧妙借助过参数化模型在初始化邻域（Lazy Training）下的动力学特性，证明输出层 Logits 导数与全网梯度的范数等价性，从而把需要做反向求导的指标成功降维为前向直接可得的代数统计。
由散度控制推导出免调参的闭式预算规则：在单纯形上引入卡方距离作为泰勒余项上界，不仅避免了拟凸优化的极点退化陷阱，还顺理成章地推导出了完全无超参数、天然适应 Top-B 预算的软权重公式。

局限与展望¶

依赖预先校验过正确性的轨迹池：论文当前实验均建立在所有候选轨迹终解均正确的假设下，未考虑包含逻辑断崖或错误结论的强噪声混杂池；若在包含事实性错误的开放候选集上运行，模型可能因无法识别伪逻辑而赋予高残差的错误轨迹过高权重。
缺乏多训练随机种子的全过程方差评估：受限于全参数 SFT 的巨大算力成本，作者主要汇报了推理时解码采样的方差，未能重复进行多次端到端全量 SFT 训练以完全隔绝训练初态随机性的微小抖动。
未来改进方向：可进一步将可学习性度量与轻量级过程检验器（Process Reward Model, PRM）联合，拓展至弱监督及自博弈强化学习的数据回放与轨迹筛选管线中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从梯度流损失衰减与 PL 条件切入推理轨迹可学习性，兼具前向代理技巧与解析卡方权重推导。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 种基座架构、4 大顶尖数学推理评测集、7 种对比基线，并附带完备的理论误差证明与微观 Token 分析。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰严谨，从 ODE 连续动力学演变自然过渡到离散前向近似，推导自洽优雅。
价值: ⭐⭐⭐⭐⭐ 在大模型推理蒸馏算力高昂的背景下，为高质量合成思维链数据的高效筛选提供了强有力的实用工程工具。