Infinite Mask Diffusion Model (IMDM)

• The paper introduces IMDM, which replaces a single mask with an infinite family to effectively capture joint token dependencies and mitigate factorization errors.
• It presents a partition-and-map construction that theoretically reduces factorization error to zero, validated by improved metrics in synthetic and real-world tasks.
• IMDM leverages advanced few-step distillation techniques like SDTT and ReDi, achieving significant improvements in text perplexity and high-fidelity image synthesis.

The Infinite Mask Diffusion Model (IMDM) generalizes discrete masked diffusion processes for generative modeling by introducing a stochastic infinite-state mask structure, enabling exact modeling of joint token dependencies within few sampling steps and scalable application to both high-dimensional text and image generation. IMDM addresses irreducible factorization errors inherent to standard masked diffusion models and achieves provable convergence and empirical advances in few-step generation, while maintaining compatibility with pre-trained masked diffusion checkpoints (Yoo et al., 11 May 2026, Conforti et al., 29 Nov 2025, Aversa et al., 2023).

1. Conceptual Foundations of Infinite Mask Diffusion

IMDM arises from the limitations of conventional masked diffusion models (MDMs) in discrete generative tasks. Standard MDMs encode information loss via a single absorbing mask state $m\notin V$ in the vocabulary $V$, leading to a forward process that systematically replaces tokens with $m$. Their reverse process leverages parallel, bidirectional prediction for generation, offering significant advantages—inference tractability, parallel decoding, simplified objectives, and conditional generation flexibility—for text and other sequential data.

However, the use of a single deterministic mask implies severe factorization errors during few-step sampling. When tokens are unmasked simultaneously, dependencies between them cannot be fully modeled: the irreversible mixing of stochasticity from categorical sampling and the deterministic masking limits expressivity. The factorization error is quantified via the conditional total correlation:

$$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$

with a nonzero lower bound for any single-masked MDM (Theorem 4.1 in (Yoo et al., 11 May 2026)).

IMDM generalizes the masking operation by introducing an effectively infinite set of stochastic mask states $M$, enabling infinitesimal stochastic "coloring" of the mask at each position and elevating the model expressivity to capture arbitrary joint dependencies between unmasked tokens in a single step (Yoo et al., 11 May 2026, Conforti et al., 29 Nov 2025).

2. The IMDM Stochastic Forward and Reverse Processes

The IMDM framework replaces the unique mask token $m$ by an infinite family $M$ of latent mask states, defining the state space $Z = V \cup M$, $V \cap M = \emptyset$. The forward process for each position embeds uniform discrete diffusion over $Z$, importing stochasticity through continuous noise. For $V$0, $V$1, and a noise schedule $V$2:

• Forward process:

$V$3

where $V$4 is uniform over $V$5 as $V$6.

• Posterior (for $V$7):

$V$8

Practical implementation perturbs the original mask embedding $V$9 using a learned MLP, taking as input a high-dimensional random noise vector $m$0, simulating the limit $m$1 while initializing to be compatible with pre-trained MDM weights (Yoo et al., 11 May 2026). This construction maintains the foundational advantages of MDMs (parallelism, bidirectionality), but crucially, IMDM can theoretically achieve zero factorization error, as formalized in Theorem 4.2 (Yoo et al., 11 May 2026).

In the continuous (image/patch) setting, as in DiffInfinite (Aversa et al., 2023), the mask diffusion model learns a semantic sketch over a compact latent space, later upsampled for high-resolution conditioning of downstream diffusion models. The "mask" here generalizes to potentially many semantic classes, but the forward diffusion retains the infinite mask principle.

3. Mitigating Factorization Error: Partition-and-Map Construction

IMDM achieves its improvement via the partition-and-map principle: for any set of positions $m$2 unmasked between steps, the infinite cardinality of $m$3 allows partitioning $m$4 into arbitrarily fine sets, each corresponding to a specific joint token assignment in $m$5 with the exact target probability. A deterministic mapping $m$6 then projects each random mask instantiation to a token vector, allowing the model to recover the true joint distribution on $m$7 for any configuration of unmaskings.

The result is that there exist parameters $m$8 and a partition of $m$9 such that the learned reverse $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$0 exactly matches the true joint $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$1, thus $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$2 (Yoo et al., 11 May 2026). Empirically, IMDM is shown to nearly saturate this bound in synthetic correlation tasks, where standard MDMs saturate their lower bound (validity $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$3 in a perfect-pair synthetic task) and IMDM achieves $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$4 validity, with the factorization error reduced by an order of magnitude.

In the mathematical foundation of discrete-state and infinite-state masking (Conforti et al., 29 Nov 2025), the masking process is made explicit on $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$5 with a countably infinite state space. The forward transition kernel is constructed as a continuous-time Markov chain (CTMC), with transitions into the mask state at rate $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$6. The backward (denoising) process relies on discrete score functions, and non-asymptotic convergence rates can be established under entropic error control on the learned scores.

