Training-Trajectory-Aware Token Selection
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 bottle-neck, 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. T3S 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.
Lay Summary
Teaching smaller artificial intelligence models to reason by copying larger models often triggers a puzzling phenomenon: the student's performance crashes sharply before slowly recovering. By analyzing the lowest point of this crash—the "bottleneck"—we discovered that the learning updates made before this point are completely unnecessary and even harmful. Digging into the data, we found this pattern is driven by a fundamental conflict between two different types of learning units, or "tokens," which cannot be learned simultaneously. One type is quickly memorized but is less relevant to actual reasoning. We discovered that learning these early tokens actively suppresses the acquisition of the other, more crucial reasoning tokens, causing the initial performance drop. To solve this, our method, Training-Trajectory-Aware Token Selection (T3S), identifies and temporarily hides these disruptive tokens around the bottleneck. This clears the optimization path, allowing the AI to focus entirely on genuine reasoning. By removing these roadblocks, T3S makes training incredibly efficient, empowering smaller models trained on just a few hundred examples to surpass their massive teachers.