Multi-Reference Self-Attention (MRSA)

• Multi-Reference Self-Attention is a mechanism that extends conventional self-attention by integrating external reference key/value pairs and masking.
• It employs a mask-logit modulated fusion strategy to selectively incorporate multi-concept cues, enhancing identity retention and text alignment.
• Empirical evaluations demonstrate that MRSA improves image fidelity and multi-concept composition with manageable computational overhead.

Multi-Reference Self-Attention (MRSA) is a mechanism developed to address the challenge of customized image generation involving multiple user-specified concepts within a single composition, without requiring retraining or fine-tuning of the base generative model. MRSA was introduced in the FreeCustom method, enabling generation based on several “one-shot” reference images—each representing a distinct concept—and a text prompt. This mechanism extends conventional self-attention in transformer-based architectures, allowing direct routing of information from both the composition feature map and the external reference latents, with strategic spatial and conceptual masking to maximize fidelity and reduce unwanted hallucination of attributes.

1. Mechanism and Extension of Standard Self-Attention

Standard self-attention (SA) in architectures such as U-Net diffusion models processes queries, keys, and values originating from the same latent feature map, limiting information exchange to intra-image context and precluding direct access to externally referenced visual concepts. MRSA augments this process by injecting key/value pairs extracted from each reference image into the SA block within the composition path. At each spatial position within the composition, the network can selectively attend either to internal features or to any location within the reference images.

This extended attention is modulated with pixel-wise binary masks $M_i$—obtained from segmentation of each reference image—and scalar weights $\omega_i$, enforcing attention only on desired regions and suppressing unrelated features. This structure directly addresses the multi-concept composition challenge, allowing the model to faithfully preserve multiple concept identities across arbitrary spatial arrangements within the synthesized output (Ding et al., 22 May 2024).

2. Mathematical Specification

At each relevant U-Net layer $\ell$ and diffusion timestep $t$:

• Let $x^{(t)} \in \mathbb{R}^{H \times W \times C}$ denote the composition-path feature map, and $x'^{i^{(t)}}$ the reference-path features for concepts $i=1 \ldots N$.
• After spatial flattening to $n = H \cdot W$, and projection via $W_q$, $W_k$, $W_v \in \mathbb{R}^{C \times d}$:

$$Q = \text{reshape}(x^{(t)} W_q) \in \mathbb{R}^{n \times d}$$

$$K = \text{reshape}(x^{(t)} W_k), \quad V = \text{reshape}(x^{(t)} W_v) \in \mathbb{R}^{n \times d}$$

For each reference image $i$: $$Q_i = \text{reshape}(x'^{i^{(t)}} W_q), \quad K_i = \text{reshape}(x'^{i^{(t)}} W_k), \quad V_i = \text{reshape}(x'^{i^{(t)}} W_v)$$ All keys and values are concatenated: $$K' = [K; K_1; \ldots; K_N] \in \mathbb{R}^{n \times d(N+1)}$$

$$V' = [V; V_1; \ldots; V_N] \in \mathbb{R}^{n \times d(N+1)}$$

A composite mask-logit matrix $M_w \in \mathbb{R}^{n \times (N+1)n}$ is constructed:

• $m_0$ is an all-ones vector for the composition path (weight 1),
• $m_i = \omega_i \cdot \text{flatten}(\text{mask}_i)$ for each reference $i$, where $\omega_i \approx 2-3$ is optimal.

Attention is computed as: $$A_{p,q} = \frac{\exp(M_{w,pq}\langle Q_p, K'_q \rangle / \sqrt{d})}{\sum_{r=1}^{(N+1)n} \exp(M_{w,pr}\langle Q_p, K'_r \rangle / \sqrt{d})}$$

$\text{MRSA}(Q, K', V', M_w) = A V'$

The output is reshaped to $H \times W \times d$ and propagated through the U-Net block.

3. Integration into Diffusion U-Net Architecture

MRSA is deployed in FreeCustom via parallel processing paths using the same base U-Net $\epsilon_\theta$ at each denoising step $t$:

• Reference path: For each input concept image $I_i$, encode and noise to $z'^i_t$, process via $\epsilon_\theta$ to extract $(K_i, V_i)$ at targeted blocks $\ell \in \Psi$ (where $\Psi$ typically is $\{5, 6\}$).
• Composition path: The generative latent $z_t$ passes through a forked U-Net $\epsilon_\theta^*$, wherein the standard SA module is replaced by MRSA in blocks $\ell \in \Psi$.

At each denoising step, reference features are updated; MRSA blocks utilize all key/value sources and mask weights, resulting in selective multi-concept feature fusion.

4. Computational Overhead

The computational cost of MRSA is characterized by:

• Standard SA on $n$ tokens: $O(n^2 d)$.
• MRSA key/value length increases to $(N+1)n$, yielding $O(n(N+1)nd) = (N+1)O(n^2 d)$ attention cost per block.
• Two U-Net “passes” per timestep (reference and composition), resulting in approximately $2\times$ the baseline cost, plus $N\times$ overhead for each replaced block.

In practical deployments (with $N=2$–$4$ concepts), this overhead is manageable: inference for $N=3$ reference images requires $\sim$58 seconds on a 3090 GPU, without any fine-tuning or preprocessing.

5. Empirical Evaluation and Ablation Analysis

Quantitative and qualitative results indicate MRSA’s efficacy for both multi-concept composition and single-concept customization:

• Multi-concept fidelity: DINOv2 similarity for FreeCustom reaches 0.7625 (vs. 0.6545 for CustomDiffusion, 0.6399 for Perfusion).
• Image quality: CLIP-IQA scores are 0.9002 (vs. 0.8921/0.8624).
• Text alignment: CLIP-T/L scores are 33.78/27.88 (vs. 29.07/23.66 and 22.14/16.17).
• User studies: Highest scores for identity consistency and alignment with prompts.

Ablation reveals that:

• Weighted mask values $\omega_i = 3$ are optimal for detail preservation; uniform weighting leads to blurring or attribute hallucination.
• Replacing only blocks $\Psi=[5,6]$ yields best trade-off between fidelity and naturalness; replacing all SA blocks degrades realism.
• Contextualizing reference images (i.e., showing interactions like a hat worn on a head) is crucial for concept retention; isolated objects underperform.

6. Practical Implementation and Controls

• Number of reference images: Only one image per concept is sufficient for customization.
• Mask extraction: Masks $M_i$ can be generated using any off-the-shelf segmentation model (e.g., GroundedSAM), with masks downsampled to match spatial dimensions.
• Mask weights: Recommended settings for $\omega_i$ are 2–3; “self” mask weight is fixed at 1.
• Layer selection: MRSA should replace self-attention only in the deepest two blocks (layers 5 and 6) for optimal results.
• Reference image composition: Incorporating context interaction in reference images, or alternatively, copy-pasting objects into plausible scenes, improves retention and synthesis quality.
• Plug-and-play: MRSA can be implemented in any transformer-based self-attention block by forking parallel U-Net passes and applying the mask-weighted MRSA routine.

7. Significance and Broader Implications

MRSA enables plug-and-play, tuning-free multi-concept composition in diffusion models, circumventing the need for iterative retraining associated with previous approaches. Its explicit mechanism for conditioned attention—leveraging targeted reference cues through masking and weighting—demonstrates substantial gains in compositionality, image quality, and controllability. A plausible implication is that such mechanisms may further generalize to other domains requiring cross-instance conditioning in generative architectures, given the empirical success of the mask-logit modulated fusion strategy (Ding et al., 22 May 2024).