MaskDM: Efficient Masked Diffusion Training

• MaskDM is a two-stage diffusion training framework that first learns a primer distribution from heavily masked images before fine-tuning on full images.
• It utilizes ViT-based backbones and carefully designed masking strategies to enhance convergence speed and reduce training computation.
• Empirical results show significant improvements, including up to 80% time savings and new FID records on datasets like CelebA-HQ.

Searching arXiv for "MaskDM" and closely related masked diffusion papers to ground the article with current references. MaskDM is a two-stage training framework for diffusion models that addresses the cost of learning full image distributions from scratch by first learning a simpler “primer distribution” from heavily masked images and then fine-tuning on unmasked images with an ordinary diffusion objective. Introduced in “Masked Diffusion Models Are Fast Distribution Learners” (Lei et al., 2023), the method is formulated for image generation, implemented with ViT-based diffusion backbones, and evaluated on CelebA, CelebA-HQ, and LSUN Church. Its central empirical claim is that masked pre-training yields substantially faster convergence, better low-data transfer, and, on CelebA-HQ $256\times256$, a FID of $6.27$, described as a new record for ViT-based diffusion models in pixel space (Lei et al., 2023).

1. Conceptual basis and definition

MaskDM starts from the observation that standard DDPM-style training asks a model to learn the full data distribution, including coarse structure and fine-grained detail, from the first optimization step. The method proposes that this is unnecessarily expensive. Instead, a diffusion model can first learn a simpler distribution that “loosely characterizes the unknown real image distribution,” then refine that initialization into the true image distribution more efficiently (Lei et al., 2023).

The paper formalizes this intermediate object as a primer distribution $p_\phi(\mathbf{x}_0)$. Let a clean image be written as

$\mathbf{x}_0=(x_0^1,\dots,x_0^N),$

and let $\tau$ index a visible subsequence. The corresponding marginal variable collection is denoted $\hat{\mathbf{x}_0^\tau}$, with marginal distribution $p(\hat{\mathbf{x}_0^\tau})$. The family $\mathcal{Q}$ consists of distributions that share the same set of marginals $p(\hat{\mathbf{x}_0})$, and a primer distribution $p_\phi(\mathbf{x}_0)$ is any member of $6.27$0 other than the true $6.27$1 (Lei et al., 2023).

The method’s guiding intuition is that heavy masking exposes the model to many marginals of the data distribution. By learning to denoise masked observations, the model approximates a distribution that preserves broad image structure and recurring patterns without yet needing to model the full joint distribution of all pixels. This suggests that masked pre-training is not an auxiliary regularizer layered onto diffusion, but a deliberate reparameterization of the early training target.

2. Two-stage training procedure

MaskDM uses a strict two-stage schedule (Lei et al., 2023).

In the masked pre-training stage, each image is heavily masked, with masking rates explored at $6.27$2, $6.27$3, and $6.27$4. The model is trained only on visible regions using a masked denoising score matching objective. The paper explicitly states that the goal is not inpainting the missing content during pre-training; the masked regions are treated as marginalized-out variables, and the model denoises the visible areas.

In the unmasked fine-tuning stage, the masked objective is discarded and the model is trained with the standard diffusion objective on full images. No masking is used in this stage. The masked pre-trained model therefore serves as an initialization that is closer to the target distribution than random weights.

This division of labor is central to the method. Pre-training is meant to learn a primer distribution; fine-tuning is meant to recover the full image distribution. The paper’s interpretation is that the first stage teaches useful visual statistics and partial image structure, while the second stage restores complete long-range dependencies and fine detail. A plausible implication is that MaskDM reduces training cost not by changing the reverse sampler or the denoiser parameterization, but by changing what the model is asked to learn early in optimization.

3. Diffusion formulation and masked denoising objective

MaskDM retains a standard DDPM-style backbone formulation (Lei et al., 2023). The forward process is written as

$6.27$5

with Gaussian transitions

$6.27$6

and closed form

$6.27$7

