Multi-Concept Customization of Text-to-Image Diffusion

In plain words

Text-to-image models can draw "a dog on a beach", but not your dog. To teach them a specific dog, chair or style from a handful of photos you normally fine-tune the model, which is slow, produces a multi-gigabyte copy per concept, and tends to make the model forget: after training on your dog, every "dog" starts looking like yours.

Custom Diffusion notices that the text prompt enters the image network only through a small set of weights, the ones that turn words into keys and values inside cross-attention. It trains only those, plus a new made-up word (written V*) that stands for your concept. While training it also shows the model ordinary real photos with similar captions, so the general word ("dog") keeps its old meaning.

Because each concept changes only those few matrices, two separately trained concepts can be merged after the fact with a small least-squares solve that takes seconds: make the merged weights reproduce each concept's outputs for its own words, while staying as close as possible to the original model on everyday captions. This lets you ask for "V1* dog wearing sunglasses in front of a V2* moongate". Training takes minutes and each concept costs about 75MB instead of a full model copy.

TL;DR

Fine-tune only W^k and W^v of every cross-attention layer of Stable Diffusion (about 5% of U-Net parameters are cross-attention, and only K/V of those are trained) plus the embedding of a new modifier token V*, for 250 steps at batch 8 and effective lr 8\times10^{-5}, mixing target images 1:1 with ~200 retrieved real LAION images whose captions resemble the target caption. Random rescaling with matching prompt words ("very small", "zoomed in") and masked loss help image fidelity. For multiple concepts either train jointly with one V*ᵢ per concept, or merge single-concept K/V matrices with a closed-form constrained least squares (~2 s). Author-reported: on par with DreamBooth in CLIP text/image alignment, lower KID, ~6 min on 2 A100 vs ~1 h on 4 A100, 75MB vs 3GB per concept.

Problem

Users want a pretrained text-to-image model to render their own concepts (pets, objects, rare categories like "moongate") from as few as four images, and to compose several such concepts in one image §1 §3.1. Fine-tuning all weights on so few images is expensive and overfits; it also causes language drift, where related words lose their previous meaning (training on "moongate" makes "moon" produce moongates; training on "V* tortoise plushy" makes every "plushy" the target) §3.1 fig. 5. Methods that fine-tune all weights, such as DreamBooth, also compose two new concepts poorly §3.2 fig. 7.

Method

Backbone: Stable Diffusion, a latent diffusion model whose U-Net receives text features through cross-attention §3.1. Diffusion training minimizes

\mathbb{E}_{\epsilon,x,c,t}\big[w_t\,\lVert \epsilon - \epsilon_\theta(x_t, c, t)\rVert\big] eq. 2

with x_t = \sqrt{\alpha_t}x_0 + \sqrt{1-\alpha_t}\epsilon §3.1.

Which weights. Following Li et al., the authors fine-tune all U-Net weights on "moongate" images and measure, per layer, the relative change \Delta_l = \lVert\theta'_l - \theta_l\rVert / \lVert\theta_l\rVert for cross-attention, self-attention and all other layers. Cross-attention changes much more than the rest while holding about 5% of the parameters §3.1 fig. 3. In single-head cross-attention with text features c \in \mathbb{R}^{s\times d} and latent image features f \in \mathbb{R}^{(h\times w)\times l}, Q = W^q f, K = W^k c, V = W^v c and

