A²-Edit: Overcoming Object Category and Mask Precision Limits for Precise Reference-Guided Image Editing

Moving an object from a reference image into another photo is essential for e‑commerce product compositing, virtual try‑on, film face‑swap, and personalized creation, but practical deployment faces two major hurdles: models typically excel only on a single object class and heavily rely on pixel‑level precise masks.

Solution – Researchers from Shanghai Jiao Tong University and Shanghai Chuangzhi Academy propose A²-Edit , a unified framework that supports arbitrary object categories and any mask precision. It combines mixed‑Transformer expert routing, mask‑annealing training (MATS), and a newly built 500K‑image multi‑category dataset (UniEdit‑500K) so that users only need a rough region to obtain identity‑preserving, structurally complete, and naturally blended edits.

Two inherent constraints of reference‑guided editing are highlighted:

Insufficient cross‑category unified modeling – different categories (clothing, portraits, vehicles, furniture, architecture) demand distinct semantic and geometric representations, yet existing methods use a single shared parameter path, causing performance drops on diverse inputs.

Heavy dependence on fine segmentation masks – real‑world masks drawn by users or generated by detectors are often coarse, leading to boundary artifacts, content leakage, structural anomalies, and background damage.

Architecture innovation : A²-Edit extends traditional mixture‑of‑experts routing from the feed‑forward network to both the attention and feed‑forward modules, embedding mixed‑Transformer blocks (MoTB) in every layer. A gating network evaluates semantic and category cues to compute expert routing weights, dynamically selecting parameter branches that match the editing target. Routing follows an Anchor‑Guided Routing (AGR) strategy: a backbone expert remains constantly active as a knowledge anchor, while auxiliary experts are chosen per‑input; only one routing pass occurs per layer, and attention and FFN share the same weights. During inference, only the backbone and a few auxiliary experts are activated, and out‑of‑distribution inputs are handled by combining experts based on feature similarity, improving generalization.

Mask‑Annealing Training (MATS) progressively relaxes mask precision across three stages: (1) learn basic editing with fine masks; (2) introduce morphological dilation and Perlin noise to simulate realistic hand‑drawn errors; (3) train directly with bounding‑box masks, forcing the model to infer pose, scale, and structure from context. This gradual training enables the model to produce natural, stable results even with coarse masks or detection boxes.

Dataset : UniEdit‑500K contains over 500,000 reference‑target image pairs covering eight major categories (clothing, portrait, animal, plant, accessories, furniture, vehicle, architecture) and 209 fine‑grained sub‑categories, providing the broad data diversity needed for the unified framework.

Experimental results : On the VITON‑HD, AnyInsertion benchmarks and the self‑built UniEdit test set, A²-Edit outperforms multiple state‑of‑the‑art image‑editing methods on both fine and coarse mask conditions. Quantitatively, it achieves the best scores on DINO‑I, CLIP‑I, LPIPS, FID, and VLM metrics; mask robustness is demonstrated by minimal degradation when switching from fine to coarse masks, unlike baselines that suffer large drops. A user study with 24 participants shows higher preference rates for A²-Edit in pairwise comparisons against all baselines. Qualitatively, the model consistently delivers semantically consistent, structurally intact results across categories, preserving clothing texture, portrait identity, and architectural coherence even with hand‑drawn coarse masks.

Ablation studies confirm the necessity of each component: removing the A²-Edit framework or limiting expert routing to the FFN reduces both quality and metrics; discarding UniEdit‑500K harms detail generation and task intent understanding; training only with fine masks causes over‑fitting to boundaries, failing under coarse masks. Incrementally adding enhanced coarse‑mask and bounding‑box training steadily mitigates these issues.

Model limitations : The current peak memory consumption is about 42 GB, posing a deployment challenge for consumer‑grade GPUs. When the provided mask region is overly large and lacks explicit textual guidance, the model may produce multiple plausible interpretations. Future work may focus on expert compression, low‑bit deployment, richer interactive prompts, and extending long‑tail category data.

Conclusion : By leveraging mixed‑Transformer expert routing, mask‑annealing training, and the extensive UniEdit‑500K dataset, A²-Edit establishes a universal reference‑guided editing solution that tolerates arbitrary object categories and mask precisions, lowering the barrier of fine segmentation while maintaining identity and detail, and opening new possibilities for virtual try‑on, e‑commerce compositing, film production, advertising, and personalized content creation.