FastComposer: Tuning-Free Multi-Subject Image Generation with Localized Attention

In plain words

Teaching an image model to draw one particular person is usually done by re-training part of the model on that person's photos. That costs time and memory for every new person. FastComposer skips the re-training. An image encoder looks at a photo of each person once, and the result is passed to the model next to the text prompt. So you can write a prompt that mentions two people and give it a photo of each, and the model draws them in a new scene without any training run for those people.

Doing this naively has two problems. When two people are in one picture, their features blur into each other, so each face takes on traits of the other. And if the photo is used from the very first drawing step, the picture copies the photo and stops following the text. The paper handles the first problem by teaching the model, during training only, to attend to the right region for each person. It handles the second by using the photos only for the later drawing steps, after the text has fixed the overall layout.

The results below are the authors' own tests. They are not independent.

TL;DR

FastComposer generates images of several people from reference photos, with no per-person fine-tuning, on top of Stable Diffusion v1-5. Its CLIP-based subject embeddings are inserted into the text conditioning, and a segmentation-mask loss on cross-attention reduces identity blending between subjects. Delayed subject conditioning keeps prompt following. The authors report much lower time and memory than fine-tuning baselines, and higher identity scores, but only on their own test sets.

Problem

Subject-driven generation that fine-tunes the model per subject is costly in compute and memory, and the paper argues this limits deployment §1. Existing methods also blend the identities of several subjects in one image §1 fig. 1 fig. 2. The paper traces this to cross-attention maps, where each subject's token spreads over all subjects in the picture instead of one region §4.2 fig. 4. A third issue is subject overfitting: conditioning on the subject from the first step makes the image reproduce the reference and ignore the prompt §4.3 fig. 2.

The setting assumes a few reference images. The optimization-based baselines use five images per subject, while FastComposer uses one randomly chosen image per subject §5.2.

Method

Subject embedding and text augmentation §4.1 eq. 2. The prompt is encoded with the CLIP text encoder ψ. Each reference subject is encoded with a CLIP image encoder φ. For every prompt token i that is the subject token i_j, the text embedding is replaced by an MLP applied to the concatenation of that text embedding and the subject embedding. Other tokens keep their text embedding. The result is the conditioning c' used in place of the original c §4.1 eq. 2. The paper gives the MLP only as a black box; see Paper vs code.

Training data §4.1 §5.1. Training pairs are built from images with captions. Noun phrases in each caption are matched to panoptic segments of the image, using a text-image similarity score and a label similarity score §5.1. Before encoding, each subject crop has its background replaced by random noise, so the model does not learn the background. At test time, natural subject photos can be used without a segmentation step §4.1.

Training objective §4.2 eq. 1 eq. 3 eq. 4. The denoising loss L_noise is the usual noise-prediction loss §3 eq. 1. A localization loss L_loc is added, computed on the cross-attention map A_{i_j} of each subject token §4.2 eq. 3. It is a balanced L1 loss: the mean attention outside the subject's mask minus the mean attention inside it. It is averaged over subjects. The total loss is L = L_noise + λ·L_loc with λ = 0.001 §4.2 eq. 4. The localization loss is applied to the five middle U-Net blocks, which the paper says contain more semantic information §4.2. Masks and this loss are used only in training §4.2.

Delayed subject conditioning §4.3 eq. 5. At sampling time the noise prediction uses the plain text conditioning c while the timestep is above α·T, and the subject-augmented conditioning c' otherwise §4.3 eq. 5. Thus the first part of sampling forms the layout from text alone, and the last α fraction refines the subjects. The authors report that α between 0.6 and 0.8 balances identity and prompt consistency §4.3 fig. 5.

Training regime §5.1. The text encoder is frozen. The U-Net, the MLP and the last two blocks of the vision encoder are trained. Ten percent of samples are trained without any condition, which enables classifier-free guidance at sampling time. Another 10% are trained text-only §5.1.

