Self-Distillation for Multi-Token Prediction

Die Arbeit stellt MTP-D vor, eine effiziente Selbst-Distillationsmethode mit einer Schleifen-Erweiterungsstrategie, die die Akzeptanzraten von Multi-Token-Vorhersage-Köpfen signifikant steigert und die Inferenzgeschwindigkeit von Large Language Models ohne wesentlichen Trainingsaufwand erhöht.

Das Wesentliche

Stell dir vor, ein großes Sprachmodell (wie ein sehr kluger KI-Assistent) ist wie ein Schreibmaschinen-Genie, das Sätze Wort für Wort schreibt.

Das Problem: Es schreibt sehr langsam. Warum? Weil es nach jedem Wort eine Pause macht, um nachzudenken: „Was kommt als Nächstes?" Erst dann tippt es das nächste Wort. Bei langen Texten dauert das ewig.

Die Forscher von Tencent haben eine Lösung namens MTP-D entwickelt. Hier ist die Erklärung, wie das funktioniert, ganz einfach und mit ein paar Bildern aus dem Alltag:

1. Das alte Problem: Der einsame Vorhersager

Früher (und bei vielen aktuellen Modellen) hatte das Genie nur einen Kopf, der das nächste Wort vorhersagte.

• Neuer Ansatz (MTP): Man gibt dem Genie mehrere Köpfe (Helfer). Der erste sagt das nächste Wort voraus, der zweite das Wort danach, der dritte das Wort nach dem nächsten usw.
• Der Vorteil: Statt auf ein Wort zu warten, kann das Genie jetzt mehrere Wörter auf einmal „vorschlagen". Wenn diese Vorschläge richtig sind, spart man sich das ständige Nachdenken und wird viel schneller.

Aber hier liegt das Problem:
Die neuen Helfer-Köpfe sind oft nicht so schlau wie der Hauptkopf. Sie machen viele Fehler. Wenn der Hauptkopf sagt: „Nein, das war falsch!", müssen wir alle vorherigen Vorschläge verwerfen und von vorne anfangen. Das nennt man eine niedrige Akzeptanzrate. Es ist, als würde ein Assistent 10 Wörter vorschlagen, aber der Chef nur das erste akzeptieren. Der Rest war Zeitverschwendung.

