🤖 machine learning

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

Dieser Beitrag stellt Learned Relay Representations (Relay) vor, eine Methode, die es Masked Diffusion Models ermöglicht, latente Informationen über einen differentierbaren pro-Token-Kanal hinweg über die Entrauschungsschritte hinweg zu propagieren, wodurch die Leistung bei Planungs- und Codierungsaufgaben verbessert und gleichzeitig die Inferenz-Latenz erheblich reduziert wird.

Ursprüngliche Autoren: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Veröffentlicht 2026-05-25

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

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

Das Problem: Die „Hard Reset"-Amnesie

Stellen Sie sich vor, Sie versuchen, ein komplexes Rätsel zu lösen, wie ein Sudoku oder das Schreiben eines Code-Stücks. Sie nutzen dabei einen sehr intelligenten Assistenten (ein KI-Modell) zur Unterstützung.

Derzeit arbeitet dieser Assistent wie eine Person mit Amnesie, die im Moment, in dem sie einen einzigen Satz beendet, alles vergisst.

Der Assistent betrachtet eine leere Seite voller Fragezeichen.
Er denkt intensiv nach, führt komplexe Kopfrechnungen durch und entscheidet sich, ein oder ein paar Buchstaben einzufügen.
Der Hard Reset: Sobald er diese Buchstaben niedergeschrieben hat, wirft er sofort alle komplexen Gedanken, Berechnungen und das „Gefühl", das er für den Rest der Seite hatte, über Bord.
Für den nächsten Schritt beginnt er völlig von vorne, betrachtet nur die Buchstaben, die er gerade geschrieben hat, und vergisst dabei völlig die reiche interne Arbeit, die er gerade geleistet hat.

Dies wird als „Hard Reset"-Problem bezeichnet. Das Paper argumentiert, dass dies ineffizient ist. Es ist wie ein Koch, der eine Suppe probiert, „Salz benötigt" aufschreibt und dann sofort den Geschmack der Brühe vergisst, bevor er das Salz hinzufügt. Er muss die gesamte Suppe jedes Mal, wenn er eine Zutat hinzufügt, von vorne probieren.

Die Lösung: Der „Relay"-Stab

Die Autoren schlagen eine neue Methode namens Relay vor.

Stellen Sie sich einen Staffellauf vor. Anstatt dass der Läufer anhält, das Rennen vergisst und von vorne beginnt, übergibt er einen Stab an den nächsten Läufer. Dieser Stab trägt die Schwungmasse, die Strategie und die Müdigkeit des vorherigen Läufers, damit der nächste genau dort weitermachen kann, wo der vorherige aufgehört hat.

In der Welt der KI ist der „Stab" ein kontinuierlicher Informationsstrom (ein verborgener Zustand), den das Modell mit sich führt.

Vor Relay: Das Modell berechnet, schreibt ein Token und löscht die Berechnung.
Mit Relay: Das Modell berechnet, schreibt ein Token und gibt eine „Notiz" für den nächsten Schritt an sich selbst weiter. Diese Notiz lautet: „Hey, ich habe über dieses spezifische Muster nachgedacht, und ich bin zu 80 % sicher bezüglich dieses nächsten Teils."

Wie es funktioniert: Das Weitergeben des Stabs lernen

Das Paper stellt einige wichtige Tricks vor, um dies zu ermöglichen:

Das „Forward-Thinking"-Training:
Normalerweise werden KI-Modelle so trainiert, dass sie nur die aktuelle Antwort richtig haben. Die Autoren haben das Modell trainiert, „forward-thinking" (zukunftsorientiert) zu sein. Sie lehrten das Modell: „Schreibe nicht nur heute den richtigen Buchstaben; schreibe eine Notiz (das Relay), die dir hilft, morgen den richtigen Buchstaben zu schreiben."
Truncated Backpropagation (BPTT):
Dies ist ein komplizierter mathematischer Begriff für „Vorausschauen zum Lernen". Stellen Sie sich vor, Sie üben eine Tanzroutine. Anstatt nur eine Bewegung zu üben, üben Sie eine Sequenz von drei Bewegungen. Wenn Sie bei der dritten Bewegung stolpern, merken Sie: „Oh, ich habe die erste Bewegung falsch gemacht, weil ich mich nicht auf die dritte vorbereitet habe."
Das Paper verwendet eine Methode namens truncated BPTT, um dem Modell während des Trainings zu ermöglichen, ein paar Schritte voraus zu schauen. Es lernt, Informationen speziell weiterzugeben, um zukünftige Schritte zu erleichtern.
Der „Soft"-Kanal:
Die weitergegebene Information ist nicht nur ein Wort oder ein Buchstabe; es ist eine „weiche" Zahl (ein kontinuierlicher Wert). Denken Sie daran wie an einen Dimmer-Schalter und nicht an ein Ein/Aus-Licht. Es ermöglicht dem Modell, subtile Hinweise und Wahrscheinlichkeiten mitzuführen, die noch nicht in eine endgültige Entscheidung umgewandelt wurden.

