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

Let $p$ be the training-prompt distribution and $\{q_{i}\}_{i=1}^{N}$ the prompts stored in the replay. When the supervised fine-tuning (SFT) demonstrations have a single identifiable generator, let $\pi_{G}$ denote this SFT data generator. Let $\pi_{R}$ be the resulting SFT reference policy, $\pi_{T}$ the model selected to provide OPD supervision, and $\pi_{\theta}$ a trainable policy initialized from $\pi_{R}$. Teacher consistency requires $\pi_{T}=\pi_{G}$(wu2026lightning). We study cross-teacher OPD, where $\pi_{T}$ is selected independently of $\pi_{G}$ or no single $\pi_{G}$ can be identified. Our method addresses this general setting through the frozen replay of Lightning OPD.

Following Lightning OPD, we sample and freeze one response $x_{i}=(y_{i1},\ldots,y_{iL_{i}})\sim\pi_{R}(\cdot\mid q_{i})$ for every prompt. Let $h_{it}=(q_{i},y_{i,<t})$ be the prefix at token $t$. The OPD teacher and the SFT reference score the same realized token, producing the cached chosen-token log-probabilities

We assume that these scores are aligned to the same response tokens. Our experiments use models with compatible Qwen-family tokenization.

Using this cache, unmodified Lightning OPD retains the original per-token advantage

and optimizes

Following standard OPD practice (agarwal2024policy; lu2025onpolicydistillation; wu2026lightning), the advantage is treated as a fixed scalar when computing parameter updates. Throughout this section, $\nabla J$ denotes the resulting advantage-weighted policy gradient rather than ordinary differentiation through the advantage. We use the same policy surrogate as Lightning OPD.

To isolate the part of the signal affected by the teacher choice, we define the teacher–reference disagreement

At initialization, $\pi_{\theta}=\pi_{R}$, so $A^{\mathrm{base}}_{it}=d_{it}$. In cross-teacher OPD, $d_{it}$ can contain both context-specific teacher evidence and recurring teacher–reference differences induced by the mismatch with the SFT data generator. The unmodified update cannot distinguish the two. The next subsection develops an operational proxy for this recurring component.

The raw disagreement $d_{it}$ can contain both context-specific teacher evidence and systematic differences in how the selected teacher and the SFT reference express a response. We make the operational assumption that disagreement associated with recurring style differences is more predictable across cached rollouts from simple lexical and coarse contextual coordinates, whereas potentially useful, reasoning-related teacher evidence is more dependent on the current context. We therefore write

where $z_{it}$ collects the token and context coordinates defined below, $b(z_{it})$ is the component that recurs at those coordinates across rollouts, and $v_{it}$ is the remaining variation. This is an operational decomposition rather than a semantic labeling of individual tokens. We use $b(z_{it})$ as a proxy for recurring style-token bias and approximate it by equally combining token- and context-based lookup predictions. We estimate these predictions by rollout-level cross-fitting and replace $d_{it}$ with its estimated residual in the training objective.

Cross-fitting is a standard sample-splitting technique for constructing held-out predictions and preventing self-fitting (chernozhukov2018double; guo2021machine). We use it only as an estimation device. We partition the cached rollouts into $K$ deterministic folds, keeping all rollouts originating from the same prompt in one fold. For a rollout in fold $k$, we estimate all statistics using only the other $K-1$ folds. Neither that rollout nor another rollout from the same prompt can therefore contribute to its lookup tables. This exclusion prevents self-fitting but does not by itself identify the predicted component as style.

To prevent longer responses from dominating the estimates, we assign each response unit total weight, distributed uniformly over its active tokens, when fitting the lookup tables. The training loss uses the same reduction by averaging active-token contributions within each response and then averaging across responses. This response-balanced reduction is shared by all methods and is not a component of style residualization.

Using the non-held-out folds, we build one lookup table for the mean disagreement of each token identity and another for each context coordinate. Rare groups are smoothed toward the global mean of the non-held-out folds, while unseen groups use that global mean directly. This prevents a rare token or context from receiving an unstable correction. Let $m_{\mathrm{tok}}^{(-k)}$ and $m_{\mathrm{ctx}}^{(-k)}$ denote the resulting smoothed lookup tables. For token $(i,t)$ in held-out fold $k$, we use the equal-weight prediction

The two lookup predictions receive equal coefficients, and averaging keeps their combination on the scale of a single disagreement estimate. Equation (7) is our concrete approximation to $b(z_{it})$ in Equation (5). At the conceptual level, both tables average the raw disagreement $d_{it}$. Because both lookup tables exclude the current rollout, $\widehat{b}_{it}$ captures recurring disagreement that is predictable from other cached rollouts rather than variation fitted specifically to the current response.

The base advantage separates into the cross-teacher disagreement and a reference-anchoring term

