PaperlayerView exact revision on arXiv

Source-linked AI summary

CAST: Credit Alignment for Cross-Tokenizer On-Policy Distillation

Jingru Yang, Jinhang Wu, Minjing Shi, Shengfeng He, Xuyan Li, Man Tan, Biye Li

arXiv:2610.13257v1cs.LGcs.CL

TL;DR

Cross-tokenizer OPD leaves token-level credit allocation unresolved even after teacher and student sequences are aligned. CAST combines uncertainty–smoothness credit assignment with selective anchoring, improving six-benchmark averages over the strongest evaluated baseline across cross-tokenizer and shared-tokenizer transfers.

  • Problem

    Cross-tokenizer OPD lacks a well-defined token-wise comparison across different segmentations and vocabularies, while chunk alignment still leaves student-token credit distribution unresolved.

    Source: §1 · p. 2

  • Method

    CAST redistributes aligned chunks’ advantages using student uncertainty and local teacher-signal smoothness, selectively anchors unaligned positions, and gives the allocation an exact KL-regularized characterization.

    Source: §1 · p. 3, §5 · p. 9

  • Results

    CAST improves the six-benchmark average over the strongest evaluated baseline by 2.45 points for Qwen3-8B to Mistral-7B, 1.32 for Qwen3-8B to Llama-3.1-8B, and 2.10 for shared-tokenizer Qwen3-1.7B.

    Source: §1 · p. 3, §5 · p. 9

  • Takeaways & Limitations

    The same credit-assignment rule applies to punctuation-defined and fixed-length chunks under shared tokenization.

    Source: §1 · p. 3, §5 · p. 9

    Evaluation covers limited model families and mainly reasoning-oriented training data; compute varies across methods, and repeated-seed coverage is four Mistral and two Llama configurations.

    Source: LIMITATIONS · p. 10

Abstract

from arXiv · show

Cross-tokenizer on-policy distillation (OPD) requires aligning teacher and student sequences before their scores can be compared. Decoded-text alignment provides such correspondence at the chunk level, but leaves a distinct problem unresolved: credit alignment. A chunk-level signal specifies how much a span should change, but not how that signal should be distributed across student tokens, while tokens without teacher correspondence remain unconstrained. We propose CAST (Chunk-aware credit Assignment with Selective Token Anchoring), a shift-or-stay framework that addresses both issues. For aligned chunks, CAST redistributes each chunk's advantage across student tokens using predictive uncertainty and local teacher-signal smoothness, emphasizing learning opportunities supported by locally consistent teacher feedback while preserving the chunk's total advantage. The resulting weights exactly solve a KL-regularized credit-allocation problem. For unaligned positions, CAST selectively anchors the student to its initial policy through a KL penalty, limiting drift where teacher supervision is unavailable. Because credit assignment is separated from chunk construction, CAST also applies to punctuation-defined or fixed-length chunks under shared tokenization. Across six mathematics, science, and code benchmarks, CAST improves the average over the strongest evaluated baseline by $2.45$ and $1.32$ points for cross-tokenizer transfer from Qwen3-8B to Mistral-7B and Llama-3.1-8B, respectively, and by $2.10$ points for shared-tokenizer transfer to Qwen3-1.7B.

1 INTRODUCTION

Cross-tokenizer OPD methods align teacher and student signals at the chunk level, but leave token-level credit allocation and unsupervised positions unresolved. CAST assigns aligned-token credit using uncertainty and teacher-signal smoothness, anchors unaligned positions, and improves benchmark averages across heterogeneous and shared-tokenizer settings.

