Skip to content

In-context Region-based Drag: Drag Any Region to Any Shape

Conference: ECCV 2026
arXiv: 2606.25907
Code: https://github.com/bcmi/ICRDrag-Region-Drag-Editing
Area: Image Generation / Diffusion Models
Keywords: Region-based Drag Editing, In-Context Learning, Attention Regularization, Diffusion Transformer, Curriculum Learning

TL;DR

ICRDrag reformulates region-based drag as an in-context learning task. By taking the source image, source region mask, and target region mask as a unified conditional input to the DiT, the target image generation is guided via two novel regularization losses: Image-Mask Attention Consistency (IMAC) and Source-Target Attention Correspondence (STAC). Additionally, the paired 287k-sample Paired Region Dataset (PRD) is constructed. It significantly outperforms existing methods in editing accuracy, visual fidelity, and user preference.

Background & Motivation

Diffusion models have shown great potential in drag-style image editing, but existing works mainly focus on point-based drag, where users specify editing intentions through several source-target point pairs. This sparse condition has inherent ambiguity: for example, dragging a source point on a face outward can mean rotating the face direction or stretching the face width; among multiple reasonable explanations, often only one matches the user's intent. Moreover, sparse point pair conditions lead to insufficient editing accuracy, making it difficult to drag the source point exactly to the target point position.

Region-based drag fundamentally solves the ambiguity problem through dense spatial conditions—source region masks and target region masks. However, only a few prior region drag methods exist: EditGAN uses masks only as a loss function rather than deeply integrating them into the generation process; RegionDrag employs a latent copy-paste strategy, resulting in unnatural boundaries and difficulty in handling complex deformations. The key challenge is: there is a lack of deep coupling mechanism between the dense structural information of the region masks and the image generation process—the masks tell you what shape the editing target is and where it is, but the generation process does not know how to utilize these structural constraints to guide image synthesis pixel-by-pixel.

The goal of this paper is to prevent masks from being "post-hoc supervision signals" or "crude feature movers," and instead let them deeply participate in the entire generation process as in-context conditions equivalent to images. The key insight is to redefine region-based drag as an in-context learning problem: the source image, source mask, and target mask constitute a unified context, and the model directly generates the target edited image in a single forward pass. Based on this, by applying two regularization constraints, image-mask consistency and source-target correspondence in the attention layers, the structural information of the masks is explicitly injected into the aggregation process of image features.

Core Idea: Pack the three conditions of drag editing (source image, source mask, target mask) into an in-context input, enabling the DiT to learn to "see image + see mask \(\rightarrow\) generate an image matching the mask constraints," and use two attention emulation losses to ensure that the mask structural information truly guides the aggregation of image features.

Method

Overall Architecture

ICRDrag is based on the Next-DiT (Lumina) architecture, modeling the region-based drag task as a conditional generation problem. The input consists of three parts: source image \(I_s\), source region mask \(M_s\) (the region to be edited, with different regions annotated with different label IDs), and target region mask \(M_g\) (the desired shape/position after editing, sharing the same label IDs as the source mask). All three and the target image \(I_g\) are mapped to the latent space through a VAE encoder, yielding \(\hat{I}_s, \hat{M}_s, \hat{M}_g, \hat{I}_g \in \mathbb{R}^{H \times W \times D}\).

Throughout training and inference, the conditional inputs \(\{\hat{I}_s, \hat{M}_s, \hat{M}_g\}\) always remain noise-free, and only the target image latent variable \(\hat{I}_g\) undergoes the noise-denoise process. The model predicts the velocity field \(\bm{v}_{\theta}(t, \hat{I}_s, \hat{M}_s, \hat{I}_g^t, \hat{M}_g)\), which is supervised by a flow-matching loss to approximate the target velocity \(\bm{u} = \hat{I}_g - \bm{\epsilon}\). During inference, starting from random Gaussian noise, denoising is performed step-by-step, and the final clean latent variable is reconstructed into the edited image \(I_e\) via the VAE decoder.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Source Image + Source Mask + Target Mask"] --> B["VAE Encoding<br/>→ Latent Condition Tokens"]
    B --> C["Modality-Specific LoRA<br/>Separate Image/Mask Feature Extraction"]
    C --> D["DiT Transformer<br/>16-layer Self-Attention + Cross-Attention"]
    D --> E["IMAC Regularization<br/>Image-Mask Attention Consistency"]
    D --> F["STAC Regularization<br/>Source-Target Attention Correspondence"]
    E --> G["Velocity Field Prediction + Flow Matching Denoising"]
    F --> G
    G --> H["VAE Decoding → Edited Image"]

