TAPO: Teaching LLMs Self-Correction by Learning from Their Own Mistakes

Forum topic · 小凯 · 2026-06-25

Summary

This forum post on zhichai.net explains TAPO (Trajectory-Augmented Policy Optimization), a training method from Alibaba Tongyi and Tsinghua/PKU researchers that improves LLM mathematical reasoning by turning the model's own wrong answers into learnable training material. Unlike standard self-distillation with GRPO, which uses KL divergence to implicitly align outputs with correct answers, TAPO constructs micro-reflective trajectories: it keeps the model's erroneous reasoning prefix up to the failure point, inserts a natural-language diagnosis, and appends a guided correction derived from correct samples in the same sampling group. Three stabilization mechanisms support training: DCS (difficulty-aware candidate selection) keeps only medium-difficulty problems, DAE (decoupled advantage estimation) prevents reward contamination between original and corrected trajectories, and OOD token suppression avoids training divergence. The reflection is internalized during training, so inference remains single-pass with no extra cost. Experiments on AIME 2024, AIME 2025, and HMMT 2025 show TAPO consistently beats GRPO baselines, improving both first-attempt accuracy and self-correction success. The paper is available at https://arxiv.org/abs/2606.18844.

> Paper: Learning from Your Own Mistakes: Constructing Learnable Micro-Reflective Trajectories for Self-Distillation > Authors: Zhilin Huang (Alibaba Tongyi + Tsinghua), Hang Gao, Ziqiang Dong, Yuan Chen, Yifeng Luo, Chujun Qin (Peking University), Jingyi Wang, Yang Yang, Guanjun Jiang (Alibaba Tongyi) > Paper link: https://arxiv.org/abs/2606.18844 > Keywords: self-distillation, reinforcement learning, mathematical reasoning, GRPO, micro-reflective trajectories

1. A Counterintuitive Observation: Why Mistakes Matter More Than Successes

Imagine learning to ride a bike. Method A: imitate someone else's posture without ever falling — you never truly understand balance. Method B: you wobble, fall, and instinctively feel which way your center of gravity shifted — the next time, you remember that feeling.

The answer is obvious, yet most current LLM training resembles Method A.

2. The Self-Distillation Dilemma: Imitating Answers Instead of Learning the Process

Self-distillation lets a model generate multiple answers, keeps the correct ones as a "teacher," and aligns the model's output distribution with them — typically via KL divergence.

The problem: KL divergence only tells the model what the answer should be, never why its previous answer was wrong. It is an implicit distribution alignment operating at the logit level. The model learns "how to output more likely-correct tokens," not "how to reflect and correct." The model never *sees* its own mistakes during training — it is pushed toward correct answers like a blind person being led.

3. TAPO's Core Breakthrough: Turning Mistakes into Teaching Material

TAPO (Trajectory-Augmented Policy Optimization) is built on one idea:

> "Your mistakes are not garbage — they are the most valuable learning material."

3.1 Micro-Reflective Trajectories: Letting Mistakes Speak

Instead of multi-turn reflection at inference (which greatly increases cost), TAPO constructs special training samples at training time. From GRPO's paired correct/incorrect samples, TAPO builds correction trajectories:

Keep the erroneous prefix up to the failure point — the trajectory starts from the model's own context, staying on-policy.
Insert a natural-language diagnosis at the exact error location — teaching *metacognition*, how to diagnose errors, not just the right answer.
Guide correction with the reference correct answer from the same sampling group.
Zero inference cost: the model internalizes the correction instinct during training; at inference it needs no explicit multi-turn reflection — a single-pass output already contains self-correction ability.

3.2 Why "Micro"?

TAPO does not write long reflective essays. It inserts diagnosis and correction only at the exact point of error, keeping trajectories compact. Benefits: high learning efficiency (no wasted "irrelevant philosophizing"), strong gradient signals near the error location, and transferable correction patterns.

4. Three Stabilization Mechanisms

Training on mistakes is a double-edged sword. TAPO addresses three failure modes:

4.1 DCS (Difficulty-aware Candidate Selection)

Inspired by Vygotsky's Zone of Proximal Development:

Too easy problems (rarely wrong) → discarded; too few error samples.
Too hard problems (almost always wrong) → discarded; error trajectories dominate.
Medium difficulty (sometimes right, sometimes wrong) → kept, the sweet spot at the model's capability boundary.

DCS automatically builds a curriculum that escalates as the model improves.

4.2 DAE (Decoupled Advantage Estimation)