The reverse process is parameterized through a noise-prediction network $6.27$8, and fine-tuning uses the usual simplified DDPM objective: $6.27$9

MaskDM introduces masking through a binary mask vector

$p_\phi(\mathbf{x}_0)$0

and positional information

$p_\phi(\mathbf{x}_0)$1

The masked clean input and masked noise are defined as

$p_\phi(\mathbf{x}_0)$2

Replacing $p_\phi(\mathbf{x}_0)$3 and $p_\phi(\mathbf{x}_0)$4 in the standard loss yields the masked denoising score matching objective: $p_\phi(\mathbf{x}_0)$5

Several properties follow directly from this construction. The model still predicts noise rather than $p_\phi(\mathbf{x}_0)$6. The training target is the masked noise $p_\phi(\mathbf{x}_0)$7, not full-image noise. The useful supervision is therefore restricted to visible support. Because masks vary, the expectation implicitly ranges over many data marginals. This is the mathematical mechanism by which the masked stage approximates the primer distribution.

4. Backbone, masking strategies, and optimization regime

MaskDM is implemented with U-ViT, a ViT-based diffusion backbone, rather than a convolutional U-Net (Lei et al., 2023). Two configurations are defined:

• MaskDM-S: U-ViT-Small, 44M parameters, depth 13, dim 512, MLP dim 2048, 8 heads.
• MaskDM-B: modified from U-ViT-Mid by removing five transformer blocks, 102M parameters, depth 12, dim 768, MLP dim 3172, 12 heads.

Patch size is $p_\phi(\mathbf{x}_0)$8 in all experiments. The authors also remove the appending convolutional blocks from U-ViT, reporting trivial impact on performance.

Three masking schemes are studied: patch-wise masking, block-wise masking, and cropping. The reported ranking is consistent: block-wise masking performs best, patch-wise masking is second, and cropping is worst. The default masking policy is:

• $p_\phi(\mathbf{x}_0)$9: $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$0 block-wise masking
• $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$1 and above: $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$2 block-wise masking

The mask rate is

$\mathbf{x}_0=(x_0^1,\dots,x_0^N),$3

where $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$4 is the number of visible variables. Under matched budget on CelebA $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$5, the best ablation is $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$6 block-wise masking at $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$7, yielding FID $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$8, compared with $\mathbf{x}_0=(x_0^1,\dots,x_0^N),$9 from scratch. By contrast, large blocks and extreme cropping perform worse. This suggests that retaining visible content distributed across the image is more useful than preserving a single local crop.

Optimization details are dataset-dependent. On CelebA $\tau$0, the best setup uses $\tau$1 block-wise masking at $\tau$2, learning rate $\tau$3, batch size 256, 150k pre-training steps, and a linear noise schedule; fine-tuning uses learning rate $\tau$4, batch size 128, 350k steps, and EMA $\tau$5. Higher-resolution settings use progressive masking schedules such as $\tau$6 or $\tau$7, with cosine schedules, warmup, and gradient clipping. The appendix states that $\tau$8 masking can be unstable under a linear noise schedule, and highlights the cosine schedule as especially important for stabilizing such high mask rates (Lei et al., 2023).

5. Empirical performance, acceleration, and transfer

The paper reports that MaskDM saves about 60% training time on lower-resolution settings and up to about 80% on CelebA-HQ $\tau$9, with gains increasing at higher resolution (Lei et al., 2023). It also reports stronger final FIDs than training from scratch under comparable or lower total compute.

Dataset	U-ViT FID	MaskDM FID
CelebA $\hat{\mathbf{x}_0^\tau}$0	2.87	2.27
LSUN Church $\hat{\mathbf{x}_0^\tau}$1	6.58	5.04
CelebA $\hat{\mathbf{x}_0^\tau}$2	12.96	6.83
CelebA-HQ $\hat{\mathbf{x}_0^\tau}$3	24.83	6.27