Key Designs

1. Modality-Specific LoRA: Resolving Image and Mask Feature Confusion

Image and mask have inherently different properties: images are rich in texture and fine details, whereas masks are sparse and only encode spatial structure. Sharing parameters to process both would contaminate the image pathway with the mask's structural representations, leading to detail loss and over-smoothed generations.

Specifically, separate LoRA modules are inserted into the Feed-Forward Networks (FFN) of the DiT for image and mask tokens, respectively. This allows each modality to learn representations tailored to its properties, sharing the backbone while avoiding cross-modality interference. This is a lightweight design—without modifying the pre-trained weights of the original DiT, only the newly added LoRA parameters are trained, which preserves the generative capacity of the foundation model and provides a dedicated process channel for mask information.

2. IMAC (Image-Mask Attention Consistency): Grounding Image Generation in Mask Structure

Core assumption: A patch in the target region of the image modality should pay attention to the same areas in the source image as its corresponding patch in the mask modality does to the source mask. In other words, "where to find visual materials" during image generation should be guided by the spatial structure defined by the masks.

Formally, let \(\mathcal{P}_g\) be the set of patches within the target region mask. For each \(p \in \mathcal{P}_g\), the token of this patch in the target image branch is taken as the query to calculate the attention map over all patches in the source image, denoted as \(\bm{A}_p^{I_s \leftarrow I_g}\). Simultaneously, the token at the same spatial location in the target mask branch is taken as the query to calculate the attention map over all patches in the source mask, denoted as \(\bm{A}_p^{M_s \leftarrow M_g}\). The IMAC loss directly minimizes the MSE between the two attention maps:

\[\mathcal{L}_{\text{imac}} = \sum_{p \in \mathcal{P}_g} \|\bm{A}_p^{I_s \leftarrow I_g} - \bm{A}_p^{M_s \leftarrow M_g}\|^2\]

Why it works: The prior that "hair regions in the mask attend to hair regions in the source mask" is very easy to learn because the mask itself consists of discrete semantic label regions. IMAC distills this "easy-to-learn structural prior" from the mask modality to the image modality, forcing the image generation process to follow the same spatial correspondence, thereby ensuring that the generated pixel content does not drift into incorrect semantic regions.

3. STAC (Source-Target Attention Correspondence): Bidirectional Mutual Attention for Correspondence Consistency

Core assumption: Related regions in the source image and the target image should mutually attend to each other. If target patch \(i\) strongly attends to source patch \(j\), then source patch \(j\) should also strongly attend to target patch \(i\) in return. This mutual information constraint ensures that the model establishes consistent correspondences between the source and target, which is crucial for spatial transformations such as object movement, scaling, and joint adjustments.

Let \(\mathbf{A}^{g \rightarrow s} \in \mathbb{R}^{|\mathcal{P}_g| \times |\mathcal{P}_s|}\) be the target-to-source attention matrix and \(\mathbf{A}^{s \rightarrow g} \in \mathbb{R}^{|\mathcal{P}_s| \times |\mathcal{P}_g|}\) be the reverse source-to-target attention matrix. Considering the product \(\mathbf{A}^{g \rightarrow s} \mathbf{A}^{s \rightarrow g} \in \mathbb{R}^{|\mathcal{P}_g| \times |\mathcal{P}_g|}\), its diagonal elements \((\mathbf{A}^{g \rightarrow s} \mathbf{A}^{s \rightarrow g})_{ii} = \sum_j \mathbf{A}^{g \rightarrow s}_{ij} \mathbf{A}^{s \rightarrow g}_{ji}\) precisely measure the sum of bidirectional attention products between target patch \(i\) and all the source patches it attends to. Maximizing the diagonal elements is equivalent to enforcing a mutual attention closed-loop. The loss function is:

\[\mathcal{L}_{\text{stac}} = -\text{Trace}(\mathbf{A}^{g \rightarrow s} \mathbf{A}^{s \rightarrow g})\]

