PaperlayerView exact revision on arXiv

Source-linked AI summary

Effective and Efficient Masked Image Generation Models

Zebin You, Jingyang Ou, Xiaolu Zhang, Jun Hu, Jun Zhou, Chongxuan Li

arXiv:2503.07197v3cs.CVcs.LG

TL;DR

Masked image generation and masked diffusion have been studied with different objectives, and the applicability of masked diffusion to image generation remained open. The paper unifies them, explores training and sampling strategies, and develops eMIGM. On ImageNet, eMIGM compares favorably with VAR, REPA, and EDM2 across the reported resolutions and settings.

  • Problem

    Masked image generation can lose information through discrete tokenization, and key design choices remain underexplored; whether masked diffusion applies to image generation is also an open question.

    Source: §1 · p. 1

  • Method

    The paper unifies masked image modeling and masked diffusion, explores training and sampling choices, and develops eMIGM with a time-interval strategy for classifier-free guidance.

    Source: §1 · pp. 1–2

  • Results

    On ImageNet 256×256, eMIGM outperforms VAR at similar NFEs and parameter counts and approaches REPA at FID 1.57 vs. 1.42; on 512×512, eMIGM-L surpasses EDM2 with fewer parameters.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported results show that eMIGM can combine strong ImageNet generation quality with reduced sampling costs and favorable comparisons to the named models.

    Source: §1 · pp. 1–2, §8 · p. 9

    The authors note that increased efficiency may increase the potential for misuse of generated images.

    Source: Impact Statement · p. 9

Abstract

from arXiv · show

Although masked image generation models and masked diffusion models are designed with different motivations and objectives, we observe that they can be unified within a single framework. Building upon this insight, we carefully explore the design space of training and sampling, identifying key factors that contribute to both performance and efficiency. Based on the improvements observed during this exploration, we develop our model, referred to as \textbf{eMIGM}. Empirically, eMIGM demonstrates strong performance on ImageNet generation, as measured by Fr\'echet Inception Distance (FID). In particular, on ImageNet $256\times256$, with similar number of function evaluations (NFEs) and model parameters, eMIGM outperforms the seminal VAR. Moreover, as NFE and model parameters increase, eMIGM achieves performance comparable to the state-of-the-art continuous diffusion model REPA while requiring less than 45\% of the NFE. Additionally, on ImageNet $512\times512$, eMIGM outperforms the strong continuous diffusion model EDM2. Code is available at https://github.com/ML-GSAI/eMIGM.

1. Introduction

The paper unifies masked image generation and masked diffusion to explore training and sampling choices, then develops eMIGM for efficient ImageNet generation. eMIGM compares favorably with VAR, REPA, and EDM2 across the reported resolutions and settings.

1. Introduction

  • Prior masked image generation lagged diffusion partly because discrete tokenization loses information, while masking schedules and loss functions remained underexplored.At limited sampling steps, MAR also fell short of VAR.Source: §1 · p. 1
  • The unified framework combines masked image modeling and masked diffusion to systematically explore training and sampling strategies.Training choices include higher masking ratios, a simple weighting function, and CFG with a mask token for unconditional generation.Source: §1 · pp. 1–2
  • The time-interval CFG strategy applies guidance only in later sampling stages, maintaining strong performance while reducing NFEs and sampling time.The authors report that early-stage guidance decreases variance and raises FID.Source: §1 · pp. 1–2
  • On ImageNet 256×256, eMIGM outperforms VAR at similar NFEs and parameter counts; eMIGM-H approaches REPA at FID 1.57 vs. 1.42.The REPA comparison uses increased NFE and model size.Source: §1 · p. 2
  • On ImageNet 512×512, eMIGM-L surpasses EDM2 while using fewer parameters.Source: §1 · p. 2
  • Larger eMIGM models improve sample quality and efficiency, with superior quality at similar training FLOPs and sampling time.Source: §1 · p. 2

2. Preliminaries

The preliminaries contrast masked image generators that predict masked tokens with masked diffusion models that progressively add and remove masks.

2.1. Masked Image Generation

  • MaskGIT predicts masked image tokens from unmasked ones, while MAR uses a diffusion model for per-token distributions to avoid discrete-tokenizer information loss.MaskGIT achieves strong performance at approximately 10 sampling steps, but its performance degrades as sampling steps increase under its default schedule.Source: §2.1 · p. 3
  • MaskGIT trains by randomly masking tokens according to a sampled ratio and decodes by unmasking the highest-confidence predictions over successive iterations.MAR instead samples its training mask ratio from a truncated Gaussian with mean 1.0, standard deviation 0.25, and range [0.7, 1.0], then uses a MaskGIT-like decoding strategy.Source: §2.1 · p. 3

