T3S — Training-Trajectory-Aware Token Selection for Continual Reasoning Distillation

Paper. Zhanming Shen, Jiaqi Hu, Zeyu Qin, Hao Chen, Wentao Ye, Zenan Huang, Yihong Zhuang, Guoshan Lu, Junlin Zhou, Junbo Zhao. Training-Trajectory-Aware Token Selection. ICML 2026. Zhejiang University · Ant Group (Inclusion AI). [arXiv]

0. The Picture in One Paragraph

When you fine-tune a strong student model (say Qwen3-8B) on chain-of-thought traces from a frontier teacher (DeepSeek-R1, QwQ), something counterintuitive happens: training loss falls monotonically, yet all evaluation metrics simultaneously collapse to a shared nadir before slowly recovering. The paper calls this the Imitation Bottleneck, caused by Imitation Shock. The root mechanism is a token-level bifurcation: surface pattern tokens (formatting, common phrases, boilerplate) anchor the gradient early and starve the harder reasoning tokens of learning signal. T3S proposes a fix that is simple to state but nontrivial to discover: mask the already-learned anchor tokens from the autoregressive loss, using a trajectory-aware confidence-shift signal to identify them. With fewer than 1,000 training examples, the resulting Qwen3-8B exceeds its DeepSeek-R1 teacher on both AIME24 (80.63% vs. 79.8%) and AIME25 (73.96% vs. 70.0%), and the same principle extends to diffusion language models (LLaDA-2.0-Mini, state-of-the-art among 16B no-think models).

1. The Problem — Imitation Shock

The standard recipe for continual reasoning distillation: sample chains-of-thought from a frontier teacher, supervised-fine-tune the student on those chains, monitor the benchmark.

The problem: on almost every setting the authors tested — across teacher models (DeepSeek-R1, QwQ), student models, dataset scales, and domains — there is a reproducible failure mode. Training loss decreases without exception. But all downstream metrics (AIME24, AIME25, MMLU-Pro, training-set accuracy) drop sharply at the same training step, hit a shared low point (the Imitation Bottleneck), and only slowly recover.

This pattern is not a random glitch. It is the same across runs. The question is why a decreasing loss produces degrading benchmarks, and why the degradation and recovery follow such a consistent shape.

2. The Discovery — Token Confidence Bifurcation

T3S isolates the failure at the token level. Define per-token confidence as the log-probability the current model assigns to the ground-truth token:

\[c_t(\theta; x, y) = \log p_\theta(y_t \mid y_{<t}, x)\]

The confidence shift between a training checkpoint $\theta_b$ and the initial model $\theta_0$ is:

\[\Delta c_t(x, y) = c_t(\theta_b;\, x, y) - c_t(\theta_0;\, x, y)\]

Tracking $\Delta c_t$ across training reveals a clean bifurcation:

Token class	$\Delta c$ trajectory	What they encode
Imitation-Anchor Tokens	Rises steeply, early	Formatting, common phrases, surface patterns
Yet-to-learn Tokens	Stays low or drops	Deep reasoning steps, problem-specific logic

Anchor tokens dominate the gradient early in training. Because the autoregressive loss treats all tokens equally, the optimizer devotes most of its capacity to tokens that are already easy — suppressing the gradient signal for the harder reasoning tokens until after the bottleneck is passed.

The critical empirical finding: static signals cannot separate these two populations. Using initial confidence alone (one time point) or local gradient geometry (one step) fails to identify which tokens are anchors. Only the change in confidence over a training window — the trajectory — reveals the bifurcation. This is what makes T3S methodologically distinct from all prior token-selection methods.

3. The Method — T3S

T3S reconstructs the training objective at the token level:

Compute $\Delta c_t$ for every token in the batch by comparing the current checkpoint to the initial model.
Identify Imitation-Anchor Tokens: those whose $\Delta c_t$ exceeds a threshold, meaning the model has already learned them and they dominate the gradient.
Mask anchor tokens from the autoregressive loss — they contribute zero gradient for this step.
Compute the loss only over yet-to-learn tokens, giving them the full gradient signal without anchor competition.

The intervention is at the loss layer only. No change to the model architecture, the optimizer, or the sampling procedure for teacher traces.

3.1 Why trajectory, not static?

The comparison to prior methods is direct:

Method	Signal used	Signal type
Rho-1	Excess loss vs. reference model	Static (single time point)
TokenCleaning / ssToken	Token-level difficulty heuristics	Static
T3S	Confidence shift $\Delta c_t$ over training	Dynamic trajectory

Ablation confirms that substituting the trajectory signal with either static alternative loses most of the gain. The two token classes only become distinguishable because of how their confidence evolves differently over time.

3.2 Extension to Diffusion Language Models

The same principle transfers to masked diffusion training. For LLaDA-2.0-Mini (a diffusion LLM), anchor tokens are those whose mask-prediction confidence has shifted most across training. T3S masks those from the diffusion loss. The paper treats this as a direct validation that the trajectory-awareness principle is architecture-agnostic.

4. Results

4.1 Super-teacher performance on AIME

With 200–1,000 training examples (far below typical fine-tuning regimes):

Model	Benchmark	Standard SFT	T3S	Teacher (DeepSeek-R1)
Qwen3-8B	AIME24	degrades	80.63%	79.8%
Qwen3-8B	AIME25	degrades	73.96%	70.0%

