A²-Edit: Precise Reference-Guided Image Editing of Arbitrary Objects and Ambiguous Masks

Conference: ECCV2026
Paper: Official paper page · PDF
Code: https://github.com/huayu-zheng/A2Edit
Area: Image Generation / Reference-Guided Image Editing
Keywords: reference-guided editing, arbitrary object insertion, Mixture of Transformers, anchor-guided routing, mask annealing

TL;DR

A²-Edit augments FLUX.1 Fill with joint attention-FFN LoRA experts, a training curriculum that progressively relaxes segmentation masks into bounding boxes, and 500,104 cross-category training samples to retain reference identity and structure under coarse masks, improving UniEdit rough-mask DINO-I from Insert Anything's 56.13 to 62.28.

Background & Motivation

Reference-guided editing is more than pasting an object into another image: the model must inherit identity, texture, or material from the reference while accommodating the target scene's pose, scale, and occlusion. Garments emphasize patterns and contours, portraits require identity preservation through non-rigid changes, and furniture or architecture depend more on geometry and perspective. Sending all these tasks through one parameter pathway can make transformations useful for one category interfere with another. Meanwhile, datasets concentrated on garments or a few object types provide limited evidence about which relationships should be shared and which need specialization.

Masks amplify this problem. A fine segmentation supplies the object boundary directly, whereas a rough stroke or bounding box only specifies approximate location. The model must then infer which pixels should belong to the replacement object and which should remain background. If training always provides accurate boundaries, enlarging the mask at inference is not merely adding noise: it removes shape information the model has learned to depend on. The paper therefore tackles category differences and imprecise spatial guidance together, rather than only improving reference features.

Core idea: share a category-adaptive expert selection between attention and FFN, and progressively train with fine masks, perturbed rough masks, and bounding boxes to shift editing decisions from supplied contours toward reference identity and scene semantics; multi-category paired data supplies the learning foundation.

Method

Overall Architecture

The core inputs are a reference image, a target image, and a user mask, with an optional text prompt. The system first extracts the reference foreground, combines SigLIP image features, VAE latents of the target, and optional T5 text features into a unified token sequence, and passes it through the FLUX.1 Fill DiT backbone. Each layer uses a Mixture of Transformers (MoT) to select lightweight LoRA experts, and a VAE decoder produces the edited image.

Two training-side components complete the approach: UniEdit-500K supplies paired examples from different categories, and the Mask Annealing Training Strategy (MATS) changes their spatial conditioning across stages. The diagram includes training and inference in one view. Mask annealing is a training curriculum, not an extra fine-mask prediction step at inference.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    DATA["UniEdit-500K<br/>Multi-Category Pairing"] --> MASK["Mask Annealing Curriculum<br/>Fine to Rough to Box"]
    MASK -->|Training conditions| ENC["Foreground Extraction and Encoding<br/>Unified Token Sequence"]
    INPUT["Reference and Target Images<br/>User Mask and Optional Text"] -->|Inference conditions| ENC
    ENC --> MOT["Joint Attention-FFN<br/>Expert Modeling"]
    MOT --> AGR["Anchor-Guided Shared Routing<br/>Layer-Wise Expert Selection"]
    AGR --> OUT["Iterative Generation and VAE Decoding<br/>Edited Image"]

Key Designs

1. UniEdit-500K Multi-Category Pairing: expose the relationship between identity preservation and deformation

The training set contains 500,104 samples across 8 major categories and 209 subcategories: garments, portraits, animals, plants, accessories, furniture, vehicles, and architecture. Pair construction differs by category instead of simply cropping two patches from one image. Garments and accessories draw on VITON-HD, DressCode, and licensed standalone-item and wearing-scene images. Portraits and animals come from VidGen-1M video frames: Grounding DINO localization, SAM segmentation, and no-reference quality filtering precede selection of two frames according to subject variation. Plants, vehicles, architecture, and furniture mainly use multi-view images with automatic pairing and manual verification.

These pairs allow reference and target to differ in pose, viewpoint, or background while still depicting the same subject. This gives the model evidence about what to copy and what may change, rather than only teaching pixel reconstruction. The paper also uses NanoBanana to supplement missing data and address imbalance in some domains, then produces detection boxes and segmentation masks for the paired images. Its licensing description distinguishes openly licensed sources from creator-authorized subsets. The main text does not quantify the synthetic fraction, so the data should not be described as entirely natural photography or uniformly licensed.