Die Ergebnisse: Schneller und intelligenter

Die Autoren testeten dies an zwei Dingen:

Sudoku-Rätsel: Sie behandelten Sudoku wie eine Planungsaufgabe. Das „Relay"-Modell löste Rätsel mit weniger Versuchen (weniger „Forward Passes") und machte weniger Fehler als das Standardmodell. Es war besser darin, das „große Ganze" im Kopf zu behalten, während es das Raster ausfüllte.
Programmieren (Software schreiben): Sie nahmen ein hochmodernes KI-Modell für Programmieren (Fast-dLLM v2) und fügten das Relay-System hinzu.
- Bessere Genauigkeit: Es schrieb besseren Code.
- Schnellere Geschwindigkeit: Es musste seine „Denkmaschine" 32 % seltener ausführen, um das gleiche Ergebnis zu erzielen.

Warum dies wichtig ist

Das Paper behauptet, dass wir durch das Stoppen des „Hard Reset" und das Erlauben, dass das Modell seine eigenen „Gedanken" wie einen Relay-Stab weiterträgt, KI-Modelle folgendermaßen verbessern können:

Intelligenter: Sie können über lange Sequenzen hinweg besser reasoning (logisch schlussfolgern), wie bei komplexem Code oder Logikrätseln.
Schneller: Sie müssen nicht jedes Mal alles von vorne berechnen, was Zeit und Energie spart.

Kurz gesagt: Das Paper lehrt KI-Modelle, ihre eigene Arbeit nicht mehr zu vergessen und stattdessen Schritt für Schritt darauf aufzubauen, genau wie ein Mensch, wenn er ein schwieriges Problem löst.

Technische Zusammenfassung: Gelernte Relais-Repräsentationen für vorausschauende diskrete Diffusionsmodelle

1. Problemstellung: Das „harte Zurücksetzen" in Maskierten Diffusionsmodellen

Maskierte Diffusionsmodelle (MDMs) erzeugen diskrete Sequenzen durch iteratives Entrauschen. Ausgehend von einer vollständig maskierten Sequenz unmaskiert das Modell schrittweise einen Bruchteil der Positionen. Allerdings besteht in der Standard-Inferenz von MDMs eine kritische Ineffizienz: Nach jedem Vorwärtsdurchlauf verwirft das Modell die reichhaltigen internen kontinuierlichen Berechnungen (verborgene Zustände), die für maskierte Positionen generiert wurden. Der nächste Schritt beginnt ausschließlich mit den neu festgelegten diskreten Tokens.

Die Autoren bezeichnen dies als das Problem des harten Zurücksetzens. Obwohl MDMs mehrere Vorwärtsdurchläufe pro Generierung durchführen, verhindert die Unfähigkeit, Zwischenberechnungen wiederzuverwenden, eine effektive Amortisierung des Denkens über die Schritte hinweg. Dies steht im Gegensatz zur rekurrenten Berechnung, bei der gezeigt wurde, dass das Entrollen eines Modells mit festen Parametern über die Schritte hinweg die Funktionsklasse erweitert, die das Modell approximieren kann, was die Leistung bei komplexen Denkaufgaben verbessert. Die zentrale Frage lautet, wie MDMs so erweitert werden können, dass sie rekurrente Berechnungen unterstützen, die reichhaltigere Informationen über sequenzielle Unmaskierungsschritte hinweg tragen, ohne Zwischenzustände zu verwerfen.

2. Methodik: Gelernte Relais-Repräsentationen (Relay)

Die vorgeschlagene Lösung, Relay, führt einen kontinuierlichen, differenzierbaren latenten Zustand ein, der über die Dekodierungsschritte hinweg weitergegeben wird und das Modell „vorausschauend" macht.