Both T3S numbers are averages over 16 runs. The student exceeds the teacher — an outcome that standard SFT cannot reach even with more data because Imitation Shock degrades the benchmark score before recovery.

Qwen3-32B with T3S approaches the performance of Qwen3-235B, the model with roughly 7× more parameters.

4.2 Consistency across settings

T3S achieves the best performance across all four combinations of training split (BOBA-200, S1K-200) and teacher model (DeepSeek-R1, QwQ), outperforming standard SFT and RRT (the parameter-level counterpart that avoids Imitation Shock by rolling back the optimizer to a pre-bottleneck state).

4.3 OOD and forgetting

On out-of-distribution reasoning benchmarks (MATH500, TheoremQA, AMC) and knowledge benchmarks (MMLU-Pro):

T3S consistently improves OOD reasoning.
T3S mitigates catastrophic forgetting of non-reasoning capabilities, relative to standard SFT.

The interpretation: by reducing anchor-token dominance, T3S prevents the model from overfitting to the teacher’s surface style, leaving non-reasoning parameters less disrupted.

4.4 dLLM results

On LLaDA-2.0-Mini:

T3S dramatically improves reasoning accuracy.
Outperforms autoregressive fine-tuning on the same model class.
Achieves state-of-the-art among 16B-scale no-think models.
Generates 15–25% fewer tokens at inference — because the model no longer overfits to surface-pattern verbosity from the teacher’s long chains.

Rho-1 and static token selectors

Rho-1 (arXiv:2404.07965) uses a reference model trained on clean data to score excess per-token loss, selecting only high-loss tokens for training. It works well on LLaMA-family models but has been shown to perform poorly — sometimes worse than full SFT — on Qwen-family models. T3S outperforms Rho-1 by using the shift in confidence rather than its absolute level, which is a more reliable signal when the student already has strong priors.

RRT (parameter-level intervention)

RRT addresses Imitation Shock at the optimizer level: it detects the onset of the bottleneck and rolls back to the pre-bottleneck checkpoint, then continues. T3S addresses the same phenomenon at the token level without any rollback or checkpoint management. The two methods are complementary in principle; in practice T3S outperforms RRT consistently.

AlphaToken / TIP

These assign token-level credit in RL-based post-training from reward model feedback. T3S operates in imitation distillation (no reward model), using the training trajectory itself as the credit signal.

6. Why It Matters

Three reasons:

It names and explains a pervasive failure mode. Imitation Shock has surely been encountered by anyone doing continual distillation — most practitioners attribute it to bad hyperparameters or insufficient data. T3S identifies the mechanistic cause at the token level and shows it is structural, not accidental.
Super-teacher performance from tiny data. The fact that a student can exceed its teacher on AIME24/25 with under 1,000 training examples — by masking some of the teacher’s output during training — reframes how we should think about data efficiency in distillation. The bottleneck is not how much teacher data you have; it is whether you are learning from the right tokens within that data.
Architecture-agnostic principle. Applying the same confidence-shift idea to diffusion LMs (LLaDA-2.0-Mini) with SOTA results suggests the token bifurcation phenomenon is not specific to autoregressive models. The trajectory signal may be a broadly applicable tool for any iterative training process over token-level objectives.

7. Limitations Worth Knowing

Token categories drift. The partition into anchor vs. yet-to-learn is not stable across training. A token that is “yet-to-learn” at step 100 may become an anchor by step 500. The current T3S formulation uses a fixed-threshold partition based on a trajectory snapshot, which does not adapt to this ongoing migration.
Static mask mismatch. Holding the mask fixed against a dynamic learning process introduces a mismatch that can create instability later in training.
No dynamic scheduling. The authors call out the natural future direction: treating token selection as an optimal control problem — designing time-dependent priority schedules that continuously re-rank tokens as training evolves. This would close the gap between the static T3S filter and the dynamic reality of learning.
Threshold sensitivity. The confidence-shift threshold that separates anchor from yet-to-learn tokens is a hyperparameter. The paper validates it across settings but does not provide a principled procedure for setting it without a validation set.

8. The Takeaway for a First Reader

If you remember three things:

Continual reasoning distillation suffers from Imitation Shock: all metrics collapse to a shared bottleneck despite decreasing loss, because surface-pattern Imitation-Anchor Tokens dominate the gradient and starve harder reasoning tokens. This is structural, not accidental.
T3S fixes it by masking anchor tokens from the AR loss, identified via a trajectory-aware confidence-shift signal $\Delta c_t$ — the change in per-token log-probability from the initial to the current checkpoint. Static signals fail; only the trajectory distinguishes the two token classes.
With under 1,000 training examples, T3S-distilled Qwen3-8B exceeds its DeepSeek-R1 teacher on AIME24 (80.63% vs. 79.8%) and AIME25 (73.96% vs. 70.0%). The same principle extends to diffusion LMs, achieving SOTA among 16B no-think models with 15–25% inference token reduction.

References

Shen, Z., Hu, J., Qin, Z., et al. (2026). Training-Trajectory-Aware Token Selection. ICML 2026. arXiv:2601.10348.
Liu, Z., et al. (2024). Rho-1: Not All Tokens Are What You Need. arXiv:2404.07965.
Related on this site: RLHF → RULER trend note — the RL post-training context this distillation work sits within.