Token Utility for On-Policy Distillation via Local Trust-Region Optimization

arXiv:2610.18883 · cs.LG, cs.AI · Submitted 2026-10-01 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Token Utility for On-Policy Distillation via Local Trust-Region Optimization".

Jane: The gist The authors derive token utility for on-policy distillation (OPD) from the objective being optimized, yielding an optimization-based construction for token utility in OPD;

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So we're looking at this paper called Token Utility for On-Policy Distillation via Local Trust-Region Optimization. It’s about how to pick which parts of the student's prediction matter most when you are using on-policy distillation.

Jane: That sounds technical, Tom. Can you tell us what the main idea is? What problem are they trying to solve with this Token Utility for On-Policy Distillation via Local Trust-Region Optimization paper?

Tom: The core thesis here is that instead of using external signals like uncertainty or disagreement to pick important positions, they derive token importance directly from the distillation objective itself. They ask, "how much can each position's loss decrease under an equal local KL budget on the student's predictive distribution?" and that answer gives them a utility.

Lu: That shifts the focus from heuristics to optimization, which is interesting because it grounds importance in what actually helps reduce the loss during training.

Meng: But how does this utility work? If you're talking about optimizing token importance, what are the actual mathematical tools they use to calculate that? I need to know how this translates into something we can actually implement on our production models.

Tom: Okay, so they define an optimization problem at each student-visited prefix. They use the Fisher geometry induced by KL divergence to find the maximum first-order improvement under a common radius epsilon, and that improvement is related to the utility st = g top t F dagger t gt.

Jane: The paper mentions this score st equals g top t F dagger t gt, where gt is the logit gradient and Ft is the categorical Fisher matrix. That sounds like a specific way to measure how much an update helps under a certain constraint.

Lalam: I see that in page three they say that for reverse-KL distillation, this score st admits a closed form: st equals Vara∼pS,t log pS t(a) pT,t(a).

Meng: A closed form is good for speed when you're running things at scale. But what about forward-KL distillation? Is the utility calculation the same there, or does it change based on the objective?

Paper summary: Tom: It definitely changes depending on the objective. For forward-KL distillation, that score st becomes related to the Pearson chi squared divergence, which is given by st FKL equals X i(q i minus p i) squared over pi.

Lu: So you have a different utility function based on whether you are doing reverse or forward KL distillation; the choice of objective dictates the resulting utility calculation. That makes sense because they need the right measure for each scenario.

Jane: And once they have this utility, what do they actually do with it? They convert this score into something called Fisher Utility Weighting, or FUW, which applies damped mean-normalized weights to the token-level OPD objective.

Tom: FUW is their continuous, mean-normalized token weighting. It’s not just a binary yes or no for a position; it’s how much that position should contribute based on the improvement scale st. They set weights proportional to the improvement scale root st.

Meng: So instead of picking positions based on whether they cause disagreement, FUW is setting their contribution level based on how much they help reduce loss under a common budget. That sounds like a more nuanced control than just selecting tokens that are different from the teacher.

Lalam: I agree with Meng, and this framework provides a reference point against which other criteria can be compared because the utility is derived directly from the objective itself.

Tom: They found that across seven mathematical reasoning benchmarks and three Qwen3 student–teacher pairs, this FUW improves over matched uniform OPD under both reverse-KL and forward-KL, with macro points ranging from one point three five to four point zero for reverse-KL and one point zero eight to two point seven one for forward-KL.

Jane: That's a solid empirical result, Tom. So what does this mean for someone just listening? It means that basing your token selection on the optimization objective can actually lead to better performance than just picking tokens based on uncertainty or how much the teacher disagrees with you.

Lu: It suggests that focusing on where the optimization effort should go, defined by this utility st, is a more principled way to allocate resources in distillation.

