Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions

Speakers: Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, Sitan Chen

Conference: ICML 2025

YouTube: https://slideslive.com/39044002

Overview

This talk, presented by Kulin Shah at ICML 2025, delves into the fundamental mechanisms and challenges of Masked Diffusion Models (MDMs), particularly concerning their approach to token ordering in language modeling. The work, a collaborative effort with Jaeyeon Kim, Vasilis Kontonis, Sham Kakade, and Sitan Chen, aims to dissect the strengths and limitations of MDMs as a burgeoning alternative to traditional auto-regressive (AR) language models. With diffusion models demonstrating compelling performance and efficiency trade-offs—such as Gemini Diffusion achieving similar performance to Flash 2 Lite with six times faster inference—understanding their core operational paradigms is paramount for advancing the field.

The central thesis explores how MDMs, by their very design, attempt to solve an exponentially larger set of sub-problems compared to AR models. While AR models predict tokens in a strict left-to-right sequence, MDMs are inherently "any-order" models, learning to predict masked tokens regardless of their position relative to already-unmasked tokens. This flexibility, while theoretically powerful, introduces significant practical hurdles, as many of these "any-order" sub-problems can be inherently difficult to learn. The research highlights a critical paradox: while MDMs struggle with tasks possessing an inherent sequential order during training, their inference can be dramatically improved through adaptive strategies that intelligently select the order of token unmasking, ultimately leading to superior performance on complex reasoning tasks.

The work is crucial for practitioners and researchers grappling with the trade-offs between model performance, inference efficiency, and training complexity in the rapidly evolving landscape of generative AI. By pinpointing where MDMs excel and where they falter, and by demonstrating the transformative power of adaptive inference, the talk provides a roadmap for developing more robust and efficient language generation systems. It underscores that while training for the "worst" (i.e., all possible permutations) can be suboptimal for certain data distributions, a "best" plan for inference can nonetheless unlock exceptional capabilities.

Background

▶ Watch: Introduction: LLMs vs. Diffusion Models (0:00)

The landscape of large language models (LLMs) has been predominantly shaped by auto-regressive (AR) models, which operate on the principle of next token prediction. These models are trained to predict the subsequent token in a sequence based on all preceding tokens, enforcing a strict left-to-right generation order. This paradigm has led to remarkable successes in various natural language processing tasks, from text generation to translation.

However, recent years have witnessed the emergence of diffusion models as a compelling alternative, particularly in generative tasks across various modalities, including language. Diffusion models, unlike their AR counterparts, typically model a reverse process that gradually denoises or unmasks corrupted data to reconstruct the original sample. For language, a popular framework is Masked Diffusion Models (MDMs).

In MDMs, the forward process is an absorbing process. It begins with an original clean sequence, say a sentence, and progressively corrupts it by masking tokens. Specifically, each token i in the sentence X0 is masked independently with a probability p(t), where t is a time step that ranges from 0 to 1. As t approaches 1, the masking probability approaches 1, meaning that at t=1, the entire sentence X1 consists of all masked tokens. This creates intermediate masked versions of the sentence, denoted as Xt.

During training, the MDM takes these partially masked sentences Xt as input and passes them through a neural network, typically a Transformer architecture. The network's objective is to predict the probability distribution over all words in the vocabulary for each of the masked tokens in Xt. The model is trained using a cross-entropy loss (also referred to as score entropy) on these masked tokens, aiming to maximize the probability of predicting the correct original word. Formally, the goal is to learn the posterior distribution P_theta(X0i | Xt) for all masked tokens i.

The inference process in vanilla MDMs starts with a completely masked sequence. It then iteratively unmasks tokens. Critically, in the standard approach, a set of masked positions is randomly sampled, and these positions are then unmasked according to the learned marginal distributions. This implies a random generation order of tokens, a stark contrast to the deterministic left-to-right order of AR models.

A crucial insight highlighted in the talk is that the MDM loss function, which learns to predict any masked token given the surrounding context and other masked tokens, is mathematically equivalent to an any-order auto-regressive loss. This means that instead of optimizing for a single, fixed permutation (like the identity permutation for left-to-right AR models), MDMs effectively optimize an expectation over all possible permutations pi of the sequence. For each permutation pi, the model aims to predict the token at position pi_i given all preceding tokens pi_1 to pi_{i-1}. This implies that MDMs are implicitly trying to solve an exponentially larger number of sub-problems compared to AR models, which only focus on one specific ordering. This inherent flexibility and the computational burden it entails form the core of the problem explored in this research.