2. Mask Annealing Curriculum: gradually remove the shortcut supplied by object contours

The first stage uses binary fine masks strictly aligned with the target object to establish reliable correspondence between the reference subject and target region. The second stage does not merely discard mask information at random. It first applies isotropic morphological dilation with a circular structuring element, extracts the outer contour, perturbs it using spatially correlated Perlin-noise displacements, clips coordinates to the image domain, and reconstructs the mask. Correlated noise produces continuous, curved drawing errors instead of the high-frequency jitter caused by independently perturbing every boundary pixel.

The third stage further relaxes the spatial hint to a bounding box, forcing the model to infer pose, scale, and local structure without a detailed contour. The purpose is not to make the mask irrelevant: it preserves the location of the edit while progressively weakening the supplied shape of the new object. The main text explains the roles of displacement magnitude, noise scale, and axis-decoupling parameters, but the corresponding equations are corrupted in the cached extraction and their numerical settings are not given. This note retains the verifiable procedure without reconstructing the missing equations.

3. Joint Attention-FFN Expert Modeling: category specialization should also change relational computation

Encoding first removes the reference background and then projects target latents, reference visual features, and optional text features into the same hidden dimension to form a token sequence. The main text abstracts this operation as token-level Fuse without detailing the exact token layout or positional encoding in the cache. It should therefore not be equated with a particular existing concatenation implementation. The important point is that expert conditioning includes both reference identity and target context, rather than selecting behavior solely from an object category name.

Conventional sparse MoE often specializes FFNs while sharing attention across categories. A²-Edit argues that this is insufficient for editing: garment texture alignment, facial non-rigid changes, and architectural structure require different relationships between tokens in the first place. It therefore introduces expert-specific LoRA increments into attention linear transformations, such as query, key, and value projections, and into FFN linear transformations in each Transformer layer. Shared backbone weights retain general generative knowledge, while routing-weighted low-rank increments from activated experts contribute to the corresponding transformations. This does not replicate a full DiT for each category.

The implementation uses 8 experts with LoRA rank 32. Although the number of experts matches the 8 major categories, the paper describes learned feature-dependent specialization, not a fixed one-to-one category-to-expert label assignment. The cached weight-combination equation also has extraction damage, so this note does not supply unverified normalization rules or a matrix multiplication order.

4. Anchor-Guided Shared Routing: retain general editing priors while selecting input-dependent assistance

An MLP gate predicts expert weights from the input features at each layer. Anchor-Guided Routing (AGR) always retains a backbone expert and selects assistants according to their weights relative to the backbone's weight; the main text also describes a highest-weighted-expert fallback. This avoids forcing all categories into one set of adaptation parameters while maintaining a stable general editing pathway instead of relying exclusively on whichever experts receive the highest scores. Expert modeling and shared routing in the diagram jointly constitute the layer's MoT, rather than two independent image generators.

Routing is performed once per layer, and attention and FFN reuse the same weights so that relationship construction and content transformation use a consistent expert combination. Ablations compare FFN-only MoE, fixed experts, Top-2, softmax weighting, and separate attention and FFN routers. One qualification matters: the method section allows assistants whose scores exceed the backbone score, whereas the ablation section says AGR's activation count is matched by 2 experts. The cache does not fully resolve selection limits or fallback boundaries. It is therefore unwarranted to claim that every input always activates exactly 2 experts, or to infer a specific load-balancing loss.

A Worked Example

Consider garment editing, one of the tasks illustrated in the paper, as a procedural example rather than an additional experiment. A user provides a reference garment, a target person image, and a rough upper-body region covering the clothing and some background. Reference foreground encoding supplies identity cues such as pattern and material, while target encoding preserves pose and surrounding context. Text can further specify the request but is not required.

After the unified tokens enter the DiT, the backbone expert supplies general inpainting ability and assistants contribute input-relevant texture and shape modeling through routing. Exposure to contour perturbations and bounding boxes during training encourages the model to interpret the loose region as an allowable generation area rather than a clothing silhouette that must be filled completely. The VAE decodes the result. This is learned behavior, not an explicit SAM step that relocalizes the garment, and it does not guarantee exact preservation of every pixel outside the mask.

Loss & Training

Routing parameters and LoRA experts are jointly optimized end to end. The main text does not specify the full training loss, optimizer, learning rate, or auxiliary regularization terms, so a familiar diffusion or flow-matching objective should not be presented as an equation reported by this paper.

