PaperlayerView exact revision on arXiv

Source-linked AI summary

Simplified and Generalized Masked Diffusion for Discrete Data

Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, Michalis K. Titsias

arXiv:2406.04329v4cs.LGstat.ML

TL;DR

Masked diffusion models have been limited by complex formulations and weak competitiveness with autoregressive models for discrete data. This paper develops a simpler continuous-time ELBO and generalized state-dependent masking framework, achieving stronger text and image-modeling results while retaining important limitations in some settings.

  • Problem

    Existing discrete masked diffusion models use complex formulations and objectives and have not matched autoregressive models in important discrete-data domains.

  • Method

    The paper derives a continuous-time masked diffusion framework with an ELBO expressed as a weighted cross-entropy integral and introduces state-dependent masking schedules.

  • Results

    Masked diffusion models outperform previous discrete diffusion proposals on GPT-2-scale text and pixel-level image modeling, with the best reported likelihood and zero-shot transfer among discrete diffusion models.

  • Takeaways & Limitations

    The simplified framework improves masked diffusion performance across text and pixel-level image modeling while providing a generalized formulation for state-dependent schedules.

  • Takeaways & Limitations

    Masked diffusions remain noncompetitive with autoregressive models on some tasks such as text8, while GenMD4 can overfit and makes zero-shot transfer and inference more challenging.

Abstract

from arXiv · show

Masked (or absorbing) diffusion is actively explored as an alternative to autoregressive models for generative modeling of discrete data. However, existing work in this area has been hindered by unnecessarily complex model formulations and unclear relationships between different perspectives, leading to suboptimal parameterization, training objectives, and ad hoc adjustments to counteract these issues. In this work, we aim to provide a simple and general framework that unlocks the full potential of masked diffusion models. We show that the continuous-time variational objective of masked diffusion models is a simple weighted integral of cross-entropy losses. Our framework also enables training generalized masked diffusion models with state-dependent masking schedules. When evaluated by perplexity, our models trained on OpenWebText surpass prior diffusion language models at GPT-2 scale and demonstrate superior performance on 4 out of 5 zero-shot language modeling tasks. Furthermore, our models vastly outperform previous discrete diffusion models on pixel-level image modeling, achieving 2.75 (CIFAR-10) and 3.40 (ImageNet 64x64) bits per dimension that are better than autoregressive models of similar sizes. Our code is available at https://github.com/google-deepmind/md4.

1 Introduction

The paper simplifies masked diffusion models through a continuous-time framework, clarifying their objectives and relationships to prior formulations. This framework improves discrete diffusion performance on text and pixel-level image modeling.

  • Motivation: Masked diffusions have lagged behind autoregressive models in important discrete-data domains because prior formulations and training objectives were unnecessarily complex.The paper positions simplification as a route toward closing this performance gap.
  • Framework: The paper studies masked, or absorbing, diffusions within a flexible continuous-time framework.This formulation is used to improve both training and understanding of the model class.
  • Objective: The continuous-time ELBO becomes a weighted integral of cross-entropy losses and can also be expressed using signal-to-noise ratio with invariance properties.This provides a simpler training objective for masked diffusion models.
  • Prior formulations: The framework unifies prior continuous-time discrete diffusion models and identifies changes that can cause expensive evaluations, high training variance, or forward-reverse inconsistency.These issues arise from modifications to the ELBO objective or model parameterization.
  • Results: Masked diffusions using the simple ELBO outperform previous proposals on GPT-2-scale text and pixel-level image modeling, with the best likelihood and zero-shot transfer among discrete diffusion models.The reported outcome spans both text and image modeling tasks.
  • Generalization: The generalized formulation permits state-dependent masking schedules and further improves predictive performance measured by test likelihoods.The schedule can depend on the current process state rather than time alone.

2 Masked Diffusion

