💬 NLP

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

Dieses Paper führt Order-Token Search ein, eine neuartige Dekodierungsmethode für Diffusion Language Models, die gemeinsam die Generierungsreihenfolge und den Token-Raum exploriert, um bestehende Baselines bei Benchmarks für mathematisches Denken und Programmierung zu übertreffen.

Ursprüngliche Autoren: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

Veröffentlicht 2026-02-06

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

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

Stellen Sie sich vor, Sie versuchen, ein komplexes Rätsel zu lösen, wie etwa eine Matheaufgabe oder eine Programmierherausforderung, aber Sie haben einen magischen Assistenten (das Diffusion Language Model), der nicht einfach von links nach rechts schreibt wie ein Mensch, der einen Buchstaben tippt. Stattdessen beginnt dieser Assistent mit einer leeren Seite voller „MASK“-Platzhalter und versucht, diese alle gleichzeitig auszufüllen, indem er rät, welche Wörter an welche Stelle gehören.

Das Problem ist, dass der Assistent ein wenig chaotisch ist. Er kann die Lücken in jeder beliebigen Reihenfolge füllen, die er möchte. Manchmal rät er das richtige Wort für die erste Lücke, bleibt dann aber stecken, weil er das falsche Wort für die zweite Lücke geraten hat. Ein anderes Mal rät er die richtigen Wörter, aber in der falschen Reihenfolge, was zu einer Sackgasse führt.

Der alte Weg: Raten und Beschneiden

Früher versuchten Menschen, dies auf zwei Hauptarten zu beheben, die beide Mängel aufwiesen:

Der „Selbstbewusste“ Ansatz: Der Assistent füllt zuerst die Lücken aus, bei denen er sich am sichersten fühlt. Es ist wie ein Wanderer, der nur auf dem Pfad geht, der am stabilsten aussieht.
- Das Gute: Er findet meist schnell eine gute Antwort.
- Das Schlechte: Wenn der „stabile Pfad“ zu einer Klippe führt (eine falsche Antwort), steckt der Wanderer fest. Er untersucht niemals andere Pfade, die vielleicht besser gewesen wären.
Der „Zufällige“ Ansatz: Der Assistent wählt die Lücken völlig zufällig aus, um sie zu füllen.
- Das Gute: Er erkundet eine riesige Vielfalt an Pfaden, sodass es sehr wahrscheinlich ist, dass er irgendwann auf die korrekte Lösung stößt.
- Das Schlechte: Er ist so zerstreut, dass er selten beim ersten Versuch den besten Pfad wählt. Es ist wie ein Wanderer, der im Kreis wandert; er wird das Gold vielleicht finden, aber er wird auch viele leere Löcher graben.

Die neue Lösung: Order-Token Search

Die Autoren dieses Papers haben eine neue Methode namens Order-Token Search eingeführt. Stellen Sie sich dies als ein Team von Entdeckern vor, die zusammenarbeiten.

Anstatt einen einzelnen Wanderer oder eine chaotische Menge auszusenden, schickt die Methode ein kleines Team (einen „Beam“) von Entdeckern aus. So arbeiten sie:

Divergierende Pfade (Die Suche): In regelmäßigen Abständen teilt sich das Team auf. Jeder Entdecker probiert eine andere Strategie aus:
- Entdecker A entscheidet sich, das erste fehlende Wort auszufüllen.
- Entdecker B entscheidet sich, das letzte fehlende Wort auszufüllen.
- Entdecker C probiert ein anderes Wort für die mittlere Stelle aus.
- Analogie: Sie erkunden sowohl das Wo des nächsten Schreibens (die Reihenfolge/Order) als auch das Was des Schreibens (das Token).
Die Bewertungskarte (Der Likelihood Estimator): Das ist der magische Teil. Das Team hat einen speziellen Richter (den Likelihood Estimator), der nicht nur auf die fertige Antwort schaut. Stattdessen betrachtet der Richter jeden Schritt, den die Entdecker gemacht haben.
- Hat der Entdecker einen logischen Zug gemacht?
- Ergibt dieser Teilsatz Sinn, wenn man berücksichtigt, was zuvor geschrieben wurde?
- Analogie: Stellen Sie sich einen Trainer vor, der einen Staffellauf beobachtet. Wenn ein Läufer früh stolpert, wartet der Trainer nicht, bis er das Rennen beendet hat, um ihn auszuschließen; er stoppt ihn sofort, weil der Schritt, den er gemacht hat, falsch war.
Das Abschneiden der Sackgassen (Pruning): Der Richter bewertet den Fortschritt jedes Entdeckers. Wenn ein Entdecker einen Pfad einschlägt, der unwahrscheinlich ist, erfolgreich zu sein (selbst wenn er noch nicht fertig ist), schneidet das Team diesen Pfad ab und konzentriert die Ressourcen auf die Entdecker, die sich auf den besten Spuren befinden.