Because $\widehat{b}_{it}$ estimates recurring teacher–reference disagreement, we residualize only $d_{it}$ and leave the reference-anchoring term unchanged. We set

The resulting advantage is

At initialization, $\pi_{\theta}=\pi_{R}$ and $A^{\mathrm{res}}_{it}=d^{*}_{it}$. We retain the Lightning OPD objective form while replacing the raw cross-teacher disagreement with its residualized counterpart

Equivalently, at the conceptual objective level, we define the fixed effective chosen-token teacher score $\widetilde{\ell}^{T}_{it}=\ell^{T}_{it}-\widehat{b}_{it}$ and use it in place of $\ell^{T}_{it}$. This score and $\widehat{b}_{it}$ are computed before training and remain fixed throughout optimization.

Table 1 presents evaluation results across five benchmarks under the two cross-teacher settings. Across both settings, Lightning OPD 2.0 consistently improves the SFT reference, whereas Lightning OPD and the two competing OPD baselines yield little or no aggregate improvement once teacher consistency is violated. In the Qwen3-4B setting, Lightning OPD 2.0 improves all five benchmarks, raising the average math and code scores over the SFT reference by $3.4$ and $3.1$ points, respectively. Compared with Lightning OPD, Lightning OPD 2.0 further improves average math and code performance by $3.1$ and $1.4$ points. It also achieves the strongest average results among all methods, outperforming IW-OPD by $3.5$ points on math and $0.2$ points on code, and TA-OPD by $6.1$ and $2.6$ points. In the Klear-Reasoner-8B-SFT setting, Lightning OPD 2.0 similarly improves all five benchmarks, with average gains of $1.0$ points on math and $3.6$ points on code over the SFT reference. Compared with Lightning OPD, it improves average math and code performance by $1.0$ and $1.4$ points, respectively, and remains stronger than IW-OPD and TA-OPD in both domains. Together, these results demonstrate that Lightning OPD 2.0 generalizes across SFT initializations and data-generation pipelines without requiring teacher consistency.

We assess whether residualization moves the cross-teacher training signal toward the signal obtained under teacher consistency. Let $\pi_{D}$ denote the model used for this post-hoc diagnostic and let $\ell^{D}_{it}$ be its chosen-token log-probability. For each scored response token in the cache, we measure the absolute deviation of the raw and residualized signals from the diagnostic reference signal

where $d_{it}=\ell^{T}_{it}-\ell^{R}_{it}$ is the original cross-teacher disagreement and $d^{*}_{it}$ is its residualized counterpart. In the Qwen3-4B-SFT setting, $\pi_{D}$ is Qwen3-8B, the provenance teacher $\pi_{G}$ that generated the SFT demonstrations. The diagnostic reference is therefore the exact teacher-consistent signal in this setting. Exact teacher-consistent scores are unavailable in the Klear-Reasoner-8B-SFT setting, so we use Klear-Reasoner-8B (kwai2026klear) as a diagnostic proxy. Neither diagnostic model contributes to correction estimation or student training.

Figure 3 compares the empirical fractions of tokens satisfying $o_{it}>\tau$ and $o^{*}_{it}>\tau$ as $\tau$ varies from $0.5$ to $3.0$ nats. The after-correction curve lies below the before-correction curve at every evaluated threshold in both settings, showing that the result is not specific to a single cutoff. At $\tau=1$, which we use as a representative reference point, residualization reduces the exceedance rate from $8.14\%$ to $3.85\%$ in the Qwen3-4B-SFT setting and from $7.19\%$ to $2.02\%$ in the Klear-Reasoner-8B-SFT setting. These changes correspond to relative reductions of $52.8\%$ and $71.9\%$, respectively. The relative reduction grows at higher thresholds, showing that residualization reduces the prevalence of large deviations from the diagnostic reference signal.

Figure 4 provides a complementary token-level view using two cached responses from the Qwen3-4B-SFT setting. For visualization, we color tokens with deviations of at most one nat in green and those above one nat in red, with darker shades denoting larger deviations. In these examples, several large raw deviations occur on discourse markers and reasoning transitions. Most fall below the reference threshold after residualization, although some large deviations remain. The examples qualitatively illustrate the threshold-sweep result without treating the visualization threshold as part of the method. Together, the aggregate and token-level evidence supports our operational interpretation of the removed component as recurring style-token bias.

Table 2 ablates the three principal design choices in Lightning OPD 2.0. The context-only and token-only variants retain one lookup at unit weight, whereas the in-sample variant retains both lookups but removes prompt-level cross-fitting. Lightning OPD serves as the uncorrected baseline, and all variants otherwise share the same cache and training configuration within each setting. The full method yields the strongest overall reported point estimates. The effects of using a single coordinate family vary across settings, while in-sample estimation is generally weaker or tied. These results support combining lexical and coarse contextual estimates and retaining cross-fitting primarily as a safeguard against self-fitting.