Fixed-Point Masked Generative Modeling

Andrea Miele, Yiming QIN, Alba Carballo-Castro, Justin Deschenaux, Pascal Frossard

Keywords: Masked generative models, Masked diffusion language models, Adaptive-depth transformers, Efficient generative modeling, Deep equilibrium models

Abstract:

Masked Generative Models (MGMs) enable parallel decoding and achieve strong performance across modalities, but require full-sequence bidirectional transformers at every step, making training costly and degrading quality under low sampling budgets.
Existing work improves efficiency via better samplers or fixed-depth backbones, but does not vary the depth of the denoiser across sampling steps.
We introduce Fixed-Point Masked Generative Models (FP-MGMs), which replace part of the denoiser with a fixed-point solver over shared attention layers to enable adaptive depth with fewer parameters.
To make it more effective for masked generation, we first introduce a cross-step consistency loss, which aligns hidden representations at nearby denoising steps and, second, three-state reuse (3SR) which warm-starts the solver using the previous solution by treating unchanged, still-masked, and newly revealed tokens differently. Together, these components define our complete training-to-inference framework for fixed-point masked generation, \emph{CoFRe}.
We also show that pre-trained MGMs can be converted into FP-MGMs with short fine-tuning, avoiding full retraining.
Across modalities, CoFRe improves the quality and cost trade-off. On OpenWebText, CoFRe reduces parameters by 38.8\%, training time by 11.5\%, and VRAM by 16.9\%, while improving generative perplexity from 830.8 to 101.8 at a budget of $96$ transformer-block forward passes, compared to MDLM. In ImageNette, CoFRe reduces training time by 48.6\% and VRAM by 50.7\%, while improving FID in all sample budgets tested. Overall, CoFRe offers a practical framework for cheaper training and stronger low-budget masked generation.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 116