1 INTRODUCTION

  • Chunk alignment makes teacher and student scores comparable, but does not determine how chunk-level credit maps to individual student tokens or constrain unaligned positions.Uniformly assigning a chunk's advantage treats all tokens as equally useful and reliable, while unaligned tokens can drift through shared parameters despite lacking direct supervision.Source: §1 · p. 2
  • CAST redistributes each aligned chunk’s advantage using student predictive uncertainty and local teacher-signal smoothness, favoring uncertain predictions with locally consistent supervision.The normalized allocation preserves the chunk’s total advantage and exactly solves a KL-regularized allocation problem.Source: §1 · pp. 2–3
  • At unaligned positions, CAST applies a selective KL penalty toward the initial policy to constrain incidental drift without competing with teacher-guided updates.For cross-tokenizer pairs, the reference is the supervised fine-tuning initialization.Source: §1 · p. 3
  • CAST’s credit rule is separate from chunk construction, so it also applies to punctuation-defined or fixed-length chunks with a shared tokenizer.Multi-rollout sampling can further improve performance by exposing the credit mechanism to multiple student trajectories per prompt.Source: §1 · p. 3
  • CAST improves the six-benchmark average over the strongest evaluated baseline by 2.45 points for Mistral-7B and 1.32 for Llama-3.1-8B, plus 2.10 points for shared-tokenizer Qwen3-1.7B.The benchmarks span mathematics, science, and code.Source: §1 · p. 3

2 RELATED WORK

Related work addresses cross-tokenizer score comparison through common spaces and decoded-text chunk alignment, while shared-tokenizer OPD methods use other distillation objectives. CAST applies token-level shift-or-stay credit after alignment and supports multiple chunk-construction schemes.

2 RELATED WORK

  • Shared-tokenizer OPD work includes GKD, reverse-KL MiniLLM, streamlined and contrastive DistiLLM objectives, KDRL, and stability methods based on multiple responses or adaptive targets.These methods assume a shared tokenizer; CAST retains its token-level credit rule with punctuation-defined or fixed-length chunks.Source: §2 · p. 3
  • Cross-tokenizer methods construct comparable score spaces using optimal transport, shared text chunks, byte-level interfaces, or token mappings; DPCA aligns on-policy sequences into text-equivalent chunks and aggregates teacher scores.The cited approaches include ULD and MultiLevelOT, ALM and byte-level distillation, MinED, and DPCA [Niu et al., 2026].Source: §2 · p. 4
  • CAST redistributes each chunk’s on-policy advantage to individual tokens using uncertainty and smoothness, while selective KL anchoring constrains drift at unaligned positions.Its shift-or-stay framework follows decoded-text alignment for heterogeneous pairs and uses fixed-length credit chunks; shared-tokenizer training here uses punctuation-defined chunks and regularizes all positions.Source: §2 · p. 4, Figure 3 · p. 4

3 METHOD

CAST first forms teacher–student alignment chunks, then assigns each chunk’s learning signal across student tokens using uncertainty and teacher-signal smoothness. It anchors unaligned predictions to the initial policy and uses multi-rollout training.

3.1 PRELIMINARIES

  • CAST uses student-rollout and teacher-retokenized log probabilities to form alignment chunks, then groups aligned student tokens into credit chunks.Credit chunks can also be punctuation-defined or fixed-length under shared tokenization.Source: §3.1 · p. 4

3.2 UNCERTAINTY–SMOOTHNESS CREDIT ASSIGNMENT

  • Within each credit chunk, CAST redistributes the teacher–student gap as token advantages while preserving the chunk’s total advantage.The allocation uniquely minimizes the stated KL-regularized objective, and the adaptive scale bounds distortion from smoothness reweighting.Source: §3.2 · p. 5
  • Token weights combine rollout-policy entropy with inverse local teacher-score variation, favoring uncertain predictions where teacher feedback is locally consistent.An adaptive floor limits inverse weights at small variations; the quantities are stop-gradient, and constant uncertainty and variation yield uniform weights.Source: §3.2 · p. 5
  • Smoothness reweighting is bounded, but uncertainty can still concentrate credit.The allocation objective’s optimality does not establish gradient or downstream effects; those are assessed separately.Source: §3.2 · p. 5

3.3 SELECTIVE TOKEN ANCHORING AND OPTIMIZATION

  • A selective KL penalty keeps unaligned sampled-token predictions close to the initial student policy, while aligned positions remain free to absorb teacher credit.The penalty vanishes when sampled-token probabilities match; for heterogeneous tokenizers, the regularizer is the anchoring loss.Source: §3.3 · p. 6
  • CAST scores and chunks n independent student responses per prompt before averaging their token objectives; shared-tokenizer training applies reference KL regularization across response positions.Unaligned positions have zero distillation advantage.Source: §3.3 · p. 6

4 EXPERIMENTS

Across cross-tokenizer and shared-tokenizer evaluations, CAST achieves higher six-benchmark averages than the strongest evaluated baselines. Component controls and training analyses examine token-credit placement, anchoring, smoothness, and training consistency.

