🤖 AI

Cornerstones or Stumbling Blocks? Deciphering the Rock Tokens in On-Policy Distillation

This paper identifies "Rock Tokens" as a persistent subset of high-loss tokens in On-Policy Distillation that consume significant optimization resources despite contributing negligibly to reasoning performance, demonstrating that strategically bypassing them can streamline the model alignment process.

Original authors: Yuxuan Jiang, Runchao Li, Shubhashis Roy Dipta, Dawei Li, Zhao Yang

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuxuan Jiang, Runchao Li, Shubhashis Roy Dipta, Dawei Li, Zhao Yang

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

Imagine you are trying to teach a talented but slightly clumsy apprentice (the Student Model) how to write perfect math proofs by having them copy a master architect (the Teacher Model).

In the world of AI, this process is called On-Policy Distillation. The apprentice writes a proof, the teacher corrects it, and the apprentice tries again. Usually, we assume that every mistake the apprentice makes is a valuable lesson. If the teacher says, "You got this wrong," the apprentice learns from it.

This paper introduces a surprising discovery: Not all mistakes are lessons. In fact, some of the most frequent "mistakes" are actually just Stumbling Blocks that waste time and energy, while a tiny few are the only real Cornerstones that hold the whole structure together.

Here is the breakdown of their findings using simple analogies:

1. The "Rock Tokens": The Stubborn Pebbles

The researchers noticed that even after the apprentice had practiced for a long time and seemed to have mastered the material, there were still specific words and symbols that kept getting "corrected" by the teacher. They called these Rock Tokens.

What are they? They aren't the complex math formulas or the deep logic. They are the "boring" stuff: spaces, new lines, brackets, punctuation, and common transition words like "So," "Wait," or "Let's."
The Problem: These tokens appear constantly (about 18% of the text). Every time the apprentice uses them slightly differently than the teacher, the computer calculates a "high loss" (a big error).
The Paradox: Even though the computer keeps screaming "ERROR!" at these tokens, the apprentice never actually learns to fix them. They remain stubbornly wrong throughout the entire training process. It's like a student who keeps writing the date wrong on every single homework assignment, and no matter how many times the teacher circles it, the student just keeps doing it the same way.

2. The Two Big Surprises

The researchers dug deeper and found two shocking truths about these "Rock Tokens":

Surprise A: The Gradient Paradox (The Noise vs. The Signal)
Usually, we think a big error means a big learning opportunity. But for Rock Tokens, the opposite is true.

The Metaphor: Imagine the teacher is trying to steer a ship. The Rock Tokens are like a massive, constant wave hitting the side of the boat. It pushes the ship hard (high "gradient"), but because the wave is so constant and predictable, the ship's rudder just ignores it. The ship keeps drifting in the same direction.
The Reality: These tokens generate a lot of "noise" in the training data. They take up a huge amount of the computer's processing power trying to fix them, but they don't actually help the model get smarter at reasoning. They are just structural habits the apprentice is too stubborn to break.

Surprise B: The "Pillar" vs. The "Stumbling Block"
The researchers asked: "If we remove these stubborn tokens, does the model crash?"

The Metaphor: Imagine a building. Most of the bricks are just filler. But a few specific bricks are Pillars—if you take them out, the roof falls. The rest are just Stumbling Blocks—if you remove them, the building stands just fine, maybe even better because it's lighter.
The Reality: They tested this by "knocking out" these tokens during the test phase.
- Result: 96% to 98% of these stubborn tokens were neutral. Removing them didn't hurt the model's math skills at all.
- The Exception: Only a tiny fraction (about 1.5% to 3.5%) were true Pillars. These were specific words like "certain," "strategic," or "initialize" that were actually critical for the logic.
- Crucial Finding: There were zero "Stumbling Blocks." Removing the stubborn tokens never made the model worse. They were just dead weight.

3. The Solution: The "Skip" Button

Since most of these "Rock Tokens" are just wasting time, the researchers tried a new strategy: Freeze the Gradients.

