🤖 AI

Trajectory-Refined Distillation

Dieses Paper stellt Trajectory-Refined Distillation (TRD) vor, eine neuartige On-Policy-Distillationsmethode, die das strukturelle Problem des „Prefix Failure“ durch die Anwendung von Trajektorien-basierten Korrekturen auf Student-Rollouts unter Lehrer-Anleitung mildert und dadurch die Genauigkeit sowie die Abdeckung des Schlussfolgerns über verschiedene Benchmarks und Modellskalen hinweg verbessert.

Ursprüngliche Autoren: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Veröffentlicht 2026-06-09

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie bringen einem Schüler (der KI) bei, wie man ein komplexes mathematisches Problem löst. Sie haben einen brillanten Lehrer (eine fortgeschrittenere KI), der die richtige Antwort kennt.

In der Standardmethode, die im Paper beschrieben wird, genannt On-Policy Distillation (OPD), sieht der Prozess so aus:

Der Schüler versucht, das Problem im Alleingang zu lösen.
Während der Schüler seine Schritte aufschreibt, beobachtet der Lehrer ihn und sagt: „Für dieses spezifische Wort, das du gerade geschrieben hast, ist dies die Wahrscheinlichkeit für das nächste Wort.“
Der Schüler versucht, sein Gehirn so anzupassen, dass es die Wahrscheinlichkeiten des Lehrers für jedes einzelne Wort widerspiegelt.

Das Problem: „Prefix Failure“ (Die falsche Abzweigung)

Das Paper identifiziert einen schwerwiegenden Fehler in dieser Standardmethode, den Prefix Failure genannt wird.

Stellen Sie sich vor, der Schüler geht einen Pfad entlang, um das Problem zu lösen. Bei Schritt 3 biegt er falsch ab. Er befindet sich nun in einer Sackgasse.

Das Dilemma des Lehrers: Der Lehrer sieht, dass der Schüler verloren ist. Der Lehrer weiß, dass der richtige Pfad mit „Warte, eigentlich...“ beginnt und dann in eine völlig andere Richtung führt.
Die Verwirrung: Aber der Lehrer ist gezwungen, auf den aktuellen Standort des Schülers (die Sackgasse) zu blicken. Daher wird der Rat des Lehrers zu einem verwirrenden Mix aus: „Bleib auf dieser Sackgasse, weil du schon hier bist“ UND „Dreh dich sofort um.“
Das Ergebnis: Der Schüler erhält ein zersplittertes Signal. Er versucht, vom Lehrer zu lernen, aber da der Schüler auf dem falschen Pfad feststeckt, wird der Rat des Lehrers fragmentiert. Der Schüler lernt, „gut darin zu sein, verloren zu sein“, anstatt zu lernen, wie er wieder auf den richtigen Weg zurückfindet. Das Paper nennt dies eine „bimodale Mischung“ – der Lehrer versucht gleichzeitig zwei entgegengesetzte Dinge zu lehren, und der Schüler wird dadurch verwirrt.

Bestehende Lösungen versuchten, dies zu beheben, indem sie lediglich die „Note“ (die Loss-Funktion) für bestimmte Wörter anpassten, aber sie behoben nicht das eigentliche Problem, dass der Schüler immer noch auf dem falschen Pfad feststeckte.

Die Lösung: Trajectory-Refined Distillation (TRD)

Die Autoren schlagen eine neue Methode vor, die Trajectory-Refined Distillation (TRD) heißt. Anstatt nur den unordentlichen Entwurf des Schülers Wort für Wort zu bewerten, ändert sie den Arbeitsablauf:

Der Entwurf: Der Schüler macht zuerst einen ersten Versuch am Problem (der „raw rollout“).
Das Umschreiben: Bevor die Bewertung erfolgt, nimmt der Lehrer diesen unordentlichen Entwurf und schreibt ihn um. Der Lehrer korrigiert die falschen Abzweigungen, korrigiert die Logik und erstellt eine saubere, perfekte Version der Lösung basierend auf dem Ausgangspunkt des Schülers.
Die Lektion: Der Schüler lernt dann aus dieser sauberen, korrigierten Version, nicht aus dem unordentlichen Original.

Die Analogie:

