# MADI: Masking-Augmented Diffusion with Inference-Time Scaling for Visual Editing

arXiv: https://arxiv.org/abs/2507.13401 · HTML: https://mlsanity.com/paper/2507.13401/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 8 of its 30 references that are in the map.
- First posted on arXiv 16 Jul 2025.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- Diffusion models / [Image generation, text-to-image and editing](https://mlsanity.com/diffusion/image-generation-text-to-image-and-editing/)

## Builds on (most cited first, 8 of 8)

- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Show-o: One Single Transformer to Unify Multimodal Understanding and Generation](https://mlsanity.com/paper/2408.12528/) (2024)
- [Muse: Text-To-Image Generation via Masked Generative Transformers](https://mlsanity.com/paper/2301.00704/) (2023)
- [OmniGen: Unified Image Generation](https://mlsanity.com/paper/2409.11340/) (2024)
- [Think before you speak: Training Language Models With Pause Tokens](https://mlsanity.com/paper/2310.02226/) (2023)
- [Fast Training of Diffusion Models with Masked Transformers](https://mlsanity.com/paper/2306.09305/) (2023)
- [[MASK] is All You Need](https://mlsanity.com/paper/2412.06787/) (2024)
- [IDEA-Bench: How Far are Generative Models from Professional Designing?](https://mlsanity.com/paper/2412.11767/) (2024)

## Abstract (arXiv)

Despite the remarkable success of diffusion models in text-to-image generation, their effectiveness in grounded visual editing and compositional control remains challenging. Motivated by advances in self-supervised learning and in-context generative modeling, we propose a series of simple yet powerful design choices that significantly enhance diffusion model capacity for structured, controllable generation and editing. We introduce Masking-Augmented Diffusion with Inference-Time Scaling (MADI), a framework that improves the editability, compositionality and controllability of diffusion models through two core innovations. First, we introduce Masking-Augmented gaussian Diffusion (MAgD), a novel training strategy with dual corruption process which combines standard denoising score matching and masked reconstruction by masking noisy input from forward process. MAgD encourages the model to learn discriminative and compositional visual representations, thus enabling localized and structure-aware editing. Second, we introduce an inference-time capacity scaling mechanism based on Pause Tokens, which act as special placeholders inserted into the prompt for increasing computational capacity at inference time. Our findings show that adopting expressive and dense prompts during training further enhances performance, particularly for MAgD. Together, these contributions in MADI substantially enhance the editability of diffusion models, paving the way toward their integration into more general-purpose, in-context generative diffusion architectures.