Pseudocode for the three parts (the authors' equations, with the code's defaults noted where they differ; see Paper vs code):

# Subject-augmented text embeddings  [§4.1] [eq. 2]
# Input:  prompt tokens (n,), subject crops (m, 3, H, W), token index per
#         subject (m,)
# Output: conditioning embeddings (n, d)

def augment_text(prompt_tokens, subject_crops, subject_token_index):
    # Encode every prompt token with the CLIP text encoder  [§3] [eq. 2]
    text_embeds = text_encoder(prompt_tokens)       # (n, d)
    # Encode each subject crop with the CLIP image encoder  [§4.1]
    subject_embeds = image_encoder(subject_crops)   # (m, d)
    # Only subject tokens change; the other tokens keep their embedding  [eq. 2]
    for j, i in enumerate(subject_token_index):
        # Fuse the subject embedding into its own token  [§4.1] [eq. 2]
        text_embeds[i] = mlp(
            concatenate(text_embeds[i], subject_embeds[j])
        )
    return text_embeds

# Training step  [§3] [§4.1] [§4.2] [eq. 1] [eq. 3] [eq. 4]
# Input:  target image, prompt, subject crops (background noised),
#         subject masks (m, H, W), token index per subject (m,)
# Output: scalar loss

def training_step(image, prompt, subject_crops, subject_masks, subject_idx):
    # Encode the target image into a latent  [§3]
    latent = vae_encode(image)
    # Sample a timestep and add Gaussian noise to the latent  [§3] [eq. 1]
    timestep = sample_timestep()
    noise = sample_gaussian(latent.shape)
    noisy_latent = add_noise(latent, noise, timestep)
    # Build the subject-augmented conditioning  [§4.1] [eq. 2]
    conditioning = augment_text(prompt, subject_crops, subject_idx)
    # Predict the noise; the paper writes a squared L2 norm here  [eq. 1]
    noise_pred, attention_maps = unet(noisy_latent, timestep, conditioning)
    loss_noise = mean_squared_error(noise_pred, noise)
    # Localize each subject token's attention on its mask  [§4.2] [eq. 3]
    loss_loc = localization_loss(attention_maps, subject_masks, subject_idx)
    # Weighted sum with lambda = 0.001  [§4.2] [eq. 4]
    return loss_noise + loc_weight * loss_loc

# Localization loss for one subject token  [§4.2] [eq. 3]
# Input:  attention (hw,) for one subject token, mask (hw,) in {0, 1}
# Output: scalar

def localization_loss_one_subject(attention, mask):
    # Mean attention inside the subject's mask  [§4.2]
    inside = mean(attention[mask == 1])
    # Mean attention outside the subject's mask  [§4.2]
    outside = mean(attention[mask == 0])
    # Push attention toward the mask: outside minus inside  [eq. 3]
    return outside - inside

# Sampling with delayed subject conditioning  [§4.3] [eq. 5]
# Input:  prompt, subject crops, token index per subject, alpha, steps
# Output: image

def sample(prompt, subject_crops, subject_idx, alpha, steps):
    # Text-only conditioning, used while the layout forms  [§4.3]
    text_cond = text_encoder(prompt)
    # Subject-augmented conditioning, used for the later steps  [§4.3]
    subject_cond = augment_text(prompt, subject_crops, subject_idx)
    # Unconditional embedding, as used in training  [§5.1]
    null_cond = null_embeddings
    latent = sample_gaussian(latent_shape)
    for step, timestep in enumerate(scheduler_timesteps(steps)):
        # Text-only while step <= (1 - alpha) * steps  [§4.3] [eq. 5]
        if step <= (1 - alpha) * steps:
            cond = text_cond
        else:
            cond = subject_cond
        # Two U-Net passes for classifier-free guidance  [§5.1]
        noise_null = unet(latent, timestep, null_cond)
        noise_cond = unet(latent, timestep, cond)
        # Guidance scale is not stated in the paper  [§5.1]
        noise_pred = noise_null + guidance * (noise_cond - noise_null)
        latent = scheduler_step(noise_pred, timestep, latent)
    return vae_decode(latent)

