Unified Masked Diffusion (UMD) Overview

• Unified Masked Diffusion (UMD) is a self-supervised framework that unifies masked autoencoding and diffusion denoising to learn robust image representations and generate high-quality visuals.
• It employs a single encoder-decoder transformer that integrates coarse patch masking with fine-grained Gaussian noise, streamlining the training process without heavy augmentations or multiple views.
• Empirical studies show UMD achieves competitive linear probing accuracy and strong generative performance, balancing computational efficiency with improved feature learning.

Unified Masked Diffusion (UMD) is a self-supervised auto-encoding objective introduced in “Unified Auto-Encoding with Masked Diffusion” that combines masked autoencoding and diffusion denoising within a single encoder-decoder transformer, with the stated goal of learning features useful for both representation learning and image generation (Hansen-Estruch et al., 2024). In this formulation, masked-patch reconstruction and diffusion denoising are treated as two corruption-and-reconstruction regimes at different granularities: coarse patch removal and fine-grained Gaussian corruption. UMD therefore reframes masked autoencoders (MAE), diffusion transformers (DiT), and intermediate masked-diffusion variants as points inside one broader denoising problem, while retaining a single backbone rather than multiple views, heavy augmentations, auxiliary encoders, or separate representation and generation modules (Hansen-Estruch et al., 2024).

1. Conceptual origin and scope

UMD arises from the observation that both generative diffusion models and self-supervised masked autoencoders are driven by reconstruction objectives under corrupted inputs. Diffusion models corrupt continuously through a scheduled Gaussian process, whereas MAE corrupts discretely by removing image patches; the paper’s main claim is that these are not fundamentally different procedures but two instances of the same reconstruction principle (Hansen-Estruch et al., 2024). The practical motivation is equally explicit: standard MAEs learn strong transferable representations but produce poor generation, especially blurry infilling, whereas DiT-style diffusion models generate well but are computationally heavy and usually yield weaker linear-probe representations unless additional self-supervised machinery is added (Hansen-Estruch et al., 2024).

UMD addresses this by introducing a single pretraining objective that preserves MAE-like representation quality while retaining DiT-like generative capability. The paper emphasizes that this is achieved without heavy data augmentations, multiple views, or additional encoders, and with improved total training time relative to prior diffusion-based methods (Hansen-Estruch et al., 2024). In this sense, UMD is not merely a hybrid of two losses; it is an augmented diffusion schedule containing a specially emphasized representation-learning step and a mixed masked/noised corruption process for the remaining steps (Hansen-Estruch et al., 2024).

Later work extends closely related unification themes in discrete diffusion. “Simplified and Generalized Masked Diffusion for Discrete Data” derives a continuous-time masked-diffusion objective as a weighted integral of masked-token cross-entropy losses and generalizes the framework to state-dependent masking schedules (Shi et al., 2024). “Effective and Efficient Masked Image Generation Models” unifies masked image generation and masked diffusion under a single masked-denoising family parameterized by the masking distribution, weighting function, and conditional parameterization (You et al., 10 Mar 2025). These developments suggest that UMD is part of a broader research trajectory in which masking, denoising, and generation are treated as variants of one reconstruction-based formalism.

2. Objective and corruption process

UMD begins from the standard diffusion forward process. Given an image $x_0$, a noisy sample at timestep $t$ is drawn as

$$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$

with $\epsilon \sim \mathcal{N}(0,I)$ (Hansen-Estruch et al., 2024). UMD modifies this standard setup in two specific ways. First, it adds a special step $t=0$ with no Gaussian corruption, effectively a noise-free but heavily masked representation-learning stage. Second, for $t \ge 1$, it applies masking not to clean images but to already noised samples, producing mixed masked-and-noised inputs $x_t^M$ (Hansen-Estruch et al., 2024).

This yields an asymmetric objective across visible and masked regions. For noisy timesteps,

