---
title: 'MADI: Masking-Augmented Diffusion Scaling'
url: https://www.emergentmind.com/topics/masking-augmented-diffusion-with-inference-time-scaling-madi
type: topic
---

# MADI: Masking-Augmented Diffusion Scaling

Searching arXiv for MADI and closely related diffusion/inference-time scaling papers.
Masking-Augmented Diffusion with Inference-Time Scaling (MADI) denotes a family of diffusion-based methods that combine masking-based training with extra inference-time computation to improve controllability, fidelity, or quality. In 2025, the acronym was used in two closely related but domain-specific settings: a visual-editing framework that couples Masking-Augmented Gaussian Diffusion (MAgD) with inference-time Pause Tokens for grounded image editing [2507.13401], and a text-style-transfer framework built on masked diffusion language models (MDMs) with seq2seq masking augmentation and verifier-based inference-time scaling during denoising [2508.10995]. Across both usages, the common design pattern is to expose the model to structured missingness during training and then exploit the iterative nature of diffusion at inference to allocate additional compute in a targeted manner.

## 1. Scope, nomenclature, and shared design pattern

The label MADI is not a single canonical algorithm. In the visual-editing literature, MADI names a framework for controllable image editing based on diffusion training with a dual corruption process and inference-time prompt expansion through Pause Tokens [2507.13401]. In the language-modeling literature, MADI refers to masked diffusion language models with inference-time scaling for text style transfer, where the scaling mechanism is a verifier-guided candidate search layered on top of the reverse diffusion process [2508.10995].

Despite the difference in modality, the two formulations share a structural template. First, masking is introduced as an auxiliary corruption or augmentation mechanism during training. Second, inference-time quality is improved by spending additional computation during sampling rather than retraining a larger model. In the visual case, the extra compute is realized through longer prompt context induced by repeated special tokens. In the language case, it is realized through classifier-free guidance (CFG), more denoising steps, and soft-value diffusion decoding (SVDD) over multiple sampled candidates.

| Variant | Training-side masking | Inference-time scaling |
|---|---|---|
| Visual MADI | MAgD with masked noisy image inputs | Pause Tokens inserted into the prompt |
| Language MADI | Target-side masking during seq2seq MDM fine-tuning | CFG and SVDD candidate selection |

A common misconception is that inference-time scaling in MADI always means increasing the number of denoising steps. The two 2025 uses of the term show otherwise: visual MADI scales conditional computation through prompt-side capacity, whereas language MADI scales through guided search over denoising candidates [2507.13401].

## 2. Visual MADI: Masking-Augmented Gaussian Diffusion

The visual-editing formulation of MADI is centered on Masking-Augmented Gaussian Diffusion (MAgD), a training strategy that overlays masking onto a standard DDPM-style forward process. For a clean image $x_0 \in \mathbb{R}^d$, the standard Gaussian corruption is
$$
q(x_t\mid x_{t-1})=\mathcal{N}\bigl(x_t;\sqrt{1-\beta_t}\,x_{t-1},\beta_t\mathbf I\bigr),\quad t=1,\dots,T,
$$
or equivalently
$$
x_t=\sqrt{\alpha_t}\,x_0+\sqrt{1-\alpha_t}\,\epsilon,\quad \epsilon\sim\mathcal{N}(0,I),\quad \alpha_t=\prod_{s=1}^t (1-\beta_s).
$$
After sampling $x_t$, MAgD applies a random spatial mask $m\in\{0,1\}^d$ with masking rate $r_{\text{mask}}$, producing
$$
\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,
$$
where $\varnothing$ is either the zero-tensor or a learned mask token embedding [2507.13401].