4.1 EXPERIMENTAL SETUP

  • Cross-tokenizer runs transfer Qwen3-8B to Mistral-7B and Llama-3.1-8B, while shared-tokenizer runs use Qwen teachers, Qwen students, and two OpenThoughts prompt sets.The protocol refers hyperparameters, compute budgets, checkpoint selection, and six-benchmark pass@1 evaluation details to Appendix C.Source: §4.1 · p. 6

4.2 MAIN RESULTS

  • 47.22 and 43.87 Avg-6 for Mistral and Llama, respectively, exceed ULD by 2.45 and 1.32 points.CAST leads ULD on five benchmarks for Mistral and four for Llama.Source: §4.2 · pp. 6–7, Table 1 · p. 7
  • 48.68 Avg-6 vs 46.58 for DistiLLM; punctuation chunks beat fixed K=6 by 1.56 points, though fixed chunks score higher on LCB (35.1 vs 34.6).Shuffled credit scores 3.20 points below CAST.Source: Table 2 · p. 7, §4.2 · p. 7, Table 3 · p. 7

4.3 DISCUSSIONS

  • 47.22 Avg-6 with selective anchoring vs 44.30 with global anchoring and 44.00 without anchoring; global anchoring also has higher entropy and lower repetition than CAST.The authors describe this pattern as consistent with global anchoring overconstraining aligned positions.Source: Table 4 · p. 8, §4.3 · p. 9, Figure 5 · p. 9
  • Removing structured token credit lowers Avg-6 from 47.22 to 45.07; shuffled weights score 44.17, with CAST improving five of six benchmarks over that control.LCB is slightly lower with CAST than with shuffled weights: 27.7 versus 27.9.Source: §4.3 · p. 8, Table 4 · p. 8
  • Across three seeds, CAST reaches 47.01 ± 0.28 Avg-6 on Mistral and 43.72 ± 0.16 on Llama, with all CAST scores above comparator highs.The Mistral and Llama seed summaries are reported in Tables 5 and 6.Source: Table 5 · p. 8, Table 6 · p. 8, §4.3 · p. 9
  • At SFT initialization, CAST raises gradient SNR from 0.111 under uniform credit to 0.256; training traces also show higher teacher endorsement and entropy, with lower repetition than DPCA.Figure 4 compares credit-rule gradient consistency; Figure 5 plots Mistral training dynamics.Source: Figure 4 · p. 8, §4.3 · p. 9, Figure 5 · p. 9
  • With anchoring and two rollouts fixed, smoothness raises LCB from 17.0 to 27.7 and Avg-6 from 45.60 to 47.22, while the other five benchmarks change by −0.20 points on average.Seventy-nine percent of redistribution occurs at teacher-variation positions.Source: Table 4 · p. 8, §4.3 · p. 9

5 CONCLUSION

CAST combines uncertainty–smoothness credit allocation for aligned chunks with selective anchoring for unaligned positions, and improves benchmark averages over the strongest evaluated baseline.

5 CONCLUSION

  • For token-level shift-or-stay cross-tokenizer OPD, CAST improves the six-benchmark average over the strongest evaluated baseline by 1.32–2.45 points for heterogeneous students and 2.10 points for Qwen3-1.7B.Source: §5 · p. 9
  • CAST redistributes aligned chunk advantages using uncertainty–smoothness credit and anchors unaligned predictions near the initial policy.The advantage-preserving rule has an exact KL-regularized characterization and extends to shared-tokenizer punctuation-defined or fixed-length chunks.Source: §5 · p. 9
  • Controls support token-specific credit placement, more consistent LM-head gradients, gains across Mistral and Llama seeds, and fewer observed collapses.Smoothness recovers code performance and broadens effective credit support; selective anchoring improves the two-rollout configuration.Source: §5 · pp. 9–10

LIMITATIONS

The evaluation is limited in model-family and training-data coverage, and training compute varies across methods.

