| license: apache-2.0 | |
| pipeline_tag: text-generation | |
| # Infinite Mask Diffusion for Few-Step Distillation | |
| This repository contains the checkpoints for the paper [Infinite Mask Diffusion for Few-Step Distillation](https://arxiv.org/abs/2605.10518). | |
| [[Project Page]](https://Ugness.github.io/official_imdm) | [[Paper]](https://arxiv.org/abs/2605.10518) | [[GitHub]](https://github.com/Ugness/IMDM) | |
| ## Overview | |
| Masked Diffusion Models (MDMs) have emerged as a promising alternative to autoregressive models in language modeling, offering the advantages of parallel decoding and bidirectional context processing. However, MDMs typically require many sampling iterations due to factorization errors stemming from simultaneous token updates. | |
| The **Infinite Mask Diffusion Model (IMDM)** introduces a stochastic infinite-state mask to mitigate the theoretical lower bound of factorization error while directly inheriting the benefits of MDMs, including compatibility with pre-trained weights. When equipped with appropriate distillation methods, IMDM surpasses existing few-step distillation methods at small step counts on benchmarks like LM1B and OpenWebText. | |
| ## Usage | |
| Please refer to the [official GitHub repository](https://github.com/Ugness/IMDM) for installation instructions and scripts for training and evaluation. | |
| ## Citation | |
| ```bibtex | |
| @inproceedings{yoo2026imdm, | |
| title={Infinite Mask Diffusion for Few-Step Distillation}, | |
| author={Yoo, Jaehoon and Kim, Wonjung and Lee, Chanhyuk and Hong, Seunghoon}, | |
| year={2026}, | |
| booktitle={ICML} | |
| } | |
| ``` |