He Kaiming's erstes Sprachmodell: 105 Millionen Parameter, kein Rückgriff auf den altmodischen Weg der GPT-Autoregressionsmethode

He Kaiming hat sich auch dem Bereich der Sprachmodelle gewidmet.

Allerdings handelt es sich diesmal nicht um das bekannte autoregressive Paradigma der „Vorhersage des nächsten Tokens“ (next token prediction), wie es hinter ChatGPT steht.

Sondern um einen anderen, in den letzten Jahren im Bildbereich sehr erfolgreichen und nun immer mehr in die Textgenerierung übernommenen neuen Ansatz: Diffusionssprachmodelle (Diffusion Language Model, DLM).

In ihrer neuesten Publikation hat das Team um He Kaiming ein neues kontinuierliches Diffusionssprachmodell vorgestellt: ELF: Embedded Language Flows.

Im Gegensatz zu vielen anderen Sprachmodellen, die noch auf Token-Ebene diffundieren, belässt ELF den gesamten Generierungsprozess im kontinuierlichen Embedding-Raum und diskretisiert erst im letzten Schritt, um die Repräsentation wieder in Tokens umzuwandeln.

Dank dieser Architektur hat ELF mit nur 105 Mio. Parametern, 45 Mrd. Trainings-Tokens und 32 Sampling-Schritten eine Reihe von führenden Diffusionssprachmodellen geschlagen.

Ein direktes Maß für die Leistung ist die Generative Perplexity auf dem OpenWebText-Datensatz, die ELF auf 24 reduziert hat.

Zur Erklärung: Die Generative Perplexity lässt ein starkes Sprachmodell die generierten Texte bewerten, um zu prüfen, wie natürlich sie klingen. Ein niedrigerer Wert bedeutet eine höhere Qualität und weniger Künstlichkeit.

Im Vergleich zu anderen Diffusionssprachmodellen hat ELF bei fast zehnmal weniger Trainings-Tokens und weniger Sampling-Schritten eine geringere Generative Perplexity erreicht.

In den letzten Jahren hat sich der Fortschritt bei Diffusionssprachmodellen fast ausschließlich auf die diskreten DLM (Discrete DLM) beschränkt.

ELF zeigt erstmals, dass ein kontinuierlicher Ansatz nicht nur funktioniert, sondern auch gute Ergebnisse liefert.

Was macht ELF?

Um ELF zu verstehen, muss man zunächst verstehen, was Diffusionssprachmodelle überhaupt tun.

Es gibt zwei Hauptansätze bei Diffusionssprachmodellen. Der erste ist die diskrete Richtung, repräsentiert durch Modelle wie MDLM und Duo, die direkt im Token-Raum diffundieren und diskrete Zufallsvariablen verarbeiten.

Der zweite ist die kontinuierliche Richtung, zu der Diffusion-LM, CDCD und DiffuSeq gehören. Diese Modelle transformieren Tokens in kontinuierliche Embeddings und entfernen Rauschen im kontinuierlichen Raum.

In früheren Studien haben diskrete Ansätze wie MDLM, LLaDA und Dream 7B den Vorsprung gehabt, da Sprache an sich diskret ist.

Das Team um He Kaiming sieht es jedoch anders: Das Problem liegt möglicherweise nicht darin, dass Sprache diskret sein muss, sondern darin, dass frühere Ansätze den kontinuierlichen Ansatz nicht konsequent verfolgt haben.

Methoden wie Diffusion-LM entfernen Rauschen im Embedding-Raum, berechnen aber bei jedem Schritt die Kreuzentropie auf Token-Ebene, was die kontinuierlichen Trajektorien an das Vokabular bindet.

Spätere Modelle wie LD4LG und Cosmos verwenden latent Diffusion, was den Rauschentfernungsprozess kontinuierlich macht, erfordert aber ein zusätzliches Decoder-Modul, um die latenten Repräsentationen in Tokens umzuwandeln.

