One-sentence Summary

Researchers from Harvard, MIT, UC Berkeley, and UT Austin demonstrate that masked diffusion models (MDMs) learn computationally intractable infilling subproblems relative to autoregressive models, yet adaptively choosing the token decoding order sidesteps these subproblems, boosting Sudoku accuracy in pretrained MDMs from under $7\%$7% to about $90\%$90% and outperforming $7\times$7×-larger teacher-forced autoregressive models.

Key Contributions

• The paper provides theoretical and empirical evidence that masked diffusion models (MDMs) face computationally intractable training subproblems compared to autoregressive models, formalized through new average-case hardness constructions based on sparse predicate observations and cavity predictions.

• This work introduces an adaptive token decoding order strategy for MDM inference, which enables the model to bypass the hard subproblems encountered during training by selecting the order of token generation dynamically.

• Results demonstrate that this adaptive inference approach raises solving accuracy on Sudoku puzzles from under 7% to about 90%, outperforming autoregressive models with seven times more parameters that were explicitly trained via teacher forcing to learn a fixed decoding order.

Introduction

Diffusion models dominate continuous generation tasks such as image and audio synthesis, but adapting them to discrete domains like text and proteins is still developing. Masked diffusion models (MDMs), which learn to reverse a random token masking process, are a promising approach but often underperform autoregressive models (ARMs) in language modeling. While MDMs offer greater flexibility at inference by decoding in any order, they face a harder training objective: predicting tokens from arbitrary unmasked contexts. The authors investigate this trade-off by proving that MDMs can encounter computationally intractable masking subproblems, even for benign data distributions. To overcome this, they propose adaptive inference strategies that intelligently select the next token to unmask, effectively bypassing these hard subproblems. Their empirical results on logic puzzles show dramatic gains, boosting Sudoku accuracy from under 7% to nearly 90% with a simple adaptive strategy, surpassing even ARMs trained with teacher forcing to learn an optimal decoding order.

Method

Masked Diffusion Models (MDMs) provide a discrete-state generalization of continuous diffusion models for sequence generation. The authors formulate both a forward masking process and a learned reverse process, then reinterpret the training objective as order-agnostic learning. They further show that this order-agnostic formulation exposes the model to computationally hard infilling problems, and propose an adaptive inference strategy that selects which positions to unmask based on model confidence, thereby bypassing these hard subproblems without any additional training.

The forward process corrupts a clean sequence $x_0 \sim p_{\mathrm{data}}$x0​∼pdata​ over a vocabulary $\{1, \ldots, m\}${1,…,m} with an additional mask token $0$0. At noise level $t \in [0,1]$t∈[0,1], each coordinate is independently either kept unchanged with probability $\alpha_t$αt​ or replaced by the mask token with probability $1 - \alpha_t$1−αt​, where the schedule satisfies $\alpha_0 \approx 1$α0​≈1 and $\alpha_1 \approx 0$α1​≈0. The reverse process is defined via the conditional distribution $q_{s|t}(x_s | x_t, x_0)$qs∣t​(xs​∣xt​,x0​) for $s < t$s<t: coordinates that are not masked remain fixed, while masked coordinates are resampled as a categorical mixture between the mask token and the true token $x_0^i$x0i​. A denoising network $p_\theta(x_t, t)$pθ​(xt​,t) is trained to predict the marginal distribution over $x_0$x0​, and the reverse transition is approximated by substituting the network output in place of $x_0$x0​. The training objective is the ELBO-based loss

In practice, the authors note that a time-embedding-free architecture is typically sufficient, since the number of masked tokens implicitly encodes the noise level. Sampling starts from the fully masked sequence and iteratively unmask positions by sampling from the learned reverse transition.

The authors then reformulate the MDM training loss to reveal its connection to masked language modeling. Let $x_0[M]$x0​[M] denote the partially masked sequence obtained from $x_0$x0​ by masking all indices in the set $M \subseteq [L]$M⊆[L]. Under mild assumptions, the loss decomposes as