$$\mathcal{L}(\theta)_{t \ge 1} = \mathbb{E}_{t,\epsilon,x_0,M} \left[ \left\| M \odot \bigl(x_0 - x_\theta(x_t^M,t)\bigr)\right\|_2^2 + \left\| (1-M) \odot \bigl(\epsilon - \epsilon_\theta(x_t^M,t)\bigr)\right\|_2^2 \right].$$

Masked patches are supervised with clean-image reconstruction $x_0$, whereas visible noised patches are supervised with noise prediction $\epsilon$ (Hansen-Estruch et al., 2024). The paper explicitly justifies this split by noting that exact Gaussian noise prediction is difficult on fully masked tokens. At the special no-noise step, the loss reduces to MAE-style masked reconstruction: $t$0 The unified objective is then

$t$1

This formulation interpolates between prior methods by changing the $t$2 sampling ratio and masking ratios: $t$3 gives MAE; $t$4 gives MaskDiT; and $t$5 gives DiT (Hansen-Estruch et al., 2024).

The conceptual role of the two parts is explicit. The $t$6 heavy-masking step promotes semantic, context-dependent inference of missing content and thereby strengthens representations, while the $t$7 noisy steps teach local denoising dynamics needed for iterative sampling (Hansen-Estruch et al., 2024).

3. Architecture and optimization regime

UMD uses a ViT-based asymmetric encoder-decoder that combines MAE and DiT design choices. For $t$8 pixel experiments, the encoder is ViT-B/4 and the decoder is a shallow 4-layer transformer. For $t$9 latent experiments, the encoder is ViT-L/2 and the decoder has 8 layers with width 1024 and 16 heads (Hansen-Estruch et al., 2024). Images or latents are patchified and linearly embedded through a convolutional projection, a CLS token is prepended, and the encoder output CLS token is used as the representation for linear probing. Following MAE, only visible patches are processed by the encoder and masked patches are inserted as learned mask embeddings in the decoder, reducing sequence length and improving efficiency (Hansen-Estruch et al., 2024).

Conditioning follows DiT conventions. Timestep information is embedded by a 256-dimensional sinusoidal/frequency embedding followed by a two-layer MLP with SiLU, and transformer blocks are modulated through AdaLN-zero. The paper states that the conditioning MLP outputs scale, shift, and gating values, with learnable positional embeddings in both encoder and decoder and a transposed convolution for reconstruction to image space (Hansen-Estruch et al., 2024). During class-conditional fine-tuning, class embeddings are added to timestep embeddings before AdaLN modulation.

On $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$0 ImageNet-1K, UMD is trained in pixel space with a cosine diffusion schedule, 1000 beta steps, and typically 125 DDIM sampling steps at inference. Pretraining uses AdamW, batch size 1024, weight decay 0.05, learning rate $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$1, $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$2, cosine decay, 40 warmup epochs, gradient clip 1.0, and 800 training epochs (Hansen-Estruch et al., 2024). Augmentation is intentionally light: random resized crop with scale $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$3 to $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$4 and random horizontal flips. The default UMD hyperparameters are

$$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$5

For class-conditional fine-tuning, the pretrained model is trained for 50 epochs on labeled unmasked images with batch size 256, learning rate $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$6, weight decay 0, $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$7, cosine schedule, 2.5 warmup epochs, EMA decay factor 0.00025, and classifier dropout 10%. Fine-tuning uses $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$8, $$q(x_t \mid x_0) = \mathcal{N}\!\left(\sqrt{\hat{\alpha}_t}x_0,\; (1-\hat{\alpha}_t)I\right) = \sqrt{\hat{\alpha}_t}x_0 + \sqrt{1-\hat{\alpha}_t}\,\epsilon,$$9, and $\epsilon \sim \mathcal{N}(0,I)$0, with DDIM 250-step sampling and typical classifier-free guidance scale $\epsilon \sim \mathcal{N}(0,I)$1 for reported metrics (Hansen-Estruch et al., 2024).

For higher-resolution experiments, UMD switches to latent diffusion. A VAE/VQGAN-style tokenizer maps $\epsilon \sim \mathcal{N}(0,I)$2 images to $\epsilon \sim \mathcal{N}(0,I)$3 latents; diffusion is then trained in latent space with a linear beta schedule, 1000 steps, and 250 DDIM sampling steps (Hansen-Estruch et al., 2024).