Im Gegensatz dazu behält ELF den gesamten Rauschentfernungsprozess im kontinuierlichen Embedding-Raum und projiziert erst im letzten Schritt t = 1 in den Token-Raum zurück.

Genauer gesagt werden bei der Trainingsphase diskrete Tokens in kontinuierliche Embeddings codiert und dann mit Rauschen versehen, um z_t zu erhalten. Das Modell kann entweder die sauberen Embeddings rekonstruieren (MSE) oder direkt die Tokens vorhersagen (CE).

Beim Inference-Prozess startet das Modell mit Gaußischem Rauschen z_0 und entfernt Rauschen kontinuierlich, bis es im letzten Schritt in den Decode-Modus wechselt und die Embeddings in Tokens umwandelt.

ELF trennt erstmals die Probleme von „kontinuierlicher Repräsentation“ und „diskreter Ausgabe“, die bisher immer als miteinander in Übereinstimmung zu bringen galten:

Die Rauschentfernung findet vollständig im kontinuierlichen Raum statt, und die endgültige Sprachgenerierung erfolgt erst im letzten Schritt der Diskretisierung.

Ohne ständige Anpassung an das Vokabular und ohne zusätzliches Decoder-Modul ist der Generierungsprozess erstmals wirklich kontinuierlich und diskret getrennt.

Dies ist auch der Schlüssel dafür, dass ELF mit weniger Sampling-Schritten und weniger Trainings-Tokens viele andere Diffusionssprachmodelle schlagen kann.

ELF ist nicht „erst diffundieren, dann decodieren“.

Bei der Implementierung löst ELF drei Probleme:

Wie werden Tokens kontinuierlich? Wie wird Rauschen im kontinuierlichen Raum entfernt? Und wie werden die Tokens am Ende wieder hergestellt?

Transformation von Tokens in kontinuierliche Embeddings

Um kontinuierliche Diffusion in der Sprache anzuwenden, muss man zunächst diskrete Tokens in kontinuierliche Repräsentationen umwandeln.

Im Paper teilt ELF die Eingabe in Token-Sequenzen und projiziert sie in den kontinuierlichen Embedding-Raum. Es gibt verschiedene Möglichkeiten, wie diese Projektion erfolgen kann.

Standardmäßig verwendet ELF einen vortrainierten T5-Encoder, um bidirektionale kontextuelle Embeddings zu erzeugen. Das Paper testet auch andere Ansätze wie gemeinsam trainierte Embeddings und zufällige Embeddings.

Wichtig zu beachten ist, dass dieser Encoder nur während der Trainingsphase verwendet wird und beim Inference-Prozess keine zusätzlichen Module hinzufügt.

Flow Matching im kontinuierlichen Embedding-Raum

Nachdem ELF kontinuierliche Repräsentationen erhalten hat, führt es Flow Matching im Embedding-Raum durch.

Einfach gesagt definiert Flow Matching eine kontinuierliche Trajektorie von Rauschen zu echten Daten:

Bei t = 0 handelt es sich um Gaußisches Rauschen;

Bei t = 1 sind es saubere Embeddings;

Alle Zwischenzustände sind lineare Interpolationen zwischen diesen beiden Zuständen, was im Paper als rectified flow bezeichnet wird.

Beim traditionellen Flow Matching prognostiziert das Netzwerk normalerweise direkt das „Geschwindigkeitsfeld“ v.

ELF folgt jedoch dem Ansatz, den das Team um He Kaiming vor sechs Monaten in der Publikation „Back to Basics: Let Denoising Generative Models Denoise“ vorgeschlagen hat –

Direkte Vorhersage der sauberen Embeddings x, also x-Prediction.

Das Trainingsziel besteht darin, die mittlere quadratische Abweichung (MSE) zwischen den vorhergesagten und den echten Embeddings zu minimieren.

Das Paper gibt zwei Gründe für die Wahl der x-Prediction:

Erstens ist sie in hochdimensionalen Repräsentationen wie 768-dimensionalen oder höheren Token-Embeddings stabiler. Zweitens stimmt sie mit dem Ziel der letzten Stufe, saubere Tokens vorherzusagen, überein.