The CelebA-HQ $\hat{\mathbf{x}_0^\tau}$4 result is the paper’s flagship number. MaskDM-B reaches FID 6.27 at 12.19 A100 days, whereas U-ViT is reported at FID 24.83 with 18.28 A100 days. The paper characterizes $\hat{\mathbf{x}_0^\tau}$5 as a new FID record for ViT-based diffusion models in pixel space (Lei et al., 2023).

Ablations show that pre-training budget continues to help. On CelebA $\hat{\mathbf{x}_0^\tau}$6, $\hat{\mathbf{x}_0^\tau}$7 block-wise masking at $\hat{\mathbf{x}_0^\tau}$8 improves from FID $\hat{\mathbf{x}_0^\tau}$9 at 50k pre-training steps to $p(\hat{\mathbf{x}_0^\tau})$0 at 100k and $p(\hat{\mathbf{x}_0^\tau})$1 at 150k. Large mask blocks degrade performance, and cropping is consistently weaker than distributed masking. This suggests that the quality of the primer distribution depends not only on how much of the image is hidden, but also on how the visible information is spatially distributed.

Transfer is a second major result. The authors report that a model pre-trained on VGGFace2 attains a 46% quality improvement when fine-tuned on a different dataset that contains only 3000 images, and still gives a 42% improvement with only 300 images (Lei et al., 2023). They also show beneficial transfer from CelebA to FFHQ and AFHQ, and even across diffusion training formulations such as DDPM, VPSDE, and VPCosine. This supports the paper’s interpretation that masked pre-training learns reusable visual structure rather than only dataset-specific shortcuts.

6. Scope, limitations, and relation to later masking-based diffusion work

MaskDM is image-specific in its presented form, and its evidence is concentrated on pixel-space image generation with U-ViT backbones. The paper also identifies several practical caveats. Mask configuration matters substantially; poor mask shapes or mask rates can hurt both speed and final quality. Very high mask rates such as $p(\hat{\mathbf{x}_0^\tau})$2 can be unstable without cosine schedules, warmup, lower learning rates, or gradient clipping. The primer distribution is a conceptual object rather than a directly estimated density, so the argument is empirical rather than theorem-driven (Lei et al., 2023).

Later work reused masking in diffusion in several distinct senses. “Di$p(\hat{\mathbf{x}_0^\tau})$3O: Distilling Masked Diffusion Models into One-step Generator” studies one-step distillation for discrete masked diffusion models and introduces token-level distribution matching on pseudo-intermediate masked states plus a token initialization strategy that mixes random image tokens with mask tokens (Zhu et al., 19 Mar 2025). “DualDiffusion: A Speculative Decoding Strategy for Masked Diffusion Models” targets masked diffusion LLMs and alternates multiple drafter steps with verifier-based remasking to improve the quality–efficiency trade-off (Goyal et al., 6 Apr 2026). “Timestep-Aware Block Masking for Efficient Diffusion Model Inference” uses a binary mask matrix $p(\hat{\mathbf{x}_0^\tau})$4 to decide which denoiser blocks to execute or bypass at each timestep in pretrained DDPM, LDM, DiT, and PixArt models (He et al., 20 Mar 2026). In a different application domain, “A Guided Unconditional Diffusion Model to Synthesize and Inpaint Radio Galaxies from FIRST, MGCLS and Radio Zoo” trains a single DDPM on mixtures of partially masked and fully masked radio images so that one network performs both inpainting and unconditional synthesis (Potevineau et al., 12 Jan 2026).

These later papers show that “masking” became a broad design pattern in diffusion research: as a pre-training mechanism, as a discrete generative process, as an inference-acceleration control signal, and as a spatial guidance mechanism. The exact name MaskDM, however, refers specifically to the two-stage primer-distribution training framework of “Masked Diffusion Models Are Fast Distribution Learners” (Lei et al., 2023). Its enduring contribution is the claim that diffusion models need not learn the full image distribution from scratch: learning many masked marginals first can be a more efficient route to strong generative performance.