Key Findings

▶ Watch: MDM training, inference, and any-order equivalence (2:00)

The research uncovers a nuanced picture of Masked Diffusion Models, revealing both their inherent challenges and their untapped potential when coupled with intelligent inference strategies.

1. MDMs Tackle Exponentially More Sub-problems, Leading to Hard-to-Learn Cases: The fundamental design of MDMs, where the loss is equivalent to an any-order auto-regressive loss, forces them to learn to predict tokens under an exponential number of possible permutations. While conceptually powerful, this leads to a situation where many of these sub-problems, especially those corresponding to "unnatural" or complex token orderings, are difficult for the model to learn effectively.

1. Theoretical Hardness for Inherently Sequential Data: The talk presents a theoretical result demonstrating this hardness. For sequences with a natural left-to-right dependency (e.g., where the last m tokens depend on the first n random tokens), auto-regressive models can easily generate samples by following the data-generating process. In contrast, MDMs struggle. The theory shows that for such distributions, with high probability over the masked input Xt, no polynomial-time algorithm can compute the posterior for any masked index i, indicating a fundamental computational barrier for MDMs on certain data structures.

1. Experimental Validation: Left-to-Right Order is Optimal for Fixed-Order Training in Language: Experiments training Transformer models on the SlimPajama dataset with various fixed permutations pi confirmed that for language, the auto-regressive (identity) permutation achieves the best validation loss for a given computation budget. As the permutation's "distance" from the left-to-right order increases, performance degrades. MDMs, by averaging over all permutations, fall "somewhere in between" these fixed-order performances, suggesting that their inherent "any-order" training can be suboptimal for tasks with strong sequential biases. This was presented as "bad news" for MDMs if only considering fixed-order training.

1. Adaptive MDM Inference Dramatically Boosts Performance: The most significant "good news" finding is the transformative impact of adaptive MDM inference. Instead of randomly unmasking tokens, the model can dynamically choose the unmasking order based on internal confidence or predictive quality. Simple rules, such as selecting the masked token with the highest predicted probability (Top Probability), the largest margin between the top two probabilities (Top Probability Margin), or the lowest prediction entropy, lead to massive performance improvements.

1. Adaptive MDM Outperforms Strong Auto-Regressive Baselines on Reasoning Tasks:

• Logic Puzzles: On logic puzzles, vanilla MDM inference yielded a mere 7% accuracy. Implementing adaptive inference, particularly using the Top Probability Margin strategy, catapulted accuracy to an impressive 90%. Crucially, this adaptive MDM even outperformed auto-regressive models that were seven times larger and had been specifically trained to decode in an optimal order (e.g., using a Sudoku solver to determine the ideal sequence). This highlights that MDMs, when guided intelligently, can leverage their parallel unmasking capabilities to excel in reasoning tasks where AR models struggle with rigid sequential decoding.
• Natural Language, Coding, and Math: The effectiveness of adaptive inference was further validated on Llama 8B models across tasks like HumanEval (coding) and various math problems, where it "almost always outperforms" vanilla inference strategies.

In summary, MDMs are presented as models that train on a wide spectrum of sub-problems with varying complexity. While some of these orders are poorly learned, adaptive inference provides a powerful mechanism to selectively avoid these poorly trained marginals during generation, thereby unlocking the true potential of MDMs for complex reasoning and generation tasks.

Technical Deep Dive

▶ Watch: Experimental setup: training with fixed permutations (4:50)

The technical exposition of the talk centers on the architecture and training of Masked Diffusion Models (MDMs), their mathematical equivalence to any-order auto-regressive loss, and the innovative adaptive inference strategies developed to overcome their limitations.

At its core, an MDM for language modeling operates on a sequence of tokens. The forward diffusion process begins with a clean sequence X0 and progressively masks tokens. For a given timestep t (typically ranging from 0 to 1), each token in X0 is independently masked with a probability p(t). As t increases, p(t) also increases, leading to more masked tokens. For instance, at t=1, all tokens are masked, resulting in X1. The intermediate state Xt is thus a partially masked version of X0.

