🤖 AI

GRACE: Gradient-aligned Reasoning Data Curation for Efficient Post-training

GRACE is a scalable, reward-model-free data curation method that improves post-training efficiency by scoring and selecting reasoning data at the step level using gradient-aligned signals, achieving full-data performance with only 5% of the dataset.

Original authors: Junjie Li, Ziao Wang, NingXuan Ma, Jianghong Ma, Xiaofeng Zhang

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Junjie Li, Ziao Wang, NingXuan Ma, Jianghong Ma, Xiaofeng Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: "The Whole is Not Equal to the Sum of Its Parts"

Imagine you are trying to teach a student how to solve a complex math problem. You give them a textbook page that shows the full solution: Step 1, Step 2, Step 3, all the way to the final answer.

The old way (Current Methods):
Most AI training systems treat that entire page as one single "good" or "bad" lesson. If the final answer is correct, the whole page gets a gold star. If the answer is wrong, the whole page gets a red X.

The Flaw: In reality, some steps on that page are brilliant insights, while others are just the student repeating what they said three lines ago, or wandering off into a confusing side-track that leads nowhere. By treating the whole page as one unit, the AI wastes time studying the boring or confusing parts and might miss the brilliant insights hidden in the middle.

The Solution: GRACE (The "Step-by-Step" Coach)

The authors created a method called GRACE (Gradient-aligned Reasoning dAta Curation). Instead of grading the whole page, GRACE acts like a super-observant coach who watches the student solve the problem one step at a time.

Here is how GRACE works, using a Hiking Analogy:

Imagine the AI is a hiker trying to reach a mountain peak (the Correct Answer). The reasoning trace is the path the hiker takes.

The "Answer Alignment" Signal (Looking at the Summit):
- GRACE asks: "Is this specific step pointing toward the mountain peak?"
- If a step moves the hiker closer to the top, it gets a high score.
- If a step makes the hiker walk in a circle or head toward a cliff, it gets a low score.
The "Trajectory Consistency" Signal (Looking at the Path Behind):
- GRACE also asks: "Does this step make sense given where the hiker just came from?"
- If the hiker is climbing a steep slope and suddenly tries to swim in a river, that's a weird jump. Even if the river is beautiful, it breaks the flow of the climb. GRACE penalizes steps that feel disconnected from the previous steps.

The Secret Sauce: The "Magic Mirror" (No Backtracking Needed)

Usually, to know if a step is good, you would have to simulate the whole training process, stop, rewind, and calculate exactly how that one step changed the AI's brain. This is like trying to measure the wind by stopping a race car, dismantling the engine, and checking the gears. It takes forever and is too slow for massive datasets.

GRACE's Innovation:
GRACE uses a clever trick called a "Representation-level Gradient Proxy."

The Analogy: Instead of dismantling the car, GRACE looks at the "exhaust fumes" (the internal signals) coming out of the engine while the car is driving forward.
By looking at these signals during a single forward pass (just reading the text once), GRACE can estimate exactly how useful a step was without ever having to "rewind" or do complex backward calculations. It's like judging a chef's skill by the smell of the food while it's cooking, rather than tasting every single bite after the meal is done.

The Results: Less Data, Better Performance

The paper tested this method on a massive dataset of math problems (MMathCoT-1M) using a vision-language model (Qwen3-VL).

The Old Way: To get great results, you usually need to feed the AI all the data (100%).
The GRACE Way:
- Using only 20% of the data (the best 1 out of 5 steps), the AI performed 8.8% better than if it had been trained on 100% of the data.
- Using only 5% of the data, the AI performed just as well as the full dataset.

Why did it get better with less data?
Training on everything is like forcing a student to read every single page of a library, including the pages with typos, boring repetition, and wrong turns. This confuses the student. GRACE filters out the "noise" and only feeds the AI the "pure gold" steps. This prevents the AI from getting confused by bad examples and helps it learn faster and smarter.

Summary

The Problem: Current AI training treats every step in a reasoning chain as equally important, wasting time on bad steps.
The Fix: GRACE scores every single step individually based on whether it helps reach the answer and fits the story so far.
The Trick: It uses a "forward-pass only" shortcut to score steps instantly, without needing slow, complex calculations.
The Outcome: You can train a smarter AI using a tiny fraction of the data, saving time and computing power while actually improving performance.

The paper concludes that the value of reasoning data isn't just about whether the final answer is right; it's about whether the journey to that answer was a constructive, logical path. GRACE finds those constructive paths and throws away the rest.

Technical Summary: GRACE – Gradient-aligned Reasoning Data Curation for Efficient Post-training

Problem Statement

Existing pipelines for curating reasoning data (Chain-of-Thought, CoT) typically operate at the sample level, assigning a single score to an entire reasoning trace based on metrics like final answer correctness or reward model scores. This approach treats every intermediate step within a trace as equally valuable. However, reasoning traces are heterogeneous; some steps directly support the final answer, while others restate content, explore irrelevant tangents, or introduce noise. Training uniformly on such traces wastes computational budget on low-value steps and dilutes the contribution of useful ones. Furthermore, current methods lack a mechanism to assess how individual reasoning steps contribute to the optimization process, often overlooking the "step-wise optimization roles" within a trace.