The Analogy: Instead of forcing the apprentice to re-learn how to write the date or use a comma every single time, the teacher says, "Okay, we agree you'll write the date your way. Let's stop wasting time correcting that and focus on the actual math."
The Result: When they stopped trying to correct these stubborn tokens:
- The model's math performance stayed exactly the same.
- The training process became 1.4 to 1.7 times faster.

The Bottom Line

The paper argues that in the current way we train AI models, we are obsessively trying to fix every single tiny mismatch between the student and the teacher.

However, this is inefficient. We are spending 18% of our effort trying to align the "boring" structural parts of the text (spaces, punctuation, common words) that the model simply refuses to change and doesn't actually need to change to be smart.

The Takeaway:
To make AI training faster and more efficient, we should stop treating every "error" as a critical lesson. We should identify these Rock Tokens (the stubborn, high-frequency structural errors) and skip them. By ignoring the "Stumbling Blocks" and only focusing on the rare "Cornerstones" (the actual logic), we can build smarter models much faster without losing any reasoning ability.

Technical Summary: Cornerstones or Stumbling Blocks? Deciphering the Rock Tokens in On-Policy Distillation

1. Problem Statement

On-Policy Distillation (OPD) has become a cornerstone of post-training pipelines for Large Language Models (LLMs), enabling token-level reasoning refinement by aligning a student model with a superior teacher using trajectories sampled from the student's own policy. While recent work in Reinforcement Learning with Verifiable Rewards (RLVR) has identified that a small subset of critical tokens (e.g., high-entropy "forking tokens") disproportionately drives reasoning gains, an analogous token-level understanding of OPD remains unexplored.

In OPD, the objective minimizes the per-token reverse KL divergence between the student and teacher policies. Intuitively, as training converges, high-loss tokens (indicating student-teacher mismatch) should diminish. However, empirical observations reveal a paradox: even after training reaches apparent saturation, a substantial subset of tokens continues to exhibit persistently high loss. The central question addressed by this work is whether these persistent high-loss tokens are indispensable "Cornerstones" (Pillars) for policy alignment or merely "Stumbling Blocks" that create unnecessary optimization redundancy.

2. Methodology

2.1 Definition and Detection of Rock Tokens

The authors define Rock Tokens as token types that exhibit persistent high-loss behavior that is stable across similar contexts, rather than being isolated spikes in long trajectories. The detection process involves three stages:

Aggregate Burden Calculation: Computing a "Rock Score" $R(v)$ for each token type $v$ , defined as the product of its mean token-level loss ( $\bar{\ell}_v$ ) and its frequency ($Freq(v)$) in student-generated trajectories.
Persistence Filtering: Identifying occurrences that remain high-loss both before and after OPD training.
Context-Consistency Filtering: Ensuring that high-loss occurrences recur in similar local contexts (using a context-window similarity metric). This distinguishes genuine structural mismatches from isolated positional noise.

The final metric, the context-aware Rock Score $R_{ctx}(v)$ , combines aggregate burden with context consistency. A cutoff $K=100$ is selected to balance selection stability (reproducibility) and loss coverage (capturing ~60% of corpus-level KL divergence).

2.2 Gradient Geometry and Persistence Analysis

To address whether Rock Tokens provide useful learning signals, the authors analyze the geometry of per-token gradients:

Magnitude vs. Frequency: Comparing the per-occurrence gradient magnitude of Rock Tokens against rare high-KL tokens.
Alignment: Measuring the cosine alignment of Rock Token gradients with a frequency-balanced global descent direction ( $G_{balanced}$ ).
Stagnation: Tracking the change in per-token KL divergence between early and late training checkpoints to determine if the loss actually decreases.

2.3 Causal Intervention: Knockout Experiments

To determine the functional necessity of these tokens, the authors employ an inference-time knockout procedure. They form a policy $\pi_{\theta}^{\setminus v}$ by forcing the logit of a candidate token $v$ to $-\infty$ at every decoding step. The causal impact is measured as the change in benchmark accuracy ( $\Delta_B(v)$ ). Tokens are categorized as:

Strong Pillars: Removal significantly degrades accuracy.
Strong Stumbling Blocks: Removal improves accuracy.
Neutral: Removal has negligible effect.