Hyperparameters

Values marked "not stated" are absent from the paper. Code values are shown only when the code sets them.

Results

Author-reported. Single-subject setup: 15 subjects from CelebA, 30 prompts each, compared with DreamBooth, Textual Inversion and Custom Diffusion §5.2. The number of images generated per prompt is not stated in the paper. The repository's evaluation command uses 4 README.md.

Single-subject results tab. 1 §5.2:

Multi-subject results, 105 pairs of the 15 subjects, 21 prompts per pair tab. 2 §5.3:

Ablation on cross-attention localization tab. 3 §5.4. The paper does not name the evaluation set; the 43.11% value matches the FastComposer row of tab. 2.

Other reported results: delayed conditioning shows that a higher subject ratio raises identity preservation and lowers editability, with 0.6–0.8 the recommended band §5.4 fig. 5. Qualitative single-subject comparisons are in fig. 6; three-subject examples are in fig. 7; cross-attention maps with and without the localization loss are in fig. 4.

Speedup and memory ratios implied by the tables (computed here from tab. 1 and tab. 2, not stated as such by the paper): single-subject total time is about 394× to 1,250× faster than the baselines, and multi-subject about 394× to 2,499×. Peak memory is 2.8× to 6.7× lower tab. 1.

Paper vs code

Code: github.com/mit-han-lab/fastcomposer at b9304a2e7d09ffe4036ebd8ed5656bb08d6a8ca9, MIT, owned by the MIT Han Lab (one of the authors is Song Han). Training, inference and single-subject evaluation code are included. Dataset construction and multi-subject evaluation code are not.

