Learning from Your Own Mistakes: Constructing Learnable Micro-Reflective Trajectories for Self-Distillation

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.

한 줄 요약

TAPO는 모델 스스로의 오류를 기반으로 학습 가능한 마이크로-반성적 추적 경로를 구성하여 자가 지도를 강화하는 새로운 정책 최적화 프레임워크이다.

핵심 기여도

TAPO는 GRPO 대비 AIME 2024, 2025, HMMT 2025에서 일관된 성능 향상을 보임.
마이크로-반성적 수정(Micro-Reflective Corrections)을 통해 모델의 오류를 진단하고 수정하는 구조적 학습 신호를 제공.
DCS, DAE, OTS를 통해 안정적인 정책 최적화를 달성.
기존 KL 기반 자가 지도 방법이 오류를 억제하는 반면, TAPO는 오류를 학습 기회로 활용.

핵심 아이디어

기존 자가 지도(self-distillation)는 KL 발산 최소화를 통해 모델의 출력 분포를 특정 목표 분포에 맞추는 방식을 사용하지만, 이는 오류의 구체적 원인을 진단하거나 수정하는 능력을 키우지 못한다. TAPO는 이 문제를 해결하기 위해 모델 스스로 생성한 정답과 오답 추적 경로를 대조하여, 오류 발생 지점에서 자연어로 진단하고 수정하는 마이크로-반성적 수정을 생성한다. 이는 모델의 실제 오류 패턴을 기반으로 하며, 오류를 억제하는 것이 아니라 학습의 출발점으로 삼는다. 학습 심리학의 ZPD(근접 발달 영역) 개념을 도입하여, 모델의 현재 능력 수준에 맞춘 학습 신호를 제공함으로써 더 효과적인 학습을 유도한다.

기술적 접근법

**Trajectory-Augmented Policy Optimization (TAPO)**: 정책 최적화 프레임워크로, 정답과 오답 추적 경로를 대조하여 마이크로-반성적 수정을 생성.
**Micro-Reflective Corrections**: 오류 발생 지점에서 자연어 진단과 수정을 삽입한 새로운 학습 경로.
**Difficulty-aware Candidate Selection (DCS)**: 모델의 능력 경계 내에서 난이도에 맞춘 학습 샘플을 선택.
**Decoupled Advantage Estimation (DAE)**: 기존 GRPO 업데이트가 왜곡되지 않도록 그룹 평균의 과도한 영향을 방지.
**OOD Token Suppression (OTS)**: 분포 외 토큰의 가중치를 낮춰 안정적인 최적화 유지.
**Pass@k 메트릭** (k=1~5)를 사용하여 AIME 2024, 2025, HMMT 2025에서 성능 평가.

주요 결과

AIME 2024에서 TAPO는 GRPO 대비 Pass@1에서 +7.2%, Pass@5에서 +12.4% 개선.
AIME 2025에서 Pass@1 +6.8%, Pass@5 +11.3%.
HMMT 2025에서 Pass@1 +5.9%, Pass@5 +9.7%.
DSR(정답률)과 ERR(오류 수정 효과성) 분석에서 TAPO는 첫 번째 추론과 오류 수정 모두에서 개선 효과를 보임.

의의 및 한계

TAPO는 기존 자가 지도 방법이 오류를 억제하는 데 집중하는 반면, 오류를 학습 기회로 전환함으로써 모델의 내재적 오류 수정 능력을 키우는 학습 심리학적 접근을 도입한 점에서 혁신적이다. 또한, 정책 최적화 과정에서 생성된 실제 오류 패턴을 기반으로 학습 신호를 구성함으로써, 모델의 분포 내 학습을 유지하는 데 성공했다. 그러나 TAPO는 오류 진단과 수정이 모두 자연어로 이루어져야 하므로, 복잡한 수학 문제에서는 진단의 정확도와 자연스러움이 한계가 될 수 있다. 또한, DCS와 DAE는 모델의 능력 수준에 따라 적절히 조정되어야 하며, 이는 추가적인 튜닝을 필요로 한다.

실용적 활용

TAPO는 수학 문제 해결, 논리적 추론, 복잡한 텍스트 생성 등에서 모델의 오류를 진단하고 수정하는 능력을 향상시키는 데 활용 가능하다. 특히, 학습 과정에서 스스로의 오류를 학습하는 방식은 교육 AI, 챗봇, 코드 생성 모델 등에서 실용적 가치가 높다.