Technische Samenvatting: Sumi – Open Uniform Diffusion Language Model vanaf nul

Probleemstelling

Hoewel diffusiemodellen zijn opgekomen als een veelbelovend alternatief voor autoregressieve (AR) taalmodellen, bestaat er een aanzienlijke kloof in de ontwikkeling van Uniform Diffusion Language Models (UDLMs) op grote schaal. In tegenstelling tot Masked Diffusion Language Models (MDLMs), die gemaskeerde tokens onomkeerbaar invullen, laten UDLMs toe dat elk token op elk moment kan worden bijgewerkt, wat theoretisch zorgt voor flexibelere generatie en zelfcorrectie.

Echter, geen enkele UDLM is tot nu toe vanaf nul voorgetraind op zowel grote parameterschalen (miljarden parameters) als grote tokenbudgetten (biljoenen tokens). Bestaande grootschalige UDLMs zijn ofwel compute-optimale checkpoints die getraind zijn op kleine tokenbudgetten, of modellen die zijn aangepast van gepre-trainde AR-architecturen (bijv. DiffusionGemma). Bijgevolg blijft het gedrag van native, grootschalige uniforme diffusie in het data-rijke regime onverkend, waarbij een helder referentiepunt ontbreeft voor het bestuderen van schalingwetten, generatiedynamiek en afwegingen tegenover gevestigde AR- en MDLM-baselines.

Methodologie

Modelarchitectuur en Trainingsdoel

De auteurs introduceren Sumi, een volledig open 7B-parameter UDLM.

• Architectuur: Sumi maakt gebruik van een tijd-agnostische bidirectionele Transformer met 36 lagen, een verborgen grootte van 4.096 en SwiGLU MLPs. Het gebruikt grouped-query attention (32 koppen, 8 KV-groepen), RMSNorm en ontkoppelde (untied) input/output embeddings. Om attention sinks te mitigeren, wordt een off-by-one softmax toegepast in de attention.
• Doel: Het model wordt getraind met de Generalized Interpolating Discrete Diffusion (GIDD) doelstelling, specifiek in de Signal-to-Noise Ratio (SNR) herparameteriseerde vorm.
• Ruis-schema (Noise Schedule): Training maakt gebruik van pure uniforme ruis met de log-SNR beperkt tot $\lambda \in [-9, 9]$.
• Tokenisatie: De OLMo 3 tokenizer wordt gebruikt met een vocabulaire van 100.278.

Trainingsregime

Sumi werd vanaf nul voorgetraind op 1,5 biljoen tokens met behulp van 288 NVIDIA H100 GPU's (totaal ~43.308 GPU-uren).

1. Pre-training (1,3T tokens): Uitgevoerd met een sequentielengte van 1.184 en een globale batchgrootte van 4.608. De datamix is zwaar gewogen naar educatieve inhoud, gefilterd en geherclassificeerd door educatieve scores afgeleid van een gedestilleerde Qwen3-32B classifier.
2. Mid-training (250B tokens): Twee fasen volgen de pre-training:
   • 130B tokens op de oorspronkelijke lengte (1.184).
   • 120B tokens op een uitgebreide lengte (4.864).
   • De mix verschuift de focus naar code, wiskunde en redeneren, bestaande uit 32,5% code, 29,7% wiskunde, 21,0% algemeen en 16,8% redeneerdata.

Inferentie-strategie

• Canvas Lengte: Het model gebruikt een vaste canvaslengte van 2.048 tokens voor alle evaluaties. Deze lengte is gekozen omdat de prestaties aanzienlijk verslechteren op substantieel kortere of langere canvases.
• Sampling: De auteurs onderzoeken confidence-based adaptive sampling versus ancestral sampling.
• Evaluatieprotocol: Alle evaluaties gebruiken de lm-evaluation-harness aangepast voor diffusie-scoring. Voor likelihood-taken worden posities voorbij het antwoord gevuld met willekeurige tokens; voor generatie-taken worden updates beperkt tot de generatielengte.

Belangrijkste Bijdragen