• MLP. The paper gives one MLP on the concatenation §4.1 eq. 2. The code runs two MLP blocks, where the first is applied to the concatenation with a residual connection to the text embedding, the second has a residual connection too, and a final LayerNorm follows. Each MLP block is LayerNorm, linear, GELU, linear model.py model.py.
• Subject embedding. The code uses the CLIP vision pooled output followed by the visual projection, giving one token per subject [code: fastcomposer/model.py#L76-L92] run_training.sh. The paper says only φ(s_j) §4.1.
• Denoising loss. The paper writes a squared L2 norm eq. 1. The code uses mean squared error over all elements model.py. This rescales L_noise relative to λ·L_loc, so the λ = 0.001 value only has the paper's meaning if the code's reduction is used.
• Localization layers. The paper says "middle 5 blocks" §4.2. The code selects five of seven named attention layers: the second and third down blocks, the middle block, and the second and third up blocks, counting from zero model.py.
• Localization loss. The code averages the balanced L1 loss over layers, heads and subjects model.py. The threshold argument is stored but never used in the loss model.py. Per-head averaging is not described in the paper, which uses single-head attention for exposition §3.
• Attention processor. For the localized layers, the code replaces fused attention with the plain processor so that scores can be read [code: fastcomposer/model.py#L297-L316]. This affects speed and memory during training.
• Subject-region loss. The paper says it applies the loss to subject regions for half the samples §5.1. The code draws one random number per batch, and masks both the prediction and the target only for the denoising loss [code: fastcomposer/model.py#L550-L560].
• Text-only samples. The paper trains 10% text-only §5.1. The shipped script does not set --text_only_prob, which defaults to 0 [code: fastcomposer/utils.py#L419-L422] run_training.sh.
• Warmup. The paper says constant learning rate 1e-5 §5.1. The script passes lr_warmup_steps (default 500) to the scheduler [code: fastcomposer/train.py#L250-L255] utils.py. Whether this changes the constant schedule depends on the diffusers version. This was not verified here.
• Delayed conditioning. Eq. 5 switches on timestep (t > αT) eq. 5. The sampler switches on step index: text-only while step <= (1 - alpha) * steps pipeline.py. At α = 0.7 and 50 steps that is steps 0 to 15, so 16 text-only steps. The shipped inference script instead passes an absolute --start_merge_step 10 run_inference.sh, which is 11 text-only steps. The Cog predictor and the single-subject evaluation use α directly, with defaults 0.7 and 0.5 predict.py run.py.
• Resolution. The paper does not state training resolution. The shipped script uses 512 and 224 for subjects run_training.sh, while the argparse default for training resolution is 256 [code: fastcomposer/utils.py#L360-L363].
• Evaluation. Prompt consistency uses CLIP ViT-L/14 in the code, as in the paper single_object_evaluation.py. The identity model is an InceptionResnetV1 trained on VGGFace2; the paper says FaceNet single_object_evaluation.py. Per-pair similarity is clipped at zero [code: evaluation/single_object/single_object_evaluation.py#L75-L76].
• Weights and data. Weights are downloaded from Hugging Face [code: README.md#L27]. The training data is a Hugging Face dataset [code: README.md#L54-L59]. Dataset construction (BLIP-2, Mask2Former, spaCy and matching) is not in the repository.

Pitfalls

• Under-specified MLP. The paper's Eq. 2 is one MLP. The code uses residual blocks, LayerNorm and a final LayerNorm model.py. An implementation that follows Eq. 2 literally will not match the released checkpoint.
• Loss scale. Eq. 1 is written as a squared L2 norm, which sums over pixels. The code uses a mean. λ = 0.001 is calibrated to the code's loss scale. Using the paper's sum form changes the balance between the two losses by the number of latent elements.
• Localization needs masks at training time only. Each subject needs a segmentation mask and a token index. The masks come from Mask2Former in the paper §5.1; the repository does not include that step.
• Localized layers may cost more. Switching those layers to the non-fused attention path (see Paper vs code) may raise training time and memory. The paper does not mention this.
• Warmup is ambiguous. The paper says constant 1e-5. The code passes a 500-step warmup to the scheduler. Check the diffusers version before assuming either behavior.
• Sampling hyperparameters are not in the paper. Guidance scale, step count and sampler are absent. The code uses guidance 5 and 50 steps, and the delayed-conditioning switch depends on the step-index convention above.
• Speedup wording. The introduction gives 300×–2500× §1 and the experiments give 300×–1200× §5.2. Neither matches the table ratios exactly. The table ratios computed in Results are the ones to use, and they depend on the rounded times in tab. 1 and tab. 2.
• Baselines are not reproducible from the paper. §5.2 refers to hyperparameters "in the appendix", but the source contains no appendix text, only a figure. Baseline settings are therefore not stated.
• DreamBooth is missing from the multi-subject table. §5.3 cites DreamBooth among the optimization methods, but tab. 2 has no DreamBooth row.
• The ablation's evaluation set is not named. tab. 3 matches the tab. 2 FastComposer row, so the multi-subject setting is likely, but this is an inference.
• Human-centric scope. Training uses only "person" segments §5.1 on a small FFHQ-based set of mostly headshots. The paper limits itself to at most three subjects in practice and notes a long-tailed count of people §6. Animals and objects are not covered §6.
• Inherited bias. The paper states that the model inherits biases from Stable Diffusion and FFHQ §6.
• Licenses not covered by this card. The paper does not state the license of the Stable Diffusion weights, the FFHQ-based training data or the released checkpoint. The repository is MIT, but that covers only the code. Check these before commercial use.
• Deepfake risk. The broader-impact paragraph is commented out (\iffalse) in the source's discussion section, so it is not in the compiled main text. It names deepfakes and consent as concerns §6.

Lineage

• Builds on: Stable Diffusion v1-5 and its latent diffusion model §3 (arXiv:2112.10752). The CLIP encoders are used unchanged except for the last two vision blocks, which are fine-tuned §5.1.
• Compared against: DreamBooth (arXiv:2208.12242), Textual Inversion, and Custom Diffusion §5.2.
• Superseded by: not known from wiki content.