4. Empirical profile on ImageNet

On $\epsilon \sim \mathcal{N}(0,I)$4 ImageNet-1K with 100-shot linear probing, UMD narrows much of the representation gap between diffusion models and MAE while remaining substantially cheaper than DiT (Hansen-Estruch et al., 2024).

Method	100-shot LP	TPU-v4-8 hours
MAE	36.6%	45
DiT	25.7%	105
MaskDiT	26.3%	72
UMD	31.8%	60

The transfer plot across STL-10, Flowers, DTD, CIFAR-100, Oxford Pets, Food101, Stanford Dogs, and ImageNet-v2 shows UMD consistently competitive with MAE and better than diffusion-only baselines, supporting the claim that the additional $\epsilon \sim \mathcal{N}(0,I)$5 step materially improves learned features (Hansen-Estruch et al., 2024).

After 50 epochs of labeled class-conditional fine-tuning on $\epsilon \sim \mathcal{N}(0,I)$6, UMD remains close to DiT and MaskDiT in FID while retaining stronger representation quality than pure diffusion baselines.

Method	FID / IS	FID-G / IS-G
MAE	34.1 / 13.4	26.8 / 18.5
DiT	21.2 / 23.0	18.9 / 46.9
MaskDiT	22.1 / 22.1	19.0 / 43.4
UMD	23.2 / 20.2	19.8 / 63.5

UMD is therefore slightly worse than DiT or MaskDiT in FID, but under classifier-free guidance it attains the highest inception score among the listed methods (Hansen-Estruch et al., 2024). The fine-tuning curves add a further practical point: UMD converges quickly in FID and IS during labeled fine-tuning and largely preserves representation quality, whereas MAE representation quality collapses as it is repurposed for generation (Hansen-Estruch et al., 2024).

At $\epsilon \sim \mathcal{N}(0,I)$7 in latent space, UMD becomes stronger still as a representation learner. At 400 epochs, MAE achieves 38.8% 100-shot accuracy, DiT trained with labels from scratch achieves 37.9%, while UMD pretrained for 400 epochs and fine-tuned for 50 labeled epochs reaches 44.6%. At 800 epochs, MAE reaches 51.1%, while UMD reaches 54.4% with FID 12.9, IS 82.0, FID-G 3.96, and IS-G 212.6 (Hansen-Estruch et al., 2024). This is one of the paper’s strongest indications that a single unified objective can support both discriminative transfer and generative quality.

5. Ablations, trade-offs, and limitations

The ablation study makes clear that UMD’s performance depends on the interaction among the $\epsilon \sim \mathcal{N}(0,I)$8 representation step, masked/noised corruption at later steps, and the dual prediction targets $\epsilon \sim \mathcal{N}(0,I)$9 and $t=0$0 (Hansen-Estruch et al., 2024). The base ablation model reports 30.9% 100-shot linear probing, with generation quality depending on which prediction is used at sampling: FID $t=0$1 and IS $t=0$2.

Several findings are especially diagnostic. Training with only $t=0$3-reconstruction gives 31.3% linear probing but poor generation (FID 26.7, IS 43.0), while training with only $t=0$4-prediction gives 31.5% linear probing but also worse generation (FID 25.0, IS 49.0). Lowering the $t=0$5 sampling ratio harms representation quality, dropping linear probing to 27.7%. Increasing masking on noisy steps harms generation, yielding 28.9% linear probing and FID 27.5/23.3 with IS 53.2/50.5 (Hansen-Estruch et al., 2024). These results support the paper’s interpretation that the noise-free masked step is the main source of representation gains, while excessive masking during noisy steps weakens the diffusion model.

The strongest raw numbers occur when noisy-step masking is removed altogether. Setting $t=0$6 gives 32.7% linear probing, FID 21.1/18.8, and IS 65.1/63.6, but sacrifices the computational advantage because all patches must then be processed on noisy steps (Hansen-Estruch et al., 2024). UMD therefore occupies a deliberate Pareto point rather than dominating every baseline in every metric.

