Rethinking Where to Edit: Task-Aware Localization for Instruction-Based Image Editing

Instruction-based image editing (IIE) models often suffer from over-editing, introducing unintended changes to regions unrelated to the desired edit. This work proposes a training-free, task-aware edit localization framework that exploits intrinsic source and target image streams within IIE models to explicitly identify edit regions and guide latent updates, thereby enhancing content consistency without compromising instruction following.

How it works

The framework operates by decomposing and propagating attention activations from the multi-modal joint attention module to derive coarse localization cues. Specifically, it first derives cross-attention submatrices, such as the cross-attention submatrix A(l)img,ca (t) ∈ R img×Ntxt, and self-attention submatrices (intra-modal interactions within each image stream). These are combined through a diffusion of the linguistic signal to obtain an attention map, which is then aggregated across transformer layers and text tokens to yield an attention-derived mask Mimg(t).

How it works (Continued)

The paper further refines these initial cues by treating them as semantic indicators and performing a subsequent refinement in the latent feature space. This involves extracting and normalizing stream-specific latent features, denoted as Fˆ(l)img (t), and constructing two feature centroids via masked average pooling, one for edited regions (C(l)img,1(t)) and one for preserved regions (C(l)img,0(t)). The final edit mask Mˆ(l)(t)[i] is then determined by assigning each spatial token i to the closest centroid based on cosine similarity.

The crucial innovation lies in the task-aware edit mask construction, which selectively leverages source and target image streams based on the editing task. The paper identifies three representative subject-centric editing primitives:

1. Subject addition: localization is primarily manifested within the target stream, i.e., Mˆ (l) (t) = Mˆ(l)tgt (t).

2. Subject removal: regions are localized within the source stream, i.e., Mˆ (l) (t) = Mˆ(l)src (t).

3. Subject replacement: this task involves coupled semantic changes spanning both streams, i.e., Mˆ (l) (t) = Mˆ(l)tgt (t) ∪ Mˆ(l)src (t).

Once the task-aware edit mask is obtained, a mask-guided latent preservation scheme is applied to constrain the evolution of target latents during denoising. This involves constructing an inverted latent Zinv (t) defined as Zinv (t) = σtZtgt(0) + (1 − σt)Zsrc, where σt is the noise schedule at timestep t. The target tokens are then updated via mask-guided latent blending: Zˆtgt(t) = Mˆ(l)(t)Ztgt(t) + 1 − Mˆ(l)(t) Zinv (t). This mechanism ensures that the evolution of target tokens is constrained to within the identified edit region, thereby enforcing localized edits while maintaining consistency in non-target regions.

The effectiveness of this approach is validated through systematic analysis on EdiVal-Bench. The study reveals that latent features provide more reliable semantic cues than attention activations, as evidenced by higher IoU scores for feature-derived masks. Furthermore, the localization strategy is inherently task-dependent, with semantic signals emerging from different image streams depending on the editing task type (e.g., addition semantics in the target stream vs. removal semantics in the source stream). The final results demonstrate that this method consistently improves content consistency in non-edited regions while maintaining strong instruction-following performance.

A systematic ablation study confirmed the robustness of the framework:

1. Attention threshold sensitivity: The feature-based assignment is robust to the choice of attention threshold, with similar IoU scores across a range of values.

2. DiT layer depth: Segmentation performance improves monotonically as layer depth increases, peaking at layer 50, suggesting that edit semantics are progressively formed and become more discriminative in deeper layers.

3. Denoising timestep: Applying latent preservation early (e.g., at t=5) significantly improves EdiVal-O, while applying it too late can lead to a semantic mismatch, indicating that the choice of timesteps must be balanced for optimal fidelity and consistency.

In summary, the proposed framework addresses over-editing by explicitly identifying edit regions through task-aware mask construction based on stream-specific semantic cues derived from latent features. This localization is then enforced via mask-guided latent blending, leading to "semantically grounded editing with minimal unintended changes.

Based on the provided scientific paper, here are specific, actionable improvements for existing AI image editing systems and what those improved systems could achieve:

The core improvement is shifting from task-agnostic editing to a task-aware localization mechanism that explicitly identifies and isolates the edit region based on the nature of the instruction.

Here are the specific enhancements:

Improving Non-Edit Region Consistency (Reducing Over-Editing):

By implementing a training-free, task-aware edit localization framework (using attention cues to derive feature centroids and task-dependent mask construction), existing models can be guided to preserve non-target regions with significantly higher fidelity. The system will stop modifying background elements, human features, or unrelated objects when the instruction only targets a specific object or region.

Task-Specific Edit Localization:

The system will dynamically adjust its localization strategy based on the editing primitive:

• For Subject Addition: It will prioritize localization signals from the target image stream to ensure new content is placed contextually within the scene, rather than drifting into unrelated parts of the source image.

• For Subject Removal: It will focus localization signals on the source image stream to ensure only existing content is erased, preventing accidental corruption of adjacent background elements.

• For Subject Replacement: It will jointly leverage cues from both streams (source for removal and target for generation) to achieve a precise cut and paste effect without introducing artifacts at the seam.

Superior Semantic Grounding via Latent Features:

The improved system will utilize latent features derived from the joint attention module as primary semantic indicators, rather than relying solely on attention maps. This leads to more robust and semantically coherent partitioning of tokens, resulting in cleaner boundaries and more complete spatial coverage for the edit region compared to methods using only attention activations.

Enforced Spatial Constraint during Denoising (Mask-Guided Latent Preservation):

The system will introduce a mechanism that selectively constrains the evolution of target image tokens during denoising. By blending between the original source latent and the modified target latent only within the task-aware edit mask, it prevents all-to-all coupling from treating all spatial tokens as equally mutable, thereby enforcing localized updates and dramatically reducing unintended changes.

Improved Handling of Complex Editing Scenarios:

The proposed framework will systematically address complex editing needs by:

• Using the union of source and target masks (for replacement) to ensure both removal and addition are perfectly synchronized.

• Implementing post-processing steps (connected component retention, hole filling, spatial expansion) to handle disjoint edit regions more effectively than current methods, ensuring full coverage even when multiple distinct edits are requested simultaneously.

The resulting improved AI system can perform the following:

Generate images with significantly higher fidelity to the intended edit instructions while maintaining near-perfect preservation of the original image's non-edited content (e.g., high-quality subject replacement without object bleeding or background alteration).

2. Execute complex, multi-step edits (like replacing a specific object) with precise spatial control over where the modification occurs, minimizing visual artifacts at boundaries.

3. Provide superior performance on specialized benchmarks like EdiVal-Bench, demonstrating enhanced content consistency and instruction following across various editing primitives compared to current state-of-the-art models (e.g., Step1X-Edit or Qwen-Image-Edit).