Warum das wichtig ist

Das Paper testete dies an schwierigen Matheaufgaben (wie den GSM8K- und MATH500-Datensätzen) und Programmieraufgaben (HumanEval).

Das Ergebnis: Der „Team von Entdeckern“ (Order-Token Search) fand konsistent häufiger die korrekten Antworten als die alten „Selbstbewussten“ oder „Zufälligen“ Methoden.
Der Vergleich: Es schnitt so gut ab, dass es Methoden ebenbürtig war oder diese sogar schlug, die eine teure, monatelange Neukalibrierung des Modells (wie diffu-GRPO) erforderten. Das bedeutet, dass man eine viel intelligentere KI erhalten kann, indem man lediglich die Art und Weise ändert, wie sie während des Tests denkt, ohne das Gehirn selbst neu trainieren zu müssen.

Eine besondere Anmerkung zu Sudoku

Das Paper hat dies auch an Sudoku-Rätseln getestet. Interessanterweise funktionierte es dort nicht gut. Die Autoren erklären, dass Sudoku strikte, globale Regeln erfordert (wie „keine wiederholten Zahlen in einer Reihe“), die die interne „Bewertungskarte“ der KI einfach nicht zu verstehen scheint. Es ist, als würde man einem Team von Wanderern eine Karte geben, die keine Klippen zeigt; egal wie gut sie explorieren, sie können nicht vermeiden, an der Kante hinunterzustürzen, wenn die Karte falsch ist. Dies deutet darauf hin, dass die KI bei manchen Aufgaben anders trainiert werden muss, nicht nur die Suchmethode.

Zusammenfassend

Das Paper zeigt, dass wir durch das Zulassen einer KI, mehrere verschiedene Wege des Schreibens (Reihenfolge) und mehrere verschiedene Wörter (Tokens) gleichzeitig zu explorieren, und indem wir ein intelligentes Bewertungssystem nutzen, um schlechte Ideen frühzeitig abzubrechen, viel bessere Ergebnisse von Diffusion Language Models erzielen können, ohne sie neu trainieren zu müssen. Es verwandelt ein chaotisches Ratespiel in eine strukturierte, effiziente Suche nach der Wahrheit.

Technisches Resümee: Verbesserung der Dekodierung von Diffusions-Sprachmodellen durch gemeinsame Suche in Generierungsordnung und Token-Raum

Problemstellung

Diffusions-Sprachmodelle (DLMs), insbesondere Maskierte Diffusionsmodelle (MDMs), bieten einen deutlichen Vorteil gegenüber autoregressiven (AR) Modellen: ordnungsunabhängige Generierung. Im Gegensatz zu AR-Modellen, die Token strikt von links nach rechts generieren, können MDMs mehrere Positionen parallel rauschen (denoise), was adaptive Generierungsordnungen ermöglicht. Aktuelle Dekodierungsmethoden versäumen es jedoch, diese Flexibilität voll auszuschöpfen.

Bestehende Dekodierungsstrategien verpflichten sich typischerweise auf eine einzelne Trajektorie pro Durchlauf, die durch heuristische Remasking-Regeln gesteuert wird:

Low-Confidence Remasking: Fixiert gierig (greedy) Tokens mit hoher Konfidenz frühzeitig. Während dies oft eine starke Genauigkeit bei Einzelproben (pass@1) liefert, beschränkt es die Suche auf eine enge Menge an Trajektorien, was zu einem langsamen Wachstum der Multi-Sample-Abdeckung (pass@k) führt und das Risiko des Feststeckens in falschen lokalen Optima birgt.
Random Remasking: Stichprobenartige Auswahl von Remasking-Mustern zur Maximierung der Trajektorien-Diversität. Während dies das pass@k verbessert, indem ein breiteres Lösungsraum exploriert wird, leidet es unter einem schwachen pass@1, da es keinen prinzipisierten Mechanismus besitzt, um die korrekte Trajektorie unter konkurrierenden Kandidaten auszuwählen.