2.2. Masked Diffusion Models

  • Masked diffusion models progressively add masks in the forward process and remove them in the reverse process, with training optimizing an upper bound on masked-token negative log-likelihood.The explicit time input is described as theoretically redundant, with empirical validation in image generation.Source: §2.2 · p. 3
  • The reverse process is specified between noise levels s and t, where 0 ≤ s < t ≤ 1.The supplied passage introduces this reverse-process characterization without providing further operational details.Source: §2.2 · p. 3

3. Unifying Masked Image Generation

The section places masked image generation and masked diffusion in a shared loss framework by aligning their masking distributions and interpreting MaskGIT’s mask ratio as scaled time.

3. Unifying Masked Image Generation

  • Removing explicit time input lets the MaskGIT objective fit the general masked-diffusion loss, with its sampled mask ratio interpreted as scaled time.The passage states that timestep dependence can be extracted as a weight coefficient outside the cross-entropy loss.Source: §3 · p. 3
  • MaskGIT, MAR, and masked diffusion models can be expressed within a unified loss after their masking distributions are aligned.MaskGIT’s expectation over its mask ratio corresponds to integrating over time, with its masked image mapped to the masked-diffusion state.Source: §3 · pp. 3–4
  • The framework identifies masking distribution, loss weighting, and conditional distribution as the components that distinguish the models.For MaskGIT and MAR, a fixed number of tokens is masked uniformly without replacement; MDM masks tokens independently with probability γ_t.Source: §3 · p. 4, Table 1 · p. 4

4. Investigating the Design Space of Training

The section explores training choices within an MDM-default formulation, including mask schedules, loss weighting, architecture, time truncation, and classifier-free guidance.

4. Investigating the Design Space of Training

  • The cosine schedule outperforms the linear schedule, while the exp schedule increases masking further but destabilizes training.The authors hypothesize that image redundancy lets cosine masking provide stronger learning signals; instability with exp is attributed as likely related to persistently large loss weights at high mask ratios.Source: §4 · p. 4
  • Setting w(t) = 1 stabilizes training and makes the exp schedule the best-performing choice.The MDM weighting w(t) = γ′_t led to unstable training, particularly with the exp schedule.Source: §4 · p. 4
  • Under the exp schedule, the MAE architecture outperforms the single-encoder transformer, so the authors adopt MAE by default.The architectures differ in whether the encoder receives masked tokens as input.Source: §4 · pp. 4–5
  • Time truncation at t_min = 0.2 can benefit performance and accelerate convergence, whereas t_min = 0.4 provides no benefit and may degrade performance.At t_min = 0.4, over 80% of image tokens are masked during training.Source: §4 · p. 5
  • CFG with Mask improves generation performance over standard CFG, and the authors adopt it as the default.The passage also reports that a fake class token negatively impacts conditional generation performance in the described MDM setting.Source: §4 · p. 5

5. Investigating the Design Space of Sampling

The sampling study tests mask schedules, diffusion samplers, and classifier-free guidance intervals to balance image quality and sampling cost. It adopts the exponential mask schedule, DPM-Solver, and a time interval for subsequent experiments.

5.1. Mask Schedule during Sampling

  • At 8 or 16 prediction steps, the exp schedule outperforms cosine, which outperforms linear, and it predicts fewer tokens early.The study adopts the exp schedule as its default for sampling.Source: §5.1 · p. 6, Figure 3 · p. 6

5.2. The Sampling Method of Diffusion Loss

  • With low mask prediction steps, DPM-Solver outperforms DDPM; at 8 steps, their FIDs are 6.6 and 10.6, respectively.DPM-Solver achieves good performance with fewer than 15 diffusion steps, while DDPM requires 100.Source: §5.2 · p. 6

5.3. Time Interval for Classifier Free Guidance

  • Applying CFG only within the interval [0.1, 0.3] achieves the best reported performance using about 60% of standard CFG's NFE.The time interval maintains performance while reducing sampling time, so the authors adopt it for subsequent experiments.Source: §5.3 · pp. 6–7
  • With 42% of function evaluations, eMIGM-H achieves performance comparable to REPA on ImageNet 256 × 256.The table caption reports this comparison for the state-of-the-art diffusion model REPA.Source: Table 2 · p. 7

6. Experiments

The experiments evaluate eMIGM on ImageNet at 256 × 256 and 512 × 512 using FID, including tests of training and inference scaling. eMIGM's reported comparisons span VAR, REPA, EDM2, and MAR.

6. Experiments

  • The evaluation measures sample quality with FID on ImageNet 256 × 256 and 512 × 512.The paper refers to Appendix D for experiment settings.Source: §6 · p. 7

