Technischer Zusammenfassung: Triplet-Block Diffusion RWKV (B3D-RWKV)

Problemstellung

Kausale Transformer-Sprachmodelle stehen vor zwei primären strukturellen Einschränkungen: der strikt sequenziellen Dekodierung, die eine Parallelisierung verhindert, und den quadratischen Attention-Kosten pro Schritt, die die Inferenz in langen Kontexten rechnerisch teuer machen. Während lineare kausale Modelle (z. B. RWKV) und diskrete Diffusionsmodelle diese Schwächen jeweils unabhängig voneinander adressieren, ist ihre Integration inhärent inkonsistent. Diffusionsmodelle erfordern bidirektionale Attention, um Token-Blöcke parallel zu entzerren, wohingegen kausale Modelle strikt unidirektional sind. Bestehende Versuche, sie zu kombinieren, erfordern oft architektonische Modifikationen, die die Effizienz des linearen Rückgrats beeinträchtigen oder ein Training von Grund auf neu notwendig machen.

Methodik

Die Autoren schlagen B3D-RWKV vor, eine Diffusionsvariante von RWKV, die eine Inferenzeffizienz von $O(L)$ mit paralleler, bidirektionaler diskreter Diffusion vereint, ohne das zugrunde liegende kausale Rückgrat zu verändern. Die Kerninnovation ist das Triplet-Block-Layout, eine Trainings- und Inferenzstrategie, die darauf ausgelegt ist, bidirektionalen Zugriff innerhalb einer strikt kausalen Architektur zu simulieren.

1. Triplet-Block-Layout (Training)

Um Diffusion auf einem kausalen Modell zu ermöglichen, wird jeder logische Generierungsblock der Größe $B$ in drei aufeinanderfolgende physikalische Blöcke innerhalb der Trainingsstichprobe erweitert:

• $b_1$ (Maskierte Kopie): Eine maskierte Version des logischen Blocks.
• $b_2$ (Verlustbehaftete maskierte Kopie): Eine identische maskierte Kopie, bei der der Entzerrungsverlust berechnet wird.
• $b_3$ (Klares Ground-Truth-Exemplar): Eine unmaskierte Version, die den rekurrenten Zustand für den nachfolgenden Block aktualisiert.

Pseudo-bidirektionaler Zugriff:
Da das kausale Rückgrat Token strikt von links nach rechts verarbeitet, hat der zum Zeitpunkt eines beliebigen maskierten Positions in $b_2$ ankommende verborgene Zustand bereits den gesamten Block $b_1$ verarbeitet. Da $b_1$ und $b_2$ dasselbe Maskenmuster teilen, werden die unmaskierten Token in $b_1$ effektiv in den verborgenen Zustand „absorbiert", bevor das Modell versucht, die maskierten Token in $b_2$ vorherzusagen. Dies gewährt dem Modell pseudo-bidirektionalen Zugriff auf seinen eigenen unmaskierten Kontext innerhalb des logischen Blocks und erfüllt die bidirektionalen Anforderungen der Diffusion, während die unidirektionale Natur des RNN/SSM-Rückgrats erhalten bleibt.

Trainingsziel:
Das Modell wird mit einem Masken-Vorhersage-Ziel trainiert. Der Verlust besteht aus:

• Cross-Entropy-Verlust ($L_{CE}$): Berechnet an überwachten Positionen, an denen Token maskiert sind.
• Vertrauensbewusster paralleler Verlust ($L_{CAP}$): In Anlehnung an LLaDA-2.0 schärft dieser Term die Wahrscheinlichkeitsverteilung an Positionen, die bereits korrekt vorhergesagt wurden, und ermutigt den Inferenz-Sampler, pro Schritt mehr Token zu commiten.

2. Inferenz: Blockweise iterative Entzerrung

Während der Inferenz generiert das Modell einen logischen Block nach dem anderen. Der Prozess umfasst:

1. Initialisierung eines Blocks mit allen maskierten Token.
2. Ausführung iterativer Entzerrungsschritte, bei denen das Modell Top-1-Wahrscheinlichkeiten für maskierte Positionen vorhersagt.
3. Commiten von Token, deren Top-1-Wahrscheinlichkeit einen Schwellenwert $\tau$ überschreitet.
4. Ein Ausweichmechanismus stellt Fortschritt sicher, indem er die $k_{min}$ vertrauenswürdigsten Token committet, falls weniger als $k_{min}$ den Schwellenwert überschreiten.
5. Sobald alle Positionen committet sind, wird der saubere Block an das Präfix angehängt, und der nächste logische Block beginnt.

