GRACE: 梯度对齐的推理数据筛选方法，实现高效后训练

# GRACE：面向高效后训练的梯度对齐推理数据筛选  
来源：https://arxiv.org/html/2605.13130  
Junjie Li  
哈尔滨工业大学（深圳），中国  
22b351018@stu\.hit\.edu\.cn  
&Ziao Wang¹  
香港浸会大学，中国  
ziaowang@hkbu\.edu\.cn  
&NingXuan Ma  
哈尔滨工业大学（深圳），中国  
2023311G27@stu\.hit\.edu\.cn  
&Jianghong Ma  
哈尔滨工业大学（深圳），中国  
香港城市大学，中国  
majianghong@hit\.edu\.cn  
&Xiaofeng Zhang²  
哈尔滨工业大学（深圳），中国  
zhangxiaofeng@hit\.edu\.cn  

###### 摘要  
现有推理数据筛选流程对整条样本评分，将所有中间步骤视为同等重要。实际上，轨迹中各步骤的贡献极不均匀，良好的推理数据选择需要逐步骤评估。我们提出GRACE，一种梯度对齐的筛选方法，将每条推理轨迹视为一系列优化事件，并通过两种互补信号对每个步骤评分：其与面向答案的梯度方向的对齐程度，以及其与先前推理轨迹的一致性。步骤级分数被聚合成样本级数值以进行子集选择，仅使用模型内部的优化信号，无需外部奖励模型或步骤标注。为了使其可扩展，GRACE引入了一种表示级梯度代理，该代理可通过单次前向传递从 token 级上游信号估计步骤级对齐。在后训练 Qwen3-VL-2B-Instruct 于 MMathCoT-1M 上时，GRACE 仅用 20% 的数据就达到完整数据性能的 108.8%，仅用 5% 的数据就保留了 100.2% 的性能，并且所选子集能有效跨模型主干迁移。  

