Hybrid Policy Distillation for LLMs

1. Problemstellung

Die Kompression von Large Language Models (LLMs) durch Wissensdistillation (Knowledge Distillation, KD) ist entscheidend für deren praktische Einsatzfähigkeit, da große Modelle hohe Rechen- und Speicheranforderungen stellen. Bestehende KD-Methoden stehen jedoch vor einem fundamentalen Dilemma, das durch drei eng verknüpfte Design-Aspekte bestimmt wird:

• Divergenz-Richtung: Die Wahl zwischen Vorwärts-KL (FKLD) und Rückwärts-KL (RKLD). FKLD fördert die Abdeckung aller Modi (Mode Coverage), führt aber oft zu übermäßig geglätteten Vorhersagen. RKLD konzentriert sich auf die wahrscheinlichsten Modi (Mode-Seeking), kann aber bei großen Lücken zwischen Lehrer und Schüler instabil werden.
• Optimierungsstrategie: Der Unterschied zwischen Verlust-basierten Ansätzen (Loss) und Belohnung-basierten Ansätzen (Reward), die unterschiedliche Gradienten-Dynamiken erzeugen.
• Datenregime: Der Konflikt zwischen Off-Policy-Daten (statische Datensätze, effizient, aber mit Train-Inference-Mismatch) und On-Policy-Daten (vom Schüler generiert, besserer Alignment, aber rechenintensiv durch Rollouts).

Bisherige Ansätze optimieren oft nur einen dieser Aspekte, was zu suboptimaler Stabilität, Effizienz oder Endleistung führt.

2. Methodik: Hybrid Policy Distillation (HPD)

Die Autoren schlagen Hybrid Policy Distillation (HPD) vor, eine Methode, die die Vorteile von FKLD und RKLD sowie von Off-Policy- und On-Policy-Daten auf Token-Ebene integriert.

Kernkonzept: Umformulierung als gewichtete Log-Likelihood
Die Autoren stellen einen einheitlichen Blickwinkel vor, der SFT, FKLD und RKLD als gewichtete Log-Likelihood-Ziele auf Token-Ebene reformuliert:
$$L(\theta) = \min_{\theta} -E_{(s_t, a_t) \sim D_\pi} [w(a_t | s_t) \log q_\theta(a_t | s_t)]$$
Dabei ist $w(a_t | s_t)$ ein Reweighting-Term, der die Diskrepanz zwischen Schüler- und Lehrer-Verteilung quantifiziert.

Der HPD-Algorithmus
HPD kombiniert zwei Mechanismen, um die inhärenten Schwächen der einzelnen Divergenzen auszugleichen:

1. Hybridisierung von Vorwärts- und Rückwärts-KL:

   • Für Experten-Token (Ground-Truth oder Lehrer-Generierung): HPD nutzt einen negativen $K_1$-Schätzer (basierend auf dem Reverse-KL-Gradienten).
     • Wenn der Schüler das Experten-Token unterschätzt ($K_1 > 0$), wird ein positiver Reweighting-Term (ähnlich FKLD) angewendet, um die Wahrscheinlichkeit zu erhöhen.
     • Wenn der Schüler das Token überschätzt ($K_1 \le 0$), wird der Vorwärts-KL-Anteil maskiert, um widersprüchliche Gradienten zu vermeiden.
   • Für vom Schüler gesampelte Token (On-Policy): Der Schüler sampelt ein Token $a_t$, das vom Experten-Token abweicht.
     • Wenn dieses gesampelte Token überschätzt wird (negativer $K_1$-Wert), wird es mit einem negativen Gewicht bestraft. Dies drückt die Wahrscheinlichkeitsmasse dieses „unvernünftigen" Tokens und verteilt sie proportional auf andere Token (insbesondere auf das Experten-Token).

2. Integration von Off- und On-Policy-Daten:

   • HPD nutzt Off-Policy-Daten als Basis.
   • Zusätzlich führt es ein leichtgewichtiges On-Policy-Sampling durch: Der Schüler sampelt ein einziges Token pro Schritt, um die Verteilung zu prüfen und unnötiges Verhalten zu unterdrücken. Dies vermeidet den hohen Rechenaufwand vollständiger Sequenz-Rollouts (wie bei MiniLLM oder GKD), behält aber die Vorteile des On-Policy-Alignments bei.