Dies erzeugt ein Explorations–Exploitations-Dilemma: Aktuelle Methoden entweder über-exploitieren lokale Konfidenz (verpassen gültige Lösungen) oder explorieren blind (scheitern an der Auswahl der besten Lösung). Das Paper argumentiert, dass das Potenzial von DLMs, diverse Generierungsordnungen und Token-Werte zu explorieren, ungenutzt bleibt, da die Standard-Dekodierung keine explizite Suche in diesem gemeinsamen Raum durchführt.

Methodik: Order-Token Search (Ordnungs-Token-Suche)

Um dies zu adressieren, schlagen die Autoren die Order-Token Search vor, einen Dekodierungsalgorithmus, der eine Beam-Suche im gemeinsamen Raum von Generierungsordnungen und Token-Werten durchführt. Der Algorithmus wechselt zwischen der Erweiterung von Kandidaten-Trajektorien und dem Beschneiden (Pruning) unvielversprechender Pfade.

1. Suchprozess (Erweiterung)

Die Order-Token Search unterhält einen Beam von $K$ teilweise entrauschten Kandidaten. Im Gegensatz zum Standard-Sampling, das einer festen Sequenz von Entscheidungen folgt, verzweigt der Algorithmus explizit über zwei Quellen der Unsicherheit in benutzerdefinierten Intervallen:

Generierungsordnung: Welche Positionen (Tokens) als nächstes unmaskiert werden sollen.
Token-Werte: Welche Tokens jenen Positionen zugewiesen werden sollen.

Um die Rechenkomplexität zu bewältigen, nutzt die Suche Block-Diffusion. Anstatt an jedem einzelnen Denoising-Schritt zu verzweigen (was einen prohibitiven Overhead verursachen würde), erweitert der Algorithmus Kandidaten an den Grenzen zusammenhängender Token-Blöcke. Innerhalb eines Blocks findet unabhängiges Denoising statt. Dies reduziert die Entscheidungspunkte von der Gesamtlänge der Sequenz auf die Anzahl der Blöcke, was eine erhebliche Diversität sowohl in der Ordnung als auch in den Token-Zuweisungen ermöglicht und gleichzeitig die Inferenzkosten handhabbar hält.

2. Beschneidungsprozess (Likelihood-Schätzung)

Die Kerninnovation liegt in der Likelihood-Schätzerfunktion, die verwendet wird, um partielle Trajektorien zu bewerten und zu beschneiden. Ein Standard-MDM-Ziel (Likelihood der Rekonstruktion der vollständigen sauberen Sequenz aus einem verrauschten Zustand) ist instabil für Zwischenzustände mit vielen verbleibenden Masken, da es das Modell zwingt, schwierige Infilling-Probleme zu lösen, für die es während des Trainings nicht optimiert wurde.

Die Autoren führen eine inkrementelle Action-Level-Bewertungsfunktion ein:

Anstatt den gesamten Zustand zu bewerten, bewertet der Algorithmus die spezifischen Denoising-Aktionen, die den Übergang von einem stärker korrumpierten Zustand $x_t$ zu einem weniger korrumpierten Zustand $x_s$ vollziehen.
Der Score $s(x_t; x_s)$ misst, wie sicher das Modell die neu enthüllten Tokens in $x_s$ vorhersagen kann, konditioniert auf die Full-Sequence-Vorhersage des Modells $x_0$ (welche als Kontext dient).
Dieser Ansatz evaluiert, ob die vergangenen Entscheidungen einer Trajektorie (Ordnungs- und Token-Wahl) durch das Modell gestützt werden, was ein Signal mit geringerer Varianz und höherer Diskriminanz für das Pruning liefert.

Der kumulative Score eines Kandidaten ist die Summe dieser inkrementellen Scores über die Suchintervalle. Der Algorithmus behält die Top- $K$ Kandidaten nach jeder Erweiterung bei und filtert so inkorrekte oder suboptimale Pfade effektiv heraus, bevor die Dekodierung fortgesetzt wird.

Wichtigste Beiträge