Training uses 4 NVIDIA A100 GPUs, batch size 8, and images at \(1024\times1024\) resolution. The total is 6000 steps: 3000 with fine masks, 1500 with augmented rough masks, and 1500 with bounding boxes. These step counts should not be described as a number of complete training-set epochs.

On a single A100 80GB, inference with 50 sampling steps averages approximately 30 seconds, versus 32 seconds for the baseline, measured over 100 evaluation samples. Peak memory is 42,074 MiB versus 39,884 MiB. This suggests limited adaptation overhead, but without variance or throughput curves it does not establish a reliable speedup or low-memory deployment.

Key Experimental Results

Main Results

Evaluation covers VITON-HD, AnyInsertion, and the authors' UniEdit test set. AnyInsertion contains 100 examples: 40 objects, 30 garments, and 30 portraits. UniEdit contains 200 test pairs, with 25 from each major category. The following rows reproduce selected rough-mask results from the paper's Table 1. Quantitative rough masks are produced by dilating fine masks; they should not be conflated with Perlin-perturbed training contours or the hand-drawn masks used for qualitative comparisons.

Dataset / Rough Masks	Method	DINO-I ↑	CLIP-I ↑	LPIPS ↓
VITON-HD	FLUX.1-Kontext	56.11	77.91	0.0972
VITON-HD	Insert Anything	60.87	78.22	0.0809
VITON-HD	A²-Edit	63.79	80.19	0.0685
AnyInsertion	FLUX.1-Kontext	54.15	74.05	0.1384
AnyInsertion	Insert Anything	60.50	75.98	0.0992
AnyInsertion	A²-Edit	61.73	77.27	0.0903

The next table selects UniEdit rough-mask results from the paper's Table 2. DINO-I and CLIP-I measure feature similarity, LPIPS measures perceptual distance, and FID measures the difference between generated and real image distributions. Scores retain the original table scales. VLM evaluation concerns boundary quality, realism, and local consistency; it is not a user preference percentage.

Method / UniEdit Rough Masks	DINO-I ↑	CLIP-I ↑	LPIPS ↓	FID ↓	VLM ↑
ACE++	56.68	73.66	0.1219	63.19	74.7
Insert Anything	56.13	73.37	0.1243	61.48	75.0
A²-Edit	62.28	77.45	0.0897	54.62	79.3

The following average expresses the main text's verbal VLM-score definition: \(n\) is the sample count, and the three component scores concern boundaries, realism, and local consistency. This is notation for the stated definition, not a new evaluation protocol. The specific evaluator model, score range, and prompts are not detailed in the available main-text cache.

\[ \mathrm{VLM}=\frac{1}{3n}\sum_{j=1}^{n}\left(s_j^{\mathrm{boundary}}+s_j^{\mathrm{realism}}+s_j^{\mathrm{consistency}}\right). \]

Ablation Study

The paper's Table 3 evaluates all variants on UniEdit; the table below retains its rough-mask columns. AE denotes the arbitrary editing framework, UE denotes UniEdit-500K, and FT, RT, and BT denote fine-mask, augmented rough-mask, and bounding-box training. The full model includes FT+RT+BT.

Configuration	DINO-I ↑	CLIP-I ↑	LPIPS ↓	FID ↓	VLM ↑
Without AE framework	56.18	73.21	0.1364	72.63	65.7
FFN-only MoE	56.11	73.69	0.1366	73.04	70.3
Without UE training	56.27	73.01	0.1347	72.84	66.3
FT only	55.65	72.92	0.1516	74.77	64.7
FT+RT	59.70	74.86	0.1073	60.33	73.0
Top-2 routing	57.42	75.90	0.1101	58.87	72.7
Fixed-2 routing	55.21	73.49	0.1406	73.13	66.0
softmax routing	60.95	76.91	0.0983	56.12	74.7
Separate attention and FFN routing	61.16	77.34	0.1003	54.90	77.3
Full A²-Edit	62.28	77.45	0.0897	54.62	79.3

Key Findings

