PaperlayerView exact revision on arXiv

Source-linked AI summary

A Unified Per-Token Gating Family for On-Policy Distillation: FKL/RKL Mixing with Multi-Channel and Bias Coefficients

Suwan Wu, Yumeng Lin, Pengcheng Yuan, Xiaolong Jiang

arXiv:2609.11768v1cs.AIcs.CLcs.LG

TL;DR

Existing per-token OPD gates use fixed signals and directions, leaving EOPD and ToDi without a direct comparison. The paper introduces a four-coefficient family that embeds aligned proxies of both methods and adds multi-channel composition and bias; exploratory sweeps favor the full family and dynamic gating, while three-seed replications yield smaller, non-significant individual effects.

  • Problem

    EOPD and ToDi fix different gating signals and directions, so their improvements and relative merits had not been directly separated or compared.

  • Method

    The paper defines λt = σ(aht + bu(x) + c + dgapt), embedding direction-aligned EOPD and ToDi proxies as one-dimensional restrictions within a shared family.

  • Results

    The full family leads matched single-channel restrictions in 33 of 36 comparable cells, while dynamic gating leads mean-matched static baselines in 19 of 26 cells.

  • Takeaways & Limitations

    The parameterization is primarily a shared coordinate system for comparing per-token gating designs in short-output classification OPD.

  • Takeaways & Limitations

    The aggregate counts summarize correlated sweeps, and none of nine targeted three-seed comparisons is individually significant at n=3.

Abstract

from arXiv · show

Per-token gating of forward/reverse KL losses has become a standard technique for on-policy knowledge distillation (OPD), but existing methods such as EOPD (Jin et al., 2026) and ToDi (Jung et al., 2025) each fix a single gating signal and a single gating direction, and the two have never been compared directly. We introduce a four-coefficient parameterization lambda_t = sigma(a * h_t + b * u(x) + c + d * gap_t) in which direction-aligned proxies of EOPD and ToDi appear as one-dimensional (1D) restrictions, and which adds multi-channel composition and an explicit bias as further degrees of freedom. On TweetEval (Barbieri et al., 2020) emotion and hate, with a Qwen3-32B teacher and a Qwen3-4B student, configurations in the full family reach higher accuracy than the matched-magnitude single-channel (entropy-only / gap-only) 1D restrictions in 33 of 36 comparable cells, and a 26-cell mean-match isolation experiment places dynamic gating ahead of effective-KL-matched static baselines in 19 of 26 cells. Because cells share training data, models, and parameter substructure, we report both counts as exploratory aggregate directional evidence rather than as independent hypothesis tests. Targeted three-seed paired replications of the nine headline comparisons singled out by that sweep -- including a third task, offensive -- are directionally consistent, but individually smaller than the single-seed estimates and not significant at n=3. We therefore present the parameterization primarily as a shared coordinate system for comparing per-token gating designs in short-output classification OPD.

1 Introduction

The paper introduces a unified four-coefficient gating family to compare entropy- and disagreement-based per-token KL mixing, while adding prompt-level signals, composition, and bias. Across exploratory sweeps, the full family generally outperforms matched single-channel restrictions and dynamic gates often outperform mean-matched static baselines, but three-seed replications temper individual effects.

  • A unified analysis framework: The family uses λt = σ(a · ht + b · u(x) + c + d · gapt), exposing four signal coefficients and an explicit bias.The signals include token entropy, prompt-level entropy, and teacher–student disagreement; the mixture is λtLRKL(t) + (1 − λt)LFKL(t).
  • A unified analysis framework: Direction-aligned proxies for EOPD and ToDi appear as one-dimensional restrictions within the shared convex-mixture family.These are structurally aligned proxies rather than faithful reproductions of the published algorithms.
  • Empirical contributions: The experiments use Qwen3-32B teachers, Qwen3-4B students, and 13 configurations per task on TweetEval emotion and hate.A third task, offensive, is added for seed-robustness replication.
  • Empirical contributions: 33 of 36 comparable cells favor the full family over matched-magnitude entropy-only or gap-only restrictions.The authors treat this as exploratory aggregate directional evidence because cells share data, models, and parameter substructure.
  • Empirical contributions: 19 of 26 mean-matched isolation experiments favor dynamic gating over static baselines with the same effective KL ratio.This comparison is designed to separate per-token structure from average KL weighting.
  • Empirical contributions: Nine targeted three-seed headline replications remain directionally consistent, but all means are smaller than single-seed estimates and none is significant at n=3.The authors therefore state empirical claims at the group level.

2 Related Work