Training alternates between conventional denoising score matching and a masked-denoising variant. With $u\sim\mathrm{Unif}(0,1)$, the masked loss is used when $u<p_{\rm MAgD}$ and $t\ge\tau_{\rm MAgD}$:
$$
\mathcal{L}_{\rm mDSM}
=\mathbb E_{x_0,\epsilon,m,t}
\bigl\|\epsilon_\theta(\tilde x_t^{\rm masked},t,c)-\epsilon\bigr\|^2,
$$
otherwise the usual DSM objective is applied:
$$
\mathcal{L}_{\rm DSM}
=\mathbb E_{x_0,\epsilon,t}
\bigl\|\epsilon_\theta(x_t,t,c)-\epsilon\bigr\|^2.
$$
The resulting per-step loss is
$$
\mathcal{L}_{\rm MAgD}=
\begin{cases}
\mathcal{L}_{\rm mDSM}, & u<p_{\rm MAgD},\;t\ge\tau_{\rm MAgD},\\
\mathcal{L}_{\rm DSM}, & \text{otherwise}.
\end{cases}
$$

The intended effect of masking noisy inputs at high-noise timesteps is explicitly described as improved locality, compositionality, and vision-language grounding. Because the model must predict noise when part of the noisy image is absent, it is trained to rely on contextual and semantic cues rather than only on dense local evidence [2507.13401]. This suggests that the method is not merely a robustness regularizer; it also functions as a representation-shaping mechanism for localized editing.

The reported implementation uses an OmniGen transformer with causal attention over interleaved image and text tokens. The mask token embedding is produced by a small two-layer MLP mapping a mask identifier to a 256-dim vector, with a pool of 1024 learnable mask embeddings. Training uses AdamW with initial learning rate $1\times10^{-4}$, cosine decay, 800-step linear warmup, batch size 128 on $4\times$ A100 GPUs, total gradient steps $\sim 4\,000$, image resolution $1024\times1024$, diffusion timesteps $T=1\,000$, masking rate $r_{\text{mask}}=0.25$, masking probability $p_{\rm MAgD}=0.5$, and noise-threshold $\tau_{\rm MAgD}=0.7$; classifier-free guidance is trained with $10\%$ prompt dropping [2507.13401].

## 3. Pause Tokens and inference-time capacity scaling in visual editing

The inference-time component of visual MADI is based on Pause Tokens, denoted $\langle\text{pause}\rangle$. These are special placeholder tokens inserted only at inference time. Given an edit prompt $p$ and optional reference image embedding $\mathrm{ref}$, the paused prompt is
$$
p_{\rm paused}=[p,\mathrm{ref},\langle\text{pause}\rangle,\dots,\langle\text{pause}\rangle],
$$
with the pause token repeated $k$ times [2507.13401].

The sampling loop is otherwise standard: initialize $z_T\sim\mathcal N(0,I)$, run the pretrained MAgD-finetuned model $\epsilon_\theta$ over $t=T,T-1,\dots,1$, and decode $z_0$. The key claim is that increasing $k$ grows prompt length, so cross-attention and self-attention attend over more tokens, and each diffusion step allocates more computation to reasoning under the conditioning. No retraining is required; Pause Tokens are treated as extra context vectors [2507.13401].

This mechanism is qualitatively different from increasing the number of diffusion steps. The denoising schedule can remain unchanged while the conditional pathway becomes more expressive. In the reported workflow, candidate pause lengths $K=\{0,8,16,32\}$ are evaluated, faithfulness is filtered by a DINO threshold, and the surviving candidate with maximal CLIP-DIR is selected. The examples provided include local edits such as “replace the net with a brick wall,” global edits such as “look like it is floating among white fluffy clouds,” and attribute edits such as “Change her hair color to white” [2507.13401].

The paper also reports that adopting expressive and dense prompts during training further enhances performance, particularly for MAgD. A plausible implication is that Pause Tokens are most effective when the model has already learned to extract structured conditioning signals from rich prompt contexts.

## 4. Language MADI: masked diffusion language models for text style transfer