Identifizierung des Trade-offs: Das Paper demonstriert empirisch, dass Standard-MDM-Dekodierungsmethoden an den entgegengesetzten Enden des Explorations–Exploitations-Spektrums liegen. Low-Confidence Remasking über-exploitiert (hohes pass@1, geringes pass@k Wachstum), während Random Remasking blind exploriert (hohes pass@k, niedriges pass@1).
Order-Token Search Algorithmus: Der Vorschlag eines neuartigen Dekodierungsalgorithmus, der im gemeinsamen Raum von Generierungsordnungen und Token-Werten sucht. Er nutzt einen stabilen, inkrementellen Likelihood-Schätzer, um partielle Trajektorien zu beschneiden, wodurch Trajektorien-Diversität in verbesserte Single-Sample-Genauigkeit umgewandelt wird.
Empirische Validierung: Die Methode zeigt konsistent über mathematische Reasoning- und Coding-Benchmarks hinweg überlegene Leistungen und erreicht oder übertrifft die Gewinne, die typischerweise durch teure Post-Training-Methoden wie diffu-GRPO erzielt werden.

Experimentelle Ergebnisse

Die Autoren evaluierten die Order-Token Search (OTS) auf GSM8K, MATH500, Countdown und HumanEval unter Verwendung von LLaDA-8B-Instruct und LLaDA-1.5 (eine RL-post-trainierte Variante).

Leistungsgewinne: OTS erzielte konsistent die höchste durchschnittliche Genauigkeit über die Benchmarks hinweg.
- Auf LLaDA-8B verbesserte OTS das pass@1 um +3,1 % auf GSM8K, +3,8 % auf MATH500, +7,9 % auf Countdown und +6,8 % auf HumanEval im Vergleich zur Low-Confidence Remasking Baseline.
- Auf LLaDA-1.5 erreichte OTS ähnliche relative Verbesserungen und erzielte eine durchschnittliche Genauigkeit von 40,4 % im Vergleich zu 37,2 % für die beste Baseline (Low-conf + MV).
Vergleich zum Post-Training: OTS erreichte oder übertraf die Performance von Modellen, die Reinforcement Learning Post-Training (diffu-GRPO) unterzogen wurden, was darauf hindeutet, dass Inference-Time-Suche ebenso effektiv wie Modell-Retraining zur Verbesserung von Reasoning-Fähigkeiten sein kann.
Ablationsstudien:
- Likelihood-Schätzer: Die spezifische inkrementelle Scoring-Funktion (Gleichung 2) war entscheidend. Varianten, die die Joint-Likelihood aller enthüllten Blöcke oder AR-Style-Likelihoods verwendeten, schnitten signifikant schlechter ab, was bestätigt, dass der maßgeschneiderte Schätzer für ein stabiles Pruning in MDMs notwendig ist.
- Suchraum: OTS übertraf "Order Search" (Suche nur nach Ordnung) und "Token Search" (Suche nur nach Tokens), welche AR-Style-Likelihoods verwendeten, was die Notwendigkeit der gemeinsamen Suche und die Diskrepanz von AR-Heuristiken zu MDMs hervorhebt.
Effizienz: OTS skaliert effektiver mit Test-Time-Compute (NFE) als Majority Voting. Bei gleichem Compute-Budget erreichte OTS eine höhere Genauigkeit (z. B. 29,3 % vs. 19,9 % auf Countdown), indem es strukturierte Suche statt unstrukturierter Stichproben nutzte.

Bedeutung und Behauptungen

Das Paper behauptet, dass die gemeinsame Ordnung-Token-Suche eine Schlüsselkomponente für die Weiterentwicklung der DLM-Dekodierung ist. Durch die explizite Exploration des Raums von Generierungsordnungen und Token-Werten sowie das Pruning basierend auf einem stabilen Likelihood-Schätzer wandelt die Methode das inhärente Multi-Trajektorien-Potenzial von Diffusionsmodellen in zuverlässige Single-Sample-Performance um.

Die Autoren positionieren diese Arbeit als Lösung für den "Exploration–Exploitation"-Engpass in der aktuellen MDM-Dekodierung. Sie behaupten, dass ihr Ansatz es DLMs ermöglicht, den lokalen Optima zu entkommen, in denen Greedy-Methoden feststecken, und die korrekten Lösungen zu wählen, die Random Sampling übersieht – und das alles ohne zusätzliches Modelltraining. Die Arbeit etabliert, dass Inference-Time-Suchstrategien, die speziell auf die parallele, ordnungsunabhängige Natur von Diffusionsmodellen zugeschnitten sind, systematische Verbesserungen in Reasoning- und Coding-Aufgaben liefern können.