Rethinking Where to Edit

Abstract

Instruction-based image editing (IIE) aims to modify an image according to a natural language instruction. Despite recent advances in diffusion transformers, existing methods often introduce unintended changes to regions unrelated to the requested edit. We attribute this limitation to the absence of an explicit mechanism for edit localization. Different editing operations (e.g., subject addition, removal, and replacement) induce distinct spatial patterns, yet existing IIE models typically perform localization in a task-agnostic manner.

To address this limitation, we propose a training-free, task-aware edit localization framework that exploits the intrinsic source and target image streams of IIE models. For each image stream, we construct feature centroids from attention-based edit cues, and then partition tokens into edit and non-edit regions based on feature similarity. Observing that effective localization is inherently task-dependent, we introduce a unified mask construction strategy that selectively leverages the source and target streams according to the editing task. We also provide a systematic analysis of our underlying insights and design choices. Extensive experiments on EdiVal-Bench demonstrate that our framework consistently improves content consistency in non-edit regions while maintaining strong instruction-following performance on top of Step1X-Edit and Qwen-Image-Edit.

Key Insight: Localization Is Task-Dependent

Modern DiT-based editors process text, target image tokens and source image tokens in one joint attention sequence. Edit semantics do not emerge uniformly across these streams. Where to look depends on what the instruction asks for.

Method

Our framework plugs into a pretrained dual-stream editor at inference time. It needs no training and no extra supervision, and its runtime and memory cost stay comparable to the base model.

1. Attention-based semantic estimation. We slice the joint attention matrix into text-to-image cross-attention and within-stream self-attention, then propagate the text signal one step through the self-attention graph. Aggregating over layers and instruction tokens gives a coarse attention mask for each image stream.
2. Feature-based semantic assignment. Attention reflects relevance more than object structure, so we treat the coarse mask as a seed. Masked average pooling over deep DiT features yields an edit centroid and a non-edit centroid, and every token is assigned to the closer one by cosine similarity.
3. Task-aware mask construction. The stream-wise masks are selected or merged according to the editing task, following the rules above, and lightly cleaned with connected-component filtering, hole filling and a small dilation.
4. Mask-guided latent preservation. At a few early denoising steps, tokens outside the mask are replaced with an inverted source latent aligned to the current noise level, while tokens inside the mask evolve freely.

Where Do Edit Semantics Emerge?

Using SAM 3 pseudo ground-truth masks on Qwen-Image-Edit, we measure how well each signal localizes the edit across denoising timesteps. Two patterns hold across tasks: deep latent features (red) localize more accurately than attention maps with or without propagation, and the informative stream switches with the task.

Results on EdiVal-Bench

We apply the framework to Step1X-Edit and Qwen-Image-Edit and evaluate on EdiVal-Bench (572 real images, nine edit types). Content consistency (CC) improves on both backbones, and instruction following (IF) is maintained or slightly improved.

Qualitative Comparison

Each case shows the edited image next to its pixel-wise difference from the input. Brighter areas mean larger changes. Our difference maps concentrate on the requested region, while the base model also alters the fox, the woman's appearance and other untouched content.

@article{he2026rethinking,
  title   = {Rethinking Where to Edit: Task-Aware Localization for Instruction-Based Image Editing},
  author  = {He, Jingxuan and Wang, Xiyu and Zheng, Mengyu and Zeng, Xiangyu and Wang, Yunke and Xu, Chang},
  journal = {arXiv preprint arXiv:2604.20258},
  year    = {2026}
}