Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

Technische Samenvatting: Training van Hybride Block Diffusion Language Models met Partiële Bidirectionaliteit

Probleemstelling
Hoge doorvoersnelheid bij generatie met lange context in grote taalmodellen (LLM's) wordt primair beperkt door geheugenbandbreedte in plaats van rekenkracht. Bij standaard autoregressieve decodering vereist elke stap het streamen van de geaccumuleerde key/value (KV) cache vanuit het geheugen, waardoor de bandbreedtevraag lineair schaalt met de contextlengte terwijl er slechts één token wordt uitgegeven. Hoewel er twee parallelle benaderingen bestaan om dit te verzachten—het verminderen van geheugentoegang via linear-time mixers (bijv. Mamba) en het vergroten van parallellisme via block-generatie—brengt het combineren hiervan een specifieke technische hindernis met zich mee.

Eerdere hybride diffusiemodellen, zoals DiffuMamba, maken gebruik van volledig bidirectionele Mamba-scans, inclusclusief een reverse-richting pass over het volledige denoising-venster. Deze reverse scan is afhankelijk van actieve tokens, wat betekent dat de interne toestanden niet alleen uit een prefix bestaan. Deze toestanden kunnen daarom niet precies worden gecached en hergebruikt over blokken, wat de efficiëntievoordelen van block diffusion tenietdoet. De kernuitdaging is het ontwerpen van een hybride architectuur die block-gewijze generatie ondersteunt met exacte caching, terwijl de voordelen van bidirectionele mixing behouden blijven.

Methodologie
De auteurs stellen BDLM Mamba-H voor, een hybride architectuur die Block Diffusion Language Models (BDLMs) combineert met Mamba- en attention-lagen. De kerninnovatie is een partiële bidirectionaliteitsmechanisme dat de reverse-richting Mamba-scan beperkt tot enkel het actieve denoising-blok, in plaats van de volledige sequentie.

Belangrijke methodologische componenten zijn:

1. Cache-gealigneerde Trainingsobjectief: Het model gebruikt een "all-block" objectief waarbij voltooide prefix-blokken worden opgeslagen als:
   • KV-states voor attention-lagen.
   • Forward-only Mamba boundary states (recurrent states) voor Mamba-lagen.
     De reverse scan is beperkt tot het huidige actieve blok ($b_t$), wat ervoor zorgt dat de prefix cache ($C_\ell(c)$) een functie blijft van de schone prefix alleen, onafhankelijk van de reverse diffusion timestep ($t$).
2. Timestep Factorisatie: Om de geldigheid van de cache te waarborgen, factoriseert de architectuur de laag-berekening in een timestep-vrije prefix operator en een timestep-geconditioneerde actieve-blok continuatie. Dit voorkomt dat de schone-prefix cache afhankelijk wordt van de specifieke reverse diffusion stap, een vereiste voor geldige hergebruik tijdens inferentie.
3. Trainingsprotocol: Het model wordt getraind met een discrete diffusion objectief over blokken. De prefix cache wordt geconstrueerd door de schone prefix te scannen, en downstream block losses trainen direct de representaties die in deze prefix states zijn opgeslagen, waardoor de training van de berekening wordt afgestemd op de cache-interface van de inference-time sampler.

Kernbijdragen

• Cache-gealigneerde Training: De auteurs formuleren een BDLM all-block objectief waarbij de gecachte objecten (KV states en Mamba prefix states) exact overeenkomen met de inference-time block caches. Dit maakt het mogelijk dat downstream block losses de prefix scan trainen die wordt gebruikt voor recurrent state hergebruik.
• Gecontroleerde Validatie en Opschaling: De auteurs voeren een uitgebreide sweep uit op 87M parameters op de DCLM dataset en schalen geselecteerde modellen naar 350M parameters met behulp van een Complete(d)P-stijl learning-rate transfer regel.
• Doorvoersnelheid bij Lange Context: De studie evalueert de inference throughput bij uitgebreide contextlengtes, waarbij wordt aangetoond dat de hybride aanpak aanzienlijk beter presteert dan zowel full-sequence baselines als pure attention-gebaseerde block diffusion modellen naarmate de contextlengte toeneemt.

Resultaten

• Validatiekwaliteit (87M Parameters): BDLM Mamba-H bereikte de beste C4-en validatie perplexiteit (61.6) onder de 87M modellen, waarmee het BDLM attention (76.5), full-sequence DiffuMamba-H (83.7) en full-sequence attention (87.6) versloeg.
• Validatiekwaliteit (350M Parameters): Op de 350M schaal bleef BDLM Mamba-H competitief met BDLM attention, met een C4-en perplexiteit van 38.2 vergeleken met 36.97 voor BDLM attention.
• Throughput Prestaties:
  • Bij 65K tokens bereikte BDLM Mamba-H 19.7x de throughput van de full-sequence DiffuMamba-H baseline.
  • Bij 262K tokens bereikte BDLM Mamba-H 3.7x de throughput van BDLM attention.
  • Hoewel full-sequence modellen sneller zijn bij korte contexten (8192 tokens) vanwege de overhead van het all-block objectief, vertonen de native BDLM modellen met caching superieure schaalbaarheid in lange-context regimes.

Betekenis en Claims
Het artikel claimt dat BDLM Mamba-H een levensvatbare architectuur vormt voor long-context diffusion. Door de reverse Mamba scan te beperken tot het actieve blok, maakt het model exacte caching over blokken mogelijk, wat het probleem van state-reuse oplost dat inherent is aan eerdere bidirectionele diffusion hybriden. De auteurs stellen dat deze aanpak hybride Mamba-attention modellen in staat stelt om competitieve validatiekwaliteit te behouden terwijl ze substantiële verbeteringen in throughput bieden voor long-context inferentie, waardoor block partial-reverse hybriden een praktische richting vormen voor het schalen van diffusion taalmodellen.

Het artikel blijft bescheiden over optimaliteit en merkt op dat de learning rate transfer rules die gebruikt zijn voor opschaling sterke kandidaten zijn gebaseerd op bestaande parametrisatiewetten, maar niet bewezen optimaal zijn voor discrete diffusion taalmodellen (dLLMs). Daarnaast erkent de auteur dat hoewel het block objectief de long-context generatie verbetert, het momenteel een trainings-throughput penalty (ongeveer 25–34% trager) met zich meebrengt vergeleken met full-sequence baselines op huidige hardwareconfiguraties.