Das Paper erwähnt auch, dass es zwar theoretisch möglich ist, zunächst die Geschwindigkeit v vorherzusagen und dann in x umzurechnen, aber dann ist es schwierig, die Gewichte zwischen Rauschentfernung und Decodierung zu teilen.

In Experimenten hat sich gezeigt, dass die Leistung der v-Prediction deutlich schlechter wird, wenn die Gewichte geteilt werden.

Zurückkehr von kontinuierlichen Embeddings zu diskreten Tokens

Die endgültige Ausgabe bei der Sprachgenerierung sind diskrete Tokens.

Deshalb projiziert ELF erst im letzten Zeitschritt (t = 1) die kontinuierlichen Embeddings wieder in den Token-Raum.

Im Gegensatz zu vielen latent Diffusion-Methoden trainiert ELF kein zusätzliches Decoder-Modul. Stattdessen wird der letzte Schritt als kontinuierliche zu diskrete Decodierung betrachtet.

Das bedeutet, dass Decoder und Rauschentferner das gleiche Netzwerk sind.

Um das Training im letzten Schritt nicht zu einfach zu machen (da die Eingabe bei t → 1 bereits sehr nahe an den sauberen Embeddings ist), fügt ELF im letzten Schritt eine Token-Ebene-Korruption hinzu, um eine gestörte Eingabe zu erzeugen.

Das gleiche Netzwerk gibt dann saubere Embeddings aus, die über eine lernbare Unembedding-Matrix W in Token-Logits projiziert werden.

Das Trainingsziel ist die standardmäßige Token-Ebene-Kreuzentropie-Verlustfunktion. Das gesamte Netzwerk teilt die gleichen Parameter und erhält zusätzlich ein binäres Modus-Token: Rauschentfernungsmodus/Decodierungsmodus.

Beim Inference-Prozess startet ELF mit Gaußischem Rauschen und entfernt Rauschen kontinuierlich, bis es im letzten Schritt t = 1 in den Decode-Modus wechselt und die endgültigen Tokens über argmax ausgibt.

Es ist erwähnenswert, dass auch die in der Bildgenerierung häufig verwendete Technik CFG (classifier-free guidance) in ELF implementiert wurde.

ELF verwendet Self-Conditioning als Bedingungssignal und wendet die Training-Time CFG an (eine Forward-Pass simuliert zwei Inference-Schritte, ohne zusätzlichen Inference-Aufwand), indem es das Verfahren aus der Bildgenerierung übernimmt.

Experimentelle Vergleiche

Im Experimentsteil beantwortet ELF eine seit zwei Jahren offene Frage:

Können kontinuierliche Diffusionssprachmodelle mithalten? Die Antwort ist: Ja, und sie gewinnen erstmals in Bezug auf Qualität, Geschwindigkeit und Trainingskosten.

Wie am Anfang erwähnt, hat ELF in der OpenWebText-Generierungsaufgabe ohne Distillation die Generative Perplexity auf 24 reduziert, indem es nur 32 Sampling-Schritte verwendet.

Frühere diskrete Diffusionsmodelle benötigten hingegen oft 1024 Schritte, um einen ähnlichen Wert zu erreichen.

Noch beeindruckender ist, dass ELF dies mit nur 45 Mrd. Trainings-Tokens erreicht hat.

Im Vergleich dazu verwenden vergleichbare Modelle normalerweise 500 Mrd. oder mehr. Das heißt, mit einer Größenordnung weniger Sampling-Schritten und einer Größenordnung weniger Trainingsdaten erzielt ELF bessere Ergebnisse.

Auch bei bedingten Generierungsaufgaben, bei denen viele Diffusionsmodelle Schwierigkeiten haben, schneidet ELF gut ab.

Ob bei der maschinellen Übersetzung von WMT14 oder bei der Textzusammenfassung von XSum, ELF übertrifft bestehende Diffusionssprachmodelle und sogar einige autoregressive Baseline-Modelle.