A$^2$-Edit: Precise Reference-Guided Image Editing of Arbitrary Objects and Ambiguous Masks

1. Problem Statement

Current reference-guided image inpainting methods face three critical bottlenecks that limit their practical application and generalization:

• Domain Specificity & Homogenization: Existing models are typically optimized for specific object categories (e.g., only garments or only portraits). They fail to handle the diverse modeling objectives required for different objects (e.g., texture consistency for clothes vs. identity preservation for faces vs. geometric rigidity for furniture). Furthermore, existing datasets suffer from limited category coverage and homogeneous distributions, hindering the learning of universal representations.
• Rigid Spatial Guidance: Most state-of-the-art methods rely on high-precision, pixel-level segmentation masks. In real-world scenarios (e.g., user sketches or automated detection boxes), masks are often coarse, ambiguous, or imprecise. When mask quality degrades, model performance drops significantly, reducing interaction flexibility.
• Lack of Unified Framework: There is no unified architecture capable of simultaneously handling arbitrary object categories and varying levels of mask precision while maintaining high fidelity and identity consistency.

2. Methodology

The authors propose A2-Edit, a unified framework designed for Arbitrary object categories and Arbitrary mask precision. The methodology consists of three core components:

A. Mixture of Transformers (MoT) Architecture

To address the challenge of diverse object categories, A2-Edit moves beyond the standard "shared parameter" approach.

• Dynamic Expert Routing: Instead of using a single pathway for all inputs, the model employs a Mixture-of-Experts (MoE) paradigm extended to the Transformer architecture.
• Dual-Module Specialization: Unlike traditional MoE which only applies to Feed-Forward Networks (FFN), A2-Edit introduces Mixture-of-Transformers Blocks (MoTB) where expert routing is applied to both the Attention and FFN modules. This allows the model to learn category-specific relational modeling (e.g., shape alignment, texture consistency) and content transformations simultaneously.
• Implementation: Each expert is a lightweight LoRA (Low-Rank Adaptation) adapter. A Backbone Expert is always active to provide universal knowledge, while Assistant Experts are dynamically selected based on the semantic characteristics of the input object via an Anchor-Guided Routing (AGR) strategy.

B. Mask Annealing Training Strategy (MATS)

To reduce reliance on perfect masks and improve robustness to coarse inputs, the authors propose a progressive training curriculum:

1. Stage 1 (Fine Mask): Training begins with high-precision segmentation masks to teach the model strict spatial alignment.
2. Stage 2 (Augmented Rough Mask): The model is trained with masks that have been morphologically dilated and perturbed with Perlin noise to simulate user drawing errors and edge extensions.
3. Stage 3 (Bounding Box): The final stage uses only bounding boxes as masks. This forces the model to infer object pose, scale, and local structure purely from context and reference semantics, rather than relying on shape priors.

C. UniEdit-500K Dataset

To support the training of such a diverse model, the authors constructed a large-scale dataset:

• Scale: 500,104 image pairs.
• Diversity: Covers 8 major categories (Garments, Portraits, Animals, Plants, Accessories, Furniture, Vehicles, Architecture) and 209 fine-grained subcategories.
• Composition: Includes both rigid and non-rigid objects, addressing the data homogeneity issues of previous datasets like VITON-HD or AnyDoor.

3. Key Contributions

1. Unified Framework (A2-Edit): A single model capable of editing arbitrary object categories with high fidelity, breaking the constraints of domain-specific editing.
2. Novel Architecture (MoT): A Mixture-of-Transformers design that dynamically routes features to specialized experts in both Attention and FFN layers, enabling superior cross-category semantic transfer and generalization.
3. Robust Training (MATS): A progressive training strategy that shifts the model from relying on geometric boundaries to developing strong contextual understanding, making it robust to ambiguous and coarse masks.
4. Comprehensive Dataset (UniEdit-500K): The most extensive dataset for this task to date, providing the necessary data diversity to train a truly universal editing model.

4. Experimental Results

The authors evaluated A2-Edit on standard benchmarks (VITON-HD, AnyInsertion) and their own UniEdit test set.

• Quantitative Performance: A2-Edit achieved State-of-the-Art (SOTA) results across all metrics (DINO-I, CLIP-I, LPIPS, FID) on both fine and rough mask conditions.
  • On the UniEdit test set (diverse categories), A2-Edit significantly outperformed baselines like FLUX.1-Fill, ACE++, and Insert Anything, particularly in the "Rough Mask" setting where other models degraded.
  • It achieved a VLM-based score of 82.3 (vs. 78.7 for the next best), indicating superior boundary quality, realism, and local consistency.
• Qualitative Performance: Visual comparisons showed A2-Edit produces sharper details, better structural integrity, and more natural blending than text-guided or mask-guided baselines. It successfully handled challenging tasks like face swapping, garment transfer, and rigid object insertion without identity distortion.
• Ablation Studies: Removing the MoT architecture, the UniEdit-500K dataset, or the MATS strategy resulted in significant performance drops, confirming the necessity of each component for cross-category generalization and mask robustness.
• User Study: In a study with 24 participants, A2-Edit was preferred over all baseline methods in the majority of comparisons.

5. Significance

A2-Edit represents a significant leap forward in generative image editing by solving the "generalization vs. precision" trade-off.

• Practical Usability: By tolerating ambiguous masks (e.g., rough sketches or bounding boxes), it lowers the barrier for user interaction, making it suitable for mobile applications and real-world workflows where pixel-perfect masks are unavailable.
• Universal Applicability: The ability to edit anything from a human face to a piece of furniture within a single unified model simplifies the deployment of image editing tools in e-commerce, virtual try-on, and creative design.
• Research Foundation: The release of the UniEdit-500K dataset and the MoT architecture provides a new foundation for future research in universal image generation and editing, moving the field away from siloed, category-specific models.