Der alte Weg: Sie schreiben eine Geschichte mit einer Logiklücke in der Mitte. Ihr Lektor liest sie und sagt: „Für das Wort ‚Katze‘ in Absatz 3 hättest du ‚Hund‘ schreiben sollen.“ Aber Sie stecken immer noch in Absatz 3 fest und versuchen herauszufinden, warum Sie überhaupt ‚Katze‘ geschrieben haben. Es ist verwirrend.
Der TRD-Weg: Sie schreiben die Geschichte mit der Logiklücke. Ihr Lektor nimmt Ihren Entwurf, korrigiert die Logiklücke und schreibt den gesamten Mittelteil um, sodass die Geschichte perfekt fließt. Dann studieren Sie die umgeschriebene Version des Lektors, um beim nächsten Mal besser Geschichten schreiben zu können.

Warum dies besser funktioniert

Behebt die Ursache: Indem der Lehrer die „falsche Abzweigung“ korrigiert, bevor der Schüler versucht zu lernen, ist der Rat des Lehrers nicht mehr verwirrend. Es ist ein klarer, einzelner Pfad.
Erforscht neue Wege: Selbst wenn der Schüler das Problem beim ersten Mal richtig löst, kann der Lehrer ihm einen anderen, kürzeren oder klügeren Weg zeigen, um das Problem zu lösen. Dies lehrt den Schüler, flexibler zu werden.
Effizienz: Das Paper fand heraus, dass die „umgeschriebenen“ Pfade oft viel kürzer und klarer waren als die ursprünglichen, ausschweifenden Versuche des Schülers, was das Training tatsächlich beschleunigte.

Die Ergebnisse

Die Forscher testeten dies an schwierigen Mathematik-Wettbewerben (wie AIME und HMMT) und Coding-Herausforderungen. Sie fanden heraus, dass:

Die neue Methode (TRD) die alten Methoden konsistent schlug.
Sie besonders gut darin war, der KI zu helfen, die schwierigsten Probleme zu lösen, bei denen sie normalerweise stecken blieb.
Es funktionierte, egal ob der Lehrer eine separate, größere KI war oder dieselbe KI, die als ihr eigenes „kluges“ Lehrer fungierte (unter Verwendung einer Technik namens „privileged information“).

Kurz gesagt argumentiert das Paper, dass man eine KI nicht effektiv lehrt, indem man lediglich ihre Fehler Wort für Wort bewertet, während sie diese noch macht. Stattdessen sollte man den Lehrer den gesamten Pfad zuerst korrigieren lassen und dann den Schüler die korrigierte Reise lernen lassen.

Technisches Resümee: Trajectory-Refined Distillation (TRD)

1. Problemstellung: Prefix-Versagen bei On-Policy-Distillation

Das Paper identifiziert eine fundamentale strukturelle Einschränkung bei der On-Policy-Distillation (OPD) und der On-Policy-Selbst-Distillation (OPSD) für Large Language Models (LLMs), die als Prefix-Versagen (Prefix Failure) bezeichnet wird.

In der Standard-OPD/OPSD generiert ein Studentenmodell einen Rollout (Trajektorie) $y_o$ , und ein Lehrer-Modell liefert dichte, pro-Token-basierte Supervision basierend auf diesem spezifischen Rollout. Die Autoren argumentieren, dass wenn der Rollout des Studenten einen falschen Denkpfad einschlägt (ein „fehlgeschlagener Präfix“), das Supervisionssignal des Lehrers strukturell kompromittiert wird:

Bimodale Lehrer-Mischung: Wenn ein Präfix falsch ist, wird die Lehrer-Verteilung zu einer Mischung aus zwei Modi: einem, der den fehlgeschlagenen Präfix fortsetzt (für die Sequenzkonsistenz), und einem anderen, der zu der korrekten Lösung umschwenkt.
Fragmentierte Gradienten: Standardmäßige Per-Token-Loss-Funktionen (KL-Divergenz) bewerten die Supervision entlang der eingefrorenen fehlgeschlagenen Trajektorie. Infolgedessen sind die Gradienten-Updates fragmentiert. Der Lehrer empfiehlt wiederholt ein Korrektur-Token (z. B. „Warte“ oder „Eigentlich“) unter der Bedingung des falschen Kontextes (des fehlgeschlagenen Präfixes), anstatt unter der Bedingung des sich entfaltenden Korrekturpfades.
Ineffektivität von Token-basierten Korrekturen: Bestehende Behebungen (z. B. Loss-Clipping, Top-K-Trunkierung, Reweighting) operieren auf der Ebene des Token-Losses. Sie versuchen, „informative“ Token auszuwählen oder „verrauschte“ zu unterdrücken, scheitern aber daran, die eigentliche Ursache zu adressieren: Der fehlgeschlagene Präfix selbst bleibt strukturell intakt. Die übermittelten Supervisionspaare weichen vom idealen Korrekturpfad ab, was verhindert, dass das Modell effektiv lernt, aus Fehlern zu rehabilitieren.

