Look Before You Select: Rethinking Vocabulary Sparsification in On-Policy Distillation

1 Introduction

• Full-vocabulary OPD captures important corrections beyond the student's most likely tokens, but backpropagating through the full vocabulary is memory-intensive, especially for long reasoning sequences.This creates an apparent trade-off between preserving full-vocabulary correction and reducing backward memory.Source: §1 · p. 1
• Sampled Token can depend strongly on the sampled token, while TopK excludes tail tokens and changes the correction through restricted-support conditioning.In the illustrated example, the sampled action incorrectly promotes “sum,” while TopK omits direct correction to “overlap” and “intersection.”Source: §1 · p. 1, Figure 1 · p. 2
• SparseOPD constructs the full-vocabulary correction before selecting logits by correction magnitude, then redistributes discarded positive and negative correction mass among retained corrections of the same sign.The resulting update backpropagates only through selected token logits.Source: §1 · p. 2
• Across six task–scale settings, SparseOPD outperforms Sampled Token and TopK and matches or exceeds Full Vocabulary.Settings span mathematical reasoning, chemistry question answering, and multimodal mathematical reasoning.Source: §1 · p. 2
• SparseOPD's complementary profiling and gradient measurements show substantially lower backward memory and close agreement with the Full Vocabulary gradient.These measurements are reported for the 4B mathematical reasoning setting.Source: §1 · p. 2

2.1 Preliminaries: Full-Vocabulary OPD

• Full-vocabulary OPD compares complete teacher and student distributions at each student-generated prefix, producing deterministic token-specific corrections whose positive and negative coefficients sum to zero.Direct differentiation retains vocabulary-sized backward intermediates, with memory cost growing with vocabulary size and the number of valid positions.Source: §2.1 · p. 3

2.2 What Existing Sparsification Changes

• Sampled-token OPD recovers the full correction only in expectation at a fixed prefix, while each realization depends on one sampled ratio and introduces sampling variance.TopK instead restricts the objective to the student's highest-probability tokens; tail coordinates receive no direct correction, and retained coefficients come from a conditional objective rather than truncating the full correction.Source: §2.2 · pp. 3–4

2.3 Full-Vocabulary Correction with Sparse Execution

• The proposed target is to approximate the full-vocabulary correction in coefficient space while restricting output-head differentiation to selected entries under an average support budget.This reframes sparsification as correction compression: construct the full correction first, then choose which token logits to retain.Source: §2.3 · p. 4, §2.2 · p. 4

3.1 Batch-global Budget Allocation

• Support is allocated across positions in proportion to correction mass, rather than uniformly, with capacity limits respected.The shared budget is distributed according to each position’s correction strength.Source: §3.1 · p. 5

3.2 Signed Selection and Compensation

• Top-C retains the largest-magnitude coefficients within each sign, while centered compensation redistributes discarded mass to preserve signed totals and zero-sum structure.Selection and compensation are complementary: support retains important coordinates and spreads residual mass across retained entries.Source: §3.2 · p. 6

3.3 Sparse Backpropagation

• After a no-gradient planning pass fixes the sparse support and coefficients, backpropagation reconstructs logits only for selected output-head rows.The full-vocabulary statistics are discarded once the sparse plan is fixed.Source: §3.3 · pp. 6–7

4.1 Experimental Setup

• The experiments use mathematics, chemistry, and multimodal reasoning, with reported evaluation spanning the listed mathematics, chemistry, and multimodal benchmarks.Mathematics and chemistry report Mean@8 accuracy; multimodal reasoning uses single-generation accuracy.Source: §4.1 · p. 7, Table 2 · p. 8

4.2 Main Results

• SparseOPD matches or exceeds Full Vocabulary in five of six task–scale settings, including Qwen3-4B mathematics (43.2 vs 41.6), but trails in 1.7B chemistry (35.8 vs 37.0).Both multimodal student scales exceed the dense reference.Source: §4.2 · p. 8

4.3 Ablation Study