6.1. Larger Models Are Training and Sampling Efficient

  • At about 0.2 seconds per image, eMIGM-L achieves FID 1.8 versus eMIGM-B's FID 2.3.Larger models also achieve better performance at the same training FLOPs, and the reported scaling tests show improved quality as training FLOPs increase.Source: §6.1 · pp. 7–8

6.2. Image Generation on ImageNet

  • FID 2.79 with 208M parameters for eMIGM-B beats VAR-d16's FID 3.30 with 310M parameters on ImageNet 256 × 256.With 478M parameters, eMIGM-L reaches FID 1.72 at 180 NFEs, outperforming the best VAR-d30's FID 1.92 with 2B parameters.Source: §6.2 · p. 8
  • FID 1.57 with 180 NFEs makes eMIGM-H comparable to REPA, which requires 425 sequential steps and self-supervised-model assistance.The authors also report that eMIGM-H outperforms Large-DiT and DiffiT.Source: §6.2 · p. 8
  • At 80 NFEs, eMIGM-L reaches FID 1.77, outperforming EDM2's FID 1.81 on ImageNet 512 × 512.At 20 NFEs, the table caption reports that eMIGM-L outperforms VAR.Source: §6.2 · p. 8, Table 3 · p. 8
  • eMIGM-L achieves competitive performance against MAR while using less than 20% of its NFE.Source: §6.2 · p. 9

7. Related Work

The related work situates eMIGM among visual generation methods and masked modeling approaches. It distinguishes its token-level guidance interval from guidance strategies operating over image noise levels.

Visual generation

  • The study identifies MaskGIT and MAR as its closest related works and unifies their masked image generation and masked diffusion approaches.Source: §7 · p. 9
  • Unlike guidance methods for continuous diffusion that operate at image noise levels, the proposed time interval applies guidance at the token level.The authors motivate this distinction by masked diffusion's irreversible token generation constraint.Source: §7 · p. 9

Masked discrete diffusion models

  • Masked discrete diffusion is reported across text, music, protein design, and image generation.Source: §7 · p. 9

8. Conclusion

The conclusion presents eMIGM as a unified framework for masked image generation and masked diffusion models. It reports comparable performance to state-of-the-art continuous diffusion models with fewer NFEs.

8. Conclusion

  • The paper unifies masked image generation models and masked diffusion models in one framework.Source: §8 · p. 9
  • The authors report that eMIGM achieves performance comparable to state-of-the-art continuous diffusion models with fewer NFEs.Source: §8 · p. 9
  • The authors state that they hope eMIGM will inspire future research in masked image generation.Source: §8 · p. 9

Impact Statement

The appendix connects MaskGIT and MDM masking formulations in expectation, then discusses mask-schedule choices; the impact statement flags potential misuse from faster generation.

Impact Statement

  • Faster generation may increase misuse potential; the authors suggest watermarking images without affecting generation quality.Source: Impact Statement · p. 9

A. Equivalence of the masking strategies of MaskGIT and MDM

  • Under a linear schedule, the MDM loss over uniformly sampled t matches MaskGIT’s uniformly sampled masked-token-count objective in expectation; the paper adopts MDM Equation (4) with an exponential schedule.Source: A. Equivalence of the masking… · p. 13

B.1. Formulations and Illustrations of Mask Schedules

  • In text generation, the linear mask schedule has the best empirical performance cited here, while cosine and exponential schedules mask more tokens during MDM’s forward process.Source: §B.1 · p. 13

B.2. Sampling Simulator Experiment

A 256-token, 16-step simulation distinguishes token-prediction patterns across mask schedules, while a separate experiment reports lower FID for guidance applied during later sampling stages than across the full interval.

B.2. Sampling Simulator Experiment

  • Across 256 tokens and 16 steps, linear predicts nearly evenly, while cosine and especially exponential schedules shift more predictions to later steps.Source: §B.2 · pp. 13–14, Figure 6 · p. 14
  • The simulator averages the number of tokens predicted per step over 10,000 runs.Source: §B.2 · pp. 13–14

C. Time Interval for Classifier Free Guidance

D. Experiment settings and results

The experiments report mask-schedule ablations and implementation choices, alongside a sampling-speed comparison between eMIGM-L and EDM2.

Image Tokenizer

Classifier-Free Guidance

  • For the main results, the authors use CFG with Mask, a constant CFG schedule, and time-interval guidance; this permits consistently high CFG without value sweeping.Source: D. Experiment settings and results · p. 16

More Mask Schedules

Sampling Speed Comparison with EDM2

More Mask Schedules

Related papers

Loading 2503.07197v3…