2. Methodik: Trajectory-Refined Distillation (TRD)

Um das Prefix-Versagen zu mildern, schlagen die Autoren die Trajectory-Refined Distillation (TRD) vor, eine Methode, die die Intervention von der Token-Ebene auf die Trajektorien-Ebene verlagert.

Kernmechanismus

Anstatt den rohen Studenten-Rollout $y_o$ direkt für die Distillation zu verwenden, führt TRD vor der Verlustberechnung einen Verfeinerungsschritt ein:

Sampling: Das Studentenmodell $\pi_\theta$ generiert einen rohen On-Policy-Rollout $y_o \sim \pi_\theta(\cdot | x)$ .
Verfeinerung (Refinement): Das Lehrer-Modell $\pi_T$ $π_{T}$ (oder das Studentenmodell selbst in der OPSD, bedingt auf privilegierten Informationen wie der Grundwahrheit $y^*$ $y^{*}$ ) generiert eine verfeinerte Trajektorie $y_r$ $y_{r}$ .
- In OPD: $y_r \sim \pi_T(\cdot | x, y_o)$ .
- In OPSD: $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$ .
- Entscheidend ist, dass der Lehrer dazu aufgefordert wird, den Rollout des Studenten geleitet durch die Referenzlösung zu revidieren, wodurch die fehlerhaften Teile des Präfixes effektiv „umgeschrieben“ werden, während die gültigen Muster des Studenten beibehalten werden.
Distillation: Die verfeinerte Trajektorie $y_r$ dient als Supervisionsziel für die standardmäßige Verteilungsanpassung (Forward KL mit vollem Vokabular).

Theoretische Begründung

TRD adressiert das Problem der „fragmentierten Gradienten“, indem es sicherstellt, dass der Supervisionskontext entlang des Korrekturpfades wächst, anstatt entlang des fehlgeschlagenen Präfixes.

Ideale Gradienten-Rekuperation: Durch die Supervision entlang von $y_r$ enthält der Kontext für jedes Token den korrigierten Präfix, was es dem Lehrer ermöglicht, ein kohärentes, mehrstufiges Korrektursignal ( $g_{ideal}$ ) bereitzustellen, anstatt ein repetitives, kontext-inkonsistentes Signal ( $g_{frag}$ ).
Exploration: Selbst wenn $y_o$ korrekt ist, setzt TRD den Studenten alternativen validen Ableitungen aus, die vom Lehrer vorgeschlagen werden, wodurch die Verteilung der korrekten Denkpfade, die der Student lernt, verbreitert wird.

Algorithmus

Der Prozess ist in Algorithmus 1 zusammengefasst:

Sample $y_o \sim \pi_\theta(\cdot | x)$ .
Generiere $y_r \sim \pi_T(\cdot | x, y_o)$ .
Update $\theta$ unter Verwendung der KL-Divergenz zwischen $\pi_\theta$ und der Lehrer-Verteilung über $y_r$ .

3. Zentrale Beiträge