The training objective of the MDM is to learn a neural network, usually a Transformer, that can predict the original unmasked token X0i given the masked sequence Xt for any masked position i. This is framed as learning the posterior distribution P_theta(X0i | Xt). The loss function used is typically cross-entropy over the vocabulary for the masked tokens, often referred to as score entropy, which aims to maximize the likelihood of predicting the correct word for each masked position. This prediction is made independently for each masked token, leveraging the full context of both unmasked and other masked tokens in Xt.

The crucial theoretical insight is that this training objective, where the model learns to predict any masked token, is formally equivalent to optimizing an any-order auto-regressive loss. An auto-regressive model typically optimizes log P(X) = sum_{i=1 to N} log P(X_i | X_{<i}), where X_{<i} denotes tokens preceding X_i in a fixed, left-to-right order. In contrast, the MDM loss is equivalent to E_pi [ log P(X | pi) ], where pi is a permutation sampled uniformly from all N! possible permutations of the sequence. For each pi, log P(X | pi) = sum_{i=1 to N} log P(X_pi_i | X_pi_1...X_pi_{i-1}). This means the MDM implicitly trains on an exponentially vast number of conditional prediction sub-problems, each corresponding to a different token ordering.

The theoretical hardness result underscores the challenge of this "any-order" training. Consider a sequence of length n+m where the first n tokens are uniformly random, and the last m tokens are deterministically dependent on the first n (e.g., X_i = X_j for some i, j). Such a distribution has an inherent left-to-right order that an AR model can easily exploit. However, the talk posits that for MDMs, with high probability over Xt, no polynomial-time algorithm can compute the posterior P_theta(X0i | Xt) for an arbitrary masked index i. This implies that the model's ability to learn certain necessary conditional probabilities is fundamentally limited by the combinatorial complexity of the task, particularly when the data has strong, non-random dependencies that are not aligned with an "any-order" approach.

To experimentally validate the impact of token ordering, the researchers trained a Transformer model on the SlimPajama dataset. They fixed various permutations pi during training and measured the log-likelihood under that specific permutation. When pi was the identity permutation, this mirrored standard auto-regressive training. As pi deviated from the identity, the validation loss increased, demonstrating that for natural language, the left-to-right order is indeed optimal for fixed-order training. The MDM, which effectively averages over all permutations, showed a performance curve that fell between the best (AR) and worst-performing fixed permutations.

The most impactful technical contribution is the introduction of adaptive MDM inference. During the iterative unmasking process, instead of randomly selecting positions to unmask, a rule-based strategy is employed to prioritize which tokens to predict. The talk highlights three such strategies:

1. Top Probability: For each masked token position i, the model predicts a probability distribution over the vocabulary. This strategy selects the position i where the maximum probability assigned to any word (max_word P(word | Xt, i)) is highest. This prioritizes positions where the model is most confident about a single word.
2. Top Probability Margin: This strategy also considers the predicted distribution. It selects the position i where the difference between the highest and second-highest predicted probabilities (max_word P(word | Xt, i) - second_max_word P(word | Xt, i)) is largest. This indicates a strong, unambiguous prediction for a single word, rather than a distribution with multiple plausible candidates.
3. Entropy: This strategy selects the position i where the entropy of the conditional distribution H(P(X0i | Xt)) is lowest. Lower entropy signifies higher confidence and a more peaked distribution, similar to the top probability strategy but capturing the overall "spread" of the distribution.

These adaptive strategies allow the MDM to dynamically avoid "poorly trained marginals"—those specific token predictions for which the model has low confidence or has learned poorly due to the inherent difficulty of that particular sub-problem. By prioritizing confident predictions, the model builds upon a stronger foundation, leading to more coherent and accurate outputs.

Experimental Setup & Results

▶ Watch: Scaling curve results: AR optimal, MDM struggles with ordering (5:50)

The experimental section of the talk provides concrete evidence for both the limitations of vanilla MDMs and the power of adaptive inference.

Fixed-Order Training Experiments