Prior work combines forward and reverse KL objectives and has explored adaptive gating at token, sequence, or batch granularity. EOPD and ToDi are the closest per-token comparators, but differ in signals, regimes, and conventions, preventing direct empirical comparison before this framework.

  • Related work: The related-work landscape spans mixed FKL/RKL objectives, per-token gating, and on-policy distillation infrastructure.EOPD belongs to an on-policy framework, whereas ToDi was proposed for offline distillation.
  • 2.1 Mixed FKL/RKL Distillation: Mixed FKL/RKL distillation balances mode-covering forward KL against mode-seeking reverse KL through a weighting parameter.Fixed mixtures use a scalar λ, while some adaptive variants operate only at sequence or batch granularity.
  • 2.2 Per-Token Gating: EOPD and ToDi: EOPD uses teacher entropy to give uncertain tokens more FKL weight, preserving distributional diversity.Its gating signal is entropy alone, and its original structure is additive rather than a convex mixture.
  • 2.2 Per-Token Gating: EOPD and ToDi: ToDi uses teacher–student disagreement to give more FKL weight where the teacher assigns more probability than the student.Its signal is computed per vocabulary entry and routes overestimated entries toward RKL.
  • 2.2 Per-Token Gating: EOPD and ToDi: EOPD and ToDi had not been directly compared because they use different parameter subspaces, scaling conventions, datasets, models, and metrics.The paper places direction-aligned proxies of both methods in one family for matched-magnitude comparison.

3 Method: Parametric Family and Aligned Restrictions of Prior Methods

The method defines a four-coefficient sigmoid gate over token entropy, prompt entropy, disagreement, and bias, then embeds entropy-only and gap-only prior-method proxies as one-dimensional restrictions. Sweep patterns suggest task-dependent coefficient preferences, but they are descriptive and seed-sensitive.

  • 3.1 Parametric Family: The family combines token entropy, prompt-level entropy, teacher–student disagreement, and bias to control per-token FKL/RKL mixing.All three signals are bounded and non-negative, with larger values indicating that the student needs more teacher-mode learning.
  • 3.2 Coefficient Patterns Observed in the Sweep: The reported coefficient patterns are descriptive single-seed observations rather than established rules.Individual cell magnitudes are seed-sensitive.
  • 3.2 Coefficient Patterns Observed in the Sweep: Negative bias appears in the best configuration for each task, shifting indifferent-token gating below λt = 0.5.The effect is clear on hate but within noise on emotion.
  • 3.2 Coefficient Patterns Observed in the Sweep: The best configuration on each task zeroes one of a or b and amplifies the other, unlike symmetric all-on baselines.The authors relate this pattern to sigmoid saturation, while noting the emotion contrast is within the sampling-SE reference band.
  • 3.2 Coefficient Patterns Observed in the Sweep: The useful direction of the gap coefficient is task-dependent, and adding it to the task-best abc configuration helps only when granularities match.Both gap directions improved a weak abc baseline in the single-seed sweep.
  • 3.3 Aligned 1D Restrictions of Prior Methods: The gap-only restriction matches ToDi’s gating direction when d < 0 but remains a structural analogue rather than a reproduction.The proxy uses a probability-difference gap instead of ToDi’s whole-vocabulary log-ratio gate and omits other algorithmic details.
  • 3.3 Aligned 1D Restrictions of Prior Methods: The entropy-only restriction preserves EOPD’s direction: higher entropy drives λt toward 0, making the mixture FKL-heavy.It is a soft, zero-bias convex relaxation rather than EOPD’s additive hard-threshold mechanism.

4 Main Results: Grid Analysis

The grid analysis evaluates a four-coefficient gating family across 13 configurations per task, varying sample-, token-, bias-, and gap-based signals. Hate shows a clear sample-level preference, whereas emotion’s small grid spread does not support a channel preference; gap effects depend on gating granularity.

  • Grid Design: The sweep varies four coefficients through abc, H, and HG configurations, covering channel balance, bias, and gap additions across 26 configuration-task cells.The protocol uses 13 OPD configurations per task and compares them with aligned single-coefficient restrictions.
  • 4.2 abc Grid: Task-Dependent on Hate, Flat on Emotion: Hate selects G4 = (0, 4, −1) with sample-level u(x) gating and ACC = 0.5586, while emotion selects G5 = (4, 0, −1) with token-level ht gating and ACC = 0.7748.The hate grid spans 3.64pp, whereas the emotion grid spans only 0.56pp.
  • 4.3 d on a Weak Baseline: Adding d to the weak G2 baseline improves accuracy in all 8 settings, with a maximum gain of +3.27pp on hate.This is a single-seed sweep across two tasks and one base configuration.
  • 4.4 d on the Task-Best abc: Granularity-Conditional: Adding d to task-best abc helps emotion but degrades hate in all four tested settings, consistent with a possible mismatch between token-level gap gating and hate’s sample-level gate.The mismatch interpretation is treated as a hypothesis because the unfavorable cell shrinks to within noise under three-seed replication.
  • 4.5 Effective-KL-Matched Static Baselines: Mean-matched static baselines trail dynamic gating in 19 of 26 cells, indicating directional evidence beyond matching the effective KL ratio.Cells share data and configuration substructure, so the count is reported as an aggregate directional summary rather than independent tests.

