💬 NLP

Lightning OPD 2.0: Mitigating Style Bias in Cross-Teacher On-Policy Distillation for Large Reasoning Models

This paper introduces Lightning OPD 2.0, a novel on-policy distillation method that mitigates style bias in cross-teacher settings by employing cross-fitted style residualization to separate stylistic differences from genuine reasoning signals, thereby enabling effective knowledge transfer even when the SFT data generator and distillation teacher are distinct models.

Original authors: Yecheng Wu, Song Han, Han Cai

Published 2026-07-31

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Yecheng Wu, Song Han, Han Cai

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 brilliant but quirky student how to solve complex puzzles. You have a "Teacher" who is a genius at math and coding, and a "Student" who is eager to learn. In the world of Artificial Intelligence, this is called distillation: the student tries to copy the teacher's thought process to become smarter. Usually, the best way to do this is to have the student practice on problems the teacher has already solved, then have the teacher grade the student's new attempts line-by-line. This is called On-Policy Distillation. It's like having a coach watch your every move and give instant feedback, rather than just waiting until the end of the game to say "good job" or "try again."

However, there's a catch. For this coaching to work perfectly, the teacher needs to be the same person who wrote the original practice problems the student studied. If the student learned from a textbook written by "Teacher A," but is now being graded by "Teacher B," things can get messy. Even if Teacher B is a genius, they might write in a different style, use different words, or think in a different rhythm. The student might get confused, thinking the teacher is correcting their math when the teacher is actually just annoyed by their handwriting. This paper tackles the problem of what happens when the "Teacher" grading the work is different from the one who wrote the original examples, and how to fix the confusion so the student can still learn effectively.

The Problem: When the Coach and the Textbook Don't Match

In the paper, the researchers noticed a frustrating glitch. They were using a technique called Lightning OPD (On-Policy Distillation) to train AI models to get better at math and coding. The setup was simple: take a model that had already learned from some examples (the "SFT reference"), let it generate new answers, and then have a super-smart "Teacher" model grade those answers to teach the student even more.

The trouble started when the "Teacher" wasn't the same model that wrote the original examples. The researchers found that when they swapped in a different, even stronger Teacher, the student's performance didn't get better; sometimes it got worse!

Why? The researchers realized the Teacher wasn't just correcting the logic (the math or the code); it was also correcting the style. Imagine a student writing a math proof. They might write, "Therefore, the answer is 42." A different teacher might prefer, "Thus, we conclude the solution is 42." If the Teacher is strict about style, they might give a low score to the word "Therefore" even though the math is perfect. The student, confused, starts changing their logic just to please the Teacher's writing habits. The researchers called this "Style Bias." The Teacher's disagreement with the student was a mix of useful corrections (fixing a wrong step) and useless noise (fixing a word choice), and the student couldn't tell the difference.

The Solution: Lightning OPD 2.0

To fix this, the team invented Lightning OPD 2.0. Their idea was to act like a detective, separating the "style" complaints from the "logic" complaints before the student learns from them.

Here is how they did it, using a playful analogy:

Imagine the Teacher and the Student are having a conversation. The Teacher says, "I don't like that word you used," and "I don't like that math step you took."

The Detective Work: The researchers realized that "style" complaints are predictable. If the Teacher hates the word "Therefore" in a math problem, they will probably hate it in every math problem, regardless of the specific numbers. It's a recurring pattern.
The Cross-Fitting Trick: To find these patterns, they split all the practice problems into groups. They looked at Group A to see what the Teacher complained about in Group B, and Group B to see what the Teacher complained about in Group A. This is called cross-fitting. It's like asking a group of friends, "What words does the Teacher always hate?" without letting the person being judged influence the answer.
Subtracting the Noise: Once they figured out the "Style Bias" (the recurring complaints about words, formatting, and rhythm), they subtracted it from the Teacher's total score.
The Result: What was left was the "Residual"—the pure, useful feedback about the actual reasoning. The student now only learns from the parts of the Teacher's feedback that actually matter for solving the problem, ignoring the stylistic grumpiness.

What They Found

The researchers tested this new method on two very different AI students: a 4-billion-parameter model (a smaller, faster learner) and an 8-billion-parameter model (a bigger, more advanced learner). They used a massive, super-smart model as the "Teacher" for both.

The results were clear:

Without the fix (Original Lightning OPD): When the Teacher was different from the original example writer, the student barely improved. The style bias drowned out the good advice.
With the fix (Lightning OPD 2.0): The student learned much faster and became significantly smarter.

Specifically, when starting with the 8-billion-parameter model, the new method helped the AI reach 82.4% accuracy on the AIME 2024 (a tough high school math competition) and 63.0% on LiveCodeBench v5 (a coding challenge). This was a big jump compared to the original method, which struggled to improve the student's score in these "mismatched" teacher scenarios.

Why It Matters

The paper suggests that you don't need to force the "Teacher" and the "Example Writer" to be the same person anymore. In the past, if you wanted to use a specific super-smart Teacher to grade your AI, you had to make sure that Teacher also wrote all your training data. That was expensive and limiting.

Lightning OPD 2.0 suggests that you can pick the best possible Teacher for grading and the best possible source for your training data, even if they are different models. By mathematically stripping away the "style" complaints, the AI can learn the reasoning without getting confused by the personality of the teacher. It's a practical way to make AI training more flexible and efficient, allowing researchers to mix and match the best tools available without worrying about them clashing.

Technical Summary: Lightning OPD 2.0

Problem Statement