Kernmechanismus

Erweiterter Zustandsverlauf: Zu jedem Inferenzschritt $k$ führt das Modell einen erweiterten Zustand $s_k = (x_{t_k}, h_k)$ , wobei $x_{t_k}$ die aktuelle teilweise maskierte Sequenz und $h_k$ ein gelernter Relaiszustand (kontinuierlicher latenter Vektor) ist.
Architektur: Das Modell besteht aus einem Backbone $f_\theta$ $f_{θ}$ , Token-Embeddings $Emb_\theta$ $E m b_{θ}$ , einem Unembedding-Head $UnEmb_\theta$ $U n E m b_{θ}$ und einem spezifischen Relais-Modul $R_\theta$ $R_{θ}$ .
- Der nächste verborgene Zustand wird berechnet als: $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$ .
- Der Relaiszustand $h_k$ wird projiziert und zu den eingebetteten Tokens addiert, bevor sie in den Backbone eintreten.
- Die Ausgabe-Logits $\ell_k$ werden über $UnEmb_\theta$ aus $h_{k+1}$ abgeleitet.
Training durch abgeschnittene BPTT: Im Gegensatz zu Standard-MDMs, die mit unabhängigen Schritten trainiert werden, wird Relay mittels Abschnittsweise Rückwärtspropagierung durch die Zeit (Truncated Backpropagation Through Time, BPTT) trainiert.
- Das Modell führt einen Rollout von $K$ Schritten (z. B. $K=2$ ) beginnend mit einer vollständig maskierten Sequenz durch.
- Während des Rollouts werden Positionen basierend auf einer Strategie (z. B. Konfidenzschwellenwerte) unmaskiert, und die festgelegten Tokens werden durch „Teacher Forcing" auf die wahren Werte gesetzt, um eine Fehlerakkumulation während des Rollouts zu verhindern.
- Der Verlust ist die Summe der Kreuzentropieverluste über die $K$ Schritte.
- Gradienten werden durch den Relaispfad ( $h_k \to h_{k+1}$ ) zurückpropagiert, wodurch das Modell nicht nur lernt, den aktuellen Schritt vorherzusagen, sondern auch Informationen in $h_k$ zu kodieren, die für zukünftige Schritte maximal nützlich sind.

Inferenz

Zum Zeitpunkt der Inferenz bleibt der Prozess hinsichtlich des Unmaskierungsplans oder der Sampling-Strategie weitgehend unverändert. Die einzige Ergänzung besteht darin, dass das Modell den Relaiszustand $h_k$ gemeinsam mit den festgelegten Tokens weiterführt. Der Relaiszustand dient als differenzierbarer Speicherkanal für Informationen, die noch nicht in diskrete Tokens kodiert wurden.

3. Hauptbeiträge

Relay-Framework: Eine allgemeine Methode zur Integration rekurrenter Berechnungen in MDMs, indem das Modell trainiert wird, einen gelernten latenten Relaiszustand über die Dekodierungsschritte hinweg mittels abgeschnittener BPTT weiterzugeben. Es ist architekturagnostisch und kompatibel mit bestehenden MDM-Beschleunigungstechniken wie Block-Diffusion und KV-Caching.
Skalierbarkeit und Anpassung: Validierung von Relay im Maßstab großer Sprachmodelle (LLM) durch vollständige Parameteranpassung von Fast-dLLM v2 (1,5B). Die Methode passt erfolgreich ein State-of-the-Art-Diskretes Sprachmodell (DLM) mit minimalen zusätzlichen Trainings-FLOPs an.
Leistungs-Latenz-Verbesserungen: Empirischer Nachweis, dass Relay die Pareto-Grenze für Leistung und Latenz vorantreibt und dabei Standard-Supervised Fine-Tuning (SFT) bei Codierungsaufgaben übertrifft, während die Inferenzlatenz signifikant reduziert wird.
Analyse des Gestaltungsraums: Umfangreiche Ablationsstudien an einer anspruchsvollen Sudoku-basierten Planungsaufgabe, die die Beiträge des Rollout-Verfahrens, des Relaiskanals und des BPTT-Trainingsziels isolieren.

4. Experimentelle Ergebnisse

Sudoku-Planungsaufgabe

Unter Verwendung des Sudoku-Extreme-Datensatzes verglichen die Autoren vier Trainingsziele:

MLM: Standardmäßige uniforme maskierte Diffusion (einzelner Schritt, kein Relais).
Rollout: Mehrstufiger Rollout mit Teacher Forcing, aber ohne Relaiskanal ( $R_\theta \equiv 0$ ).
Relay (sg): Beinhaltet den Relaispfad, verwendet jedoch Stop-Gradients (keine BPTT durch das Relais).
Relay: Vollständige Methode mit $K=2$ BPTT durch das Relais.

Ergebnisse:

Genauigkeits-NFE-Grenze: Relay erzielte den besten Kompromiss zwischen Genauigkeit und Anzahl der Vorwärtsbewertungen (NFE).
Auswirkung der Komponenten: Der Übergang von MLM zu Rollout verbesserte die Leistung. Das Hinzufügen des Relaiskanals (Relay sg) führte zu einem signifikanten Sprung. Schließlich lieferte die Aktivierung von BPTT (Relay) die besten Ergebnisse und zeigte, dass das Modell lernt, das teilweise Brett unter aggressiverem Unmaskieren konsistent zu halten.
Effizienz: Bei einem festen Konfidenzschwellenwert legte Relay pro Vorwärtsdurchlauf mehr Zellen fest und behielt dabei die Legalität bei (keine Verletzungen von Zeilen/Spalten/Boxen) und erreichte die gleiche Genauigkeit in weniger Gesamtvorwärtsdurchläufen im Vergleich zu Baselines.

Anpassung vortrainierter Modelle (Fast-dLLM v2)

Die Autoren passten Fast-dLLM v2 (1,5B) unter Verwendung von Relay auf einer Mischung aus Code- und Mathematikdatensätzen an.

Codierungsleistung: Auf den Benchmarks HumanEval und MBPP übertraf Relay das Standard-SFT.
- HumanEval: Relay erreichte 42,1 % Genauigkeit (vs. 38,4 % für SFT) mit 88,3 NFE (vs. 130,7 für SFT).
- MBPP: Relay erreichte 46,6 % Genauigkeit (vs. 43,9 % für SFT) mit 78,8 NFE (vs. 84,8 für SFT).
Latenz: Relay reduzierte die Inferenzlatenz um bis zu 32 % im Vergleich zu Standard-SFT, während gleichzeitig die Genauigkeit verbessert wurde.
Speicher-Overhead: Profiling auf einer A100-GPU zeigte, dass Relay zwar den Speicher für lebende Aktivierungen während des zweiten Vorwärtsdurchlaufs erhöht, der Spitzen-Speicherverbrauch jedoch nahezu identisch mit dem von Vanilla-SFT bleibt (ca. 20,1 GiB vs. 21,2 GiB), da der Engpass der Rückwärtsdurchlauf der Vokabularprojektion ist und nicht die Speicherung des Relaiszustands.

5. Bedeutung und Behauptungen

Die Arbeit behauptet, dass Relay die Natur der iterativen nicht-autoregressiven Generierung grundlegend verändert, indem es MDMs ermöglicht, Zwischenberechnungen in kontinuierlicher Form zu akkumulieren.

Vorausschauend: Die Methode ermöglicht es Modellen, Zwischenrepräsentationen explizit für zukünftige Unmaskierungsentscheidungen zu optimieren, anstatt jeden Schritt als ein neues Vorhersageproblem zu behandeln.
Vorantreiben der Pareto-Grenze: Durch die Reduzierung der Anzahl der Vorwärtsdurchläufe, die erforderlich sind, um eine Zielgenauigkeit zu erreichen, verbessert Relay die Leistung-Latenz-Grenze und macht diskrete Diffusionsmodelle effizienter für komplexe Denkaufgaben und Generierungsaufgaben.
Kompatibilität: Der Ansatz wird als Plug-in-Erweiterung präsentiert, die mit bestehenden State-of-the-Art-DLMs und Beschleunigungstechniken (z. B. Block-Diffusion, KV-Caching) funktioniert, ohne dass architektonische Überholungen erforderlich sind.

Die Autoren schließen, dass Relay einen bedeutenden Schritt in Richtung nicht-gieriger, vorausschauender Generierung darstellt und ein kontinuierliches Substrat für latentes Denken bietet, das weiter für Interpretierbarkeit und Anwendung in anderen Modalitäten (z. B. Bilder, Moleküle) erforscht werden kann.