LIMITATIONS

  • Evaluation covers a limited set of model families and mainly reasoning-oriented training data, so domain gains may vary with student capacity or prompt composition.Broader model and data coverage remains future work.Source: LIMITATIONS · p. 10
  • Training compute varies across methods, constraining matched-compute comparisons.Matched-compute evaluation remains future work.Source: LIMITATIONS · p. 10
  • Repeated-seed performance covers four Mistral and two Llama configurations, while direct tests of credit allocation’s effects on full-model optimization and training collapse remain future work.Source: LIMITATIONS · p. 10

A ADDITIONAL METHOD DETAILS

The method constructs credit chunks from aligned token groups under heterogeneous tokenization or from punctuation or fixed-length partitions under a shared tokenizer.

A.1 CHUNK CONSTRUCTION

  • Under heterogeneous tokenization, decoded-text alignment matches contiguous teacher and student token groups after Unicode and tokenizer-specific whitespace normalization.A greedy two-pointer scan extends the shorter side until the decoded strings match; excessively large or unreliable matches are discarded, leaving aligned positions A and unaligned positions U.Source: §A.1 · p. 13
  • After teacher-score projection, contiguous aligned student tokens are regrouped into fixed-length K-token credit chunks, using K = 6 unless stated otherwise.A shorter residual chunk is allowed, and credit-chunk boundaries need not match alignment-chunk boundaries.Source: §A.1 · p. 13
  • With a shared tokenizer, response tokens have identity correspondence and punctuation-defined chunks split after commas, periods, exclamation marks, and question marks.The fixed-length variant instead partitions responses into consecutive K = 6 token chunks, with a shorter final chunk when needed.Source: §A.1 · p. 13

B.1 TEACHER SCALE AND TRAINING DISTRIBUTION

The supplementary analyses show that sampling, anchoring, teacher scale, prompt distribution, and token-level credit allocation have distinct effects across the reported settings.

B.1 TEACHER SCALE AND TRAINING DISTRIBUTION

  • For Qwen3-1.7B, the 8B teacher reaches 48.68 versus 48.05 for the 30B-A3B teacher; for Qwen3-4B, the 30B-A3B teacher reaches 64.58 versus 62.33 with the 8B teacher.With the 4B student, v2 raises averages to 64.05 with the 8B teacher and 64.93 with the 30B-A3B teacher; the larger-teacher gains concentrate on the three AIME benchmarks.Source: Table 7 · p. 13, §B.1 · p. 13, Table 8 · p. 14

B.2 ADDITIONAL SAMPLING ANALYSIS

  • At n=2, selective anchoring raises Avg-6 by 3.22 points versus no anchoring, and selective anchoring beats global anchoring by 2.92 points; n=4 lowers Avg-6 to 45.05.With anchoring fixed, n=2 raises Avg-6 from 46.07 to 47.22 versus n=1, while n=4 reduces LCB by 7.8 points relative to n=2.Source: §B.2 · p. 15

B.3 TOKEN-LEVEL CREDIT ANALYSIS

  • Content words and space tokens receive 1.53× and 1.58× their uniform credit shares, while punctuation, sub-word continuations, and digit/operator tokens receive 0.79×, 0.54×, and 0.30×.In the illustrated six-token chunk, CAST assigns 45.8%, 28.6%, and 25.6% to the three high-entropy positions, while near-deterministic digit and operator tokens receive effectively no share.Source: Figure 6 · p. 15, §B.3 · p. 15

B.4 TRAINING DYNAMICS: PROTOCOL AND ANALYSIS

Training diagnostics compare CAST with alternative credit and anchoring choices, then test whether its smoothness signal reflects teacher structure rather than projection geometry. CAST shows more consistent sampled gradients, while smoothness attribution concentrates mainly on teacher-variation positions.

B.4 TRAINING DYNAMICS: PROTOCOL AND ANALYSIS

  • Over the final 50 updates, DPCA’s teacher-endorsed fraction was 0.19 versus 0.40 for CAST, while response entropy was 0.22 versus 0.43, gradient norm 1.86 versus 0.55, and repetition 0.51 versus 0.44.Global anchoring had a higher six-benchmark average than CAST’s diagnostic comparison? No: it scored 44.30 versus 47.22 despite favorable entropy and repetition statistics.Source: §B.4 · p. 16

B.5 GRADIENT DIAGNOSTIC PROTOCOL

  • CAST had the highest SNR and training–held-out mean-gradient cosine at both fixed checkpoints; relative to uniform credit, SNR rose from 0.111 to 0.256 and from 0.093 to 0.182.At SFT initialization, mean held-out rollout alignment rose from 0.279 to 0.438; these are sampled head-gradient diagnostics, not downstream performance results.Source: §B.5 · p. 16