In practice, multi-head attention is extracted from the 7th-11th transformer layers (the authors' experimental analysis revealed that middle layers are best suited for applying such constraints; early layers focus on local content localization, late layers focus on spatial refinement, and the middle layers integrate the source-target context). The bidirectional attention matrix product trace is calculated after averaging over multiple heads. Both IMAC and STAC use softmax-normalized attention maps.

4. Two-Stage Curriculum Training Strategy: Progressive Learning from Full to Sparse Masks

In real-world applications, users might only annotate a few areas to be edited (incomplete masks). Compared to full masks where all semantic regions are segmented out, the model faces sparser information: edits may involve changes outside the explicitly annotated regions, and the model must infer the collateral adjustments needed in unannotated regions.

In the first stage, the model is trained with full region masks for 60,000 steps (batch size 2, lr \(1 \times 10^{-4}\)), allowing it to learn the basic "mask \(\rightarrow\) image" mapping capability under information-rich conditions. In the second stage, it switches to incomplete masks and is trained for another 2,000 steps (batch size 1, lr \(5 \times 10^{-5}\)): 1-5 key regions are randomly sampled and kept from the full mask while the rest of the regions are filled with a gray value, and dilation operations are randomly applied to the sampled regions to enhance robustness against imprecise user inputs. This easy-to-hard curriculum design smoothens the optimization process and avoids training instability caused by starting directly with incomplete masks.

A Complete Example

Take the facial drag in Figure 2 as an example. The user provides: (1) source image—a frontal face; (2) source mask—three blue annotations for hair (label 1), face (label 2), and arm (label 3); (3) target mask—the shape after stretching the hair region to the left side and slightly rotating the face (red annotations, label IDs kept consistent with the source mask).

System workflow: The three inputs are encoded by VAE and passed into the image LoRA and mask LoRA to extract dedicated features. In the 7th-11th layers of the DiT, the IMAC regularization ensures that when a target image patch belonging to the "hair" region aggregates source image features, its attention distribution is consistent with the mask modality's attention distribution where "hair patch attends to source mask hair region." Thus, the generated hair texture originates from the source hair region without shifting onto the face. Concurrently, the STAC regularization ensures that if a facial patch in the source image is attended to by a target image patch, that source patch will also attend back to the target patch, forming a bidirectional lock. After denoising and decoding, the output displays a face slightly rotated, hair stretched according to the mask shape, and other unannotated regions (e.g., the background) automatically inferred with reasonable collateral changes.

Loss & Training

The total loss is a weighted sum of three terms: \(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{flow}} + \lambda_1 \mathcal{L}_{\text{imac}} + \lambda_2 \mathcal{L}_{\text{stac}}\). Where \(\mathcal{L}_{\text{flow}}\) is the standard flow-matching loss (the predicted velocity field approximates the negative of \(\bm{\epsilon} - \hat{I}_g\), while the actual target velocity is \(\hat{I}_g - \bm{\epsilon}\)). The timestep \(t\) is sampled from a LogNormal(0, 1) distribution, which is standard practice in Rectified Flow (Esser et al., 2024), focusing training on intermediate signal-to-noise ratio regions compared to uniform sampling.

IMAC and STAC, as auxiliary losses, are computed only on patches covered by the target region mask (\(\mathcal{P}_g\) and \(\mathcal{P}_s\)) and are not applied to the background to avoid meaningless constraints interfering with generation quality. \(\lambda_1\) and \(\lambda_2\) are hyperparameters determined through experimental tuning.

Key Experimental Results

Main Results

PRD Benchmark Quantitative Comparison (Table 1). The PRD benchmark contains 1,000 manually verified samples with real target images as ground truth, allowing full-reference metrics.

Method MSE \(\downarrow\) LPIPS \(\downarrow\) SSIM \(\uparrow\) MD(RegionDrag) \(\downarrow\) MD(DragLoRA) \(\downarrow\)
DragDiffusion 0.0937 0.1836 0.5993 5.17 26.79
SDE-Drag 0.1017 0.2018 0.5849 7.97 44.04
DiffEditor 0.0959 0.1949 0.6071 23.45 31.19
FastDrag 0.0962 0.2049 0.5862 6.01 31.22
Inpaint4Drag 0.0972 0.1923 0.6102 4.24 23.62
GoodDrag 0.0902 0.1761 0.6094 3.70 18.01
DragLoRA 0.0933 0.1870 0.5974 4.62 23.96
RegionDrag 0.0977 0.1944 0.6076 8.02 43.05
ICRDrag 0.0735 0.1610 0.6284 3.66 22.34