2.4 Selective Distillation (Freezing Gradients)

To evaluate the impact of Rock Tokens on the training process, the authors introduce a window-aware token reweighting mechanism. They freeze the gradients of Rock Tokens and their associated local high-divergence windows (setting the weight $\lambda=0$ ) during OPD training. This "Just Not Train" strategy allows them to isolate whether these tokens are essential functional anchors or unproductive optimization burdens.

3. Key Findings

3.1 Identity and Pervasiveness

Composition: Rock Tokens are primarily syntactic and structural scaffolding, including LaTeX/math delimiters, Markdown/whitespace structures, discourse markers (e.g., "So", "Wait"), and digits. They mark boundaries in reasoning rather than content-bearing words.
Prevalence: Despite representing only ~6% of the vocabulary, Rock Tokens account for approximately 18% of the total token frequency in generated outputs.

3.2 The Gradient Paradox

Alignment: Contrary to the assumption that they are noise, Rock Tokens exhibit non-trivial positive directional alignment with the global optimization target. They provide gradients that point toward the correct optimization path.
Stagnation: Despite this alignment, Rock Tokens remain stagnant throughout training. Their KL divergence does not significantly decrease (median $\Delta KL \approx 0$ ), unlike rare high-KL tokens which show substantial reduction.
Cause: The authors hypothesize that adaptive optimizers (like Adam) suppress the effective learning rate for these ubiquitous tokens due to heavy-tailed class imbalance, trapping the student in pre-existing structural habits despite the presence of correct directional signals.

3.3 Causal Necessity (Pillars vs. Stumbling Blocks)

Rarity of Pillars: Only a tiny fraction (1.5%–3.5%) of Rock Tokens act as "Strong Pillars" whose removal degrades reasoning performance.
Absence of Stumbling Blocks: No "Strong Stumbling Blocks" were found; removing Rock Tokens never improved performance.
Orthogonality: The causal importance of a token (Pillarhood) is orthogonal to standard metrics like entropy, frequency, or residual KL. Existing importance taxonomies fail to predict which tokens are causally necessary.
Buffer Zone: The student model appears to maintain a "buffer zone" of persistent residuals to protect a few critical semantic anchors, rather than relying on the entire high-loss set.

3.4 Training Efficiency

Selective Freezing: By freezing the gradients of Rock Tokens (and their local windows) during training, the authors achieved a 1.4× to 1.7× increase in training wall-clock speed.
Performance Impact: This acceleration resulted in negligible loss in downstream reasoning performance (measured on AIME, HMMT, and MATH benchmarks). This suggests that a significant portion of the optimization bandwidth spent on Rock Tokens is redundant.

4. Significance and Claims

The paper challenges the conventional wisdom that aligning every high-loss token is indispensable for learning in On-Policy Distillation. The authors claim:

Identification of a Hidden Dynamic: They identify "Rock Tokens" as a persistent high-loss phenomenon that unveils a hidden aspect of OPD optimization dynamics, distinct from the "forking tokens" observed in RLVR.
Inefficiency in Current Pipelines: A substantial portion of training signals in OPD provides significantly lower-than-expected contributions to the model's ultimate reasoning performance.
Efficiency-Performance Frontier: The work demonstrates that strategically bypassing these "stumbling blocks" (via gradient freezing) can significantly streamline the alignment process. This offers a more efficient paradigm for large-scale model distillation without compromising reasoning integrity.

The paper concludes that while Rock Tokens are not entirely disposable (as a small subset acts as necessary pillars), non-uniform gradient allocation can effectively bypass computational redundancy, optimizing the efficiency-performance frontier in on-policy alignment.

5. Limitations

The authors note that their empirical evaluations rely predominantly on competitive mathematical reasoning, and the role of Rock Tokens may differ in open-ended generation or coding tasks. Additionally, the study utilizes a specific 30B-to-4B distillation pairing, leaving the generalizability across different architectural families or scales as an open question. Finally, the mitigation strategy employs a strict binary gradient-freeze mask, which the authors suggest could be refined in future work with dynamic reweighting schemes.