Source-linked AI summary
Learning from Your Own Mistakes: Constructing Learnable Micro-Reflective Trajectories for Self-Distillation
Zhilin Huang, Hang Gao, Ziqiang Dong, Yuan Chen, Yifeng Luo, Chujun Qin, Jingyi Wang, Yang Yang, Guanjun Jiang
TL;DR
Self-distillation typically aligns token-level distributions but offers limited insight into the model’s specific reasoning errors. TAPO constructs micro-reflective correction trajectories from the model’s own mistakes and achieves consistent improvements over GRPO under the same training steps, strengthening both first-pass reasoning and error correction.
Problem
Existing self-distillation methods provide implicit token-level distributional supervision without insight into the underlying reasoning process or controllable corrective guidance.
Method
TAPO constructs micro-reflective trajectories that preserve the learner’s valid prefix, diagnose its first error, and insert corrected reasoning from the model’s own solutions.
Results
TAPO achieves consistent improvements over GRPO under the same number of training steps across AIME 2024, AIME 2025, and HMMT 2025.
Takeaways & Limitations
TAPO strengthens both first-pass reasoning and error-correction effectiveness, indicating that its training signal transfers beyond surface-level reflection behavior.
Abstract
from arXiv · showhide
Self-distillation improves reasoning in large language models by using the model's own rollouts as training signal, typically through implicit logit-level alignment that minimizes KL divergence toward a privileged target distribution. However, because this supervision is generated via uncontrolled sampling, it provides no diagnostic insight into the model's specific errors or corrective guidance for its individual failure patterns. Consequently, the model learns to imitate a privileged distribution rather than receiving fine-grained corrections that pinpoint where and why its reasoning fails. In this paper, we propose Trajectory-Augmented Policy Optimization (TAPO), which advances self-distillation from implicit distributional alignment to explicit trajectory construction. During RL training, the model produces both correct and incorrect rollouts to the same query, and TAPO leverages this contrastive structure to construct micro-reflective corrections, new training trajectories that retain the model's erroneous reasoning up to the point of failure, then insert a natural-language diagnosis and corrected reasoning guided by a correct reference from the same sampling group. Since each trajectory is anchored in the learner's own prefix and solutions, the corrective signal preserves the model's on-policy distribution to a greater extent than the position-wise alignment imposed by KL-based methods. To integrate these trajectories, TAPO introduces difficulty-aware candidate selection at the model's capability boundary and decoupled advantage estimation to prevent gradient contamination. Experiments on AIME 2024, AIME 2025, and HMMT 2025 show that TAPO achieves consistent improvements over GRPO under the same number of training steps. Further analysis demonstrates that TAPO strengthens both first-pass reasoning and error-correction effectiveness.
1 Introduction
TAPO advances self-distillation from token-level distributional alignment to explicit micro-reflective trajectory construction based on the model’s own errors and correct solutions. It combines difficulty-aware selection, decoupled advantage estimation, and OOD token suppression, improving reasoning and error correction while retaining standard single-pass inference.
- Motivation: Existing self-distillation minimizes token-level KL divergence to a privileged target distribution but provides no insight into the underlying reasoning process.The targets may come from an external teacher or the same model conditioned on a verified correct solution.
- TAPO: TAPO constructs micro-reflective training trajectories that preserve erroneous reasoning until failure, then add a diagnosis and corrected reasoning from the same sampling group.This advances self-distillation from implicit distributional alignment to explicit trajectory construction.
- TAPO: Difficulty-aware Candidate Selection targets problems within the model’s capability boundary and produces an emergent curriculum as the model improves.The selection mechanism is designed to provide useful corrective supervision without focusing on problems beyond the model’s current capability.
- TAPO: Decoupled Advantage Estimation prevents inflated group rewards from contaminating the base GRPO update, while OOD Token Suppression stabilizes optimization around corrective tokens.Both mechanisms address challenges introduced when constructed trajectories are integrated into advantage-based reinforcement learning.
- Inference: TAPO applies micro-reflective construction only during training and uses standard single-pass generation without thinking mode at inference.The intended outcome is internalized error diagnosis and recovery rather than surface-level reflection patterns at test time.
- Results: Across AIME 2024, AIME 2025, and HMMT 2025, TAPO consistently improves over GRPO under the same training steps and strengthens both first-pass reasoning and error correction.Direct Solution Rate and Effective Reflection Rate analyses support these two improvements and indicate transfer to general reasoning capability.
2 Related Work
Related work develops model reasoning through verifiable-reward reinforcement learning, self-distillation, iterative revision, and error-driven correction. These approaches also highlight the importance of keeping training data aligned with the model’s own distribution.
- Reinforcement learning with verifiable rewards: RLVR improves language-model reasoning by training on verifiable rewards, building on STaR and ReST’s iterative self-training with self-generated correct solutions.DeepSeek-R1 further showed that RL training can elicit sophisticated chain-of-thought reasoning.
- Distillation and error-driven correction: OPD, OPSD, and ROSD learn from model-generated outputs through token-level KL-divergence supervision, progressively replacing external teachers with privileged self-supervision and error-focused reflection.These methods minimize KL divergence toward a target distribution external to the student’s current reasoning at each token position.
- Iterative revision and self-correction: Self-Refine, Reflexion, and SCoRe improve outputs through critique, textual reflection, or corrected responses generated in iterative or multi-turn processes.SCoRe specifically trains self-correction with multi-turn reinforcement learning after an initial response.
- Distribution alignment: Training on data that deviates from the model’s distribution can cause instability and performance degradation, motivating measures of in-distribution versus out-of-distribution alignment.DDT formalizes this issue with a centered log-likelihood criterion, while off-policy RL data introduces bias requiring correction.
3 Method
TAPO constructs micro-reflective corrective trajectories from the model’s own correct and incorrect rollouts, preserving valid reasoning before the first error and inserting diagnosis-driven corrections. It combines difficulty-aware construction with decoupled advantage estimation and token-level suppression to maintain distributional proximity and stabilize training.
- Micro-Reflective Trajectory Construction: TAPO transforms contrastive information from GRPO response groups into learnable corrective trajectories through candidate selection, trajectory synthesis, and reward evaluation.Construction is performed within each GRPO training iteration.
- Candidate Selection: Reflective reconstruction applies only to Zone of Proximal Development problems containing at least npos correct and nneg incorrect responses.This excludes mastered and beyond-capability problems from trajectory construction.
- Trajectory Synthesis: Each micro-reflective trajectory preserves the learner’s valid reasoning prefix through the first error, then inserts an error diagnosis and corrected reasoning.This design maintains distributional proximity while intervening only from the first error point onward.
- Decoupled Advantage Estimation: Decoupled advantage estimation prevents constructed trajectories from contaminating original-rollout advantages, avoiding length penalties that can cause response length collapse, policy entropy collapse, and performance degradation.Jointly appending constructed trajectories to the original GRPO group produces advantage contamination because incorrect samples are typically longer than correct ones.
- Training Stabilization: OOD Token Suppression down-weights out-of-distribution corrective tokens while preserving learning signal from mathematical reasoning portions and reducing gradients from meta-cognitive transition phrases.Natural tokens receive weights near 1.0, while corrective transition tokens are selectively attenuated.
4 Experiments
Across cold-start and direct-training settings, TAPO improves benchmark performance over GRPO and OPSD, while ablations and training analyses support its trajectory construction and optimization design. TAPO also strengthens both first-pass reasoning and error correction, with more stable optimization behavior.
- Benchmark Results: 62.50% Pass@1 on AIME 2024 makes cold-start TAPO the strongest reported method, exceeding Cold-start + GRPO by 9.58 points and Coldstart + OPSD by 4.79 points.On AIME 2025, TAPO matches GRPO at 46.88% while exceeding OPSD by 3.55 points; on HMMT 2025, it reaches 31.46%, surpassing GRPO by 2.71 points.
- Benchmark Results: 63.12% Pass@1 gives direct-training TAPO an 8.54-point improvement over both GRPO and OPSD on AIME 2024.The passage also reports stronger inference-time performance for Qwen3-8B-Instruct w/ Thinking than the non-thinking base model across all three benchmarks.
- Reasoning and Reflection: Cold-start TAPO achieves higher Direct Solution Rate than Cold-start GRPO across all benchmarks, indicating stronger first-pass reasoning without reflection.Direct Solution Rate measures correctly answered problems solved directly, without intermediate reflection or self-correction.
- Reasoning and Reflection: TAPO achieves higher Effective Reflection Rate than GRPO across all benchmarks, supporting improved error-correction capability alongside stronger first-pass reasoning.The ERR improvement is most pronounced on AIME 2024 and more modest on AIME 2025 and HMMT 2025.
- Micro-Reflective Design: Full reconstruction underperforms TAPO across all benchmarks, showing that preserving the learner’s valid prefix is an important component of micro-reflective trajectory construction.The ablation replaces the entire solution with a complete regeneration rather than correcting from the error point onward.
- Optimization Ablations and Dynamics: Removing OTS degrades performance across all benchmarks, while TAPO’s OTS weighting and decoupled advantage estimation are associated with lower gradient norms and more stable optimization than OPSD.Cold-start+TAPO and Qwen3+TAPO maintain the lowest gradient norms, whereas Cold-start+OPSD has the highest; OPSD also exhibits the highest and most variable policy-gradient loss.
5 Conclusion
TAPO advances self-distillation from implicit distributional alignment to explicit trajectory construction. It constructs micro-reflective corrections from the model’s own error patterns and integrates them through difficulty-aware ZPD candidate selection and decoupled advantage estimation.
- 5 Conclusion: TAPO advances self-distillation from implicit distributional alignment to explicit trajectory construction.The method treats errors as material for constructing corrective trajectories rather than merely suppressing them.
- 5 Conclusion: TAPO constructs micro-reflective corrections from the model’s own error patterns.These corrections leverage the model’s errors as training information.
- 5 Conclusion: TAPO integrates its constructed corrections through difficulty-aware ZPD candidate selection and decoupled advantage estimation.The passage identifies both mechanisms as components of TAPO’s integration procedure.