5 Comparison with Aligned 1D Restrictions of Prior Methods

The four-coefficient family is compared with matched-magnitude entropy-only and gap-only restrictions, which are structurally aligned proxies rather than reproductions of EOPD and ToDi. Across comparable cells, multi-coefficient configurations generally lead, although shared data and substructure make the counts exploratory directional evidence.

  • Comparison setup: The gap-only and entropy-only comparisons retain only d or −a, respectively, while marking degenerate matched-magnitude restrictions as N/A.Gap-only follows ToDi’s direction when d < 0, and entropy-only reverses the EOPD sign convention because λ_t weights RKL rather than FKL.
  • Comparison scope: The comparison reuses the family’s accuracies but separately trains and evaluates the gap-only and entropy-only cells under an identical overall training setup.The restrictions differ from the published methods in signal granularity, gradient treatment, functional form, and coefficient tuning.
  • Aggregate findings: 19 of 26 mean-matched isolation cells favor dynamic gating over static baselines with the same effective KL ratio.No configuration falls below the negative 1σ sampling-SE band, but the cells are correlated and the band does not represent across-seed variability.
  • Aggregate findings: 33 of 36 comparable cells favor the full family over matched single-channel restrictions, including 14 of 16 hate cells and 19 of 20 emotion cells.The three losses include one hate comparison outside the reference band, attributed to a granularity mismatch; per-cell magnitudes are not stable across seeds.
  • Configuration patterns: Multi-coefficient configurations, especially H and HG5, attain the strongest sweep accuracies and largest gains over aligned 1D restrictions.The result supports multi-channel composition within this implementation, not a faithful reproduction or direct validation of the published algorithms.
  • Interpretation: The 36-cell aggregate is treated as directional evidence rather than independent hypothesis testing because cells share data, models, and overlapping configuration substructure.This dependence limits interpretation of the count as a collection of separate statistical tests.

6 Seed Robustness and a Third Task

Three-seed paired replications across emotion, hate, and offensive reduced the single-seed effect estimates, while most comparisons remained directionally positive but individually nonsignificant.

  • Design: The nine replications were selected from the single-seed sweep and included offensive alongside emotion and hate, covering restriction and isolation comparisons.
  • Results: Every three-seed mean was smaller than its single-seed estimate by roughly 1.5–3.1×, so the paper adopts the three-seed means as more reliable effect sizes.
  • Uncertainty: None of the nine comparisons was individually significant at n = 3, with p ∈[0.29, 0.51], so claims remain at the group level.
  • Results: Eight of nine three-seed replicated gaps remained directionally positive across both comparison axes and all three tasks.The lone negative was hate HG4,d+2, which moved from −1.78pp to −0.72±0.92pp and inside the reference band.
  • Uncertainty: Dynamic configurations showed greater seed variability than mean-matched static baselines, separating training variability from test-set sampling uncertainty.Across-seed variability was 0.2–1.8pp, while dynamic configurations had 0.9–1.8pp versus 0.20pp for static baselines on emotion and offensive.

7 Conclusion

The paper presents a four-coefficient gating family as a shared coordinate system for comparing per-token designs, with aggregate advantages scoped to short-output classification OPD. Its conclusions remain group-level because the comparisons use proxies, single-seed sweeps, limited tasks and models, and grid-searched coefficients.

  • Conclusion: The four-coefficient family places direction-aligned EOPD and ToDi proxies in one space, with multi-coefficient gating ahead of matched 1D restrictions in 33 of 36 cells.The paper reports headline magnitudes of roughly 0.5–1.7pp with wide three-seed intervals.
  • Conclusion: Dynamic gating led effective-KL-matched static baselines in 19 of 26 cells, while the three per-task headline pairs remained positive across seeds.
  • Scope: The evidence supports group-level directional claims rather than per-cell advantages because the full grid and isolation experiment used one seed and headline replications were nonsignificant at n = 3.
  • Scope: The aligned entropy-only and gap-only comparators are family-internal proxies, not reproductions of the published EOPD and ToDi algorithms.
  • Scope: The conclusions are limited to short-text, short-output TweetEval classification with one Qwen3-32B/4B pair and should not be extrapolated to long-form generation, reasoning, or other architectures.
  • Practical limitation: Because coefficient selection relies on manual design and grid search without a principled automatic criterion, the family is intended for analysis and comparison rather than turnkey deployment.

