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

Abstract: 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 LLM which incorporates Binary Encoding and Index Shuffling. Our scaling analysis reveals that MDM-Prime-v2 is 21.8$\times$ 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.

• The paper introduces binary encoding and index shuffling to achieve optimal subtoken granularity and a tighter variational bound for diffusion language models.
• It establishes compute-optimal scaling laws, reducing validation perplexity and significantly improving compute efficiency compared to prior models such as ARM.
• Empirical results and spectral analyses demonstrate enhanced generalization, robustness, and effective parameter utilization across diverse benchmarks.

MDM-Prime-v2: Compute-Optimal Scaling of Diffusion LLMs with Binary Encoding and Index Shuffling

Introduction

MDM-Prime-v2 presents a substantial advancement in the study of discrete diffusion LLMs (MDMs) by addressing critical inefficiencies in previous MDM architectures. The work rigorously analyzes the influence of subtokenization granularity and distribution on the variational training bound, identifies BPE-induced entropy limitations, and proposes a theoretically-grounded subtokenizer design. The two main innovations—binary encoding and index shuffling—significantly enhance the model's likelihood estimation capability and compute optimality, achieving new state-of-the-art perplexity and generalization in the diffusion paradigm.

Figure 1: Overview of MDM, MDM-Prime, and the MDM-Prime-v2 enhancements, comparing tokenization granularity and the introduction of binary encoding and index shuffling.

Methodology

Limitations of Prior MDM-Prime

MDM-Prime augmented the standard MDM framework by mapping each token to a sequence of $\ell$ sub-tokens via an invertible function $f_\ell$, thus allowing token granularity to be tuned and introducing intermediate “partially-masked” latent states. Despite empirical improvements in perplexity, MDM-Prime lacked a theoretical prescription for optimal granularity $\ell$ and did not address the entropy structure induced by conventional BPE tokenizers. These deficiencies led to suboptimal likelihood modeling—particularly due to highly non-uniform sub-token distributions.

Tightening the Variational Bound: Optimal Subtoken Granularity

The paper derives a fundamental result: the variational upper bound on log-likelihood monotically tightens as $\ell$ increases, reaching its theoretical optimum at $\ell = \lceil\log_2 V\rceil$, where $V$ is the vocabulary size. This specifies that binary encoding of token indices is optimal for subtokenization under the variational framework, as proved in the main results.

Figure 3: Loss envelopes and isoFLOP/isoloss curves showing model scaling across MDM, ARM, and MDM-Prime-v2.

Index Shuffling: Maximizing Sub-token Entropy

The second key insight is that the informativeness of each sub-token—measured by entropy—critically governs model performance. Typical BPE tokenizers produce sub-token distributions with low entropy due to index correlations (frequent tokens clustered at low indices). Binary encoding alone is insufficient: the authors show that maximizing $q_\alpha(y_t)$ entropy in the forward process is necessary for a tighter bound and superior generative modeling. To decorrelate indices and achieve near-uniform entropy, the paper introduces random index shuffling prior to binary encoding:

Figure 5: Entropy analysis before and after index shuffling, illustrating high-entropy sub-token distributions after shuffling.

This operation is static, requires no FLOPs at train/test time, and is implemented via lookup at preprocessing.

Resulting Subtokenizer

By composing index shuffling with binary encoding, the unified subtokenizer achieves (1) monotonic tightening of the variational bound and (2) maximal sub-token entropy, yielding MDM-Prime-v2. Empirically, both the training NLL and validation perplexity decrease strictly with increasing $\ell$ only when index shuffling is used (see Fig. 4 in the paper).

Scaling Analysis and Empirical Results

Compute-Optimal Frontier and Scaling Laws

Comprehensive scaling studies are conducted on C4 and OWT, sweeping model parameters $N$ and training tokens $D$ under fixed compute ($C$). Using Chinchilla-style scaling law fits, the analysis identifies the compute-optimal allocations for ARM, MDM, and MDM-Prime-v2. Crucially, MDM-Prime-v2 consistently establishes a lower minimal validation perplexity at a fixed compute than either ARM or MDM, regardless of scale.

Figure 7: Isoloss curves for ARM, MDM, and MDM-Prime-v2; showing the superior efficient frontier for MDM-Prime-v2.

For example, on compute-optimal OWT experiments (at $2.89 \times 10^{20}$ FLOPs):

• MDM-Prime-v2 achieves PPL of 7.77, outperforming ARM (12.99), MDM (18.94), and MDM-Prime (13.41).
• Compute efficiency is improved by 21.8$\times$ versus ARM.

Zero-shot Generalization

On standard zero-shot PPL and commonsense reasoning benchmarks (LAMBADA, WikiText, PTB, LM1B, AG News, ArXiv, SciQ, SocialIQA, etc.), MDM-Prime-v2 at 1.1B parameters consistently surpasses comparably-sized ARMs (e.g., GPT-Neo, OPT, Pythia, BLOOM, TinyLLaMA) and MDMs:

Figure 2: Compute-optimal parameter and token allocations for MDM-Prime-v2 under increasing FLOP budgets.

Figure 10: Conditional generation sample by MDM-Prime-v2 (1.1B) on long-context summary task, demonstrating improved coherence.

MDM-Prime-v2 achieves highest accuracy in 6/8 reasoning tasks (notably on SciQ and McTaco), affirming its transfer and generalization properties.

Architectural and Analytical Insights

Attention and Spectral Analysis

Investigation of attention scores and projection matrix spectra reveals that MDM-Prime-v2 alleviates the "attention sink" effect observed in MDMs, exhibiting richer routing and higher stable rank (10.0 vs 8.3), which correlates with enhanced expressivity.

Figure 12: Comparison of attention score patterns between MDM and MDM-Prime-v2, highlighting the reduction in attention sinks in MDM-Prime-v2.

Figure 4: Singular value distributions for query-key-value projections, showing heavier tails in MDM-Prime-v2 indicative of increased model capacity usage.

Robustness

Performance remains invariant with respect to model width/depth trade-offs, embedding table design, and random seeds for index shuffling, indicating the robustness of the proposed techniques.

Figure 6: Analysis of performance invariance to architectural and shuffling design choices.

Implications and Theoretical Significance

MDM-Prime-v2 is the first order-agnostic discrete diffusion LLM to surpass ARMs in the compute-optimal regime. Its joint advances in subtokenizer theory and practical scaling unlock a new scaling law for MDMs: compute-optimality is driven more strongly by increased data ($\hat{b}$) than by model parameters—a notable divergence from ARM scaling. This enables more efficient and robust LLM pretraining in data-abundant settings, and potentially narrows or inverts the longstanding performance gap vs. ARMs in high-compute regimes.

Practically, the subtokenizer construction is static, efficient, and architecture-agnostic, facilitating straightforward adoption in large-scale pretraining and ablation studies.

Limitations and Future Directions

Despite the advancements, open questions remain regarding optimal learned subtokenizer strategies (beyond random shuffling), improving inter-token modeling at high mask rates, and explicit integration with post-training regimes for specialized tasks. Developing entropy-maximizing deterministic subtokenizer algorithms, and extending these principles to multimodal and multilingual settings, represent promising avenues for subsequent research.

Conclusion

MDM-Prime-v2 establishes a new compute-optimal scaling law for discrete diffusion LMs, driven by a theoretically-motivated design for subtokenization that resolves prior bottlenecks in entropy and likelihood estimation. Its robust empirical outperformance across scaling regimes and benchmarks demonstrates that order-agnostic generative models can be made competitive, or even dominant, in the large-scale language modeling landscape.