AttnRouter: Per-Category Attention Routing for Training-Free Image Editing on MMDiT

2026-05-02

AI Summary

This work addresses the performance degradation in multimodal diffusion Transformers (MMDiT) caused by prompt misalignment during image editing. To this end, the authors propose KVInject, an attention manipulation mechanism, coupled with AttnRouter, a category-aware routing strategy. Implemented in the Qwen-Image-Edit-2511 model, this approach enables efficient, training-free structure-preserving editing by fusing key/value (K/V) features in a single forward pass and leveraging a CLIP zero-shot classifier to automatically select optimal attention operations. The study identifies that K/V injection within specific layers during early denoising steps (S0–7) constitutes a critical subcircuit for effective editing. Experiments show a 6.4% improvement in editing quality when using ground-truth labels, with automatic routing recovering 98% of this gain; optimal performance is achieved with an injection weight α in the range [0.3, 0.5].

Abstract

We study training-free image editing on Qwen-Image-Edit-2511, a 60-block multi-modal diffusion transformer (MMDiT) that concatenates noise and source-image tokens within a single attention stream. We make three contributions. (i) We introduce KVInject, a single-forward attention manipulation that alpha-blends source-half key/value projections into the noise-half within a localized layer/step band. KVInject is simpler than the classical two-pass MasaCtrl recipe and avoids the prompt-mismatch failure mode that disables MasaCtrl on MMDiT (composite score drops 31% versus baseline). (ii) We show that no single attention operation dominates across edit types, motivating AttnRouter, a per-category routing table that dispatches edits to the operation that best preserves source structure for that type. With ground-truth categories the router improves the CLIP-T+DINO-I composite by 6.4% over the editing baseline; an automatic CLIP zero-shot classifier closes 98% of this gap despite only 55% category accuracy. (iii) Through layer-, step-, and alpha-band ablations we localize the editing-effective attention sub-circuit: K/V injection in early denoising steps (S0-7) recovers nearly all of the gain of full-step injection, while injection in early (L0-15) or late (L45-60) layer bands fails to drive editing entirely; alpha in [0.3, 0.5] is a stable sweet spot. We also report negative results that highlight what does not transfer from the UNet folklore: simple K/V rescaling never beats baseline and aggressive variants collapse generation entirely (composite 0.084). We release code, pre-computed routing tables, and a 100-sample stratified subset of ImgEdit-Bench used in all ablations.

Problem

• training-free image editing
• MMDiT
• attention routing
• per-category editing
• KV injection

Innovation

• KVInject
• AttnRouter
• training-free image editing
• MMDiT
• attention routing