Identifikation des Prefix-Versagens: Das Paper formalisiert „Prefix Failure“ als strukturelle Ursache für die Degradation bei der On-Policy-Distillation und erklärt, warum Interventionen auf Token-Ebene nicht ausreichen, um aus Denkfehlern zu rehabilitieren.
Korrektur auf Trajektorien-Ebene: TRD wird als das erste Optimierungsdesign auf Trajektorien-Ebene vorgeschlagen, das explizit die Wiederherstellung fehlgeschlagener Präfixe anstrebt und dabei die On-Policy-Support-Beschränkung respektiert, ohne prohibitiven Rechenaufwand zu verursachen.
Verbesserung in zwei Regimen: TRD verbessert die Leistung in zwei Szenarien:
- Fehlerkorrektur: Es mildert das Prefix-Versagen ab, indem es problematische Präfixe vor der Distillation korrigiert.
- Exploration: Es erweitert die Reasoning-Abdeckung, indem es Studenten auch dann alternativen validen Ableitungen aussetzt, wenn der ursprüngliche Rollout korrekt war.
Anwendbarkeit auf Selbst-Distillation: Die Methode lässt sich natürlich auf OPSD erweitern, wobei der Student als sein eigener Lehrer fungiert, der durch privilegierte Informationen (Ground Truth) bedingt ist, was die Notwendigkeit eines separaten Lehrer-Modells eliminiert.

4. Experimentelle Ergebnisse

Die Autoren evaluierten TRD über fünf Wettbewerbs-Mathematik-Benchmarks (AIME24/25, HMMT25, BeyondAIME, AMOBench) und Code-Generierungs-Benchmarks (HumanEval+, MBPP+, LiveCodeBench) unter Verwendung von Qwen3-Modellen (1.7B, 4B, 8B).

Wichtigste Erkenntnisse

Überlegene Leistung: TRD übertrifft konsistent vorangegangene Baselines (Forward/Reverse KL, mit/ohne Clipping/Trunkierung) sowohl in OPD- als auch in OPSD-Settings.
Signifikante Gewinne bei schwierigen Aufgaben: Die Verbesserungen sind bei den schwierigsten Benchmarks am ausgeprägtesten.
- AMOBench: TRD erreichte eine relative Verbesserung von ~50 % im Pass@16 für Qwen3-8B unter dem OPSD-Setup im Vergleich zum Basismodell.
- AIME24/25: TRD zeigte konsistente Gewinne, insbesondere bei kleineren Modellen (z. B. +4,6 % auf AIME24 für Qwen3-1.7B).
Stabilität: Im Gegensatz zu einigen Baselines, die auf bestimmten Benchmarks an Leistung verlieren (z. B. Reverse KL mit Leistungsabfall auf AIME24), fiel TRD nie unter die Leistung des Basismodells.
Effizienz: Im OPSD-Setting waren die verfeinerten Trajektorien $y_r$ signifikant kürzer (die mediane Länge wurde um den Faktor ~9 komprimiert) als die rohen Rollouts $y_o$ , was die Trainingszeit (Wall-Clock Time) trotz des zusätzlichen Sampling-Schritts um ~60 % reduzierte.
Code-Generierung: TRD erreichte oder übertraf die Baselines auf HumanEval+ und MBPP+, wobei die Gewinne auf LiveCodeBench begrenzt waren, was darauf hindeutet, dass die Verfeinerungsfähigkeit des Lehrers durch die Aufgabenkomplexität begrenzt sein könnte.

5. Bedeutung und Behauptungen

Das Paper behauptet, dass TRD einen notwendigen Wandel in der Post-Training-Strategie für LLMs darstellt. Die Autoren argumentieren:

Interventionen auf Token-Ebene sind unzureichend: Die strukturelle Natur des Prefix-Versagens erfordert Korrekturen auf Trajektorien-Ebene.
On-Policy-Support bleibt erhalten: TRD erreicht diese Korrekturen, ohne die On-Policy-Beschränkung zu verletzen (im Gegensatz zur Verwendung einer festen Experten-Policy), was es mit Standard-Recheneffizienz-Paradigmen kompatibel macht.
Selbstverbesserung ist praktikabel: Der Erfolg von TRD im OPSD-Setting zeigt, dass Modelle ihre eigenen Denkpfade unter Verwendung privilegierter Informationen (Ground Truth) effektiv verfeinern können, ohne ein separates, größeres Lehrer-Modell zu benötigen, was einen leichtgewichtigen Weg zur Selbstverbesserung bietet.

Die Arbeit legt nahe, dass zukünftige Distillationsmethoden die strukturelle Integrität der Denk-Trajektorie berücksichtigen müssen, nicht nur die statistische Ausrichtung einzelner Token.