\text{Attention}(Q,K,V) = \text{Softmax}\Big(\frac{QK^T}{\sqrt{d'}}\Big)V eq. 3

Since the text enters only through W^k and W^v, only those two matrices of every cross-attention layer are updated §3.1 fig. 4 model.py.

Text encoding. If the concept has a name ("moongate") it is used as the caption. For a personal instance of a category a new modifier token is prepended ("photo of a V* dog"); V* is initialized from a rarely occurring token embedding and optimized together with the K/V matrices §3.1. The rare token is token-id 42170 (second concept: 47629) in Stable Diffusion's CLIP tokenizer app. C. It was chosen by counting all 49408 tokens in 200K LAION-400M captions and taking an alphabetic token with roughly 5 to 10 occurrences that is not a substring of another token app. E.

Regularization set. 200 real images from LAION-400M whose captions have CLIP text-feature similarity above 0.85 to the target prompt are used, with their captions, as regularization during fine-tuning §3.1. Target images are oversampled so that target and regularization samples appear in equal proportion app. E.

Augmentation. Target images are randomly resized and the prompt gets a matching phrase; the loss is computed only on the valid image region [§3.2, "Training details"]. Exact recipe app. E: with probability 1/3 resize to 1.2 to 1.4× and add "zoomed in" or "close up"; otherwise resize to 0.4 to 1.0×, and if the ratio is below 0.6 add "far away" or "very small" and propagate the loss only inside the image region. The start-token embedding is detached during fine-tuning app. E.

# Custom Diffusion single-concept fine-tuning  [§3.1] [§3.2] [app. E]
# Input: target_images (4 to ~135), category word, pretrained SD,
#        reg_set: ~200 (image, caption) pairs retrieved from LAION-400M
# Output: fine-tuned K/V of all cross-attention layers + V* embedding

def custom_diffusion_finetune(target_images, category, sd, reg_set):
    # new token V*, initialized from rare token id 42170  [app. C]
    v_star = add_token(sd.tokenizer, init_from_id=42170)
    target_caption = "photo of a " + v_star + " " + category
    # trainable: W_k, W_v of every cross-attention layer, plus V*  [§3.1]
    params = []
    for layer in sd.unet.cross_attention_layers:
        params += [layer.W_k, layer.W_v]
    params += [embedding_row(v_star)]
    # lr 1e-5 scaled by total batch 8 -> 8e-5  [app. E]
    optimizer = make_optimizer(params, lr=8e-5)  # AdamW in code
    for step in range(250):
        # equal share of target and regularization samples  [app. E]
        batch = sample_balanced(target_images, reg_set, batch_size=8)
        images, captions, masks = [], [], []
        for image, caption, is_target in batch:
            mask = ones_like_latent()
            if is_target:
                # random rescale plus size words in the prompt  [app. E]
                image, caption, mask = resize_augment(image, target_caption)
            images.append(image)
            captions.append(caption)
            masks.append(mask)
        latents = sd.vae.encode(images)
        t = uniform_timesteps(len(batch), 1000)
        noise = gaussian_like(latents)
        noisy = sqrt(alpha[t]) * latents + sqrt(1 - alpha[t]) * noise
        # start-token K/V detached from the graph  [app. E]
        pred = sd.unet(noisy, t, sd.text_encoder(captions))
        # loss only on valid (non-padded) latent pixels  [§3.2]
        loss = masked_mean(squared_error(pred, noise), masks)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    return params


def resize_augment(image, caption):
    # 1 in 3: zoom in 1.2-1.4x and crop back to 512  [app. E]
    if random() < 1 / 3:
        ratio = uniform(1.2, 1.4)
        image = center_region_crop(resize(image, ratio), 512)
        caption = choice(["zoomed in ", "close up "]) + caption
        return image, caption, ones_like_latent()
    # otherwise shrink 0.4-1.0x and paste on a blank canvas  [app. E]
    ratio = uniform(0.4, 1.0)
    image, region = paste_at_random_position(resize(image, ratio), 512)
    if ratio < 0.6:
        caption = choice(["far away ", "very small "]) + caption
    # loss only inside the pasted region  [app. E]
    return image, caption, latent_mask(region)

Multiple concepts, joint training. Merge the training sets of all concepts, give each concept its own modifier token V*ᵢ initialized from a different rare token, and train K/V plus all V*ᵢ together §3.2; 500 steps for two concepts §3.2, "Training details".

Multiple concepts, closed-form merge. Given pretrained matrices W_0 and per-concept fine-tuned matrices W_n (n = 1..N) for each layer and each of K and V, solve §3.2

\hat{W} = \arg\min_W \lVert W C_{\text{reg}}^\top - W_0 C_{\text{reg}}^\top\rVert_F \quad\text{s.t.}\quad W C^\top = V, C = [c_1 \cdots c_N]^\top,\qquad V = [W_1 c_1^\top \cdots W_N c_N^\top]^\top eq. 4

C \in \mathbb{R}^{s\times d} stacks the text features of the s target words of all concepts (all captions per concept flattened and concatenated); C_{\text{reg}} \in \mathbb{R}^{s_{\text{reg}}\times d} holds text features of about 1000 randomly sampled regularization captions §3.2. With W_0 \in \mathbb{R}^{o\times d} app. B, Lagrange multipliers give, assuming C_{\text{reg}} is non-degenerate and a solution exists,

\hat{W} = W_0 + v^\top d,\quad d = C(C_{\text{reg}}^\top C_{\text{reg}})^{-1}, \quad v^\top = (V - W_0 C^\top)(d\,C^\top)^{-1} eq. 5

Derivation app. B eq. 7 eq. 8 eq. 9: the Lagrangian is L = \tfrac12\lVert W C_{\text{reg}}^\top - W_0 C_{\text{reg}}^\top\rVert - \operatorname{trace}(v(WC^\top - V)) with v \in \mathbb{R}^{s\times o}; setting the gradient to zero yields W = W_0 + v^\top C(C_{\text{reg}}^\top C_{\text{reg}})^{-1}, and substituting into WC^\top = V gives v^\top. The merge takes about 2 s when the single-concept models exist §3.2.

# Closed-form multi-concept merge, per layer and per K / V  [eq. 4] [eq. 5]
# W0: (o, d) pretrained; W_list[n]: (o, d) fine-tuned for concept n
# target_feats[n]: (s_n, d) text features of concept n's target words
# reg_feats: (s_reg, d) text features of regularization captions

def merge_projection(W0, W_list, target_feats, reg_feats):
    # stack all target words of all concepts  [eq. 4]
    C = concat_rows(target_feats)                        # (s, d)
    # values each fine-tuned model assigns to its own words  [eq. 4]
    V = concat_cols(
        [W_n @ c_n.T for W_n, c_n in zip(W_list, target_feats)]
    )                                                    # (o, s)
    # d = C (C_reg^T C_reg)^{-1}  [eq. 5]
    gram = reg_feats.T @ reg_feats                       # (d, d)
    D = solve(gram, C.T).T                               # (s, d)
    # v^T = (V - W0 C^T)(D C^T)^{-1}  [eq. 5]
    v_t = (V - W0 @ C.T) @ inverse(D @ C.T)              # (o, s)
    # W_hat = W0 + v^T d  [eq. 5]
    return W0 + v_t @ D

Compression (optional). The difference W - W_0 of fine-tuned K/V matrices has steeply decaying singular values, so it can be stored as a truncated SVD. "Top k%" keeps singular values up to the rank where their cumulative sum reaches k% of the total; top 60% gives 5× smaller storage with similar metrics, top 1% lowers image alignment app. C fig. 14 fig. 15. Enforcing a low-rank update during fine-tuning gave sub-optimal results app. C.

Hyperparameters

Results

Author-reported. Ten target datasets (2 scenes, 2 pets, 6 objects; 4 to 135 training images), Stable Diffusion backbone. Text-alignment is CLIP text-image similarity (V* removed from the prompt), image-alignment is CLIP image similarity to the target images, KID (\times10^3) is against 500 retrieved real images of the related category. Single-concept: 20 prompts × 50 samples per dataset; multi-concept: 8 prompts, 400 images per pair over five pairs §4 §4.1 §4.2 app. D.

Main comparison tab. 1

The pretrained model's KID is 22.27, identical to Textual Inversion, which does not change the model tab. 7.

Human preference (Amazon Mechanical Turk, 800 responses per questionnaire; % preferring Ours over the baseline) §4.3 tab. 2

Ablations (single-concept, same setup) §4.4 tab. 3

Choice of V* (means over nine datasets) app. C tab. 5

Random init shifts the plain category word toward the target images more than rare-token init fig. 18.

MS-COCO FID (50 DDPM steps, scale 6; pretrained model 16.35) tab. 6

CustomConcept101 (101 concepts; Ours and DreamBooth both use generated regularization images here; DDPM 200 steps) app. A tab. 4

(Bold as printed in tab. 4; note 0.763 > 0.757 in the multi-concept text-alignment column.)

Compute §4.1: Ours ~6 min on 2 A100; Ours (w/ fine-tune all) 20 min on 4 A100; Textual Inversion 20 min on 2 A100; DreamBooth ~1 h on 4 A100; batch 8 for all. Storage 75MB per concept vs 3GB for DreamBooth.

Per-dataset numbers are in tab. 7 (single-concept) and tab. 8 (multi-concept, per target). Ours beats baselines on the average metric for every pair except Table+Chair, where most methods are comparable §4.2 tab. 8.

Paper vs code

Checked against adobe-research/custom-diffusion at 7c19c9a, CompVis-based path (train.py, src/model.py, src/finetune_data.py, src/composenW.py).

• Steps. Paper: 250 single / 500 joint §3.2. Config files: 300 finetune_addtoken.yaml and 550 finetune_joint.yaml. The diffusers recipe in the README uses 250 / 500 README.md README.md.
• Retrieval. Paper: captions with CLIP text similarity > 0.85 §3.1. Code queries the LAION-400M kNN service with the target text (aesthetic_weight=0.1), requests 2× the needed count and keeps the first downloadable results; no similarity threshold is applied retrieve.py retrieve.py.
• Shrink range. Paper: 0.4 to 1.0× app. E. Code samples the new size uniformly from size // 3 to size, about 0.33 to 1.0× finetune_data.py.
• Prompt phrases are prepended, not appended as §3.2 says: code builds "a far away photo of a V* dog", "very small …", "zoomed in …", "close up …" finetune_data.py.
• Mask. The valid-region mask is on the 64×64 latent grid and is shrunk by one latent pixel on each side of the pasted image finetune_data.py; the masked loss is a per-sample mean over masked latents model.py.
• "Detach the start token" app. E is implemented in every cross-attention layer by detaching K and V at sequence position 0 model.py. In the text encoder, every token embedding except the modifier tokens is detached custom_modules.py.
• Merge regularization set. Paper: ~1000 random captions §3.2. Code takes the first 200 lines of a regularization-caption file and uses the features of all tokens after the start token up to and including the end token (skipping the first 4 if the caption begins with "photo of a") composenW.py composenW.py.
• Merge target set. Paper: all training captions per concept §3.2. Code uses just two strings per concept, "V*ₙ {category}" and "a photo of V*ₙ {category}" (style: "painting in the style of …"), and drops duplicate target-feature rows before solving composenW.py. The solve uses an LU factorization of C_{\text{reg}}^\top C_{\text{reg}} and an explicit inverse of dC^\top composenW.py.
• Joint training batches. Code draws one batch per concept dataset (each mixing target and regularization images) and sums the two losses model.py train.py; the paper only says the datasets are combined §3.2.
• Compression keeps singular values up to, but excluding, the one at which the cumulative share first exceeds the ratio compress.py.

Pitfalls

• Batch 8 means 2 GPUs × 4. The learning rate 8\times10^{-5} comes from scaling the base 10^{-5} by the total batch app. E train.py. Running on one GPU with the same config quietly halves the effective batch and the learning rate; the README suggests doubling steps on 1 GPU for the diffusers trainer README.md.
• Eq. 2 and eq. 7 print unsquared norms, but the derivation in app. B (eq. 8) is the gradient of the squared Frobenius norm. Implement the merge as squared least squares.
• Transpose in eq. 4. With W \in \mathbb{R}^{o\times d} and C \in \mathbb{R}^{s\times d}, WC^\top is o\times s, so the constraint needs V as o\times s (columns W_n c_n^\top). The printed V = [\dots]^\top suggests s\times o. Code uses the o\times s form composenW.py.
• Merge conditioning. C_{\text{reg}}^\top C_{\text{reg}} (d\times d, d = 768) must be invertible, so you need well over 768 regularization token features; dC^\top (s\times s) must also be invertible, so target rows must be linearly independent. Duplicate rows (the same word in several captions) make it singular, which is why the code removes them composenW.py. The paper only states the non-degeneracy assumption §3.2.
• Merging needs V* embeddings too. Each concept also has its own trained token embedding; the merged model must load all of them, and target features must be computed with them in place composenW.py.
• The merge is per layer and per K/V. Apply eq. 5 separately to every cross-attention W^k and W^v §3.2.
• Generated vs real regularization. Retrieved real images matter for KID: generated ones give 34.70 vs 20.96 tab. 3 and saturation artifacts for the plain category word, which the authors attribute to their higher learning rate than DreamBooth app. C. Yet the later CustomConcept101 results use generated regularization for both Ours and DreamBooth app. A; do not compare those numbers with tab. 1.
• Retrieval depends on an external service (LAION kNN API) in the code retrieve.py; the 0.85 threshold in the paper is not reproduced by it.
• Text/image-alignment trade-off. Text-alignment falls and image-alignment rises with training steps app. D fig. 20; step count is the main knob. Even with regularization, "photo of a {category}" drifts toward the target images fig. 24.
• Composition failures. Similar categories (teddybear + tortoise plushy) and cat + dog together fail; the pretrained model already struggles and their attention maps overlap §4.4 fig. 11 app. C fig. 19. Three concepts work only for some combinations fig. 16.
• Joint beats optimization on most multi-concept metrics tab. 1 fig. 7; use the closed-form merge when per-concept models already exist and 2 s matters more than quality.
• Sequential training forgets the first concept §4.2.
• Licensing. The CompVis-based code is under the non-commercial Adobe Research License (the diffusers training scripts carry MIT modifications); reimplement from the paper for commercial use.

Lineage

• Builds on: Stable Diffusion / latent diffusion (arXiv:2112.10752) as the backbone; DDPM (arXiv:2006.11239); the weight-change analysis follows Li et al., "Few-shot image generation with elastic weight consolidation" (NeurIPS 2020) §3.1.
• Compared against: DreamBooth (arXiv:2208.12242), which fine-tunes the whole U-Net with generated prior images and a fixed rare token; Textual Inversion (arXiv:2208.01618), which only learns a new token embedding §4.
• Applications: image editing with Prompt-to-Prompt (arXiv:2208.01626) §4.4 fig. 10; style learning with "A painting in the style of V* art" §4.4 fig. 9.