Technische Zusammenfassung: Draft-OPD

Problemstellung

Speculative Decoding (SD) beschleunigt die Inferenz von Large Language Models (LLM), indem ein leichtgewichtiges Draft-Modell Token-Blöcke vorschlägt, die anschließend parallel von einem größeren Zielmodell verifiziert werden. Während aktuelle Mainstream-Methoden (z. B. EAGLE-3, DFlash) diese Draft-Modelle durch Supervised Fine-Tuning (SFT) auf festen Trajektorien trainieren, die vom Zielmodell generiert wurden, stößt dieser Ansatz an eine Leistungsgrenze.

Die Autoren identifizieren eine fundamentale Offline-zu-Inferenz-Mismatch:

1. Trainingsphase: Das Draft-Modell lernt aus statischen Trajektorien, bei denen jedes Präfix vom Zielmodell generiert wurde.
2. Inferenzphase: Das Draft-Modell wird an Blöcken evaluiert, die es selbst vorgeschlagen hat. Die Akzeptanzlänge (die Anzahl der pro Verifizierungsschritt akzeptierten Token) wird durch die vom eigenen Policy induzierten Zustände des Draft-Modells bestimmt, nicht durch die statischen Zustände, die während des SFT gesehen wurden.

Folglich erreicht fortgesetztes SFT schnell ein Plateau, und das Draft-Modell verbessert seine Akzeptanzlänge auf Testdaten nicht weiter. Darüber hinaus ist eine direkte Anwendung standardmäßiger On-Policy Distillation (OPD) auf Draft-Modelle schwierig, da:

• Nur-Draft-Rollouts (bei denen das Draft-Modell Sequenzen unabhängig generiert) instabil sind und oft in repetitive Ausgaben degenerieren.
• Ziel-unterstützte Rollouts (bei denen das Zielmodell das Draft-Modell korrigiert) hochwertige Sequenzen produzieren, aber das „On-Policy"-Signal eliminieren. Das Zielmodell verwirft die Fehler des Draft-Modells, sodass das Draft-Modell nur aus der Verteilung des Zielmodells lernt und die spezifischen Fehler, die die spekulative Akzeptanz begrenzen, nicht adressiert.

Methodik: Draft-OPD

Um diese Lücke zu schließen, schlagen die Autoren Draft-OPD vor, ein On-Policy-Distillationsframework, das speziell für trainingsbasierte Draft-Modelle entwickelt wurde. Die Methode ermöglicht es dem Draft-Modell, aus seinen eigenen Inferenzzeit-Zuständen zu lernen, einschließlich sowohl erfolgreicher Vorschläge als auch der spezifischen Fehler, die zu Verifizierungsfehlern führen.

Das Framework besteht aus drei koordinierten Komponenten:

1. Rollout mit Fehlerpositions-Sammlung

Anstatt das Draft-Modell zu zwingen, vollständige Sequenzen allein zu generieren, verwendet Draft-OPD ziel-unterstützte Rollouts, um stabile, hochwertige Fortsetzungen sicherzustellen.

• Während des Speculative Decoding schlägt das Draft-Modell einen Block von Token vor, und das Zielmodell verifiziert diese.
• Das System erfasst die Startposition jedes vorgeschlagenen Blocks als Anker.
• Dieser Prozess erzeugt ein gültiges Rollout (unter Einhaltung der Zielverteilung), während die spezifischen Zustände bewahrt werden, in denen das Draft-Modell gehandelt hat.

2. Replay zur Berechnung der Log-Wahrscheinlichkeiten

Um das während des ziel-unterstützten Rollouts verlorene „On-Policy"-Signal wiederzugewinnen, spielt das System den Draft-Prozess von den aufgezeichneten Ankern aus neu ab.

• Ausgehend von einem Ankerzustand (dem vom Zielmodell verifizierten Präfix) simuliert das System den Vorschlag des nachfolgenden Blocks durch das Draft-Modell erneut.
• Es werden Log-Wahrscheinlichkeiten sowohl für das Draft-Modell (Schüler) als auch für das Zielmodell (Lehrer) auf diesen vom Draft generierten Präfixen berechnet.
• Entscheidend ist, dass dies dem System ermöglicht, die Verteilung des Zielmodells auf Token zu bewerten, die während der initialen Verifizierung abgelehnt wurden, wodurch die spezifischen Fehler erfasst werden, die das Draft-Modell gemacht hat.

3. Akzeptanzbewusstes Distillationsziel

Das Trainingsziel behandelt akzeptierte und abgelehnte Token unterschiedlich, um die spekulative Akzeptanz zu optimieren:

