MDM-Prime-v2: Binaire Codering en Indexherschikking Maken Compute-optimaal Schalen van Diffusion Language Models Mogelijk

MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models

March 17, 2026

Auteurs: Chen-Hao Chao, Wei-Fang Sun, Junwei Qua, Chun-Yi Lee, Rahul G. Krishnan

cs.AI

Samenvatting

Gemaskerde diffusiemodellen (MDM) vertonen superieure generalisatie wanneer ze worden aangeleerd met een gedeeltelijk maskeringsschema (Prime). Deze methode zet tokens om in sub-tokens en modelleert het diffusieproces op sub-tokenniveau. Wij identificeren twee beperkingen van het MDM-Prime raamwerk. Ten eerste ontbreken er instrumenten om de hyperparameterkeuze voor de tokengranulariteit in de subtokenizer te sturen. Ten tweede constateren we dat de functievorm van de subtokenizer de waarschijnlijkheidsschatting aanzienlijk verslechtert in combinatie met veelgebruikte Byte-Pair-Encoding (BPE) tokenizers. Om deze beperkingen aan te pakken, bestuderen we de strakheid van de variational bound in MDM-Prime en ontwikkelen we MDM-Prime-v2, een gemaskerd diffusie-taalmodel dat Binaire Codering en Indexherschikking incorporeert. Onze schaalanalyse toont aan dat MDM-Prime-v2 21,8 keer rekenzuiniger is dan autoregressieve modellen (ARM). In compute-optimale vergelijkingen behaalt MDM-Prime-v2 een perplexiteit van 7,77 op OpenWebText, wat beter presteert dan ARM (12,99), MDM (18,94) en MDM-Prime (13,41). Bij het opschalen van het model naar 1,1B parameters, toont ons model verder superieure zero-shot nauwkeurigheid op diverse common sense-redeneertaken.

English

Masked diffusion models (MDM) exhibit superior generalization when learned using a Partial masking scheme (Prime). This approach converts tokens into sub-tokens and models the diffusion process at the sub-token level. We identify two limitations of the MDM-Prime framework. First, we lack tools to guide the hyperparameter choice of the token granularity in the subtokenizer. Second, we find that the function form of the subtokenizer significantly degrades likelihood estimation when paired with commonly used Byte-Pair-Encoding (BPE) tokenizers. To address these limitations, we study the tightness of the variational bound in MDM-Prime and develop MDM-Prime-v2, a masked diffusion language model which incorporates Binary Encoding and Index Shuffling. Our scaling analysis reveals that MDM-Prime-v2 is 21.8times more compute-efficient than autoregressive models (ARM). In compute-optimal comparisons, MDM-Prime-v2 achieves 7.77 perplexity on OpenWebText, outperforming ARM (12.99), MDM (18.94), and MDM-Prime (13.41). When extending the model size to 1.1B parameters, our model further demonstrates superior zero-shot accuracy on various commonsense reasoning tasks.