A On-Policy Distillation Background and Infrastructure

The paper uses on-policy distillation in which the student generates responses and the teacher supplies token probabilities during training. Its implementation combines Qwen3 models, vLLM, Megatron-LM, and TweetEval classification evaluation.

  • Background: On-policy distillation has the student generate responses during training while the teacher computes per-token target probabilities on the rollout tokens.This addresses the train–test distribution shift associated with offline distillation on fixed teacher-generated responses.
  • Infrastructure: The implementation uses a Qwen3-4B student, Qwen3-32B teacher, vLLM generation and teacher serving, ZMQ message passing, and Megatron-LM backpropagation.
  • Training: Training uses learning rate 10−6, batch size 72, 100 steps for emotion, 200 steps for hate and offensive, and rkl_ratio = 0.5 or E[λt] for static baselines.
  • Evaluation: Evaluation uses greedy decoding with a 32-token cap and strict JSON label output, producing 1–3-token classification responses scored by accuracy.
  • Data: The experiments cover TweetEval emotion, hate, and offensive, with test sizes of 1421, 2970, and 860 examples respectively.

B Three-Seed Replication Details

The replication analysis pairs both sides of each comparison within three seeds and reports uncertainty through paired differences, t-tests, and wide confidence intervals. Seed variability was especially high for several headline dynamic configurations, while Figure 5 probes gap-signal behavior across task and granularity settings.

  • Seed variability: Across-seed standard deviations ranged from 0.20–1.82pp and were largest for configurations producing the biggest single-seed gaps, indicating optimistic selection effects.Mean-matched static baselines were the most stable on emotion and offensive, at SD = 0.20pp.
  • Statistics: Each replication computes paired per-seed differences, their mean and SD, a two-sided paired t-test, and a t-based 95% CI with n = 3.The resulting CI half-width is roughly 2.5 SDs of the paired differences, so intervals are wide by construction.
  • Gap analysis: Figure 5 compares adding the gap coefficient to a weak G2 baseline versus stacking it on each task’s best abc configuration.The appendix presents this as a decomposition of the single-seed sweep’s gap-coefficient effect.
  • Gap analysis: The gap signal was positive in all eight weak-G2 settings, but stacking on task-best abc was positive in 3/4 emotion settings and 0/4 hate settings.The caption labels emotion as token-level and matched, while hate is sample-level and mismatched.
  • Reporting: The reported comparisons use held-out accuracy, with the appendix also listing best-observed single-seed values within the explored sweep rather than per-method hyperparameter optimization.

C Additional Sweep Views

The appendix provides two supplementary views of the single-seed sweep: a decomposition of the gap coefficient and a best-observed-accuracy comparison.

  • Additional Sweep Views: Figure 5 decomposes the effect of adding gap coefficient d to the weak G2 baseline or to each task’s best abc configuration.This supplements the Section 4.4 analysis.
  • Additional Sweep Views: Table 4 compares the best-observed accuracy of each side within the explored sweep.It offers a partial alternative to the matched-magnitude comparison in Section 5.
  • Additional Sweep Views: These views are referenced from the main text but are not needed to follow its primary argument.

D Design Choices Not Available to the Aligned Restrictions

The parameterization adds a bias coefficient and permits coefficient signs unavailable to the aligned gap-only and entropy-only restrictions. These freedoms make additional design choices observable, including a reversal of the entropy-gating direction on emotion.

  • Design Choices Not Available to the Aligned Restrictions: The family adds two structural freedoms absent from the gap-only and entropy-only 1D restrictions: an explicit bias and flexible coefficient signs.
  • Bias coefficient c: Best configurations use c = −1 on hate and emotion, although the emotion contrast lies within the sampling-SE reference band.The contribution of nonzero c is therefore descriptive rather than conclusive for these tasks.
  • Sign of coefficient a: Emotion’s best configuration G5 uses a > 0, so high entropy receives more RKL rather than the more-FKL direction fixed by EOPD’s a < 0 rule.This regime is unavailable to EOPD’s original entropy-only formulation.
Loading 2609.11768v1…