Masked diffusion progressively converts discrete tokens into a special mask state and learns a reverse process that reconstructs the original data. Its continuous-time formulation provides tractable marginals, reversals, and masking schedules.

  • Process definition: The forward process independently replaces discrete tokens with a special mask state, while the learned reverse process turns masked sequences back into coherent data.The discrete state space is augmented with one additional mask state.
  • Discrete-time process: Each discrete-time transition preserves the current state with probability 1−β_i and jumps to the mask state with probability β_i.The transition matrix has size (m+1) × (m+1), including the mask category.
  • Terminal behavior: The forward process is expected to approach a point mass at the mask state at sufficiently large terminal time.This requires α_T to become very small or zero.
  • Continuous-time limit: The continuous-time process is obtained by setting β_i=β(t(i))/T and taking T→∞, yielding compatible marginal and transition distributions.The resulting process supports closed-form quantities for arbitrary times.
  • Masking schedules: Masking schedules are designed with α_0=1 and α_1 near zero, and the literature includes linear, geometric, and cosine choices.The cosine schedule slows reverse unmasking early and can reduce simultaneous conflicting unmaskings.
  • Time reversal: For the conditioned time reversal, a masked token either jumps to its original state or remains masked, while an unmasked token stays unchanged.The jump probability is determined by the masking schedule between the two times.

3 Model and Objective

The model approximately reverses the masked forward process with a neural network predicting clean-token distributions. Its ELBO simplifies in continuous time to a cross-entropy-based objective whose schedule dependence is characterized through signal-to-noise ratio.

  • Reverse model: The generative model approximately reverses forward transitions using a reverse distribution p_θ(x_s|x_t).The reverse model is conditioned on the current corrupted state and time.
  • Parameterization: The neural network uses mean-parameterization to predict a probability vector for the clean data, with the mask output forced to zero.The model predicts the mean or clean-token distribution rather than masked data.
  • Boundary distributions: The model specifies a clean-data conditional and a terminal prior concentrated near the mask state when α_1 is near zero.These components complete the reverse model and its boundary conditions.
  • ELBO: The ELBO is a lower bound on the log marginal likelihood and decomposes into KL divergence terms for the reverse process.The objective is defined for the discrete-time diffusion model before taking the continuous-time limit.
  • Simplified objective: The reverse-process KL terms reduce to cross-entropy between predicted logits and the clean data, and the discrete objective forms a lower-bounded Riemann sum of a continuous integral.The resulting simple masked diffusion model is called MD4.
  • Continuous-time objective: The continuous-time limit produces an ELBO equal to −L_∞ after sending the initial evaluation time toward zero.The boundary reconstruction term vanishes in this limit.
  • Schedule dependence: Expressing α_t through log-SNR shows that the loss depends on schedule endpoints, while uniform time sampling still makes the schedule affect estimator variance.The schedule changes variance even when its endpoint SNR values determine the loss dependence.

4 Sampling

MD4 uses ancestral sampling for generation, while cosine schedules and discretization grids substantially improve sample quality by reducing conflicts during iterative unmasking.

  • Ancestral sampling produces slightly higher sample quality than Euler discretization for MD4 generation.
  • FID on ImageNet 64×64 improves from 70 to 17 with cosine scheduling at T = 256 sampling steps.
  • Cosine discretization slows early unmasking, reducing conflicts when multiple tokens are revealed in one step.The same cosine grid can also be used when models are trained with a linear schedule, yielding comparable sample quality.
  • The schedule improvement is attributed to information redundancy: revealed tokens make remaining tokens more predictable.
  • Sampling findings developed on images also translate well to text, while top-p sampling, classifier-free guidance, and predictor-correctors remain future improvements.

5 Relation to Existing Work

The paper relates masked diffusion to CTMCs, score-based formulations, and any-order autoregressive models, identifying computational and consistency trade-offs in prior approaches.

  • The masked diffusion framework unifies several existing continuous-time discrete diffusion models and identifies their objective or parameterization changes.
  • Continuous-Time Markov Chains (CTMC): Forward and reverse masked diffusion transitions can be expressed with CTMC transition-rate matrices Q(t) and R_x0(t).
  • Continuous-Time Markov Chains (CTMC): Rate-matrix ELBO formulations can require N prediction-model evaluations, whereas the analytic cross-entropy objective requires one pass and has lower variance.The doubly stochastic estimator used to reduce the former computational burden has significantly higher variance.
  • Score parameterization: Score parameterization can break compatibility between learned reverse and forward processes when its constraint is not enforced.The paper's constraint-enforcing parameterization yields more stable training and better results.
  • Any-order autoregressive models: The continuous-time reverse process is equivalent to an any-order autoregressive model after ordering tokens by their unmasking times.The masking schedule remains an additional design degree of freedom because it changes the distribution of unmasking times.

6 Generalization to State-dependent Masking Schedules

