Training-Trajectory-Aware Token Selection

TL;DR

This paper introduces T3S, a token selection method that improves continual distillation by addressing token-level confidence bifurcation, leading to significant performance gains in reasoning benchmarks.

Contribution

The paper proposes a novel training-trajectory-aware token selection technique that reconstructs the training objective at the token level to enhance distillation effectiveness.

Findings

T3S improves AR and dLLM performance across multiple models.

Qwen3-8B surpasses DeepSeek-R1 on reasoning benchmarks with few examples.

T3-trained LLaDA-2.0-Mini exceeds its AR baseline, achieving state-of-the-art results.

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. 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.