If corrected and original trajectories share one advantage computation, the raw error trajectory's reward contaminates the gradient. DAE decouples them:

Original trajectories: task reward only, standard GRPO update.
Corrected trajectories: independently computed advantages, keeping the correction signal pure.

Like regular homework vs. error-correction exercises with separate grading.

4.3 OOD Token Suppression

When the model starts generating tokens never seen in training, divergence can occur. This mechanism down-weights anomalous tokens, keeping training stable within the model's comfort zone.

5. Experimental Results

TAPO was evaluated on three high-difficulty math competition benchmarks: AIME 2024, AIME 2025, HMMT 2025.

Key conclusions:

At the same number of training steps, TAPO consistently outperforms GRPO and traditional self-distillation on all three benchmarks.
First-attempt accuracy improves — TAPO raises initial reasoning quality, not just error correction.
Self-correction improves: when the first solution is wrong, the success rate of self-repair is higher.

Why is the first-attempt gain important? Micro-reflective trajectories let the model "rehearse" error scenarios in training — like a pilot who trains in a simulator: not only better failure handling, but better initial procedure, because of a deeper understanding of what causes errors.

6. Comparison with Other Methods

vs. GRPO

| Dimension | GRPO | TAPO | |------|------|------| | Signal source | Correct trajectories only | Correct trajectories + error-trajectory corrections | | Learning mode | Imitate correct distribution | Learn diagnosis/repair from mistakes | | Alignment | Implicit KL alignment | Explicit trajectory construction | | Inference cost | Single-pass | Single-pass (reflection internalized at training) | | Metacognition | Weak | Strong |

vs. Multi-turn Reflection (Self-Refine, CoT)

| Dimension | Multi-turn reflection | TAPO | |------|------|------| | Inference | Multi-turn generation | Single-pass | | Cost | High (multiple LLM calls) | Low (one call) | | Error handling | Explicit multi-turn checking | Implicitly internalized |

TAPO's philosophy: move "reflection ability" from inference to training. Learn to reflect during training; be instinctively correct at inference.

Complementarity with HiVA

HiVA solves how multi-agent systems organize; TAPO solves how a single agent learns from its mistakes. An HiVA agent network could train each agent with TAPO — yielding agents with both collaboration (topology evolution) and self-correction (micro-reflection) abilities.

7. Analogies

Driving: Traditional self-distillation is a coach who silently steers you back to the right route. TAPO is a coach who, at the moment of a mistake, says: "You turned too early because you didn't notice the oncoming car. Next time, slow down, check the mirror, then turn." Complete diagnosis + correction feedback at the instant of error.
Medical training: Instead of just imitating a senior doctor's records, the intern is told: "Your reasoning was right, but you missed the link between symptom three and symptom one. Next time, prioritize disease X for this combination."
Gaming: Traditional RL learns statistically which action sequences survive; TAPO generates a "death replay": "You moved right at second 3, but the enemy predicted it. Next time, feint or use cover."

8. Why Learning from Mistakes Is Hard — and Valuable

> Wrong samples have far higher information density than correct ones.

A correct answer may just mean the model "took a familiar path." A wrong answer exposes the fuzzy zones at the model's knowledge boundary. TAPO doesn't merely train on wrong samples (which would teach more errors); it (1) preserves the error prefix, (2) inserts natural-language diagnosis, (3) guides the correction reasoning, and (4) internalizes it into single-pass output — cultivating expert-level intuition that avoids errors before they happen.

9. Limitations and Outlook

Limitations: 1. Quality dependence on the diagnosis in constructed trajectories — an inaccurate "correction generator" teaches wrong correction patterns. 2. Computational overhead: building micro-reflective trajectories (error-point identification, diagnosis generation) costs more than pure GRPO. 3. Task scope: validated mainly on mathematical reasoning; for open-ended tasks, "error" and "correction" are less well defined.

Directions: cross-task transfer (code generation, logical reasoning); combining micro-reflection with external tools (calculators, code interpreters) for verification; human-in-the-loop correction; and integration with HiVA-style multi-agent architectures.

10. One-Sentence Takeaway

> TAPO shows that the highest form of learning is not memorizing correct answers, but making mistakes part of you — like a reflex that snaps back on course at the very moment you drift.

Reference [1] Huang, Z., et al. Learning from Your Own Mistakes: Constructing Learnable Micro-Reflective Trajectories for Self-Distillation. https://arxiv.org/abs/2606.18844.