• Replacing correction-magnitude ranking with student-probability ranking lowers average accuracy by 2.0 points; removing compensation lowers it by 1.4, and fixed per-position allocation by 0.9.The ablations hold the remaining components and total support budget fixed on Qwen3-4B mathematical reasoning; compensation preserves aggregate signed correction mass, while global allocation redistributes capacity toward stronger corrections.Source: §4.3 · p. 9, Table 3 · p. 9

4.4 Analysis and Discussion

• SparseOPD is closest to Full Vocabulary on all four gradient-fidelity metrics, with relative ℓ2 error about one-third of TopK’s (0.15 vs 0.45).Cosine similarity and norm ratio are near one, and SparseOPD has the highest sign agreement.Source: §4.4 · p. 10, Table 4 · p. 10
• At 8K sequence length, SparseOPD reduces backward memory by 70.5%, with larger relative savings at longer measured sequences.Full-vocabulary statistics are computed without an autograd graph and discarded after planning; backward time improves more modestly because the Transformer backbone still requires full backpropagation and planning adds overhead.Source: §4.4 · p. 10, Table 5 · p. 10

• SparseOPD differs from sampled or restricted-support approaches by forming the full-vocabulary correction before compressing it for sparse backpropagation.The cited prior approaches reduce cost through vocabulary sampling or support restriction (Anshumann et al., 2025).Source: §5 · p. 10

6 Conclusions

• SparseOPD selects entries by correction magnitude and compensates discarded correction mass to compress updates derived from the complete teacher and student distributions.The conclusion characterizes sparse execution as preserving much of the training benefit of full-vocabulary supervision across mathematics, chemistry QA, and multimodal reasoning.Source: §6 · p. 10
• The paper reports close agreement with the full-vocabulary gradient and substantially lower peak backward memory in the measured setting.The overall conclusion is that SparseOPD preserves full-vocabulary correction while reducing backpropagation memory.Source: §6 · pp. 10–11

Reproducibility statement

• The reproducibility statement locates datasets, model configurations, metrics, and benchmark results in Section 4, with training and efficiency procedures in Appendix B.It also points to Sections 2 and 3 for SparseOPD and comparison-baseline specifications.Source: Reproducibility statement · p. 11

A More Related Work

• Related work reduces vocabulary-distillation cost through sampling or restricted support, whereas SparseOPD forms the deterministic full correction before compressing it for sparse differentiation.The cited research addresses supervision on on-policy trajectories or vocabulary approximation; SparseOPD focuses on computational realization of token-level full-vocabulary supervision.Source: A More Related Work · p. 17

• The section specifies data splits, student adaptation, optimization, and evaluation protocols for the main accuracy experiments.Source: B Additional Experimental Details · p. 17
• Efficiency measurement is handled in a separate study.Source: B Additional Experimental Details · p. 17
• The section’s scope is experimental detail rather than results: it describes protocols used for accuracy experiments and efficiency measurement.Source: B Additional Experimental Details · p. 17

B.1 Training Data and Model Adaptation

• Training uses DAPO14k for mathematics, a fixed seed-42 90%/10% train/held-out split of Chemistry L3 in SciKnowEval V2, and the original MMR1-Math-RL-Data-v0 release for multimodal tasks.The multimodal validation split is excluded from training.Source: §B.1 · p. 17
• Each method starts from the same student checkpoint within a model–task pair and trains LoRA parameters; multimodal vision towers remain frozen and bias parameters are not trained.Shared settings and task-specific settings are listed in Tables 6 and 7.Source: §B.1 · pp. 17–18, Table 6 · p. 18, Table 7 · p. 18
• All tasks use the student’s native Qwen chat template with thinking mode disabled, with task-specific prompt instructions for mathematics and chemistry.Mathematics prompts request step-by-step solutions and boxed answers; chemistry uses the same prompt construction for training and evaluation.Source: §B.1 · p. 18, Table 8 · p. 19

B.2 Optimization and Distillation