## 1 引言  
大规模推理数据集已成为后训练大型语言模型和视觉-语言模型[30 (https://arxiv.org/html/2605.13130#bib.bib17),35 (https://arxiv.org/html/2605.13130#bib.bib18)]的基石。标准使用方式是在整个推理轨迹上监督模型，将每一步都视为同等有价值的目标。实际上，轨迹中各步骤的贡献极不均匀：有些直接支持最终答案，而另一些则重复之前的内容、探索无关的分支或引入噪声。均匀地训练它们会浪费预算在低价值步骤上，并稀释有用步骤的贡献。随着推理语料库增长到数百万条轨迹[19 (https://arxiv.org/html/2605.13130#bib.bib41)]，且每次后训练耗费数百 GPU 小时，这一成本已变得显著。因此，选择训练内容至关重要[39 (https://arxiv.org/html/2605.13130#bib.bib20)]，而对于推理数据，这意味着逐步骤评估，而非对整个轨迹排序。  

现有数据筛选方法通过基于正确性[29 (https://arxiv.org/html/2605.13130#bib.bib21)]、奖励模型[20 (https://arxiv.org/html/2605.13130#bib.bib22)]或样本级影响力[32 (https://arxiv.org/html/2605.13130#bib.bib5),16 (https://arxiv.org/html/2605.13130#bib.bib23),10 (https://arxiv.org/html/2605.13130#bib.bib24)]选择样本，来提升训练效率，但这些方法都作用于整条轨迹的粒度。因此，一条最终答案正确但中间步骤质量差的轨迹[27 (https://arxiv.org/html/2605.13130#bib.bib25)]，会被视为与精心推理的轨迹同等有价值。这突显了一个基本局限：当前方法缺乏评估单个推理步骤如何贡献于优化的机制。  

参见图注  
图 1：GRACE 的动机与实证效果。左图：推理轨迹被视为一系列优化事件，每个步骤都会引起一个更新方向，其效用取决于与目标方向及演化轨迹的对齐程度。右图：跨基准的下游性能雷达图。GRACE 仅使用训练数据的一小部分即可达到全部数据或更好的性能。  

在本工作中，我们从优化角度重新审视推理数据。我们不将推理轨迹视为静态的监督目标，而是将其建模为一系列优化事件，其中每个推理步骤都会产生一个局部训练信号，影响朝最终答案的梯度方向。从这个角度看，推理数据的效用不仅取决于外部属性（如正确性或长度），还取决于其中间步骤是否建设性地支持优化。基于这一观点，我们提出 GRACE（梯度对齐推理数据高效后训练筛选），一种通过估计步骤级优化效用进行细粒度数据筛选的方法。GRACE 不修剪或重写推理轨迹，而是根据两个互补标准为每个步骤分配效用分数：(i) 其与面向答案的优化方向的对齐程度，以及 (ii) 其与累积推理轨迹的一致性。这些信号捕捉了每个步骤的任务驱动贡献和轨迹感知贡献。随后，这些步骤级分数被聚合成样本级效用分数，使得在保持样本级训练简单性的同时实现有效的子集选择。图 1 (https://arxiv.org/html/2605.13130#S1.F1) 展示了 GRACE 的动机，并提供了实证证据，表明优化感知的筛选方法可以在使用更少训练样本的情况下保持强性能。  

一个关键挑战是，真正的步骤级梯度在大规模计算上难以处理。为每个推理步骤朴素地计算梯度需要将轨迹分解为多个训练实例并执行多次反向传播。为解决这一问题，GRACE 引入了一种表示级梯度代理，利用 token 级上游信号近似步骤引起的优化方向。该代理使得从单次前向传递中高效估计步骤级对齐成为可能，从而使优化感知的筛选对大规模 CoT 数据集可行。  

我们通过在 MMathCoT-1M[19 (https://arxiv.org/html/2605.13130#bib.bib41)] 上后训练 Qwen3-VL-2B-Instruct[26 (https://arxiv.org/html/2605.13130#bib.bib6)]，并在涵盖数学推理和通用视觉问答的多样化多模态基准套件上评估模型，来评估 GRACE。GRACE 一致地识别出高价值子集：仅用 20% 的筛选数据训练即超过完整数据性能，在基准上平均达到完整数据结果的 108.8%，而仅用 5% 的数据则保留了 100.2%。此外，所选子集能有效跨模型主干迁移，表明所提出的基于优化的信号捕捉了超出特定模型配置的内在数据价值。  

我们的贡献有三方面：  
1. 引入了推理数据的优化视角，将推理轨迹视为一系列优化事件，并强调步骤级对齐在有效学习中的作用。  
2. 提出了 GRACE，一种推理数据筛选方法，通过聚合源自面向答案对齐和轨迹一致性的步骤级优化信号，用于样本级子集选择。  
3. 开发了一种表示级梯度代理，无需逐步骤参数空间梯度计算即可实现步骤级对齐的可扩展估计。  

## 2 方法  
参见图注  
图 2：GRACE 筛选流程。(1) 给定输入及其推理轨迹，GRACE 识别每个步骤和答案的 token 集合。(2) 一个固定的评分模型在单次前向传递中提取 token 级上游信号，并按 token 集合分组。(3) 分组信号被平均为梯度代理，并通过答案对齐和轨迹对齐进行评分。(4) 步骤分数被聚合为样本值，用于排序和 top-ρ 子集选择。  

在本节中，我们介绍 GRACE。我们首先定义推理数据筛选问题，然后推导步骤级优化效用，引入其可扩展的表示级代理，并描述样本级子集选择。整体流程如图 2 (https://arxiv.org/html/2605.13130#S2.F2) 所示。  

### 2.1 问题形式化  
考虑推理数据集 D = {z_i}_{i=1}^N，其中每个样本为 z_i = (x_i, s_i, a_i)，x_i 表示输入，s_i = (s_{i,1}, s_{i,2}, ..., s_{i,K_i}) 表示推理步骤序列，a_i 表示最终答案。令 f_θ 表示模型，T_{i,k} 和 T_{i,ans} 分别表示步骤 s_{i,k} 和答案段落的 token 位置。对于任意 token 集合 T，定义平均 token 级损失 L(θ; z_i, T) = (1/|T|) Σ_{t∈T} L_t(θ; z_i)。相应地，步骤损失和答案损失分别为 L_{i,k} = L(θ; z_i, T_{i,k}) 和 L_i^{ans} = L(θ; z_i, T_{i,ans})，推理轨迹和答案上的完整损失为 T_{i,full} = ∪_{k=1}^{K_i} T_{i,k} ∪ T_{i,ans}，L_full(θ; z_i) = L(θ; z_i, T_{i,full})。  

标准后训练最小化 L_SFT(θ; D) = (1/|D|) Σ_{z_i∈D} L_full(θ; z_i)。 (1)  
我们的目标是选择紧凑子集 S ⊂ D，预算 |S| = ⌈ρ|D|⌉，其中 ρ ∈ (0,1) 是选择比例，使得在 S 上后训练能够保持或改善与完整数据集训练相比的下游性能。为此，GRACE 为每个样本分配标量值 V(z_i)，并选择排名最高的子集：S = {z_i ∈ D | rank_V(z_i) ≤ ⌈ρ|D|⌉}。 (2)  

关键问题是如何为推理数据定义 V(z_i)。GRACE 通过估计每个推理步骤的优化效用并将这些步骤级信号聚合成样本级数值来解决。由于这些数值是从模型内部信号计算得出的，评分模型应提供稳定的表示。遵循先前工作[33 (https://arxiv.org/html/2605.13130#bib.bib2),11 (https://arxiv.org/html/2605.13130#bib.bib3)]，我们通过预热初始模型 f_θ0 在 D 的 γ 比例子集上获得评分模型 f_θ，并在数据评分期间保持其固定。  

### 2.2 步骤级优化效用  
我们基于推理步骤在优化目标中的贡献来定义其效用。为清晰起见，在讨论单个样本 z = (x, s, a) 时省略样本索引 i。令 T_k 表示步骤 s_k 的 token 集合，并记 L_k(θ; z) = L(θ; z, T_k)。我们采用标准的一阶影响视角[22 (https://arxiv.org/html/2605.13130#bib.bib1),9 (https://arxiv.org/html/2605.13130#bib.bib26)]。考虑步骤 s_k 引起的微小更新：θ' = θ - η ∇_θ L_k(θ; z)， (3)  
其中 η > 0 是学习率。令 L_tar(θ; z) 表示指定所需优化方向的目标损失；在式 (3) 下，其一阶变化近似为：  
L_tar(θ'; z) - L_tar(θ; z) ≈ -η ⟨∇_θ L_k, ∇_θ L_tar⟩ (4)  
= -η ||∇_θ L_k|| ||∇_θ L_tar|| cos(∇_θ L_k, ∇_θ L_tar)。  

这表明，当一个步骤引起的梯度方向与目标梯度方向对齐时，该步骤局部有益。由于步骤长度和梯度尺度可能在不同推理片段间变化很大，我们关注归一化的方向分量：  
A_k^{tar} ≜ cos(∇_θ L_k, ∇_θ L_tar)。 (5)  
详见附录 B (https://arxiv.org/html/2605.13130#A2)。  

目标方向的不同选择对应不同的步骤效用概念。在 GRACE 中，我们考虑两个互补目标：  
(1) 面向答案的目标。我们将 L_tar 实例化为 L_ans，其中 L_ans(θ; z) = L(θ; z, T_ans) 是答案片段上的损失，T_ans 表示其 token 集合。这给出了：  
A_k^{ans} ≜ cos(∇_θ L_k, ∇_θ L_ans)。 (6)  
它衡量该步骤是否支持优化最终答案。虽然面向答案的目标捕捉了步骤是否支持最终答案，但它没有刻画该步骤是否与前面的推理过程一致。  

(2) 轨迹一致性目标。推理步骤形成一个有序轨迹，而非独立的监督信号。对于有前文上下文的步骤，我们共同定义历史参考方向及其相应的对齐分数：  
A_k^{hist} ≜ cos(∇_θ L_k, r_k)， r_k ≜ Normalize( Σ_{j<k} ∇_θ L_j )。 (7)  
这里 r_k 是通过归一化所有前一步梯度的和构造的累积方向，代表步骤 s_k 之前集体优化的方向。轨迹一致性 A_k^{hist} 衡量当前步骤的梯度与累积历史方向的余弦相似度，确保推理过程逻辑上连贯。  

步骤级分数。结合这两个信号，我们定义步骤 s_k 的效用分数为：  
Score_k = 
\begin{cases} 
A_k^{ans}, & k = 1, \\
\alpha A_k^{ans} + (1-\alpha) A_k^{hist}, & k > 1,
\end{cases}  (8)  
其中 α ∈ [0,1] 平衡面向答案和轨迹一致性目标。该公式确保第一个步骤仅根据其与答案的相关性进行评分，而后续步骤则通过当前任务相关性（A_k^{ans}）和轨迹连贯性（A_k^{hist}）进行评估。  

### 2.3 表示级梯度代理  
步骤级梯度 ∇_θ L_k 涉及整个参数空间，计算成本过高：为每个步骤单独反向传播需要将每个推理步骤视为独立数据点，导致每个样本进行 K 次反向传播。为解决这一问题，我们引入表示级梯度代理，从前向传播中估计 ∇_θ L_k 的方向，无需显式的参数空间梯度计算。  

核心思想源于标准 Transformer 层的反向传播分析。考虑单一注意力层，令 h ∈ ℝ^{d} 为特定位置处输入序列第 t 个 token 的隐藏状态，即 h = H_t。该隐藏状态经仿射变换后产生该 token 上的 logits ℓ = W h + b，其中 W ∈ ℝ^{V×d} 是语言建模头，b ∈ ℝ^V 是偏置。令 v ∈ ℝ^V 为真实 token 在 logits ℓ 处的梯度，其分量由损失函数 L_t(θ; z) 确定。通过应用链式法则，该 token 处损失相对于隐藏状态 h 的梯度为 ∇_h L_t = W^⊤ v。 (9)  

基于此，我们推导表示级梯度代理。考虑 token 集合 T（如步骤 k 的 token 集合），梯度 ∇_θ L_k 相对于注意力层输出的分量与 token 级梯度 ∇_h L_t (t∈T) 密切相关。我们的关键见解是，这些梯度之间的角度可以通过 token 级信息充分估计，而无需完整的参数空间梯度。我们通过式 (10) 中的代理余弦相似度来近似 cos(∇_θ L_k, ∇_θ L_tar)（命题 1 中形式化）：  
ĝ_k = Normalize( Σ_{t∈T_k} ∇_h L_t )， ĝ_tar = Normalize( Σ_{t∈T_{tar}} ∇_h L_t )， \widehat{cos}(∇_θ L_k, ∇_θ L_tar) = cos(ĝ_k, ĝ_tar)。 (10)  

**命题 1**（表示级代理）。令 ∇_h L_t 为式 (9) 在 token 位置 t 处针对隐藏状态 h_t 的梯度。在宽带注意力层和不相关偏差的简化假设下，cos(∇_θ L_k, ∇_θ L_tar) 可由式 (10) 定义的 cos(ĝ_k, ĝ_tar) 近似。  
证明见附录 A (https://arxiv.org/html/2605.13130#A1)。  

我们指出，此代理是对真实梯度的近似——梯度 ∇_θ L 包含式 (9) 仅捕捉的投影分量，并且由于表示级梯度聚合跨越 token 集合，方向信息得以保留，但尺度信息丢失。由于我们仅需要方向对齐，因此尺度信息的缺失是可接受的。与需要对每个推理步骤进行单独反向传播的自然梯度相比，该代理在单一前向传播中计算隐藏状态梯度，并仅在与语言建模头交互时进行一次反向传播，无论解析出的 token 集合数量如何。  

### 2.4 样本级分数与子集选择  
结合步骤级分数和表示级梯度代理，我们获得代理的步骤级分数：  
Ŝcore_k = 
\begin{cases} 
\widehat{A}_k^{ans}, & k = 1, \\
\alpha \widehat{A}_k^{ans} + (1-\alpha) \widehat{A}_k^{hist}, & k > 1,
\end{cases}  (12)  
其中  
\widehat{A}_k^{ans} ≜ cos(ĝ_k, ĝ_ans)， \widehat{A}_k^{hist} ≜ cos(ĝ_k, r̂_k)， r̂_k ≜ Normalize( Σ_{j<k} ĝ_j )。 (13)  

然后，我们通过合并定义步骤级分数的聚合来定义样本级分数 V(z)： (14)  
V(z) = (1/|S|) Σ_{s_k∈S} Ŝcore_k，  
其中 S 是推理步骤的集合。  

此设计承认一个关键属性：如果一个样本包含许多低质量步骤（无论其答案段落的对齐如何），它将获得较低的样本级分数。相反，由一致高对齐步骤组成的轨迹将产生高分。使用排序公式（式 (2)），我们按 V(z) 降序排列数据集，根据预算 ρ 保留前 ρ 个样本，形成筛选后的训练子集 S。  

## A 表示级梯度代理的证明  
命题 1：令 ∇_h L_t 为式 (9) 在 token 位置 t 处针对隐藏状态 h_t 的梯度。在宽带注意力层和不相关偏差的简化假设下，cos(∇_θ L_k, ∇_θ L_tar) 可由式 (10) 定义的 cos(ĝ_k, ĝ_tar) 近似。  

证明：首先，考虑一个简化的单层线性模型以激励该代理。令对 token 集合 T 的损失为 L_T = (1/|T|) Σ_{t∈T} L_t，其中 L_t 是 token 位置 t 处的交叉熵损失。对于单个 token t，令 ℓ_t = W h_t ∈ ℝ^V 为 logits，其中 W ∈ ℝ^{V×d}。设 v_t = ∂L_t/∂ℓ_t 为 logits 处的导数向量。通过链式法则，关于隐藏状态 h_t 的梯度为 ∇_h L_t = W^⊤ v_t。 (A.1)  

现在考虑相对于 W 的梯度。对于 L_t，关于 W 的梯度是 ∇_W L_t = v_t h_t^⊤。（为简洁起见省略偏差。）因此，对于 token 集合 T，关于 W 的梯度为 ∇_W L_T = (1/|T|) Σ_{t∈T} v_t h_t^⊤。  

类似地，对于目标集合 T_tar，我们得到 ∇_W L_tar = (1/|T_tar|) Σ_{t∈T_tar} v_t h_t^⊤。  

参数空间的弗罗贝尼乌斯内积 ⟨∇_W L_T, ∇_W L_tar⟩_F = Tr((∇_W L_T)^⊤ ∇_W L_tar) 可写为：  
⟨∇_W L_T, ∇_W L_tar⟩_F = (1/|T||T_tar|) Σ_{t∈T} Σ_{t'∈T_tar} (v_t^⊤ v_{t'})(h_t^⊤ h_{t'})。 (A.2)  

这在 v_t 和 h_t 之间引入了交叉项。然而，在宽网络以及 h_t 和 v_t 分量近似不相关的假设下（该假设在随机初始化以及有限训练后非极端定向对齐下保持，且在 Transformer 的表示级广泛成立），交叉项平均趋于零，从而得到：  
⟨∇_W L_T, ∇_W L_tar⟩_F ≈ (1/|T||T_tar|) Σ_{t∈T} Σ_{t'∈T_tar} (v_t^⊤ v_{t'}) · (特定假设下的平均点积)。  

为了获得仅基于 token 级梯度的余弦，我们注意到 Σ_{t∈T} ∇_h L_t = Σ_{t∈T} W^⊤ v_t = W^⊤ (Σ_{t∈T} v_t)。由于 W 是固定的，Σ_{t∈T} ∇_h L_t 与 Σ_{t∈T} v_t 的方向相同。因此，我们可以直接使用 Σ_{t∈T} ∇_h L_t 来估计参数空间梯度之间的余弦。此外，由式 (A.1) 得到：  
||Σ_{t∈T} ∇_h L_t|| = ||W^⊤ Σ_{t∈T} v_t|| ≤ ||W||_F · ||Σ_{t∈T} v_t||，  
但方向信息保留。由于余弦仅取决于方向，我们直接使用归一化的表示级和：  
ĝ_k = Normalize( Σ_{t∈T_k} ∇_h L_t )， ĝ_tar = Normalize( Σ_{t∈T_{tar}} ∇_h L_t )，  
并定义 \widehat{cos}(∇_θ L_k, ∇_θ L_tar) = cos(ĝ_k, ĝ_tar) = (ĝ_k^⊤ ĝ_tar) / (||ĝ_k|| ||ĝ_tar||)。 (A.3)  

该途径在假设梯度在所有 token 上具有足够各向异性分布且聚合保留了主导方向时，近似了真实梯度方向余弦。完整的链式法则会涉及更深的层，但跨层的梯度组合在宽网络假设下可分解，且有中心极限定理支持表示级方向的近似成立。因此，我们采用式 (A.3) 作为计算效率高的代理。 ∎  

**备注 1**。式 (A.2) 中的交叉项表明该代理是一种近似，而非精确相等。然而，对于此处考虑的大规模模型，我们的实证结果表明，该代理捕捉到的相对方向差异足以进行有意义的信号对齐，且筛选性能一致。当假设不成立时（例如，极端窄网络），该代理的质量可能会下降。  

## B 步骤级对齐公式的推导  
我们提供从一阶泰勒展开到步骤效用的完整推导。给定在小数据集上进行少量训练步骤后的模型参数 θ，考虑仅从步骤 s_k 的损失更新参数：θ' = θ - η ∇_θ L_k(θ; z)，其中 η > 0。在 θ 处应用一阶泰勒展开得到：  
L_tar(θ'; z) = L_tar(θ; z) + ⟨∇_θ L_tar(θ; z), θ' - θ⟩ + O(||θ' - θ||^2) (15)  
= L_tar(θ; z) - η ⟨∇_θ L_tar(θ; z), ∇_θ L_k(θ; z)⟩ + O(η^2)。 (16)  

因此，  
L_tar(θ'; z) - L_tar(θ; z) = -η ⟨∇_θ L_k(θ; z), ∇_θ L_tar(θ; z)⟩ + O(η^2)。 (17)  

在小步长下，忽略高阶项即得式 (4)。一阶项表明，对步骤效用最重要的量是梯度余弦。