• Objective: To understand how different token orderings impact model performance in language modeling.
• Model: A Transformer architecture was used as the base model.
• Dataset: The SlimPajama dataset, a large-scale natural language corpus, was used for training.
• Methodology: The researchers fixed a specific permutation pi for token prediction during training. This allowed them to simulate training an auto-regressive model under a specific, potentially non-standard, order. They then estimated the log-likelihood under this permutation pi (i.e., P(pi_i | pi_1...pi_{i-1})).
• Baselines/Comparisons:
• Auto-regressive (AR) Training: Corresponds to using the identity permutation (pi = left-to-right order).
• Masked Diffusion Model (MDM) Training: Represented as an expectation over all possible permutations.
• Variations: Permutations with increasing "distance" from the left-to-right order were also explored.
• Metrics: Validation loss was plotted against computation budget (scaling curve).
• Results: The auto-regressive models (trained with the identity permutation) consistently achieved the best validation loss for a given computation budget. As the fixed permutation pi diverged further from the natural left-to-right order of language, the model's performance (validation loss) became progressively worse. The MDM's performance, being an average over all permutations, fell "somewhere in between" these extremes. This result empirically supports the theoretical claim that for language, which has an inherent left-to-right structure, training on a fixed, natural order is optimal, and the MDM's "any-order" training can be a disadvantage.

Adaptive Inference Experiments

• Objective: To demonstrate the effectiveness of adaptive strategies during MDM inference compared to vanilla (random) unmasking.
• Model: Masked Diffusion Models were used for these experiments, with Llama 8B being specifically mentioned for some tasks (implying a Llama-based MDM architecture).
• Tasks & Datasets:

1. Logic Puzzles: Specific logic puzzles were used to test reasoning capabilities.
2. Natural Language, Coding, and Math Tasks: These included HumanEval (a benchmark for code generation) and various math-type tasks.

• Inference Strategies Compared:
• Vanilla MDM Inference: Tokens are unmasked by randomly sampling positions.
• Adaptive MDM Inference (Rules):
• Top Probability: Unmask the token whose most likely assignment has the highest probability.
• Top Probability Margin: Unmask the token where the difference between the highest and second-highest predicted probabilities is greatest.
• Entropy: Unmask the token with the lowest entropy (most confident prediction) in its conditional distribution.
• Baselines for Logic Puzzles:
• Vanilla MDM inference.
• Auto-regressive models that were seven times larger than the MDM and were specifically "trained to decode in optimal order" (e.g., using a Sudoku solver to determine the ideal generation sequence). These AR models had strictly more information during training.
• Metrics: Accuracy for logic puzzles, and general performance (implied accuracy/score) for HumanEval and math tasks.
• Headline Results:
• Logic Puzzles:
• Vanilla MDM accuracy: 7%.
• Adaptive MDM accuracy (using Top Probability Margin): 90%.
• The choice of adaptive strategy mattered significantly: Top Probability yielded 19% accuracy, while Top Probability Margin jumped to 90%.
• Crucially, the adaptive MDM outperformed the 7x larger auto-regressive models that were trained with optimal decoding orders.
• Natural Language, Coding, and Math Tasks (on Llama 8B): Adaptive inference "almost always outperforms" vanilla inference strategies, indicating its broad applicability and benefit across diverse language-related challenges.

These experimental results conclusively demonstrate that while the "any-order" training of MDMs can lead to poorly learned marginals for certain prediction sequences, intelligent, adaptive inference strategies can effectively navigate these weaknesses, leading to vastly superior performance, even surpassing larger and more informed auto-regressive counterparts on tasks requiring complex reasoning.

Practical Implications

▶ Watch: Adaptive inference achieves high accuracy on logic puzzles (8:00)

The findings presented in "Train for the Worst, Plan for the Best" carry significant practical implications for anyone involved in developing, deploying, or researching advanced AI/ML models, especially in the domain of generative language.

For practitioners and model builders, the primary takeaway is that the raw training paradigm of Masked Diffusion Models, while offering theoretical flexibility, does not automatically translate to optimal performance. The "any-order" training exposes MDMs to an exponentially vast number of sub-problems, many of which are suboptimal or even theoretically hard to learn for data with inherent sequential structure like natural language. This suggests that simply scaling up MDMs with vanilla training and inference might hit a ceiling for certain tasks. Instead, the real power of MDMs, particularly for complex reasoning or generation, is unlocked through adaptive inference strategies. This means that when deploying an MDM, developers should not settle for random unmasking; investing in and implementing intelligent token selection mechanisms during generation is crucial.