Hauptbeiträge

1. B3D-RWKV-7.2B: Das erste Diffusions-basierte, lineare RNN-Sprachmodell, das im 7B-Maßstab mit einem Masken-Vorhersage-Ziel trainiert wurde. Es integriert parallele Token-Auswahl in das RWKV-7-Rückgrat, ohne dessen ursprüngliche Parameter zu verändern.
2. Triplet-Block-Framework: Ein universelles Trainingsrezept, das strikt kausale Rückgrate (die strikte Kausalität und nach vorne propagierende Zustands-Eigenschaften erfüllen) in Block-Diffusionsmodelle umwandelt. Dieser Ansatz vermeidet die architektonische Komplexität einer Umwandlung des Rückgrats in einen bidirektionalen Block.
3. Effizienz und Leistung: Die Methode erzielt signifikante Durchsatzverbesserungen gegenüber der Standard-Decodierung bei kausalen Modellen, während sie wettbewerbsfähige Schlussfolgerungsfähigkeiten beibehält.

Experimentelle Ergebnisse

Das Modell wurde an einem 8-Aufgaben-Suite evaluiert, der allgemeines Wissen (MMLU, ARC, PIQA, RACE) und Schlussfolgerungs-/Matheaufgaben (GSM8K, MATH, GPQA) umfasst.

• Genauigkeit: B3D-RWKV-7.2B erreicht eine Genauigkeit, die mit dem RWKV-7-Baseline und anderen Diffusionsmodellen ähnlicher Skalierung (z. B. LLaDA-8B, Dream-7B) vergleichbar ist. Es übertrifft die Baselines insbesondere bei Schlussfolgerungsaufgaben wie ARC-Challenge und RACE, was der pseudo-bidirektionalen Wahrnehmung zugeschrieben wird, die das Schlussfolgern verbessert.
• Durchsatz: Das Modell zeigt eine durchschnittliche 1,6-fache Beschleunigung der Dekodierungsrate im Vergleich zur RWKV-7-Baseline über Kontextlängen von 1K bis 512K Token. Mit angepassten Sampling-Parametern ist eine 2,02-fache Beschleunigung mit einem leichten Qualitätsverlust erreichbar.
• Mathematisches Schlussfolgern: Die Arbeit verzeichnet einen leichten Leistungsabfall bei komplexen Matheaufgaben (z. B. MATH) im Vergleich zur kausalen Baseline. Die Autoren führen dies auf die Natur der parallelen Dekodierung zurück, die Schwierigkeiten mit der hochstrukturierten, schrittweisen Verifikation hat, die für symbolische Matheaufgaben erforderlich ist, bei denen keine Teilpunkte vergeben werden.

Bedeutung und Behauptungen

Die Arbeit behauptet, einen effizienten Weg zu bieten, vortrainierte kausale Sprachmodelle ohne architektonische Änderungen in Diffusions-Sprachmodelle zu transformieren. Durch die Nutzung des Triplet-Block-Layouts vereint B3D-RWKV erfolgreich die Inferenzeffizienz von $O(L)$ linearer Modelle mit den parallelen Dekodierungsfähigkeiten der Diffusion.

Die Autoren positionieren die Arbeit als strukturelle Lösung für die Inkompatibilität zwischen kausalen und Diffusions-Paradigmen. Sie räumen Einschränkungen ein, darunter die aktuelle Validierung an einem einzigen Rückgrat (RWKV-7), die 3-fachen Kosten der physikalischen Sequenzlänge während des Trainings und die Verwendung eines relativ kleinen SFT-Datensatzes ohne RL-Ausrichtung, was zu Genauigkeitsrückgängen bei bestimmten Aufgaben beitragen könnte. Der Kernbeitrag bleibt jedoch der Nachweis, dass eine hochdurchsatzige, parallele Diffusionsgenerierung auf strikt kausalen, linearen Rückgraten durch Datenlayout statt durch architektonische Neugestaltung erreichbar ist.