ICRDrag leads comprehensively across all metrics, with MSE reduced by 18.5% compared to the best baseline GoodDrag, LPIPS reduced by 8.6%, and SSIM improved by 0.019. Notably, RegionDrag (the only directly comparable region drag method) is weaker than strong point-drag baselines like GoodDrag across all metrics, whereas ICRDrag pushes the performance of region-based drag to a new upper bound.

User Study (Table 2, preference votes from 50 participants on DragBench). DragBench lacks ground truth and is evaluated on three subjective dimensions.

Method Realism \(\uparrow\) Fidelity \(\uparrow\) Region Accuracy \(\uparrow\)
GoodDrag 0.2876 0.2120 0.1276
RegionDrag 0.2442 0.2314 0.3518
ICRDrag 0.4682 0.5566 0.5206

ICRDrag wins with an overwhelming advantage: its realism is 1.6 times that of the second place, its fidelity is 2.4 times, and in terms of region accuracy, GoodDrag achieves very low accuracy (0.1276) due to using only point pair conditions, while ICRDrag achieves a 48% improvement compared to RegionDrag.

Ablation Study

Configuration MSE \(\downarrow\) LPIPS \(\downarrow\) Description
Full ICRDrag 0.0735 0.1610 Full model, IMAC + STAC
w/o IMAC — — Edited results deviate from target masks, with structural distortion and boundary misalignment
w/o STAC — — Fine-grained details of source image are lost; textures and identity features are altered or lost during dragging

Quantitative results are only shown for the full model, while quantitative ablation data are in the supplementary material. Qualitative results (Figure 8) clearly demonstrate the respective roles of the two losses: without IMAC, the image content cannot strictly align with the target mask shape; without STAC, the texture details of the source image cannot be maintained.

Key Findings

  • IMAC and STAC have a clear division of labor: IMAC is responsible for "editing accuracy" (making generated content conform to the target mask), while STAC is responsible for "content fidelity" (preserving the details of the source image). They complement each other and are both indispensable.
  • Division of labor of attention maps across different layers: Early layers (1-6) model source content localization, middle layers (7-11) integrate source-target context (hence IMAC/STAC are applied here), and late layers (12-16) perform spatial refinement—this layered characteristic provides the empirical basis for the authors' choice of loss-application layers.
  • Strong cross-dataset generalization: Realistic photos from Adobe Stock can be edited reasonably without any fine-tuning, demonstrating that ICRDrag learns a general "mask \(\rightarrow\) image" mapping capability rather than a specific pattern of the PRD distribution.
  • Difficult scenarios: In non-rigid editing scenarios with large topological changes, occlusions, and human limb repositioning, ICRDrag still generates reasonable results, whereas point-drag baselines clearly fail in these cases.

Highlights & Insights

  • The framing of "In-context learning for drag" is highly ingenious: It liberates region-based drag from traditional approaches like iterative optimization, GAN latent-space editing, or feature copy-pasting, transforming it into a conditional generation problem. Essentially, it "generativizes" the drag task—producing results in a single forward pass, which is faster and more stable than iterative methods.
  • IMAC is essentially a cross-modality attention distillation: The mask modality easily learns correct spatial correspondences (as masks are discrete label regions). Distilling this "easy-to-learn" structural prior into the image modality via L2 loss avoids the image modality having to search for attention distributions from scratch. This idea can migrate to any task requiring "using structural information of a simple modality to constrain the generation of a complex modality," such as using depth maps to constrain novel view synthesis or edge maps for image-to-image translation.
  • The trace maximization formula of STAC is simple and elegant: The trace of the bidirectional attention product naturally captures the mutual attention closed-loop—requiring no extra forward or backward passes, but simply extracting and computing from existing attention matrices with zero additional inference overhead. This concept of "imposing structural constraints on attention matrices" can be transferred to other tasks requiring bidirectional correspondence (e.g., bidirectional image matching, video inter-frame correspondence).
  • Practical value of the two-stage curriculum training: The first stage uses full masks to learn foundational capabilities, and the second stage uses only 2,000 steps (3.2% of the total steps) to adapt to sparse masks. This division of labor strategy of "doing the heavy lifting on easy tasks first, then fine-tuning on hard tasks" is highly efficient and can be extended to other generative tasks handling sparse or noisy conditions.