This expression shows that the MDM objective is a linear combination of losses over all possible infilling masks: the optimal predictor must solve every masking problem, i.e., predict the posterior marginal of each token conditioned on an arbitrary subset of visible context. The authors refer to this as order-agnostic training, since the model must be equally capable at every masking pattern. In contrast, autoregressive models (ARMs) only learn a restricted set of $L$L infilling problems, each predicting the next token given all preceding ones, which the authors term order-aware training. This distinction is central to the subsequent analysis: order-agnostic training is strictly more demanding, and the authors show that many of the masking subproblems encountered during MDM training are computationally intractable.

To make this hardness precise, the authors introduce a family of data distributions called latents-and-observations (L&O) distributions. In such a distribution, a sequence contains $N$N latent tokens sampled from a prior, and $P$P observation tokens each generated as a (possibly randomized) function of a subset of the latents via observation functions $\mathcal{O}_1, \ldots, \mathcal{O}_P$O1​,…,OP​. Order-aware training is computationally tractable for these distributions: by permuting the sequence so that the latent tokens appear first, each token can be predicted efficiently given its predecessors. However, order-agnostic training encounters hard masking problems. Concretely, in an example with sparse predicate observations, where each observation token is the output of a fixed predicate $g$g applied to a $k$k-element subset of latents, the authors prove that for a range of masking probabilities, predicting any of the masked latent tokens given the visible context is computationally intractable for polynomial-time algorithms. This hardness holds in the average case and even when the algorithm knows all distribution parameters, establishing that the difficulty is inherent to the masking problem itself rather than to a particular training procedure. The authors therefore argue that a large fraction of the subproblems in order-agnostic training are computationally hard, and empirically verify that MDMs trained on such data exhibit performance imbalances: small errors on easy subproblems but large errors on hard ones.

Given this hardness, a natural question is whether MDMs can still be made effective during sampling. The authors show that the answer is yes: the reverse process of an MDM admits multiple valid sampling paths. Whereas vanilla inference randomly selects which positions to unmask at each step, the authors observe that with an ideal model that perfectly solves all masking problems, any unmasking order yields the same final distribution, because the chain rule factorizes the likelihood identically regardless of the order. In practice, since the model is not equally accurate on all subproblems, different sampling paths lead to different likelihood estimates. This motivates adaptive MDM inference: instead of selecting the set of positions to unmask randomly, the authors propose using an oracle $\mathcal{F}(\theta, x_t)$F(θ,xt​) that strategically chooses the next positions to decode, thereby avoiding hard masking problems.

The authors introduce two concrete oracle designs. The first, Top-K probability, ranks masked positions by the maximum predicted probability over the vocabulary, i.e., $\max_{j} p_\theta(x^i = j | x_t)$maxj​pθ​(xi=j∣xt​), and unmask the $K$K positions with the highest certainty. This is a simple and effective proxy for confidence, but it can be misleading when the model is uncertain between two or more high-probability values. To address this, the authors propose Top-K probability margin, which instead ranks positions by the difference between the two most probable values at each position, $|p_\theta(x^i = j_1 | x_t) - p_\theta(x^i = j_2 | x_t)|$∣pθ​(xi=j1​∣xt​)−pθ​(xi=j2​∣xt​)∣. This margin-based score provides a more reliable estimate of uncertainty when the model assigns similar probabilities to multiple tokens, while degenerating to the Top-K criterion when a single value dominates. Both oracles are purely inference-time modifications: they require no additional training and can be applied to any pretrained MDM. The authors demonstrate that pairing a pretrained MDM with an effective ordering oracle allows the model to sidestep the hard subproblems identified earlier, substantially improving likelihood matching despite the model's imperfect performance on difficult masking tasks.

Experiment

The experiments first establish that masked diffusion models (MDMs) suffer from a hardness imbalance across the many masking subproblems they train on: order-agnostic training creates computationally harder prediction tasks than left-to-right autoregressive training, and the model's final error is concentrated on these harder subproblems. To address this, the authors introduce adaptive MDM inference, which selects easier subproblems at generation time. This strategy consistently outperforms vanilla inference on L&O-NAE-SAT, text benchmarks (including LLaDA 8B), and logic puzzles, where Top-K probability margin notably solves Sudoku at 89.49% versus 6.88% for vanilla inference. Furthermore, adaptive MDM inference outperforms autoregressive models, even when those models are trained with explicit ordering information, and shows greater robustness when generalizing to harder puzzles, suggesting that the order-agnostic training of MDMs combined with adaptive inference provides a more flexible and effective approach to sequence generation.