The paper also reports that removing AdaLN harms both representation learning and generation badly, and that latent UMD does not achieve the same wall-clock speedups over DiT as pixel-space UMD because VAE preprocessing adds overhead (Hansen-Estruch et al., 2024). Two further caveats are explicit. First, the corruption scheme remains hand-designed: coarse patch masking and fine Gaussian noise coexist rather than emerging from a learned continuous corruption family. Second, on $t=0$7, UMD still does not fully match MAE’s linear-probe accuracy or DiT’s best FID (Hansen-Estruch et al., 2024). The authors therefore describe future work on a learned gating mechanism that could interpolate more smoothly between coarse and fine corruption.

6. Subsequent developments and broader interpretations

Later masked-diffusion research broadens the unification program along several axes. In discrete data modeling, MD4 shows that masked diffusion can be written as a simple weighted masked-token cross-entropy objective and extended to state-dependent masking schedules, effectively generalizing the masked-diffusion family beyond a single global schedule (Shi et al., 2024). In image generation, eMIGM shows that masked image generation models such as MaskGIT and MAR, and masked diffusion models, can be unified within one training and sampling framework parameterized by the masking distribution, weighting function, and per-token conditional decoder (You et al., 10 Mar 2025). In language modeling, OeMDM treats generation order itself as part of the masked-diffusion parameterization, giving a single framework that can express standard MDMs, autoregressive models, and block diffusion (Hong et al., 2 Feb 2026).

A second line of work extends unification across modalities. “The Design Space of Tri-Modal Masked Diffusion Models” trains a single masked discrete diffusion model from scratch on text, image-text, and audio-text data using one transformer backbone and one shared token stream (Bethune et al., 25 Feb 2026). “Omni-Diffusion” presents a mask-based discrete diffusion model that unifies understanding and generation across text, speech, and images in an any-to-any multimodal setting (Li et al., 6 Mar 2026). “Lavida-O” extends unified masked diffusion to image-level understanding, object grounding, image editing, interleaved generation, and $t=0$8 text-to-image synthesis through an Elastic Mixture-of-Transformers architecture and planning/self-reflection loops (Li et al., 23 Sep 2025). These papers do not reproduce the exact image-space UMD objective of (Hansen-Estruch et al., 2024), but they extend the same general thesis: masked denoising can serve as a common substrate for multiple objectives, modalities, and generation interfaces.

A third line of work studies theory and systems aspects of masked diffusion rather than the specific MAE–diffusion bridge. An energy-minimization interpretation shows that masked diffusion models can be understood as minimizing equivalent transport energies under suitable schedules (Chen et al., 17 Sep 2025). Self-conditioning adaptation shows that carrying forward previous clean-state predictions can turn repeated masked-token inference into iterative refinement without extra denoiser evaluations during sampling (Cardei et al., 28 Apr 2026). Progressive UnMAsking (PUMA) aligns training-time masks with inference-time unmasking policies and reports approximately 2.5× pretraining speedup at 125M scale (Kim et al., 10 Feb 2026). At the same time, “Scaling Beyond Masked Diffusion LLMs” cautions that masked diffusion is important but not sufficient as a universal organizing principle for discrete diffusion language modeling, since perplexity can be informative within a family but misleading across families (Sahoo et al., 16 Feb 2026).

Taken together, these later papers suggest that Unified Masked Diffusion is best understood not only as the specific image objective proposed in (Hansen-Estruch et al., 2024), but also as a broader research direction concerned with collapsing masked reconstruction and diffusion-style denoising into shared objectives, shared schedulers, shared backbones, or shared multimodal token spaces. Within that broader landscape, the 2024 UMD proposal remains notable for making the MAE–DiT connection explicit and for showing, in a single encoder-decoder transformer, that one pretraining objective can support both linear probing and class-conditional generation (Hansen-Estruch et al., 2024).