MDM-Prime-v2: Binäre Kodierung und Index-Vermischung ermöglichen rechenoptimales Skalieren von Diffusions-Sprachmodellen
MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models
March 17, 2026
Autoren: Chen-Hao Chao, Wei-Fang Sun, Junwei Qua, Chun-Yi Lee, Rahul G. Krishnan
cs.AI
Zusammenfassung
Maskierte Diffusionsmodelle (MDM) zeigen eine überlegene Generalisierung, wenn sie mit einem partiellen Maskierungsschema (Prime) erlernt werden. Dieser Ansatz wandelt Tokens in Sub-Tokens um und modelliert den Diffusionsprozess auf Sub-Token-Ebene. Wir identifizieren zwei Einschränkungen des MDM-Prime-Frameworks. Erstens fehlen uns Werkzeuge, um die Hyperparameter-Auswahl der Token-Granularität im Sub-Tokenizer zu steuern. Zweitens stellen wir fest, dass die Funktionsform des Sub-Tokenizers die Likelihood-Schätzung in Kombination mit häufig verwendeten Byte-Pair-Encoding (BPE) Tokenizern erheblich verschlechtert. Um diese Einschränkungen zu adressieren, untersuchen wir die Tightness der variationellen Schranke in MDM-Prime und entwickeln MDM-Prime-v2, ein maskiertes Diffusions-Sprachmodell, das Binäre Kodierung und Index-Vermischung integriert. Unsere Skalierungsanalyse zeigt, dass MDM-Prime-v2 21,8-mal recheneffizienter ist als autoregressive Modelle (ARM). In rechenoptimalen Vergleichen erreicht MDM-Prime-v2 eine Perplexität von 7,77 auf OpenWebText und übertrifft damit ARM (12,99), MDM (18,94) und MDM-Prime (13,41). Bei einer Erweiterung der Modellgröße auf 1,1 Mrd. Parameter zeigt unser Modell zudem eine überlegene Zero-Shot-Genauigkeit bei verschiedenen Common-Sense-Reasoning-Aufgaben.
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.