MADI: Masking-Augmented Diffusion Scaling

• The paper introduces a framework that integrates masking during training with extra computation at inference to boost controllability and fidelity.
• Visual MADI leverages pause tokens for extended prompt context, while language MADI employs masked diffusion paired with verifier-guided candidate selection.
• Empirical evaluations show notable improvements, including enhanced CLIP-DIR metrics in image editing and up to +2.8 BLEU gains in text style transfer.

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 (Kadambi et al., 16 Jul 2025), and a text-style-transfer framework built on masked diffusion LLMs (MDMs) with seq2seq masking augmentation and verifier-based inference-time scaling during denoising (Padole et al., 14 Aug 2025). 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 (Kadambi et al., 16 Jul 2025). In the language-modeling literature, MADI refers to masked diffusion LLMs 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 (Padole et al., 14 Aug 2025).

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 (Kadambi et al., 16 Jul 2025).

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 (Kadambi et al., 16 Jul 2025).

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 $$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,$$0:

$$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,$$1

otherwise the usual DSM objective is applied:

$$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,$$2

The resulting per-step loss 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,$$3

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 (Kadambi et al., 16 Jul 2025). 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 $$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,$$4, cosine decay, 800-step linear warmup, batch size 128 on $$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,$$5 A100 GPUs, total gradient steps $$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,$$6, image resolution $$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,$$7, diffusion timesteps $$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,$$8, masking rate $$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,$$9, masking probability $$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).$$0, and noise-threshold $$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).$$1; classifier-free guidance is trained with $$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).$$2 prompt dropping (Kadambi et al., 16 Jul 2025).

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

The inference-time component of visual MADI is based on Pause Tokens, denoted $$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).$$3. These are special placeholder tokens inserted only at inference time. Given an edit prompt $$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).$$4 and optional reference image embedding $$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).$$5, the paused prompt is

$$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).$$6

with the pause token repeated $$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).$$7 times (Kadambi et al., 16 Jul 2025).

The sampling loop is otherwise standard: initialize $$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).$$8, run the pretrained MAgD-finetuned model $$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).$$9 over $x_t$0, and decode $x_t$1. The key claim is that increasing $x_t$2 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 (Kadambi et al., 16 Jul 2025).

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 $x_t$3 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” (Kadambi et al., 16 Jul 2025).

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 LLMs for text style transfer

The language formulation of MADI builds on masked diffusion LLMs operating directly on discrete token sequences. Let $x_t$4 be a length-$x_t$5 sentence, and let $x_t$6 denote the special mask token. Continuous time is indexed by $x_t$7, with $x_t$8 as the mask retention probability. The forward process independently keeps or masks each token:

$x_t$9

The joint corruption factorizes across positions, while the reverse process is defined through $m\in\{0,1\}^d$0 for $m\in\{0,1\}^d$1 and uses a neural denoiser $m\in\{0,1\}^d$2 to approximate the intractable posterior $m\in\{0,1\}^d$3 (Padole et al., 14 Aug 2025).

A simplified variational bound yields a continuous-time masked-language-modeling objective:

$m\in\{0,1\}^d$4

For seq2seq style transfer, fine-tuning uses an input concatenation

$m\in\{0,1\}^d$5

where $m\in\{0,1\}^d$6 is the source sentence in style $m\in\{0,1\}^d$7 and $m\in\{0,1\}^d$8 is the target sentence in style $m\in\{0,1\}^d$9. The forward masking process is applied only to the target half $r_{\text{mask}}$0, leaving $r_{\text{mask}}$1 intact. At each training step, $r_{\text{mask}}$2 is sampled, $r_{\text{mask}}$3 is set, and each target token is independently replaced by $r_{\text{mask}}$4 with probability $r_{\text{mask}}$5. The transformer denoiser $r_{\text{mask}}$6 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 (Padole et al., 14 Aug 2025).

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 $r_{\text{mask}}$7, $r_{\text{mask}}$8, $r_{\text{mask}}$9, weight decay $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$0, 1K warmup steps, and inverse-square-root decay (Padole et al., 14 Aug 2025).

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 $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$1. For a candidate denoised sequence $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$2 at an intermediate step, the reward is defined as cosine similarity to the source sentence embedding:

$$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$3

The soft value at time $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$4 is approximated by

$$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$5

where $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$6 is sampled once to approximate the expectation (Padole et al., 14 Aug 2025).

At each reverse step from $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$7 to $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$8, the model computes the usual denoising kernel, draws $$\tilde x_t^{\rm masked}=m\odot \varnothing +(1-m)\odot x_t,$$9 independent candidates, evaluates a soft value for each candidate, and keeps the maximizing candidate:

$\varnothing$0

This is described as approximately sampling from a modified distribution

$\varnothing$1

in a greedy, reward-maximizing fashion (Padole et al., 14 Aug 2025).

SVDD can be combined with classifier-free guidance. The conditional and unconditional logits are combined as

$\varnothing$2

with reported CFG scale $\varnothing$3. The verifier uses the MPNet-base-v2 sentence-embedding model. Denoising steps are varied over $\varnothing$4 and candidate counts over $\varnothing$5, with $\varnothing$6 described as a good trade-off in practice (Padole et al., 14 Aug 2025).

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 $\varnothing$7 cost per step, CFG sacrifices some diversity, and very large $\varnothing$8 or $\varnothing$9 can over-focus on style at the expense of fluency (Padole et al., 14 Aug 2025).

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 $u\sim\mathrm{Unif}(0,1)$0 CLIP-I, $u\sim\mathrm{Unif}(0,1)$1 DINO, $u\sim\mathrm{Unif}(0,1)$2 CLIP-DIR, and $u\sim\mathrm{Unif}(0,1)$3 MLLM; MAgD-no-scale $u\sim\mathrm{Unif}(0,1)$4, $u\sim\mathrm{Unif}(0,1)$5, $u\sim\mathrm{Unif}(0,1)$6, and $u\sim\mathrm{Unif}(0,1)$7; MAgD-with-scaling $u\sim\mathrm{Unif}(0,1)$8, $u\sim\mathrm{Unif}(0,1)$9, $u<p_{\rm MAgD}$0, and $u<p_{\rm MAgD}$1 (Kadambi et al., 16 Jul 2025). On Complex-Edit, OmniGen has MLLM aggregate $u<p_{\rm MAgD}$2, while MAgD reports $u<p_{\rm MAgD}$3; on IDEA-Bench, OmniGen reports $u<p_{\rm MAgD}$4 MLLM success rate and MAgD reports $u<p_{\rm MAgD}$5 (Kadambi et al., 16 Jul 2025).

In language style transfer, evaluation covers WikiLarge for complex-to-simple simplification using BLEU, SARI, and LENS, and Bible prose for PUB$u<p_{\rm MAgD}$6BBE and PUB$u<p_{\rm MAgD}$7ASV 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 $u<p_{\rm MAgD}$8 grows. Applying SVDD with $u<p_{\rm MAgD}$9 on top of CFG yields up to $$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,$$00 BLEU on WikiLarge and $$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,$$01 to $$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,$$02 BLEU on Bible test sets. Ablations report that increasing $$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,$$03 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 $$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,$$04 helps up to approximately $$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,$$05–$$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,$$06 and then plateaus or degrades; increasing $$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,$$07 improves performance with diminishing returns after approximately $$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,$$08 candidates (Padole et al., 14 Aug 2025).

Setting	Reported result	Source
Emu-Edit	MAgD-with-scaling improves CLIP-DIR from $$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,$$09 to $$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,$$10 over OmniGen	(Kadambi et al., 16 Jul 2025)
Complex-Edit	MLLM aggregate rises from $$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,$$11 to $$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,$$12	(Kadambi et al., 16 Jul 2025)
WikiLarge	CFG+SVDD yields up to $$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,$$13 BLEU	(Padole et al., 14 Aug 2025)
Bible prose	CFG+SVDD yields $$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,$$14 to $$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,$$15 BLEU	(Padole et al., 14 Aug 2025)

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 (Wang et al., 1 Mar 2025). UMF, by contrast, casts the unmasking trajectory of masked diffusion LLMs as a Monte Carlo Tree Search over deterministic partial-unmasking actions, emphasizing low-variance branching under a fixed Number of Function Evaluations budget (Misaki et al., 4 Feb 2026). 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.