For infrastructure teams and those concerned with model deployment, the talk highlights a promising efficiency trade-off. Diffusion models like Gemini Diffusion are shown to achieve similar performance to auto-regressive models but with significantly faster inference (e.g., six times faster). This speedup is a major advantage for real-time applications or large-scale deployments, potentially reducing computational costs and latency. However, this efficiency gain is often contingent on employing adaptive inference. Implementing adaptive strategies might introduce a slight increase in complexity during the inference step (e.g., needing to compute confidence scores for all masked tokens before selecting the next one), but the demonstrated performance boost far outweighs this overhead. Infra teams should be prepared to support dynamic, adaptive decoding pipelines rather than static, pre-determined ones.

Model architects and researchers should consider the implications for future MDM design. The current "any-order" training approach might be too broad. Future work could explore ways to:

1. Prioritize Training: Develop training objectives that selectively emphasize or improve the learning of "important" permutations (e.g., those closer to the left-to-right order for language) while still retaining some of the flexibility of any-order generation.
2. Identify and Mitigate Hard Sub-problems: Integrate mechanisms during training to identify which specific orderings or conditional probabilities are poorly learned and focus additional capacity or regularization on them.
3. Hybrid Approaches: Combine the strengths of AR and MDM models, perhaps using an AR component to bootstrap initial tokens or guide the overall structure, with an MDM then filling in details in parallel.

The talk also sheds light on the trade-offs and limitations. While adaptive inference dramatically improves performance, the speaker acknowledges that it does not guarantee sampling from the correct data distribution if the underlying marginals are fundamentally poorly trained. Adaptive inference is a heuristic to avoid poorly trained marginals, not to fix them. This means that while it's a powerful tool, it doesn't absolve the need for better training methods that might inherently improve the quality of all marginals, or at least the critical ones. The fact that simple adaptive rules (like Top Probability Margin) yield such large gains also suggests there's ample room for more sophisticated, learned adaptive strategies.

In essence, the practical implication is a call to action: embrace the efficiency of MDMs, but be strategic about their application. Their training might be "for the worst" in terms of covering all permutations, but intelligent "planning for the best" during inference is key to realizing their full potential, especially for challenging reasoning tasks where traditional AR models might be constrained by their sequential nature.

Key Takeaways

• MDMs' "Any-Order" Training Paradox: Masked Diffusion Models (MDMs) inherently train on an exponentially larger set of sub-problems (all possible token permutations) compared to auto-regressive (AR) models, which only focus on a fixed order.
• Struggles with Inherent Sequentiality: This comprehensive training strategy can be detrimental for tasks like language modeling, which possess a strong, natural left-to-right sequential dependency. Theoretical and experimental results show that MDMs can struggle to learn certain "hard" permutations, leading to poorly trained conditional marginals.
• Vanilla Inference is Suboptimal: Standard MDM inference, which randomly unmasks tokens, often performs poorly because it can inadvertently rely on these poorly trained marginals, resulting in low accuracy (e.g., 7% on logic puzzles).
• Adaptive Inference is a Game-Changer: By intelligently selecting the order of token unmasking during inference (e.g., based on prediction confidence via "Top Probability Margin"), MDMs can dramatically improve performance. This strategy allows the model to avoid poorly learned sub-problems and prioritize confident predictions.
• Outperforming Strong AR Baselines: Adaptive MDM inference enabled the model to achieve 90% accuracy on logic puzzles, significantly outperforming vanilla MDM and even larger, optimally-trained auto-regressive models (which were 7x larger and had access to optimal decoding orders).
• Broad Applicability and Future Directions: The benefits of adaptive inference extend to natural language, coding (HumanEval), and math tasks. Future research should focus on developing even more sophisticated adaptive strategies and exploring better trade-offs between training flexibility and inference efficiency.

About the Speaker(s)

The talk was presented by Kulin Shah, who introduced the work on understanding token ordering in masked diffusion models. This research is a collaborative effort with Jaeyeon Kim, Vasilis Kontonis, Sham Kakade, and Sitan Chen. While specific titles and affiliations for all authors were not detailed within the transcript, their collective contribution highlights a strong interdisciplinary team working at the forefront of generative AI research, particularly in the domain of diffusion models for language.