Paper summary: Meng: For deployment, I'm looking at how this would run. They mention they use a candidate set Ct of top sixteen from both student and teacher distributions with a size of at most thirty-two for training loss calculation. That sounds like manageable overhead compared to evaluating every single token position during the entire process.

Tom: Exactly, it keeps the candidate set small and controlled while still using this utility to decide which ones get the weight. The paper notes that several independently motivated criteria, like student entropy or KL disagreement, concentrate their importance on nearly the same positions as this derived utility at a fifty percent selection fraction <ref:2610.18883#pg1,concentrate their importance on nearly the same positions as>.

Jane: That's a very strong hint from the results. It suggests that many of those other methods are actually pointing in the right direction when you look at where they are most important, but FUW provides a single score derived from the training objective.

Lalam: And I can also tell you that this framework accounts for where those other criteria diverge, meaning it tells you exactly where they stop agreeing.

Tom: The paper concludes that they’ve created an optimization-based construction for token utility in on-policy distillation by specifying the distillation objective and deriving its local utility under a common predictive distribution budget to characterize or prioritize student-visited positions.

Jane: So, to wrap up, this Token Utility for On-Policy Distillation via Local Trust-Region Optimization paper shows that you can derive a single, closed-form score from the training objective itself that consistently outperforms existing methods on various benchmarks and Qwen3 pairs.

Lu: The implication is that instead of choosing between different heuristic signals like uncertainty or mismatch direction, you can use the mathematical structure of the distillation loss to automatically determine token importance.

Meng: For a production system, this means we could replace some of those more complex selection mechanisms with this single score derived from the objective, which simplifies our pipeline and gives us a more robust way to handle token weighting at scale.

Tom: Yeah, it's about moving from choosing what looks good on paper to picking what mathematically maximizes the loss reduction under a specific constraint. That’s the mechanics here.

Conclusion: Tom: So we’re wrapping up this look at the paper on Token Utility for On-Policy Distillation via Local Trust-Region Optimization.

Jane: It boils down to deriving a token importance score directly from the distillation objective itself, instead of picking signals like uncertainty or disagreement.

Lu: The authors use a local trust-region optimization approach, using Fisher geometry induced by KL divergence to find how much loss decreases under an equal local budget.

Meng: So what’s the payoff here? What does this utility actually mean when we try to deploy it at scale?

Tom: It means they turn that math into something called Fisher Utility Weighting, which is a continuous weighting for each token position.

Jane: And the results show this method outperforms uniform OPD on seven mathematical reasoning benchmarks and three Qwen3 student-teacher pairs.

Lu: They see macro point improvements of one point three five to four point zero in reverse KL distillation and one point zero eight to two point seven one in forward KL distillation when using this utility, compared to the uniform baseline.

Meng: That’s a solid number, but what about the limitations? Does this work if we’re dealing with really massive models or complex reasoning tasks?

Tom: The paper notes that they evaluate it across those specific benchmarks and pairs. It focuses on mathematical reasoning, so we should keep that in mind for other domains.

Jane: The implication for us is that instead of guessing which tokens are important based on external heuristics, we can use the objective itself to guide our token weighting.

Lu: It offers a principled way to characterize student-visited positions by looking at how they impact the loss under a specific distributional budget constraint.

Tom: And it’s competitive with other selective methods because it uses a single closed-form score derived right from the training objective.

Jane: So, moving forward, this suggests we can ditch some of those separate heuristics and use this utility to automatically prioritize where our distillation effort should go next.

Lu: The idea is that by following that optimization utility, we get a more accurate map of which tokens actually matter for improving the model's performance in a structured way.

Narayanan Sadagopan, Yilan Chen, Vijay Lingam

AWS AI Labs

cs.LG, cs.AI

Submitted: 2026-10-01

Updated: 2026-10-01

Comments: 10 pages, Accepted to Pre-to-Post workshop at NeurIPS 2026

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 86/100

The gist: The gist The authors derive token utility for on-policy distillation (OPD) from the objective being optimized, yielding an optimization-based construction for token utility in OPD; this framework