Methodology: GRACE

The authors propose GRACE (Gradient-aligned Reasoning dAta Curation for Efficient post-training), a method that shifts the curation granularity from the sample level to the step level. GRACE views a reasoning trace not as a static supervision target, but as a sequence of optimization events.

1. Step-Level Optimization Utility

GRACE defines the utility of a reasoning step based on its contribution to optimizing a target objective. It scores each step $k$ using two complementary signals derived from the model's internal optimization dynamics:

Answer-Oriented Alignment ( $A_{ans}^k$ ): Measures the cosine similarity between the gradient induced by the current step ( $g_k$ ) and the gradient induced by the final answer segment ( $g_{ans}$ ). This determines if the step constructively supports the final objective.
Trajectory Consistency ( $A_{hist}^k$ ): Measures the cosine similarity between the current step's gradient and a historical reference direction ( $r_k$ ), which is a weighted average of gradients from preceding steps. This ensures the step maintains coherence with the evolving reasoning trajectory.

The final step-level score is a weighted combination:
$\text{Score}_k = \alpha A_{ans}^k + (1-\alpha) A_{hist}^k$
where $\alpha$ balances the two signals. These step scores are aggregated to compute a sample-level value $V(z)$ , which is used to select the top- $\rho$ subset of data.

2. Representation-Level Gradient Proxy

A primary challenge is that computing true step-level gradients requires decomposing traces and performing repeated backward passes, which is computationally intractable at scale. To address this, GRACE introduces a representation-level gradient proxy:

Instead of computing parameter-space gradients ( $\nabla_\theta L$ ), the method operates on the upstream signals at the representation interface (the hidden states before the output projection).
For each token $t$ , an upstream signal $u_t = W_{out}(p_t - y_t)$ is computed, where $p_t$ is the predicted probability and $y_t$ is the ground truth.
Step-level gradients are approximated by averaging these token-level signals ( $g_k = \frac{1}{|T_k|} \sum u_t$ ).
Crucially, this proxy allows the estimation of step-level alignment using only a single forward pass per sample, avoiding per-step backpropagation and the need for external reward models or step annotations.

Key Contributions

Optimization-Based Perspective: The paper reframes reasoning data curation by modeling traces as sequences of optimization events, highlighting that step-level alignment with the target objective and trajectory is critical for effective learning.
GRACE Framework: A novel curation method that aggregates step-level optimization signals (answer alignment and trajectory consistency) for fine-grained subset selection.
Scalable Proxy: The development of a representation-level gradient proxy that enables efficient, step-level data valuation without the computational cost of per-step gradient computation.

Experimental Results

The authors evaluated GRACE by post-training Qwen3-VL-2B-Instruct on the MMathCoT-1M dataset (a large-scale multimodal mathematical reasoning corpus) and testing on diverse benchmarks including mathematical reasoning (MathVista, MathVision) and general visual question answering (HallusionBench, MMBench, MME).

Performance vs. Data Volume:
- Using only 20% of the curated data, GRACE achieved 108.8% of the full-data performance (averaged across benchmarks), significantly outperforming the full dataset and all baselines.
- Using only 5% of the data, GRACE retained 100.2% of the full-data performance.
Comparison with Baselines: GRACE consistently outperformed heuristic selectors (Random, Longest, Stepmax) and state-of-the-art gradient-based data selection methods (LESS, ICONS, CADC). For instance, at 20% data, GRACE (108.8%) surpassed CADC (102.6%) and LESS (101.7%).
Transferability: Subsets selected by GRACE using a Qwen3-VL-2B scoring model transferred effectively to other backbones (Qwen2.5-VL-3B, LLaVA-1.5-7B, Qwen3-VL-8B), consistently surpassing their respective full-data baselines. This suggests the captured data value is intrinsic to the reasoning structure rather than model-specific.
Efficiency: The proxy extraction process is lightweight, requiring only forward passes. The authors report that collecting signals for the candidate pool took 6.5 single-node hours and 5.7GB of storage, avoiding the expensive backward passes required by gradient-projection methods.

Significance and Claims

The paper claims that the value of reasoning data depends not merely on external quality indicators (correctness, length, or preference scores) but on whether intermediate steps constructively support the optimization trajectory.

GRACE demonstrates that:

Step-level granularity is essential: Treating reasoning traces as indivisible units overlooks the varying utility of intermediate steps.
Internal signals are sufficient: High-quality data curation can be achieved using only the model's internal optimization signals (gradients/proxies) without external reward models or human annotations.
Efficiency and Performance: Optimization-aware curation allows models to match or exceed full-data performance using a fraction of the training data, substantially reducing post-training costs (GPU hours) while maintaining or improving downstream capabilities.

The authors note limitations, including the reliance on a fixed scoring model configuration and the current focus on multimodal CoT data, suggesting future work could extend these principles to pure language reasoning or reinforcement learning-based post-training.