GenMD4 generalizes masked diffusion by making masking schedules depend on token values, enabling selected tokens to unmask earlier while retaining a valid continuous-time Markov formulation.

  • State-dependent masking schedules make a token's unmasking probability depend on both time and token value.This targets settings where some tokens are more significant and should be revealed earlier.
  • The approach is motivated by revealing more significant tokens earlier, but learned schedules can overfit dataset statistics and weaken zero-shot transfer.
  • The generalized forward process assigns each possible token value its own masking function α_t,i and defines transitions through a categorical Markov process.
  • The state-dependent process remains a valid continuous-time Markov chain through consistent marginal and transition relationships.
  • GenMD4 extends the MD4 loss to state-dependent rates, recovering MD4 when α_t is a scalar schedule multiplied by an all-ones vector.
  • Token-dependent schedules use α_t,i = 1 − t^w_i, where w_i controls how quickly tokens of value i transition to the mask state.The model learns the token-specific parameters w_i > 0 through ELBO optimization.

7 Experiments

The experiments evaluate masked diffusion on text and pixel-level image modeling, finding strong likelihood and transfer results alongside schedule- and scale-dependent trade-offs.

  • Text: MD4 outperforms previous discrete diffusion models across all five zero-shot tasks at small scale and on four of five tasks at medium scale.It also exceeds GPT-2 on all tasks except LAMBADA at small scale, where it is second best.
  • Text: MD4 models converge faster and achieve better final likelihoods than prior methods on OpenWebText.The evaluation uses zero-shot perplexity on benchmark datasets, with training curves shown on an OpenWebText validation subset.
  • Text: On Text8, MD4 outperforms previous discrete and continuous diffusion models and closely related AO-ARMs, but remains behind an autoregressive transformer and Discrete Flow.GenMD4 achieves significantly better BPC than MD4 on this small-vocabulary benchmark.
  • Text: MD4 is outperformed by same-sized autoregressive models on Hellaswag, with AR-small reaching 30% accuracy in 50k steps versus MD4-small in 200k steps.At medium scale, AR reaches 37% in 270k steps compared with MD4 in 1 million steps.
  • Pixel-level image modeling: MD4 establishes a new state of the art among discrete diffusion models for pixel-level images, surpassing the best reported AR result on CIFAR-10 and matching strong larger or continuous baselines on ImageNet 64×64.Images are represented as discrete pixel tokens without using pixel-value ordinal structure; the reported comparisons are in bits per dimension.

8 Conclusion

The paper presents a simplified continuous-time framework for masked diffusion, including a weighted cross-entropy ELBO and state-dependent masking schedules. It reports improved text and image modeling, while identifying remaining limitations in competitiveness, transfer, and inference.

  • Contributions: The framework expresses the masked-diffusion ELBO as a weighted time integral of cross-entropy losses and extends it to state-dependent masking schedules.The generalized formulation is called GenMD4.
  • Results: MD4 outperforms existing discrete and continuous diffusion models on text, while pixel-level image results outperform similar-sized autoregressive models and approach continuous diffusion likelihoods.GenMD4 provides further likelihood improvements over the state-independent formulation.
  • Limitations: Masked diffusions remain noncompetitive with autoregressive models on some tasks, including text8.The authors conjecture that autoregressive models better use capacity because they learn one generation order.
  • Limitations: GenMD4 can overfit dataset statistics, reducing zero-shot transfer effectiveness, and state-dependent schedules make inference more challenging.These limitations constrain the practical use of the generalized formulation.

A Discrete-time derivation

The discrete-time derivation divides the diffusion interval into T steps and defines a forward transition matrix over vocabulary tokens plus a mask state. Its marginal transition is formed by composing the stepwise matrices and taking a continuous-time limit.

  • Time discretization: Time is divided into T intervals with endpoints s(i) = (i − 1)/T and t(i) = i/T.The derivation uses these intervals to define discrete forward transitions.
  • Forward process: The forward transition matrix Q_i has dimension (m+1)×(m+1), where m is vocabulary size and the additional state is the mask.It preserves tokens with probability 1 − β_i, masks tokens with probability β_i, and otherwise has zero transition probability.
  • Forward process: The state x_t is represented as a one-hot vector of length m+1, and the forward conditionals specify probabilities over the possible states at each time step.The extra category corresponds to the absorbing mask state.
  • Marginal distribution: The marginal forward distribution is obtained by composing the stepwise transition matrices as Q̄_i = ∏_{j=1}^i Q_j.The product captures the cumulative noising process from the initial state to time t(i).
  • Continuous-time limit: As T approaches infinity, the cumulative product converges to a continuous-time marginal transition Q̄(t).The limiting expression is obtained from the accumulated masking schedule.