Key concepts

On-Policy Distillation (OPD): A technique that transfers knowledge from a strong teacher model to a student model while keeping the student's own training distribution intact. The student generates the data, and the teacher provides supervision only at the specific positions (tokens) the student actually visits during training.
Token Utility ($s_t$ or $st$): A mathematical measure derived from optimization that quantifies how much a specific token's loss can decrease if it is allowed an equal local change in the student's predictive distribution. It is calculated using the logit gradient and the categorical Fisher matrix, representing the maximum first-order improvement under a common distributional budget.
Fisher Utility Weighting (FUW): A continuous, mean-normalized token weighting scheme derived from the objective-based utility. Instead of deciding if a position contributes (like TIP), FUW uses the magnitude of its utility to set how much each position should contribute to the total loss during training, proportional to the improvement scale.

Terminology

Distillation — 1× in this paper · explained in 2 episodes
On-Policy Distillation — 1× in this paper · explained in 3 episodes

Summary

The gist The authors derive token utility for on-policy distillation (OPD) from the objective being optimized, yielding an optimization-based construction for token utility in OPD; this framework provides a complementary perspective to recent notions based on uncertainty, disagreement, reasoning role, support compatibility, or mismatch type.

How it works

The core idea is to define importance through the optimization problem itself: how much can each position’s loss decrease under an equal local KL budget on the student’s predictive distribution (Page 1). This leads to the derivation of a common optimization framework for token utility in OPD, where comparing positions under a common distributional budget yields st = g⊤t F†t gt (Page 2). The solution to the Fisher-constrained steepest-descent problem yields the local token optimization utility as s = g⊤F†g (Page 3).

Key Derivations

The utility is derived differently depending on the distillation objective:

For reverse-KL OPD, it reduces exactly to student-weighted log-ratio variance for reverse-KL (Page 3). The closed form is given by s = Varv∼pS log pS(v) / pT (v) (Page 3).
For forward-KL OPD, the utility becomes the Pearson χ2 divergence and is given by sFKL = Xi(qi − pi)2 / pi (Page 17).

Practical Implementation

The derived utility is converted into Fisher Utility Weighting (FUW), which applies damped, mean-normalized weights to the token-level OPD objective (Page 2). For training, the authors use a candidate set defined as Ct = Top16(pS,t) ∪ Top16(pT,t), Ct ≤ 32 (Page 6). The per-position training loss is calculated using this candidate set and renormalized distributions, and the resulting weights are set as wt = 1 for every position if the utility is positive, or zero otherwise (Page 10).

Comparative Analysis

The framework is evaluated across seven mathematical-reasoning benchmarks and three Qwen3 student–teacher pairs, showing that FUW improves over matched uniform OPD under both reverse-KL (1.35–4.0 macro points) and forward-KL (1.08–2.71 macro points) (Page 1). Furthermore, the analysis reveals that several of them concentrate importance on largely the same positions as the derived utility at a 50% selection fraction, while diverging at stricter fractions and for direction-specific criteria (Page 2). The framework is shown to be robust against variations in scoring or weighting rules because high Bρ measures agreement in effective high-importance support, not equality of the underlying weight vectors (Page 21).

Conclusion

The paper concludes that the framework provides an optimization-based construction for token utility in OPD: specify the distillation objective, derive its local utility under a common predictive-distribution budget, and use that utility to characterize or prioritize student-visited positions (Page 10). This method successfully improves aggregate performance over uniform weighting under both reverse- and forward-KL distillation. The analysis is local and first-order: it evaluates predictive distributions independently at fixed student-visited prefixes, whereas shared transformer parameters couple positions during training (Page 9).

--- Page 1 ---

Preprint

TOKEN UTILITY FOR ON-POLICY DISTILLATION VIA

LOCAL TRUST-REGION OPTIMIZATION

