English

Demystifying MaskGIT Sampler and Beyond: Adaptive Order Selection in Masked Diffusion

Machine Learning 2026-04-23 v2 Probability Machine Learning

Abstract

Masked diffusion models have shown promising performance in generating high-quality samples in a wide range of domains, but accelerating their sampling process remains relatively underexplored. To investigate efficient samplers for masked diffusion, this paper theoretically analyzes the MaskGIT sampler for image modeling, revealing its implicit temperature sampling mechanism. Through this analysis, we introduce the "moment sampler," an asymptotically equivalent but more tractable and interpretable alternative to MaskGIT, which employs a "choose-then-sample" approach by selecting unmasking positions before sampling tokens. In addition, we improve the efficiency of choose-then-sample algorithms through two key innovations: a partial caching technique for transformers that approximates longer sampling trajectories without proportional computational cost, and a hybrid approach formalizing the exploration-exploitation trade-off in adaptive unmasking. Experiments in image and text domains demonstrate our theory as well as the efficiency of our proposed methods, advancing both theoretical understanding and practical implementation of masked diffusion samplers.

Keywords

Cite

@article{arxiv.2510.04525,
  title  = {Demystifying MaskGIT Sampler and Beyond: Adaptive Order Selection in Masked Diffusion},
  author = {Satoshi Hayakawa and Yuhta Takida and Masaaki Imaizumi and Hiromi Wakaki and Yuki Mitsufuji},
  journal= {arXiv preprint arXiv:2510.04525},
  year   = {2026}
}

Comments

23 pages, fixed cleveref-related issue

R2 v1 2026-07-01T06:18:34.851Z