Across all tested (N, P) configurations, adaptive MDM inference markedly improves observation token prediction accuracy compared to vanilla MDM inference. Vanilla inference performs near or below the 75% naive guessing baseline, while adaptive inference consistently achieves accuracy levels above 88%, with the best results exceeding 93%. Adaptive inference outperforms vanilla inference by substantial margins for every (N, P) pair, with gaps ranging from about 15 to 26 percentage points. Vanilla inference scores between 62.84% and 78.06%, sometimes falling below the 75% naive guessing baseline, whereas adaptive inference always stays above 88%. The accuracy gap widens as the problem size increases, with the largest relative improvement observed at (100, 200) where adaptive nearly reaches 89% while vanilla drops to under 63%.

On Sudoku puzzles, adaptive masked diffusion model inference, particularly the Top-K probability margin strategy, achieves the highest accuracy, surpassing both autoregressive models trained with or without ordering information and vanilla masked diffusion inference. This improvement is notable given the masked diffusion model uses far fewer parameters than the autoregressive baselines. Top-K probability margin adaptive inference reaches 89.49% accuracy, outperforming the strongest autoregressive baseline trained with ordering information at 87.18%. The adaptive Top-K probability margin method substantially outperforms vanilla masked diffusion inference (6.88%) and Top-K probability sampling (18.51%). The best masked diffusion model uses only 6M parameters, compared to 42M for the autoregressive model, while still achieving higher accuracy.

Adaptive MDM inference methods significantly outperform vanilla MDM and ARMs on the Zebra puzzle, achieving over 98% accuracy with a smaller parameter count. Both Top-K and Top-K probability margin strategies show strong gains, with Top-K slightly ahead, while even the stronger ARM baseline with ordering information remains below these adaptive MDM approaches. Adaptive MDM strategies exceed 98% accuracy, far surpassing vanilla MDM at 76.9% and ARMs even when trained with ordering information. The best adaptive MDM method uses Top-K probability, while Top-K probability margin achieves nearly identical accuracy, both with only 19M parameters. ARMs require 42M parameters and still underperform adaptive MDMs, including the variant given supervised order information.

Adaptive inference strategies for the LLaDA 8B model consistently outperform vanilla inference on coding and math benchmarks. The Top-K probability margin strategy shows the strongest gains, particularly on harder code completion and mathematical reasoning tasks, while both adaptive methods remain competitive on general knowledge and story reasoning. Both adaptive strategies outperform vanilla inference across all evaluated tasks, with the margin variant achieving the largest improvements on multi-line and split-line code generation. Top-K probability margin beats plain Top-K on challenging tasks like mathematical reasoning and split-line code, though not on MMLU, suggesting better handling of ambiguous token probabilities.

On hard Sudoku puzzles that require unseen strategies or backtracking, the adaptive MDM inference strategy based on Top-K probability margin achieves the highest accuracy, substantially outperforming ARMs trained with order information. MDMs with adaptive inference also show greater robustness to the distribution shift toward harder puzzles, despite using far fewer parameters. Top-K probability margin adaptive MDM inference outperforms ARM with ordering by a significant margin on hard Sudokus. Random MDM inference performs poorly, demonstrating the critical role of adaptive inference strategies for MDMs. The best MDM achieves higher accuracy than the ARM while using only about one-seventh of the parameters.

Across multiple benchmarks, including observation token prediction, Sudoku, Zebra puzzles, and coding/math tasks, adaptive masked diffusion model (MDM) inference consistently outperforms vanilla MDM and often autoregressive models (ARMs), with the Top-K probability margin strategy delivering the strongest gains. Adaptive inference achieves high accuracy even with significantly fewer parameters, showing robustness to harder problem instances and better handling of ambiguous token probabilities. These improvements highlight the critical role of adaptive inference strategies for MDMs.