Algorithmus-Ablauf:
Für jeden Zustand $s_t$ im Offline-Datensatz:

1. Berechne den Reverse-KL-Lücke ($k_1$) für das Experten-Token.
2. Sampel ein alternatives Token $a_t$ vom Schüler und berechne dessen Lücke ($k'_1$).
3. Passe die Gewichte ($w^*_t$ für Experten, $w_t$ für gesampelte Token) dynamisch an, basierend auf den Vorzeichen von $k_1$ und $k'_1$.
4. Optimiere das Ziel mit den gewichteten Log-Likelihoods.

3. Wichtige Beiträge

• Einheitliche Sichtweise: Die Autoren etablieren eine theoretische Grundlage, die verschiedene KD-Methoden als Varianten eines gewichteten Log-Likelihood-Ziels erklärt. Dies ermöglicht ein systematisches Verständnis und die Entwicklung neuer Ansätze.
• HPD-Algorithmus: Die Entwicklung einer Methode, die Forward- und Reverse-KL-Signale hybridisiert und dabei die Effizienz von One-Hot-Supervision (wie bei SFT) bewahrt, aber durch dynamisches Reweighting die Verteilungsanpassung verbessert.
• Effizienz ohne Kompromisse: HPD erreicht On-Policy-ähnliche Stabilität und Performance, ohne die teuren vollständigen Rollouts zu benötigen, die für On-Policy-Distillation typisch sind.

4. Ergebnisse

Die Methode wurde auf drei verschiedenen Aufgabenbereichen evaluiert: Mathematisches Schlussfolgern (lange Generierung), Dialog (Personalisierung) und Code-Generierung.

• Mathematisches Schlussfolgern (Off-Policy):

  • HPD übertrifft konsistent SFT, SeqKD, RKLD und JSD auf Benchmarks wie AIME, AMC und GPQA.
  • Besonders beeindruckend ist die Leistung von kleineren Modellen (z. B. Qwen 2.5 3B), die durch HPD Fähigkeiten erreichen, die mit deutlich größeren Modellen vergleichbar sind (Verbesserung von ~28% auf ~40% im Durchschnitt).
  • HPD verhindert den „Entropy Collapse" (Zusammenbruch der Entropie), der bei SFT häufig auftritt, und sorgt für eine stabile Konvergenz.

• Dialog und Personalisierung (Off-Policy):

  • Auf Benchmarks wie MT-Bench und Arena-Hard zeigt HPD die beste Leistung, insbesondere bei der Erhaltung von Mehr-Turn-Dialogfähigkeiten und Kontextverständnis.

• Code-Generierung (Off-Policy):

  • HPD erzielt die höchste durchschnittliche Pass-Rate auf HumanEval und MBPP und zeigt dabei eine höhere Robustheit über verschiedene Modellfamilien hinweg als andere KD-Methoden.

• On-Policy Szenario:

  • Als Initialisierung für On-Policy-Distillation (OPD) dient HPD als starkes Fundament. Modelle, die mit HPD initialisiert und dann mit OPD weiter trainiert werden, übertreffen die reinen SFT+OPD-Pipelines signifikant.

5. Bedeutung und Ausblick

HPD adressiert das zentrale Problem der Wissensdistillation bei LLMs: den Zielkonflikt zwischen Stabilität, Effizienz und der Qualität der Verteilungsanpassung.

• Praktische Relevanz: Die Methode ermöglicht das Training kleiner, effizienter Modelle, die die Leistung großer Lehrermodelle in komplexen Aufgaben (Mathematik, Code, Dialog) nahezu erreichen, ohne den enormen Rechenbedarf von Reinforcement Learning oder vollständigen On-Policy-Rollouts.
• Zukunftsperspektive: Die Autoren schlagen vor, HPD bereits in früheren Trainingsphasen (Mid-Training oder Pre-Training) einzusetzen, um das Potenzial für die Skalierung von LLMs weiter zu untersuchen.

Zusammenfassend bietet HPD einen robusten, recheneffizienten und leistungsfähigen Rahmen für die Kompression von LLMs, der die Grenzen zwischen Off-Policy- und On-Policy-Lernen sowie zwischen verschiedenen Divergenz-Maßen überbrückt. Der Code ist öffentlich verfügbar.