Limitations & Future Work

  • Limitations not explicitly discussed by the authors: The PRD dataset is constructed based on the OpenVid video dataset. Video frame pairs naturally guarantee scene consistency and lighting stability between the source and target images—this implies that ICRDrag might perform poorly in drag scenarios with dramatic changes in scene lighting or drastic background modifications, which is not evaluated with diagnostic experiments in the paper. Additionally, masks must be provided by users via tools like SAM or drawn manually, which imposes a higher entry barrier than point-drag.
  • Lack of inference efficiency comparison: As a single-forward generation method based on DiT, ICRDrag's inference speed is theoretically superior to iteratively optimized point-drag methods (e.g., DragDiffusion, GoodDrag), but the paper does not provide quick quantitative measurements of inference time.
  • Relationship to DragFlow: ICRDrag is a training-dependent method (requiring training on PRD), while concurrent work DragFlow is a training-free method (relying on region-level affine supervision on DiT feature layers)—their application scenarios and cost structures are completely different, a trade-off not discussed in the paper.
  • Future directions: (1) Extending ICRDrag to video drag-style editing, utilizing temporal consistency across video frames to improve editing stability; (2) Combining LLMs for text-guided region-based drag ("drag the mug on the left to the right") to lower the cost of manual mask annotation; (3) Exploring the transfer of IMAC/STAC regularization concepts to other DiT-based controllable generation tasks (e.g., subject-driven image generation, virtual try-on).
  • vs RegionDrag: RegionDrag performs copy-paste operations in the latent space and moves features by converting region masks into point pairs, which is essentially still a point-based drag variant. ICRDrag inputs masks as first-class conditions to the DiT and injects structural constraints into the generation core via attention regularization. The disadvantage of ICRDrag is its need for training, whereas RegionDrag is training-free; the advantage is significantly higher editing quality, especially in complex deformations and boundary blending.
  • vs GoodDrag / DragDiffusion: Representative point-drag methods like GoodDrag and DragDiffusion rely on iterative optimization (motion supervision + point tracking). Their fundamental limitation lies in the ambiguity caused by the sparseness of point-pair conditions. ICRDrag eliminates this ambiguity from the source by using region masks. Their relationship is not "different methods for the same task," but rather "region-based drag defines a more precise input space, leading to a higher upper bound."
  • vs DragFlow: DragFlow is also a training-free DiT-based region drag method, replacing point-wise supervision with region-level affine supervision. ICRDrag and DragFlow have orthogonal technical routes—one is training-dependent in-context learning, and the other is training-free feature constraint. Future research could combine both directions: using DragFlow's feature constraint ideas to design stronger training-free baselines, or using ICRDrag's attention regularization ideas to enhance DragFlow's generation quality.
  • vs In-Context Learning Series (Prompt-to-Prompt, Self-Supervised ICL): ICRDrag's in-context design shares the same lineage as vision ICL works—concatenating multiple related inputs to form a context, letting the model implicitly infer the task. The difference is that ICRDrag's "context" consists of heterogeneous image-mask multimodal information rather than homogeneous samples (like image pairs), and it ensures effective context utilization through explicit attention regularizations (rather than relying solely on implicit learning).

Rating

  • Novelty: ⭐⭐⭐⭐ Formulating region drag as an in-context learning task is a clever perspective shift. The designs of IMAC and STAC attention regularizations have clear intuitive support and elegant formulations. However, the overall framework is built upon the existing Next-DiT architecture, with core contributions leaning towards loss function designs and training strategies.
  • Experimental Thoroughness: ⭐⭐⭐⭐ Quantitative/qualitative/user study comparisons were thoroughly conducted on the PRD benchmark and DragBench, covering both point-drag and region-drag baselines, accompanied by analysis of difficult cases and cross-dataset transfer tests. However, the complete quantitative results for the ablation study are relegated to the supplementary material, with only qualitative demos in the main text.
  • Writing Quality: ⭐⭐⭐⭐ The problem definition is clear (Section 3 defines signs and tasks), and the method unfolds step-by-step (Framework \(\rightarrow\) IMAC \(\rightarrow\) STAC \(\rightarrow\) Curriculum Training). Figures are sufficient (Figure 2's three sub-figures cover the framework/IMAC/STAC). However, the details of PRD dataset construction are scattered between the method and experiments, introducing slight discontinuities during reading.
  • Value: ⭐⭐⭐⭐ Established a new performance benchmark for region-based drag (comprehensively beating RegionDrag and all point-drag methods). The PRD dataset of 287k samples is an independent contribution to the community. The attention regularization of IMAC/STAC has cross-task migration potential. Practical applications are restricted by the extra cost of manual user-defined region masks.