The language formulation of MADI builds on masked diffusion language models operating directly on discrete token sequences. Let $x_0=[x_0^0,\dots,x_0^{L-1}]$ be a length-$L$ sentence, and let $\langle\mathrm{mask}\rangle$ denote the special mask token. Continuous time is indexed by $t\in[0,1]$, with $\alpha_t:=1-t$ as the mask retention probability. The forward process independently keeps or masks each token:
$$
q_{t|0}(x_t^i|x_0^i)=
\begin{cases}
\alpha_t & \text{if } x_t^i=x_0^i,\\
1-\alpha_t & \text{if } x_t^i=\langle\mathrm{mask}\rangle.
\end{cases}
$$
The joint corruption factorizes across positions, while the reverse process is defined through $q_{s|t}(x_s^i|x_t)$ for $0\le s<t\le1$ and uses a neural denoiser $p_\theta(x_0|x_t)$ to approximate the intractable posterior $q_{0|t}$ [2508.10995].

A simplified variational bound yields a continuous-time masked-language-modeling objective:
$$
L_{MDM}(\theta)=
\int_0^1 \frac{\alpha'_t}{1-\alpha_t}\;
\mathbb E_{x_t\sim q_{t|0}(\cdot|x_0)}
\left[
-\sum_{i:x_t^i=\langle\mathrm{mask}\rangle}
\log p_\theta(x_0^i|x_t)
\right]dt.
$$

For seq2seq style transfer, fine-tuning uses an input concatenation
$$
x^{inp}=[y^1,\dots,y^{L_2},\langle sep\rangle,x_0^1,\dots,x_0^{L_1}],
$$
where $y$ is the source sentence in style $s_1$ and $x_0$ is the target sentence in style $s_2$. The forward masking process is applied only to the target half $[x_0]$, leaving $y$ intact. At each training step, $t\sim\mathrm{Uniform}(0,1)$ is sampled, $\alpha_t=1-t$ is set, and each target token is independently replaced by $\langle\mathrm{mask}\rangle$ with probability $1-\alpha_t$. The transformer denoiser $NN_\theta$ predicts the original target tokens at masked positions via cross-entropy. No explicit time embeddings are required, because the masking schedule is said to be known to the network implicitly by the fraction of masked tokens [2508.10995].

The denoiser is a 12-layer Transformer encoder with hidden size 768, 12 heads, FFN 3072, and approximately 167M total parameters. Optimization uses AdamW with $\beta_1=0.9$, $\beta_2=0.95$, $\epsilon=10^{-5}$, weight decay $0.01$, 1K warmup steps, and inverse-square-root decay [2508.10995].

## 5. Verifier-based inference-time scaling in language MADI

The inference-time scaling mechanism in language MADI is soft-value diffusion decoding (SVDD), a derivative-free verifier-based procedure using an off-the-shelf sentence embedding model $\mathcal E(\cdot)$. For a candidate denoised sequence $\hat x_0$ at an intermediate step, the reward is defined as cosine similarity to the source sentence embedding:
$$
R_{\mathcal E,t}(\hat x_0,y):=
\cosine(\mathcal E(\hat x_0),\mathcal E(y))
=
\frac{\langle \mathcal E(\hat x_0),\mathcal E(y)\rangle}
{\|\mathcal E(\hat x_0)\|\cdot\|\mathcal E(y)\|}.
$$
The soft value at time $t$ is approximated by
$$
v_t(x_t)\approx R_{\mathcal E,t}(\hat x_0(x_t),y),
$$
where $\hat x_0(x_t)\sim p_\theta(x_0|x_t)$ is sampled once to approximate the expectation [2508.10995].

At each reverse step from $x_t$ to $x_{t-\Delta t}$, the model computes the usual denoising kernel, draws $M$ independent candidates, evaluates a soft value for each candidate, and keeps the maximizing candidate:
$$
m^\star=\arg\max_m v_{t-\Delta t}(x_{t-\Delta t}^{(m)}).
$$
This is described as approximately sampling from a modified distribution
$$
p_\alpha(x)\propto p_\theta(x)\cdot e^{R(x)/\alpha}
$$
in a greedy, reward-maximizing fashion [2508.10995].

SVDD can be combined with classifier-free guidance. The conditional and unconditional logits are combined as
$$
L=\gamma L_{\rm cond}+(1-\gamma)L_{\rm uncond},
$$
with reported CFG scale $\gamma=1.4$. The verifier uses the MPNet-base-v2 sentence-embedding model. Denoising steps are varied over $T\in\{8,16,64\}$ and candidate counts over $M\in\{1,2,4,8,16\}$, with $M=4$ described as a good trade-off in practice [2508.10995].

The paper emphasizes that CFG and SVDD represent complementary inference-compute axes: “logit guidance” versus “verifier search.” It also states explicit trade-offs: SVDD adds a linear $M\times$ cost per step, CFG sacrifices some diversity, and very large $T$ or $\gamma$ can over-focus on style at the expense of fluency [2508.10995].

## 6. Empirical results, related methods, and interpretation

The visual and language MADI variants are evaluated on different tasks and with different success criteria, but both are presented as improving control without requiring a larger base model. In visual editing, benchmarks include Emu-Edit, MagicBrush, Complex-Edit, and IDEA-Bench, with metrics CLIP-I, DINO, CLIP-T, CLIP-DIR, and MLLM score. On Emu-Edit, the reported numbers are: OmniGen $0.820$ CLIP-I, $0.882$ DINO, $0.122$ CLIP-DIR, and $8.00$ MLLM; MAgD-no-scale $0.873$, $0.927$, $0.126$, and $8.40$; MAgD-with-scaling $0.869$, $0.894$, $0.134$, and $8.43$ [2507.13401]. On Complex-Edit, OmniGen has MLLM aggregate $7.53$, while MAgD reports $8.62$; on IDEA-Bench, OmniGen reports $57.41\%$ MLLM success rate and MAgD reports $74.07\%$ [2507.13401].

In language style transfer, evaluation covers WikiLarge for complex-to-simple simplification using BLEU, SARI, and LENS, and Bible prose for PUB$\to$BBE and PUB$\to$ASV using BLEU, ROUGE-L, METEOR, and BERTScore. Fine-tuned MDM without guidance is reported to already match or slightly exceed autoregressive baselines in BLEU but to produce too-short outputs on WikiLarge. Adding CFG steadily improves BLEU, SARI, LENS, and Bible metrics, especially as $T$ grows. Applying SVDD with $M=4$ on top of CFG yields up to $+2.8$ BLEU on WikiLarge and $+0.6$ to $1.0$ BLEU on Bible test sets. Ablations report that increasing $T$ helps up to approximately 16 steps and then plateaus or slightly drops without CFG, while with CFG it scales robustly even to 64 steps; higher $\gamma$ helps up to approximately $1.4$–$1.6$ and then plateaus or degrades; increasing $M$ improves performance with diminishing returns after approximately $8$ candidates [2508.10995].

| Setting | Reported result | Source |
|---|---|---|
| Emu-Edit | MAgD-with-scaling improves CLIP-DIR from $0.122$ to $0.134$ over OmniGen | [2507.13401] |
| Complex-Edit | MLLM aggregate rises from $7.53$ to $8.62$ | [2507.13401] |
| WikiLarge | CFG+SVDD yields up to $+2.8$ BLEU | [2508.10995] |
| Bible prose | CFG+SVDD yields $+0.6$ to $1.0$ BLEU | [2508.10995] |

The broader significance of MADI is clearer when placed next to adjacent inference-time scaling methods for diffusion. ReMDM introduces remasking during the backward process so that already generated discrete tokens can be masked again and corrected, and it treats increasing the number of denoising steps as an explicit compute-quality axis [2503.00307]. UMF, by contrast, casts the unmasking trajectory of masked diffusion language models as a Monte Carlo Tree Search over deterministic partial-unmasking actions, emphasizing low-variance branching under a fixed Number of Function Evaluations budget [2602.04344]. These comparisons indicate that MADI belongs to a wider shift in diffusion research from fixed samplers toward compute-adaptive decoding.

A further distinction concerns what “masking” means. In visual MADI, masking is a training-time corruption over spatially structured image inputs. In language MADI, masking is the native discrete corruption used by MDMs, specialized to seq2seq style transfer and supplemented by a semantic verifier. This suggests that the unifying content of MADI is not a single corruption operator, but the combination of masked conditioning and post-training inference-time scaling.

Source: https://www.emergentmind.com/topics/masking-augmented-diffusion-with-inference-time-scaling-madi