Narayanan Sadagopan1∗ Yilan Chen2† Vijay Lingam1

1AWS AI Labs 2UC San Diego

ABSTRACT

Token-selection methods for on-policy distillation (OPD) prioritize student-visited positions using diverse signals, including uncertainty, teacher–student disagreement, and mismatch direction. We instead derive token importance from the distillation objective itself: how much can each position’s loss decrease under an equal local KL budget on the student’s predictive distribution. The answer is the utility st = g⊤t F†t gt, where gt is the logit gradient and Ft is the categorical Fisher matrix. For reverse-KL it equals the student-weighted variance of the student–teacher log-probability ratio, and for forward-KL it equals the Pearson χ2 divergence, so the appropriate utility follows from the choice of objective. On shared training checkpoints, several independently motivated criteria, including student entropy, KL disagreement, and TIP, concentrate their importance on nearly

the same positions as the derived utility at a 50% selection fraction, while diverging at stricter fractions and for direction-specific criteria. We instantiate the

utility as Fisher Utility Weighting (FUW), a continuous, mean-normalized token

weighting. Across seven mathematical reasoning benchmarks and three Qwen3

student–teacher pairs, FUW improves over matched uniform OPD under both reverse-KL (1.35–4.0 macro points) and forward-KL (1.08–2.71 macro points),

and is competitive with existing selective methods using a single closed-form score

derived from the training objective.

--- Page 2 ---

Preprint

equal-budget comparison separates the amount of disagreement from how effectively a local update can resolve it. We formulate an independent local trust-region problem at each student-visited prefix. Using the Fisher geometry induced by KL divergence, the maximum first-order improvement under a common radius ε is √2εst, where st = g⊤t F†t gt. Here, gt is the gradient of the token-level distillation objective with respect to the student logits and Ft is the categorical Fisher matrix. For reverse-KL OPD, the general score admits the closed form st = Vara∼pS,t log pS(a) pT,t(a). Thus,

the variance of the log-probability ratio is not introduced as a disagreement heuristic. It is the optimization utility induced by an equal local distributional budget constraint. We convert this score into Fisher Utility Weighting (FUW), which applies damped, mean-normalized weights to the token-level OPD objective. Unlike selection schemes such as TIP and TIDE, which decide whether a position contributes, FUW uses the utility’s magnitude to set how much each position contributes, with weights proportional to the improvement scale √st. Since the utility is derived from the objective rather than chosen, it also provides a reference against which existing criteria can be compared. We find that several of them concentrate importance on largely

the same positions as the derived utility, and that the framework accounts for where they diverge. Our contributions are

threefold:

--- Page 3 ---

1 INTRODUCTION

On-policy distillation (OPD) transfers capabilities from a stronger teacher while retaining the student’s state distribution (Agarwal et al., 2024; Gu et al., 2024). The student generates the training trajectories, and the teacher supplies dense distributional supervision at the prefixes that the student actually visits. This avoids the state-distribution mismatch of offline distillation and has become an effective recipe for improving reasoning models. Yet dense supervision does not imply that every supervised position is equally useful. Some student and teacher distributions are already closely aligned; elsewhere they disagree, but the correction may require a large change in the student’s predictive distribution. Uniform OPD, which weighs all token positions the same, does not distinguish between these cases. Recent work addresses this heterogeneity using properties of the visited positions. Entropy-based methods emphasize uncertain decisions (Jin et al., 2026). TIP (Xu et al., 2026) combines student uncertainty with teacher–student divergence; DEAR (Xiao et al., 2026) separates decision tokens from supporting evidence tokens; Token Teachability (Wang et al., 2026) asks whether teacher corrections lie within compatible student support; and TIDE (Yu et al., 2026) distinguishes student-excess from student-deficit probability mass. These approaches operationalize token importance through properties such as uncertainty, disagreement, reasoning role, support compatibility, or mismatch direction. We take a different approach, defining importance through the optimization problem itself: how much can each position’s expected distillation objective improve under an equal local budget on the student’s predictive distribution?