On-policy distillation (OPD) has emerged as a stable, cost-effective alternative to reinforcement learning for post-training large reasoning models. It leverages a high-capability teacher to provide dense, token-level supervision on student-generated trajectories. However, a critical prerequisite for effective OPD, identified in prior work (Lightning OPD), is teacher consistency: the model providing OPD supervision must be the same model that generated the demonstrations used for the initial Supervised Fine-Tuning (SFT) reference policy.

In practice, this condition is frequently violated. SFT data often have mixed or undocumented provenance, or practitioners may wish to use a specific SFT dataset generated by one model while distilling with a different, potentially stronger teacher model. This cross-teacher setting introduces a mismatch where raw teacher–reference disagreement contains not only useful, context-specific reasoning signals but also a recurring component associated with stylistic differences (wording, formatting, reasoning cadence). When uncorrected, this "style-token bias" accumulates across rollouts, systematically interfering with the reasoning signal and causing distillation performance to degrade, even when the new teacher is stronger than the original SFT generator.

Methodology: Lightning OPD 2.0

To address this, the authors introduce Lightning OPD 2.0, a method that mitigates style bias through cross-fitted style residualization. The core insight is that style-related disagreement is predictable across cached rollouts based on simple lexical and coarse contextual coordinates, whereas reasoning-related disagreement is more dependent on the specific problem state.

The method proceeds as follows:

Decomposition of Disagreement: The raw teacher–reference log-probability difference ( $d_{it} = \ell^T_{it} - \ell^R_{it}$ ) is modeled as the sum of a recurring style bias component $b(z_{it})$ and a residual variation $v_{it}$ (representing context-specific evidence).
$d_{it} = b(z_{it}) + v_{it}$
Observable Coordinates: To estimate the recurring component $b(z_{it})$ $b (z_{i t})$ , the method uses two types of coordinates:
- Token Identity: Captures lexical preferences.
- Coarse Context: Defined by the normalized response position and the reference-policy surprisal (a measure of how atypical a token is under the SFT reference).
Cross-Fitted Estimation: The cached rollouts are partitioned into $K$ $K$ folds. For a held-out fold $k$ $k$ , the method constructs two lookup tables using only the other $K-1$ $K - 1$ folds:
- A token-lookup table estimating mean disagreement per token identity.
- A context-lookup table estimating mean disagreement per context coordinate pair.
- These tables are smoothed to handle rare groups.
Residualization: For each token in the held-out fold, the estimated style bias $\hat{b}_{it}$ is calculated as the average of the token and context lookup predictions. This estimate is subtracted from the raw disagreement to produce a residualized signal:
$d^*_{it} = d_{it} - \hat{b}_{it}$
Optimization: The residualized disagreement $d^*_{it}$ replaces the raw disagreement in the Lightning OPD objective. The student policy is updated using the advantage derived from this residualized signal, effectively removing the recurring style bias while preserving the context-specific reasoning evidence.

Key Contributions

Formulation of Cross-Teacher OPD: The paper formally defines the setting where the SFT data generator and the OPD teacher are selected independently, demonstrating that naive application of OPD in this setting leads to substantial performance degradation.
Identification of Style-Token Bias: The authors identify a recurring component of teacher–reference disagreement that is predictable across rollouts and associated with stylistic differences, explaining why raw cross-teacher supervision can be misleading.
Lightning OPD 2.0 Algorithm: The introduction of a cross-fitted style residualization technique that estimates and subtracts this bias using token and context lookup tables. This allows the SFT reference and distillation teacher to be chosen independently without rebuilding the SFT pipeline.

Experimental Results

The method was evaluated on mathematical reasoning and code generation benchmarks using two distinct cross-teacher settings:

Setting 1: SFT reference from Qwen3-4B (generated by Qwen3-8B) with an OPD teacher of Qwen3-30B-A3B-Thinking-2507.
Setting 2: SFT reference from Klear-Reasoner-8B (trained on DeepSeek-R1 distilled data) with the same Qwen3-30B OPD teacher.

Key Findings:

Performance Gains: Lightning OPD 2.0 consistently outperforms the baseline Lightning OPD and other OPD variants (IW-OPD, TA-OPD) in cross-teacher settings.
- In the Qwen3-4B setting, it improved average math reasoning by 3.1 points and code generation by 1.4 points over the baseline Lightning OPD.
- In the Klear-Reasoner-8B setting, it improved average math reasoning by 1.0 point and code generation by 1.4 points.
State-of-the-Art: Starting from Klear-Reasoner-8B-SFT, Lightning OPD 2.0 achieved 82.4% on AIME 2024 and 63.0% on LiveCodeBench v5.
Mechanism Analysis: Post-hoc analysis showed that the residualization significantly reduced the deviation of the training signal from a diagnostic reference signal (representing the ideal teacher-consistent signal). For instance, in the Qwen3-4B setting, the rate of tokens with large deviations (>1 nat) dropped from 8.14% to 3.85% after correction.
Ablation: Removing either the token lookup, context lookup, or cross-fitting resulted in lower performance, confirming the necessity of combining lexical and contextual estimates with sample-splitting to prevent self-fitting.

Significance and Claims

The paper claims that Lightning OPD 2.0 establishes a practical approach to cross-teacher on-policy distillation. By relaxing the strict requirement of teacher consistency, it allows practitioners to:

Select the most suitable model for generating SFT demonstrations (to maximize the quality of the reference policy) independently of the model selected for subsequent distillation (to maximize the final post-training result).
Utilize existing SFT datasets without the substantial cost of regenerating demonstrations for every candidate teacher.

The authors position this work as a solution to a practical bottleneck in scaling reasoning models, enabling more flexible and efficient post-training pipelines where the SFT generator and the OPD teacher are decoupled. The method is presented as a robust framework that achieves state-of-the-art performance on reasoning benchmarks while maintaining the stability and cost-efficiency of offline OPD.