• For every method, the student generates one on-policy completion and the frozen teacher scores that completion by teacher forcing, while the methods use distinct distillation objectives.Full Vocabulary uses reverse KL over the full vocabulary; TopK uses the student’s Top-20 support with within-support renormalization; Sampled Token uses a sampled-action log-probability difference without a baseline. SparseOPD uses a batch-global budget of 20N, with at most 32 retained coordinates per sign at each position.Source: §B.2 · pp. 18–19

B.3 Evaluation Protocol

• Mathematics and chemistry use Mean@8, averaging binary correctness across eight independently generated answers per question and then across questions.Answer extraction and grading are consistent across methods within each benchmark.Source: §B.3 · p. 19
• Multimodal benchmarks follow LMMs-Eval, using one generation per question and exact-match answer grading.Each evaluation split is evaluated in full.Source: §B.3 · p. 19

B.4 Efficiency Measurements

• Profiling uses full-parameter updates on Qwen3-4B-Base with Qwen3-4B as the frozen teacher, and each method processes the same data.This profiling study is separate from the LoRA accuracy experiments.Source: §B.4 · p. 19, Table 9 · p. 19
• Forward time covers loss-producing forward calls, while backward time includes corresponding autograd calls, recomputation, and communication but excludes the optimizer.The reported stage time is the larger of the two GPU totals.Source: §B.4 · p. 19
• The efficiency table reports total peak Torch allocations by stage before baseline subtraction, while the main memory result reports extra memory.Step time covers the complete update window, including vocabulary-sparsification planning.Source: Table 10 · p. 20
• 70.5% lower peak extra backward memory than Full Vocabulary at 8K sequence length, with a 14.5% increase in complete-update time.The reported time outside forward and backward includes input preparation and optimizer work, so it cannot all be attributed to token selection.Source: §B.4 · p. 20
• SparseOPD’s backward-time reduction is limited because selected hidden and weight rows must still be materialized, dense parameter gradients accumulated, and the full Transformer backbone backpropagated through.The time outside forward and backward also includes input preparation and optimizer work.Source: §B.4 · p. 20

B.5 Parameter Sensitivity Analysis

• At the default budget, SparseOPD achieves strong average accuracy across training seeds with a small standard deviation on 1.7B mathematical reasoning.The authors say this consistency suggests the main result is not driven by a particularly favorable seed.Source: §B.5 · p. 20, Table 11 · p. 21
• Across tested average support budgets K ∈ {16, 20, 24}, larger budgets improve average accuracy, with K = 24 achieving the best result.MATH500 and Minerva remain relatively stable, while AMC23 benefits more from increased support; overall performance remains robust.Source: §B.5 · p. 21

B.6.1 Mathematics

• The mathematics subsection includes example problems on polynomial perfect squares, prime divisibility, expected area, common roots, and integer constraints.These are presented as example questions rather than sensitivity results.Source: §B.6.1 · p. 22

B.6.2 Chemistry

• The chemistry examples include solution pH, identifying a compound from spectral clues, and counting chemically distinct proton signals after a reaction sequence.The continued examples ask about reaction-rate changes with pH and interpreting duplicated alpha-proton signals.Source: §B.6.2 · p. 23, Chemistry (continued) · p. 24

Chemistry (continued)

• Further chemistry examples describe a peptidic compound with duplicated alpha-proton peaks and matching-mass LC-MS peaks, with the listed answer identifying a diastereoisomer mixture.The passage gives the observations and answer but does not report an evaluation result.Source: Chemistry (continued) · p. 24

B.6.3 Multimodal Geometry

• The multimodal geometry examples include a square-and-circle shaded-region probability question and an area-ratio problem involving a quadrilateral intersecting a square.The subsection also includes a route-distance question.Source: §B.6.3 · p. 25

Multimodal Geometry (continued)

• Continued multimodal geometry examples ask for a rectangle side length from its diagonal and an angle, and include a problem involving opposite rays and an angle bisector.Source: Multimodal Geometry (continued) · p. 26