2. Die Lösung: MTP-D (Der „Lehrmeister"-Effekt)

Die Forscher haben eine clevere Methode namens Selbst-Distillation (MTP-D) erfunden.

Stell dir vor, der Hauptkopf ist ein erfahrener Lehrer, und die Helfer-Köpfe sind seine Schüler.

• Das alte Training: Die Schüler lernten nur aus dem Buch (den richtigen Antworten). Sie waren oft unsicher und machten Fehler.
• Das neue Training (MTP-D): Die Schüler schauen dem Lehrer genau über die Schulter. Sie sehen nicht nur das richtige Wort, sondern auch, welche Wörter der Lehrer am wahrscheinlichsten für richtig hielt, bevor er sich entschied.
• Die Technik: Die Forscher lassen die Schüler die „Gedanken" des Lehrers kopieren (die Wahrscheinlichkeiten der Top-Wörter), aber sie lassen den Lehrer selbst in Ruhe (damit er nicht verwirrt wird).

Das Ergebnis: Die Schüler werden fast so schlau wie der Lehrer. Sie machen viel weniger Fehler. Wenn sie jetzt mehrere Wörter auf einmal vorschlagen, stimmt der Chef viel öfter zu.

3. Der „Schleifen-Trick" (Looped Extension)

Jetzt wird es noch genialer. Was, wenn wir nicht nur 4 Helfer haben wollen, sondern 16?
Normalerweise wäre das Training von 16 Köpfen extrem teuer und schwierig.

Die Forscher nutzen einen Trick namens „Looped Extension" (Schleifen-Erweiterung):

• Stell dir vor, du hast eine Gruppe von 4 gut trainierten Schülern.
• Du kopierst diese Gruppe und setzt sie als neue Gruppe dahinter.
• Weil die ersten 4 schon so gut sind, sind die neuen 4 (die Kopien) schon halbwegs fit.
• Du trainierst sie nur noch ein wenig weiter, um sie perfekt aufeinander abzustimmen.

Die Analogie: Es ist wie eine Fabrik-Kette. Wenn du eine gute Produktionslinie hast, kannst du einfach eine identische Kopie dahinterstellen. Die neue Linie muss nicht bei Null anfangen, sie startet mit dem Wissen der alten. So kannst du die Geschwindigkeit massiv steigern, ohne das ganze System neu zu erfinden.

4. Was bringt das alles?

• Geschwindigkeit: Durch diese Methode wird die KI bis zu 2,5-mal schneller (in manchen Tests sogar noch mehr). Das bedeutet, du bekommst Antworten viel schneller.
• Qualität: Die KI wird nicht dumm. Der Hauptkopf (der Chef) bleibt genau so schlau wie vorher.
• Kosten: Es kostet kaum mehr Rechenleistung beim Training, aber bringt einen riesigen Geschwindigkeitsvorteil beim Nutzen.

Zusammenfassung in einem Satz

Die Forscher haben eine KI so trainiert, dass ihre „Vorschlags-Abteilungen" dem Chef genau so gut denken wie er selbst, und sie haben eine Methode gefunden, diese Abteilungen wie eine Reihung von Kopien zu vermehren – was die KI extrem schnell macht, ohne ihre Intelligenz zu verlieren.

Kurz gesagt: Sie haben aus einem langsamen, vorsichtigen Schreiber ein schnelles, koordiniertes Team gemacht, das im Takt arbeitet.

Technische Zusammenfassung

Problemstellung

Mit dem Skalieren von Large Language Models (LLMs) wird die Inferenzeffizienz zu einem kritischen Engpass. Die meisten Modelle nutzen das Paradigma der Next-Token Prediction (NTP), bei dem Token sequentiell generiert werden, was zu hoher Latenz und Rechenkosten führt, insbesondere bei langen Sequenzen.
Multi-Token Prediction (MTP) wurde als Lösung vorgeschlagen, um mehrere zukünftige Token parallel vorherzusagen und so die Inferenz zu beschleunigen. Bestehende Ansätze (wie bei DeepSeek-V3) stoßen jedoch auf zwei Hauptprobleme:

1. Begrenzte Akzeptanzraten: Die von den MTP-Köpfen generierten Token werden vom Hauptkopf (Main Head) oft nicht akzeptiert. Da die kumulative Akzeptanzrate exponentiell mit der Anzahl der vorhergesagten Token abfällt, führt eine geringe Einzelakzeptanz zu einem Zusammenbruch des Beschleunigungseffekts.
2. Schwierigkeiten beim gemeinsamen Training: Das gleichzeitige Trainieren mehrerer MTP-Köpfe mit dem Hauptkopf ist komplex. Oft führt dies zu einem „Wippen-Effekt" (Seesaw effect), bei dem die Verbesserung der MTP-Köpfe die Leistung des Hauptkopfes verschlechtert, was in der Praxis inakzeptabel ist.

Methodik: MTP-D und Looped Extension

Die Autoren schlagen MTP-D (Self-Distillation for Multi-Token Prediction) vor, eine einfache, aber effektive Methode, die diese Probleme adressiert, ohne die Hauptleistung zu beeinträchtigen.

1. Self-Distillation im Pre-Training

MTP-D nutzt eine selbstdistillierende Strategie während des Pre-Trainings, um die MTP-Köpfe an den Hauptkopf anzupassen:

• Gradient-Entkopplung (Gradient-Detached): Der Hauptkopf dient als Lehrer (Teacher), die MTP-Köpfe als Schüler (Students). Um zu verhindern, dass Gradienten vom Hauptkopf zurückfließen und dessen Optimierung stören, werden die Logits des Hauptkopfes für den Distillationsverlust „stop-gradient" (abgekoppelt).
• TopN-Logits-Auswahl: Da moderne LLMs sehr große Vokabulare haben (z. B. 122.880 Token), ist eine vollständige Distillation über das gesamte Vokabular rechenintensiv und numerisch instabil. Die Autoren wählen stattdessen nur die TopN (z. B. 10.000) wahrscheinlichsten Token des Hauptkopfes aus, um den KL-Divergenz-Verlust (Kullback-Leibler) zu berechnen. Dies reduziert den Overhead erheblich und konzentriert sich auf die relevanten Wahrscheinlichkeitsverteilungen.
• Verlustfunktion: Der Gesamtverlust besteht aus dem herkömmlichen Cross-Entropy-Verlust (gegenüber den Ground-Truth-Token) und dem zusätzlichen KL-Divergenz-Verlust, der die Verteilung der MTP-Köpfe an die des Hauptkopfes angleicht.

2. Looped Extension Strategy (Schleifen-Erweiterung)

Um die Anzahl der MTP-Köpfe kosteneffizient zu erhöhen (Skalierung), schlagen die Autoren eine looped extension-Strategie vor:

• Anstatt neue Köpfe von Grund auf zu trainieren, werden bereits trainierte Gruppen von MTP-Köpfen (z. B. Köpfe 1 bis 4) kopiert, um neue Gruppen (Köpfe 5 bis 8) zu initialisieren.
• Anschließend erfolgt ein Weiter-Pre-Training (Continued Pre-Training) nur auf den neuen Köpfen, während der Hauptmodell und die bereits trainierten Köpfe eingefroren (frozen) bleiben.
• Dies nutzt die inhärente strukturelle Konsistenz der MTP-Architektur aus und ermöglicht eine Skalierung auf bis zu 16 Köpfe mit minimalem zusätzlichem Rechenaufwand.

Wichtige Beiträge

1. MTP-D Framework: Ein neuartiges Self-Distillation-Framework, das die Akzeptanzraten der MTP-Köpfe signifikant erhöht, während die Leistung des Hauptkopfes erhalten bleibt und die Trainingskosten nur marginal steigen.
2. Looped Extension: Eine Strategie zur kosteneffizienten Skalierung der Anzahl der MTP-Köpfe durch Wiederverwendung trainierter Gewichte und Weiter-Pre-Training.
3. Systematische Validierung: Umfassende Experimente auf sieben Benchmarks, die zeigen, dass die Methode nicht nur die Akzeptanzraten verbessert, sondern auch die Skalierbarkeit von MTP-Architekturen aufzeigt.

Ergebnisse

Die Experimente wurden auf 2B-Dense- und 10B-MoE-Modellen (Mixture of Experts) durchgeführt:

• Steigerung der Akzeptanzraten: MTP-D mit 4 Köpfen erhöht die Akzeptanzraten der MTP-Köpfe um durchschnittlich 7,5 % im Vergleich zu herkömmlichem MTP, bei vergleichbarer Hauptkopf-Leistung.
• Inferenz-Beschleunigung:
  • Bei 1 Kopf ergibt sich eine Beschleunigung von ca. 14 %.
  • Bei 4 Köpfen steigt die Beschleunigung auf 22,9 %.
  • Durch die Looped Extension auf 16 Köpfe wird eine maximale Beschleunigung von 220,4 % (im Vergleich zur 1-Kopf-Konfiguration) erreicht.
• Skalierbarkeit: Während herkömmliches MTP bei der Erweiterung auf 8 Köpfe (ohne Training) in der kumulativen Akzeptanzrate auf unter 1 % einbricht, hält MTP-D Werte von über 26 % aufrecht und ermöglicht so eine stabile Skalierung.
• Daten-Effizienz: Die Looped Extension benötigt nur sehr wenig zusätzliche Daten (70B Token im Vergleich zu 350B für das Initial-Training), um signifikante Geschwindigkeitsgewinne zu erzielen.

Bedeutung und Ausblick

Das Paper zeigt einen praktischen Weg auf, um die Inferenzgeschwindigkeit von LLMs drastisch zu erhöhen, ohne die Modellqualität zu opfern.

• Praktische Anwendbarkeit: Die Methode ist leicht in bestehende Pre-Training-Pipelines integrierbar und erfordert keine komplexen Architekturrevolutionen.
• Skalierungspotenzial: Die Ergebnisse belegen, dass MTP-Architekturen durch Self-Distillation und Looping deutlich besser skalieren als bisher angenommen, was für zukünftige sehr große Modelle und lange Kontexte entscheidend ist.
• Zukünftige Arbeit: Die Autoren planen, die Methode auf Post-Training-Szenarien (Fine-Tuning) zu übertragen und die Skalierbarkeit auf noch größere Modelle und diverse Datensätze zu validieren.

Zusammenfassend stellt MTP-D einen effizienten Hebel dar, um das „Token-by-Token"-Paradigma zu überwinden und die nächste Generation schnellerer und leistungsfähigerer LLMs zu ermöglichen.