Suppose that every position is allowed the same infinitesimal change in the student’s next-token

distribution. Which position permits the largest decrease in its expected distillation loss? This∗Corresponding author.

†Work done during an internship at AWS AI Labs.

--- Page 4 ---

--- Page 5 ---

2 THEORY

The KL–Fisher expansion and the natural-gradient solution to a Fisher-constrained steepest-descent problem are classical results (Amari & Nagaoka, 2000; Kakade, 2001). Our contribution is to view token importance in on-policy distillation through the lens of local optimization: comparing positions under a common distributional budget yields objective-specific utilities in closed form. The trust-region serves only to derive these per-position utilities; it is not imposed as a constraint during training.

Improvements for AI systems

textbfProcessor Allocation via Fisher Utility Weighting (FUW): The improved system can allocate optimization effort based on how much can each position’s loss decrease under an equal local KL budget on the student’s predictive distribution, which is quantified by the utility st = g⊤t F†t gt. This leads to a continuous, mean-normalized token weighting that outperforms uniform OPD, as shown by FUW improving macro mean accuracy over matched uniform OPD across all three Qwen3 student–teacher pairs.
textbfObjective-Specific Utility Adaptation: The system can dynamically derive the appropriate utility based on the distillation objective, stating that the appropriate utility follows from the choice of objective (e.g., for reverse-KL it equals the student-weighted variance of the student–teacher log-probability ratio, and for forward-KL it equals the Pearson χ2 divergence). This allows for superior performance in both reverse-KL (1.35–4.0 macro points) and forward-KL (1.08–2.71 macro points).
textbfPrecise Token Prioritization: The system can prioritize tokens by the derived utility rather than arbitrary heuristics, as the utility measures the maximum first-order improvement available at token position t after normalizing out the common maximum distributional budget. This results in a ranking where high Bρ measures agreement in effective high-importance support, not equality of the underlying weight vectors, allowing for more targeted knowledge transfer.
textbfContextual Robustness: The system can maintain performance consistency across different model pairs and training trajectories by utilizing the derived utility to ensure that shared effective support does not imply identical weighting behavior, leading to a robust framework where the same qualitative effective-support structure appears at both ρ = 0.2 and ρ = 0.5 in cross-pair robustness checks.
textbfAdaptability to Forward Distillation: The system can effectively implement token utility for forward-KL OPD by using the corresponding objective-specific utility, which is found to be the Pearson χ2 (pT pS) divergence, improving both Macro and pooled Avg@8 on all three model pairs compared to uniform forward-KL distillation.

Abstract

Token-selection methods for on-policy distillation (OPD) prioritize student-visited positions using diverse signals, including uncertainty, teacher-student disagreement, and mismatch direction. We instead derive token importance from the distillation objective itself: how much can each position's loss decrease under an equal local KL budget on the student's predictive distribution? The answer is the utility s t=g t F t g t, where g t is the logit gradient and F t is the categorical Fisher matrix. For reverse-KL it equals the student-weighted variance of the student-teacher log-probability ratio, and for forward-KL it equals the Pearson χ squared divergence, so the appropriate utility follows from the choice of objective. On shared training checkpoints, several independently motivated criteria, including student entropy, KL disagreement, and TIP, concentrate their importance on nearly the same positions as the derived utility at a 50% selection fraction, while diverging at stricter fractions and for direction-specific criteria. We instantiate the utility as Fisher Utility Weighting (FUW), a continuous, mean-normalized token weighting. Across seven mathematical reasoning benchmarks and three Qwen3 student-teacher pairs, FUW improves over matched uniform OPD under both reverse-KL (1.35-4.0 macro points) and forward-KL (1.08-2.71 macro points), and is competitive with existing selective methods using a single closed-form score derived from the training objective.