B.6 SMOOTHNESS ATTRIBUTION

  • 79% of smoothness-only reweighting fell on teacher-variation positions, versus 21% at projection-flat interiors; shared-tokenizer analyses also found comparable local teacher structure.Multi-student alignment chunks cover 19% of tokens, and the floor is active in 36% of credit chunks; shared-tokenizer median variation is 0.145 versus 0.094 cross-tokenizer, with variation above the floor in 76% versus 64% of chunks.Source: §B.6 · pp. 16–17, Figure 7 · p. 17

C TRAINING AND EVALUATION DETAILS

The appendix specifies training compute, evaluation, and checkpoint selection for the reported experiments. Configurations share a six-benchmark pass@1 protocol, but checkpoint selection differs between point estimates and seed summaries.

C TRAINING AND EVALUATION DETAILS

  • Training uses a prompt batch size of 64, learning rate 10^-6, and maximum prompt and response lengths of 1,024 and 16,384 tokens; heterogeneous pairs use K=6 and CAST defaults to n=2 rollouts.Runs use eight H100 GPUs and 250 optimizer steps; the teacher scores student-generated responses.Source: C TRAINING AND EVALUATION DETAILS · p. 17, Table 9 · p. 18
  • Default CAST generates approximately 355M response tokens and 369M actor tokens over training, with approximately 30 hours of wall-clock time per run.The n=1 configurations use approximately half the generation and teacher-scoring budget; n=4 uses twice the n=2 budget.Source: C TRAINING AND EVALUATION DETAILS · pp. 17–18
  • Point estimates select the highest Avg-6 checkpoint among evaluated seed–checkpoint pairs, whereas Table 5 averages each seed’s best checkpoint; benchmark-based selection can make reported averages optimistic.For Llama, Table 1 uses the highest Avg-6 among three per-seed checkpoints, while Table 6 summarizes all three.Source: C TRAINING AND EVALUATION DETAILS · p. 18
  • Evaluation uses pass@1 across six benchmarks, with 32 samples per AIME problem and four per problem elsewhere; three-seed means and standard deviations are reported for selected configurations.Other comparisons use single runs; each AIME problem represents 3.3, 3.3, or 6.7 percentage points for 2024, 2025, or 2026, respectively.Source: C TRAINING AND EVALUATION DETAILS · p. 18

D CASE STUDY

The case studies contrast answer completion and unfinished responses across four selected 2024 AIME I problems. CAST completes three shown calculations, while the lottery case is solved by both CAST and ULD.

Case study 1: Grid paths and repeated direction sequences

  • On the grid-path problem, CAST completes the count as 147 + 147 = 294, while DPCA and ULD repeat alternating directions without returning a count.SFT continues reconsidering individual moves.Source: D CASE STUDY · p. 19

Case study 2: Enclosing sphere and unfinished calculations

  • On the enclosing-sphere problem, CAST reaches r^2 = 657/64 and returns p + q = 721; the other response tails shown do not provide an extracted final answer.The analysis contrasts CAST’s completed answer with intermediate algebra in the baseline tails.Source: D CASE STUDY · p. 19

Case study 3: Lottery probability and shared success

  • On the lottery problem, CAST and ULD both return the correct sum 116, while SFT and DPCA provide no extracted answer.Both successful responses give the conditional probability as 1/115.Source: D CASE STUDY · p. 20

Case study 4: Logarithmic system and repeated reconsideration

  • On the logarithmic-system problem, CAST and ULD conclude with 25, while SFT repeatedly revisits an invalid value and DPCA continues substitutions without returning xy.The analysis describes the baseline tails as unresolved calculations.Source: D CASE STUDY · p. 20

E PERFORMANCE ACROSS TRAINING SEEDS

Across evaluated seeds, CAST’s selected Avg-6 exceeded the selected controls for Mistral, and its highest reported Llama checkpoint exceeded ULD’s. The appendix also reports credit concentration and collapse counts for diagnostic runs.

E PERFORMANCE ACROSS TRAINING SEEDS

Related papers

Loading 2610.13257v1…