• On UniEdit rough masks, the full model improves DINO-I by 6.15, reduces FID by 6.86, and reduces LPIPS by 0.0346 relative to Insert Anything. Agreement across these measures is more persuasive than semantic similarity alone.
• Training only with fine masks yields rough-mask LPIPS of 0.1516; adding RT reduces it to 0.1073, and adding BT reduces it further to 0.0897. The final addition also adds a training stage and steps. Without a step-matched control, the entire gain cannot be attributed to bounding boxes themselves.
• Separate versus shared routing differs by only 0.28 in FID, while LPIPS falls from 0.1003 to 0.0897. Coordinating relational modeling with content transformation appears to help local detail rather than producing a dramatic improvement on every metric.
• Full-model UniEdit DINO-I rises from 61.71 with fine masks to 62.28 with rough masks, but VLM falls from 82.3 to 79.3. Coarser masks do not improve every quality dimension, and these results do not imply that larger mask errors are always better.
• In a user study with 24 participants, A²-Edit receives 69.0% preference against Insert Anything. This complements automatic metrics, but the main text does not fully specify vote counts or confidence intervals.

Highlights & Insights

• Expert specialization extends from FFNs into attention projections, targeting the relational nature of reference editing. Matching a texture across different poses is not merely a matter of transforming token content independently.
• MATS turns coarse-mask robustness into an actionable training curriculum. The transferable lesson is not one particular noise generator, but identifying a shortcut in the conditioning and gradually removing it after basic correspondence has been learned.
• Paired data, category specialization, and a spatial curriculum are complementary. Increasing category coverage alone does not prove that specialists emerge; the joint-routing ablations provide more direct evidence for that interpretation.

Limitations & Future Work

• The main text has no dedicated failure-case or limitations section. The following are reader assessments of the reported evidence, not failure modes the authors have already demonstrated.
• Arbitrary categories are not yet equivalent to strictly unseen-category generalization. UniEdit testing uses the same 8 major categories covered by training, without a clear held-out-category protocol or complete subcategory breakdown. Cross-category and cross-source holdouts would strengthen the claim.
• Quantitative evidence for arbitrarily coarse masks mainly involves dilation. Displaced masks, missing coverage, holes, and disconnected regions lack separate degradation curves and should be evaluated alongside real hand-drawn inputs.
• The UE ablation is described as using the untrained initial model without that data, not training on an equal-sized alternative dataset. It supports the value of domain training but cannot independently establish that this dataset's category diversity is superior to any equally large alternative.
• The custom test set contains only 200 pairs, and FID uncertainty is not reported. VLM protocol details, AGR edge cases, the full loss, and the synthetic data fraction are insufficiently specified in the available main-text cache. More complete reproduction details and repeated-run statistics would help separate robust gains from evaluation variation.
• Peak memory of 42,074 MiB and approximately 30 seconds per generation remain deployment barriers. Sampling acceleration or expert compression could be investigated while retaining identity and coarse-mask benefits; the reported configuration should not be called real-time editing.

Related Work & Insights

• Versus AnyDoor and MimicBrush: all use reference images for object-level editing. A²-Edit further emphasizes cross-category parameter specialization and degraded spatial guidance. Its advantage is supported in these experiments, not established for every editing task.
• Versus ACE++ and Insert Anything: these explore unified or in-context insertion, while A²-Edit adds joint attention-FFN experts and a mask curriculum. Backbone, data, and training also differ, so main comparisons should be read alongside ablations rather than as isolated causal evidence for MoT.
• Versus FFN-only MoE: conventional experts primarily specialize content transformations; this work also conditions relational computation and shares selection between the two modules. The idea could transfer to multi-domain inpainting or personalized generation, but would need compute-matched validation in each new task.
• Versus text-guided editing: the main text places qualitative cross-paradigm comparisons with Qwen-Image-Edit-2509 and FLUX.2-klein-4B in supplementary material. The current cache provides no corresponding numerical results, so they should not be counted as quantitatively defeated baselines.

Rating

• Novelty: 4/5. Joint attention-FFN LoRA experts, anchor-guided routing, and the mask curriculum form a task-specific design, although their foundational components are not entirely new.
• Experimental Thoroughness: 4/5. Multiple test sets, routing and curriculum ablations, and a user study are provided, but strict unseen-category testing, matched training budgets, and uncertainty analysis remain missing.
• Writing Quality: 3/5. The problem and workflow are clear, but routing activation details, the loss, and parts of the evaluation protocol need elaboration; extraction damage in cached equations should not be treated as a typesetting fault in the original paper.
• Value: 4/5. Reducing the mask-precision burden has practical relevance, and multi-category paired data is useful, although deployment costs and generalization boundaries require further validation.