4. Training Regimes and Few-Step Distillation

IMDM training employs the standard evidence lower bound (ELBO) or denoising loss; in the Rao-Blackwellized form the per-step objective coincides with MDMs:

$$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$7

Model initialization uses pre-trained MDM checkpoints by construction. For sampling acceleration, IMDM leverages advanced few-step distillation regimens:

• Self-Distillation Through Time (SDTT): Uses a long-step teacher to train a few-step student via progressive KL loss across forward timesteps.
• Rectified Discrete Flow (ReDi): Constructs a rectified coupling $$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$8 for more accurate student training, replacing the empirical data distribution in the loss by the teacher-corrected joint.
• Combined SDTT+ReDi: Applies SDTT pre-distillation followed by ReDi for maximal error minimization (Yoo et al., 11 May 2026).

Distillation can be applied with no change to the pipeline except replacing MDM layers with IMDM. Empirical results show that IMDM, when equipped with these distillation techniques, surpasses baseline single-mask approaches in perplexity and correlation metrics over both synthetic and large-scale text tasks (LM1B, OpenWebText).

5. Empirical and Theoretical Validation

Summary Table: Token Correlation Task (1-Step Generation)

Method	Validity (%)	Token Entropy (bits)	Factorization Error
MDM+ReDi	≈50	≈0.69	≈0.693 (log 2)
IMDM+ReDi	97.7	≈0.69	0.082

Validity: Probability Mass on Joint Correct Sequences.

Language Modeling Benchmarks

For text (LM1B, OpenWebText), IMDM consistently produces lower generative perplexity than MDM in the extreme few-step regime ($$\mathrm{TCe}(Z_s|Z_t) = \mathbb{E}_{Z_t}[\mathrm{KL}(q(Z_s|Z_t) \parallel \prod_{\ell} q(Z_{s,\ell}|Z_t))]\,,$$9 steps). For instance, with SDTT and ReDi combined on GPT-2 Small, generative perplexity at 1 step drops from $M$0 (MDLM) to $M$1 (IMDM), and at 4 steps from $M$2 to $M$3. On OpenWebText, IMDM’s generative perplexity at 2 steps is $M$4, compared to $M$5 for baseline MDLM (Yoo et al., 11 May 2026).

For discrete-state convergence, the non-asymptotic convergence theorem for IMDM on $M$6 yields a total variation bound that scales linearly (up to logarithmic terms) with the dimension $M$7, and depends on the entropic approximation error of the learned score (Conforti et al., 29 Nov 2025). This advances theoretical understanding over bounded-support score functions and finite-state models.

6. Extensions to Image Synthesis and Parallel Generation

In high-resolution image synthesis, notably histopathology whole-slide images, IMDM serves as the conceptual basis for hierarchical generative architectures such as DiffInfinite (Aversa et al., 2023). There, mask diffusion models generate coarse semantic tissue sketches in latent space, which are upsampled and used for local conditioning of high-fidelity image diffusion models operating in VQ-VAE latent spaces. The parallel random patch (RP) diffusion sampling procedure enables scalable, seamless generation of extremely large images:

• Each denoising step acts on random, potentially overlapping latent patches, avoiding tiling artifacts and supporting efficient parallelization.
• Mask information guides class-conditional feature updates, and classifier-free guidance further improves sample fidelity/diversity tradeoffs.

Empirical results show that DiffInfinite achieves improved precision (0.98), recall (0.44), and higher authenticity (0.86-0.98) in large-scale image synthesis relative to prior baselines. Downstream classification and segmentation experiments show that models trained on synthetic data generated by mask-based diffusion techniques can approach or match those trained on real data, and native augmentation via synthetic data can improve overall generalization (Aversa et al., 2023).

7. Limitations, Complexity, and Practical Considerations

IMDM introduces minimal computational overhead beyond a small MLP for noise injection atop the mask embedding. Compatibility with pre-trained MDMs is attained by initializing the perturbation MLP to zero. The injected noise dimension $M$8 must be sufficiently large (e.g., $M$9 for LM1B) to ensure adequate coverage of the infinite state space; performance degrades for small $m$0.

Potential limitations include:

• The practicality of partition-and-map matching for the joint posteriors is bounded by model capacity and the success of the distillation pipeline—zero factorization error does not arise automatically.
• When the number of reverse steps increases, the factorization-error bound that defines IMDM's advantage vanishes, and the model converges to standard MDM performance.
• For truly infinite discrete spaces, as in the formalization on $m$1, some truncation or importance weighting is necessary to manage the state space, and finite second moment conditions must be satisfied for convergence guarantees (Conforti et al., 29 Nov 2025).

IMDM provides a theoretically grounded, empirically validated framework for fast, high fidelity, diffusion-based sequence and image generation, particularly effective in the regime of few sampling steps, and is extensible to hybrid discrete-continuous and arbitrarily high-dimensional domains (Yoo et al., 11 May 2026, Conforti et al., 29 Nov 2025, Aversa et al., 2023).