Training-Trajectory-Aware Token Selection

TL;DR

This work identifies Imitation Shock as a universal failure mode in continual reasoning distillation, where performance collapses even as training loss falls, due to a bottleneck at which two token groups—Imitation-Anchor Tokens and yet-to-learn tokens—exhibit antagonistic training dynamics. It introduces Training-Trajectory-Aware Token Selection (T3S), which uses trajectory signals around the Imitation Bottleneck to mask anchor tokens in AR distillation and to target yet-to-learn tokens in diffusion-style LMs, thereby clearing the optimization path and accelerating reasoning transfer. The approach yields consistent gains across autoregressive and diffusion-language-model settings, reduces inference tokens by 15–25%, and improves robustness to out-of-distribution reasoning and forgetting. The results demonstrate that token-level, trajectory-informed targeting can unlock efficient, scalable reasoning transfer from teacher models with hundreds of examples, and that anchor suppression is a critical lever for stable continual distillation.

Abstract

Efficient distillation is a key pathway for converting expensive reasoning capability into deployable efficiency, yet in the frontier regime where the student already has strong reasoning ability, naive continual distillation often yields limited gains or even degradation. We observe a characteristic training phenomenon: even as loss decreases monotonically, all performance metrics can drop sharply at almost the same bottleneck, before gradually recovering. We further uncover a token-level mechanism: confidence bifurcates into steadily increasing Imitation-Anchor Tokens that quickly anchor optimization and other yet-to-learn tokens whose confidence is suppressed until after the bottleneck. And the characteristic that these two types of tokens cannot coexist is the root cause of the failure in continual distillation. To this end, we propose Training-Trajectory-Aware Token Selection (T3S) to reconstruct the training objective at the token level, clearing the optimization path for yet-to-learn tokens. T3 yields consistent gains in both AR and dLLM settings: with only hundreds of examples, Qwen3-8B surpasses DeepSeek-R1 on competitive reasoning benchmarks, Qwen3-32B approaches Qwen3-235B, and T3-trained LLaDA-2.0-Mini exceeds its AR baseline, achieving state-of-the-art performance among all of 16B-scale no-think models.

Figures (12)

• Figure 1: Imitation Shock under DeepSeek-R1 distillation on BOBA-200. Although training loss decreases monotonically (a), training-set answer accuracy and multiple benchmarks (AIME24/25, MMLU-Pro) drop sharply to a shared minimum stage and then recover, revealing an Imitation Shock.
• Figure 2: Imitation Shock is universal across settings. We observe the same "crash then recover" trajectory across (a) different teachers, (b) different datasets, (c) larger-scale datasets, (d) different student backbones, and (e) different training domains. See Appendix \ref{['sec:appendix-universal-shock']} for detailed setups and metrics.
• Figure 3: Word clouds at the Imitation Bottleneck (identified by training-set accuracy). Token sizes are proportional to the magnitude of confidence change relative to the base model.
• Figure 4: Anchor tokens and delayed learning dynamics on BOBA-200. (a) Under the full distillation objective, anchor tokens increase monotonically, while the remaining tokens crash early and recover only after anchors stabilize. (b) When training only the remaining tokens (excluding anchor tokens from the loss), their confidence rises monotonically, reversing the "crash-then-recover" pattern and indicating suppressive coupling from anchor-token learning.
• Figure 5: Checkpoint-wise one-step intervention on Imitation-Anchor Tokens. Panel (a) reports $\Delta \mathcal{L}_{\text{other}}$ after one gradient step that optimizes only anchor tokens at each checkpoint. Panel (b) shows the corresponding anchor-token loss $\mathcal{L}_{\text{anchor}}$ at that checkpoint.