B Continuous-time derivation

The section derives the continuous-time masked diffusion process, its reverse transitions, and a simplified training objective. It also contrasts this objective with prior formulations and identifies numerical and computational issues in alternative approaches.

  • Continuous-time process: The continuous-time forward process is obtained by taking the infinite-step limit of the discrete-time masked Markov process.The transition matrix between arbitrary times solves the Kolmogorov forward equation, with a matrix-exponential solution.
  • Reverse process: The analytic forward process permits closed-form computation of the marginal transition distribution and the time reversal conditioned on the original token.The reverse transition can be represented as a categorical distribution or an explicit matrix whose masked-state row describes unmasking behavior.
  • Training objective: Proposition 3 rewrites the continuous-time objective using the approximate reverse transition-rate matrix obtained from the model prediction.The reformulation connects the objective to transition-rate expectations while absorbing terms independent of θ into a constant.
  • Comparison with prior objectives: The prior transition-rate loss requires N neural-network evaluations for multidimensional data, whereas the proposed loss requires one evaluation for a single stochastic estimate.Monte Carlo rewriting of the prior loss produces a doubly stochastic estimator with large variance.
  • Consistency and numerical issues: The score parameterization used by Lou et al. can violate consistency between the learned reverse model and the forward process.The section also notes that finite endpoint masking introduces a support mismatch unless the prior treatment is modified.

I.1 Derivations and time continuous limit

This appendix extends the masked-diffusion derivations to multiple tokens and state-dependent rates. It gives unbiased gradient procedures for both neural-network and schedule parameters.

  • State-dependent extension: For N tokens with state-dependent rates, the masked diffusion factorizes across tokens while the masking probability of each token can depend on the full state.The full negative lower bound uses weights determined by the state-dependent rates.
  • Reverse transitions: The reverse transition for a masked token matches the corresponding forward-process reversal, while an unmasked token remains fixed.The derivation treats separately the cases xt = m and xt ≠ m.
  • Reverse parameterization: The reverse model predicts an m + 1 dimensional categorical distribution through a neural network that receives the full masked state and time features.The mask output is constrained because the original data tokens cannot themselves be mask states.
  • Gradient estimation: An unbiased neural-network gradient is obtained by sampling one time point and one noisy state from q(xt|x0) to estimate the integral and expectation.The schedule parameters are harder because they also determine the sampling distribution q(xt|x0).
  • Gradient estimation: RLOO with two samples is used to estimate the REINFORCE term for schedule parameters, which can otherwise have high variance.The direct gradient component is straightforward, whereas the distribution-dependent term requires the variance-reduction procedure.

J Experimental Details

The experiments use continuous-time training with discrete-time reverse sampling across text and image benchmarks. The section documents model architectures, optimization settings, evaluation procedures, and reported benchmark artifacts.

  • Common setup: All experiments train with a continuous-time loss and generally generate samples using a 1000-timestep discrete-time reverse process.An exponential moving average factor of 0.9999 is used for evaluation and generation.
  • Text experiments: OpenWebText models use standard 12-layer transformers with GPT-2-compatible dimensions and tokenizer settings.Text chunks have length 256, training lasts 1 million steps, and the batch size is 512.
  • Text evaluation: Hellaswag predictions are made by comparing likelihood estimates for padded question–answer sequences using 32 Monte Carlo samples.The answer option with the highest estimated likelihood is selected.
  • Image experiments: Image experiments use U-Net plus self-attention architectures, with 2 million training iterations and dataset-specific optimization settings.CIFAR-10 and ImageNet 64×64 use different learning rates, batch sizes, weight decay factors, and TPU allocations.
  • Reported evaluations: The section reports OpenWebText perplexity and ImageNet FID results through Tables 5–6 and corresponding figures.Additional unconditional ImageNet samples and MD4 Medium text samples are provided for qualitative inspection.

K.5.2 MD4-M unconditional sample 2: 1024 tokens

The sample and ablation material compares generated text under different masking schedules and discretization choices. It emphasizes that cosine-style unmasking improves qualitative sample coherence and that the JAX implementation has a sampling-related caveat.

  • Conditional text samples: Cosine-schedule sampling slows unmasking near t = 1 and tends to produce more consistent and meaningful conditional text than uniform unmasking.This comparison is illustrated with MD4-M samples using linear and cosine training or sampling schedules.
Loading 2406.04329v4…