• Akzeptierte Token: Für Token, die die Verifizierung bestehen, verwendet die Methode die Forward KL-Divergenz ($D_{KL}(P_{target} || P_{draft})$). Dies ermutigt die Draft-Verteilung, die Zielverteilung dort abzudecken, wo sie bereits übereinstimmen, und verstärkt so zuverlässiges Verhalten.
• Abgelehnte Token: Für Token, die die Verifizierung nicht bestehen, verwendet die Methode die Reverse KL-Divergenz ($D_{KL}(P_{draft} || P_{target})$). Dies bestraft das Draft-Modell dafür, dass es Token eine hohe Wahrscheinlichkeit zuweist, die das Zielmodell ablehnt, und lehrt das Draft-Modell effektiv, seine eigenen hochkonfidenten Fehler zu vermeiden.
• Gewichtung: Abgelehnte Token werden exponentiell basierend auf ihrer Position innerhalb des Blocks gewichtet ($w_k = \gamma^{k-1}$). Fehler, die früher in einem Block auftreten, werden stärker gewichtet, da sie den gesamten nachfolgenden Suffix ungültig machen und einen größeren Einfluss auf die gesamte Akzeptanzlänge haben.

Hauptbeiträge

1. Identifikation einer Einschränkung: Die Arbeit zeigt, dass offlines SFT für Draft-Modelle aufgrund der Diskrepanz zwischen Trainingszuständen (vom Ziel generiert) und Inferenzzuständen (vom Draft generiert) schnell ein Plateau erreicht.
2. Analyse von OPD-Herausforderungen: Sie erklärt, warum Standard-OPD für Draft-Modelle versagt: unabhängige Rollouts sind instabil, während naive ziel-unterstützte Rollouts das On-Policy-Fehlersignal auslöschen.
3. Draft-OPD-Framework: Sie stellt ein neuartiges Framework vor, das ziel-unterstützte Rollouts mit Fehlerpositions-Replay kombiniert und effizientes Nachtrainieren auf Verifizierungszeit-Fehler ermöglicht.
4. Akzeptanzbewusstes Ziel: Sie schlägt ein hybrides KL-Divergenz-Ziel vor, das zwischen akzeptierten und abgelehnten Token unterscheidet, um spezifisch die Fehler zu adressieren, die den spektralen Geschwindigkeitsvorteil begrenzen.

Experimentelle Ergebnisse

Die Autoren evaluierten Draft-OPD an der Qwen3-Familie (4B, 8B und 30B-A3B-Thinking) über mathematisches Schlussfolgern (GSM8K, MATH-500, AIME), Code-Generierung (MBPP, HumanEval, SWE-bench Lite) und allgemeinen Chat (MT-Bench).

• Leistungssteigerungen: Unter einem angepassten Trainings-FLOPs-Budget übertraf Draft-OPD EAGLE-3 und DFlash konsistent.
  • Es wurde eine über 5-fache verlustfreie Beschleunigung für Thinking-Modelle erreicht.
  • Es verbesserte die durchschnittliche Akzeptanzlänge und den Geschwindigkeitsvorteil um 23 % gegenüber EAGLE-3 und um 13 % gegenüber DFlash.
• Durchsatz: In SGLang-Deployments-Benchmarks zeigte Draft-OPD konsistente Durchsatzgewinne über alle Parallelitätsstufen (1 bis 32) hinweg, mit durchschnittlichen Verbesserungen der Akzeptanzlänge von ca. 11,2 % gegenüber DFlash.
• Ablationsstudien:
  • Die einfache Anwendung von SFT auf OPD-generierten Daten (ohne den Replay-Mechanismus) führte zu keiner Leistungsverbesserung, was die Notwendigkeit des On-Policy-Fehlersignals bestätigt.
  • Die Verwendung einer einzigen KL-Richtung (alle-forward oder alle-reverse) war dem gemischten, akzeptanzbewussten Ziel unterlegen.
  • Die zufällige Auswahl von Replay-Ankern war weniger effektiv als die Auswahl von Ankern basierend auf Fehlerpositionen, was die Bedeutung der Fokussierung des Trainings auf Fehlerzustände unterstreicht.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass Draft-OPD einen bedeutenden Schritt vorwärts bei der Optimierung trainingsbasierter Draft-Modelle darstellt. Durch den Wechsel von statischer, offliner Überwachung zu dynamischer, on-policy Distillation, die explizit Verifizierungszeit-Fehler adressiert, durchbricht die Methode das Leistungsplateau des traditionellen SFT.

Die Autoren betonen, dass ihr Ansatz die verlustfreie Natur des Speculative Decoding bewahrt (die Ausgabe-Verteilung des Zielmodells bleibt erhalten), während die Effizienz erheblich verbessert wird. Sie positionieren diese Arbeit als praktische Lösung zur Beschleunigung von „Thinking"-Modellen und komplexen Schlussfolgerungsaufgaben, bei denen Standard-Draft-Modelle oft aufgrund der hohen Unsicherheit langer Schlussfolgerungsspuren Schwierigkeiten haben. Die Arbeit schließt, dass das Nachtrainieren von Draft-Modellen auf ihren eigenen Inferenzzeit-Fehlern eine kritische Richtung zur Verbesserung des Speculative Decoding ist und einen Weg zu höheren Akzeptanzlängen und schnellerer Inferenz ohne Kompromisse bei der Generierungsqualität bietet.