1. Eerste Native Grootschalige UDLM: Sumi is de eerste UDLM die vanaf nul is voorgetraind op een grote parameterschaal (7B) en een groot tokenbudget (1,5T), wat een baseline biedt voor het bestuderen van native uniforme diffusie.
2. Open Release: De auteurs stellen de modelgewichten, checkpoints en het volledige trainingsrecept vrij, inclusclusief een volledige specificatie van de datamix over publiekelijk beschikbare corpora.
3. Exploratieve Inferentie-probes: Het artikel presenteert directionele analyses op generatiedynamiek, waaronder canvaslengte-gevoeligheid, de volgorde van token-commitment, parallelle decodering mogelijkheden en het potentieel voor zelfcorrectie.

Resultaten

Benchmark Prestaties

Sumi werd geëvalueerd op 13 benchmarks binnen vier categorieën en vergeleken met open AR-modellen (Falcon-7B, Llama 2-7B, OLMo-7B) die getraind zijn op vergelijkbare tokenbudgetten.

• Algemene Kennis & Coding: Sumi presteert competitief en behaalt de hoogste scores onder de modellen die onder hetzelfde protocol zijn geëvalueerd op MMLU, TruthfulQA, HumanEval en MBPP. Dit wordt toegeschreven aan de educatie- en code-zware datamix.
• Redeneren & Wiskunde: Sumi is competitief met Llama 2-7B en presteert beter dan Falcon-7B op taken zoals GSM8K en ARC, hoewel het achterblijft bij grotere of anders getrainde modellen (bijv. LLaDA-8B, Llama 3-8B).
• Common Sense: Sumi presteert aanzienlijk minder goed op common sense benchmarks (PIQA, HellaSwag, WinoGrande). De auteurs schrijven dit gat toe aan de educatie-zware datamix, wat de common sense prestaties heeft doen verslechteren, hoewel zij opmerken dat het gat waarschijnlijk groter is dan alleen de datacompositie kan verklaren.

Observaties tijdens de Inferentie

• Canvas Gevoeligheid: De vloeiendheid van de generatie is zeer gevoelig voor de canvaslengte. Sumi genereert vloeiend binnen een taakafhankelijke band (gecentreerd rond 2.048), maar verslechtert scherp buiten dit bereik, met name op korte canvases.
• Commitment Order: Hoewel de trainingsdoelstelling order-agnostisch is, induceert confidence-based sampling een zelfgeorganiseerde commitment order. Het model heeft de neiging om bepaalbare tokens vroegtijdig vast te leggen (commit) en onzekere tokens uit te stellen, een structuur die niet aanwezig is bij ancestral sampling.
• Parallelle Decodering: Op coding-taken (HumanEval, MBPP) ondersteunt het model bescheiden parallelle decodering (het vastleggen van maximaal 4 tokens per stap) met minimaal verlies aan nauwkeurigheid. Echter, meerstaps rekenkunde (GSM8K) blijft order-gevoelig en verslechtert onmiddellijk bij parallelle stappen.
• Zelfcorrectie: Een expliciet revisiebudget (het toestaan van extra denoising-passes) levert geen zelfcorrectie op. Hoewel het model vastgelegde tokens overschrijft, zijn deze bewerkingen voornedoeldelijk "round trips" ($A \to B \to A$) in plaats van gerichte verbeteringen, waardoor de uiteindelijke nauwkeurigheid onveranderd blijft.

Betekenis en Claims

Het artikel positioneert Sumi als een fundamentele bron voor de gemeenschap om native uniforme diffusie op schaal te bestuderen. De auteurs benadrukken dat hun inferentie-probes "directioneel in plaats van conclusief" zijn, bedoeld om toekomstig onderzoek te stimuleren naar slecht begrepen aspecten van UDLMs, zoals:

• De intrinsieke generatiegedragingen van uniforme diffusie versus gemaskeerde diffusie.
• De condities waaronder zelfcorrectie bereikt kan worden.
• De impact van datacompositie op de specifieke sterktes en zwaktes van diffusiemodellen.

De auteurs stellen expliciet dat Sumi wordt vrijgegeven als een voorgetraind basismodel zonder instructie-tuning of veiligheidsfiltering, bedoeld om onderzoek te ondersteunen in plaats van directe inzet. Zij waarschuwen dat het model de standaardrisico's overneemt die geassocieerd worden met basis LMs, waaronder het potentieel voor schadelijke outputs en feitelijke onjuistheden.