eDiff-I: Text-to-Image Diffusion Models with an Ensemble of Expert Denoisers

In plain words

Drawing an image from a text prompt with a diffusion model starts from pure static and removes the noise step by step. Early steps decide the rough layout and what the picture is about, and the text prompt matters most there. Late steps only add fine detail, and the text hardly matters any more.

The authors use one network for all of those steps. They argue that a single network has to do two very different jobs, so they train several specialist networks instead, each for a stretch of the noise levels. To keep training cheap, they first train one shared network, then copy it and fine-tune each copy on its own stretch. Running the specialists costs the same per step as running one network.

They also feed the model three kinds of text and image signals, so the picture can follow the wording in a different way. A reference image can pass its style along. A final trick lets a user paint a rough region on a canvas for a word in the prompt, to say where that object should appear. The paper reports that this beats earlier large text-to-image models on a standard benchmark. Those numbers are the authors' own.

TL;DR

Text-to-image cascade (64 px base, then 256 px and 1024 px super-resolution) whose base and first super-resolution stages are ensembles of expert denoisers, each trained on a range of noise levels and branched from a shared model in a binary tree §4.1 fig. 5. Conditioning uses T5-XXL, CLIP L/14 text and CLIP L/14 image embeddings with independent dropout §4.2. The paper's headline is a zero-shot COCO FID-30K of 6.95 at 256 px in its best configuration, against 7.23 for Parti and 7.27 for Imagen tab. 1. All results are author-reported. Training data is about one billion text-image pairs, partly proprietary, and no code or weights are released §5 §5.1.

Problem

Large text-to-image diffusion models are strong but expensive, and the usual way to scale them is to make one denoiser bigger, which slows down sampling because the denoiser runs at every step §1. The authors also observe that the denoiser's use of the text prompt changes over the course of sampling: at high noise the prompt drives the layout, and at low noise it is largely ignored §4 fig. 3 fig. 4. They want more capacity without more cost per step, and they want the conditioning to be more controllable §1.

Method

Cascade §4.2 fig. 5:

1. Base model generates 64×64 images.
2. SR256 upsamples to 256×256.
3. SR1024 upsamples to 1024×1024. It is trained on random 256×256 crops and applied at 1024×1024 app. A §5.

Denoiser and noise schedule §3 eq. 1 eq. 2:

• The denoiser is preconditioned as in the EDM formulation: a skip term plus a scaled network output F_θ, with σ_data = 0.5 eq. 1.
• Training noise levels follow a log-normal distribution, ln σ ~ N(−1.2, 1.2²). The loss weight λ(σ) cancels the output scale of F_θ §3.
• Sampling starts from N(0, σ_max² I) and solves the generative ODE from σ_max down to about zero eq. 2 §3.

Expert ensemble, binary-tree branching §4.1 tab. 5:

1. Train one model on the full noise distribution p(σ).
2. Split p(σ) into 2^l intervals of equal area at level l. The paper gives no exact σ edges.
3. Initialise each child from its parent and fine-tune it on its own interval for a short run.
4. The authors grow only the leftmost and rightmost leaves at each level. The rightmost leaf covers high noise and the leftmost covers low noise. All intermediate intervals are merged into one extra expert.
5. The final base ensemble has three experts: low-noise leaf, high-noise leaf and a middle expert §4.1. Other configurations (2-expert, 3-expert, with the SR256 ensemble) are listed in §5.1.

Conditioning §4.2 app. A:

• T5-XXL text embeddings (113 tokens), CLIP L/14 text embeddings (77 tokens), and CLIP L/14 image embeddings (1 token). All are precomputed for the dataset.
• Global conditioning: the projected pooled CLIP text and CLIP image embeddings are added to the time embedding. Pooled T5 is not used.
• Cross-attention: after each self-attention block, a cross-attention block attends over the concatenated tokens plus a learnable null token.
• Independent dropout on each embedding. A dropped embedding is zeroed. Dropping all three gives the unconditional branch used for classifier-free guidance §4.2.

# Denoiser with EDM preconditioning  [§3] [eq. 1]
# x: noisy image (B, 3, H, W); sigma: (B,); cond: embeddings
def denoise(net, x, sigma, cond):
    # Combined noise scale, sigma_data = 0.5  [§3]
    scale = sqrt(sigma**2 + sigma_data**2)
    # Network sees the scaled input and ln(sigma)/4  [§3] [eq. 1]
    raw = net(x / scale, cond, log(sigma) / 4)
    skip = (sigma_data / scale)**2 * x
    return skip + sigma * sigma_data / scale * raw

# Training loss for one expert  [§3] [§4.1]
# The expert only sees sigma inside its own interval  [§4.1]
def expert_loss(net, images, cond, interval):
    # Sample ln(sigma) ~ N(-1.2, 1.2^2), restricted to the interval  [§3] [§4.1]
    sigma = sample_lognormal_in(interval, mean=-1.2, std=1.2)
    noisy = images + sigma * randn_like(images)
    denoised = denoise(net, noisy, sigma, cond)
    # Weight lambda(sigma) = (sigma* / (sigma * sigma_data))^2  [§3]
    weight = (sqrt(sigma**2 + sigma_data**2) / (sigma * sigma_data))**2
    return weight * squared_error(denoised, images)

# Independent condition dropout  [§4.2] [tab. 7]
# Rates: T5 0.25, CLIP text 0.2, CLIP image 0.9
def drop_conditions(t5, clip_text, clip_image):
    # Each embedding is zeroed on its own  [§4.2]
    if random() < 0.25:
        t5 = zeros_like(t5)
    if random() < 0.2:
        clip_text = zeros_like(clip_text)
    if random() < 0.9:
        clip_image = zeros_like(clip_image)
    return t5, clip_text, clip_image

# Sampling by solving the generative ODE  [§3] [eq. 2]
# sigma_max and sigma_min values are not stated in the paper
def sample(models, cond, sigmas):
    # Start from the prior N(0, sigma_max^2 I)  [§3]
    x = randn(shape) * sigmas[0]
    for sigma, sigma_next in pairs(sigmas):
        # Pick the expert whose interval contains sigma  [§4.1] (rule implied)
        net = expert_for(models, sigma)
        # DEIS multistep solver: order 6 base, order 3 SR  [tab. 7]
        # Step formula not given in this paper
        x = deis_step(net, x, sigma, sigma_next, cond)
    return x

Paint-with-words §4.3 fig. 6 (training-free):

• The user selects a phrase and paints a binary mask. Masks are bilinearly downsampled to each cross-attention layer's resolution.
• Build an input attention matrix A (N_i × N_t). Each text token's column is its phrase mask, flattened, and zero if the token is in no painted phrase.
• Add w·A to the attention logits before the softmax: softmax((QKᵀ + wA)/√d_k) V.
• Schedule for w: w = w′ · log(1 + σ) · max(QKᵀ). w′ is user-specified §4.3.

# Paint-with-words cross-attention  [§4.3]
# queries (N_i, d_k); keys, values (N_t, d_k); masks per phrase
def cross_attention(queries, keys, values, bias, w_user, sigma):
    # Attention logits between image and text tokens  [§4.3]
    logits = queries @ keys.T              # (N_i, N_t)
    # Strength grows with noise and scales with max(QK^T)  [§4.3]
    strength = w_user * log(1 + sigma) * max(logits)
    # Bias A is built from the user masks  [§4.3]
    weights = softmax((logits + strength * bias) / sqrt(d_k))
    return weights @ values

Hyperparameters

Branching schedule (base model) tab. 5 app. B. Each row is initialised from the row named in "from":

SR branching tab. 6: M(0,0) from scratch for 2M; M(1,0) and M(1,1) from M(0,0) for 300K each.

Experiment-specific schedule §5.1: the 4-expert ensemble splits the 500K-iteration baseline, trains 50K per child, then splits again at 550K and trains for another 50K, giving 600K steps per expert. The baseline is trained to 800K for a fair sample count.

Results

Author-reported. Zero-shot FID-30K on MS-COCO 2014 validation, 256 px, 30K captions §5 Evaluation tab. 1. Lower is better.

Config definitions §5.1: a = baseline base + baseline SR256; b = a with the ensemble SR256 (E¹₀, E¹₁); c = 2-expert base (E⁹₅₁₁ and its complement) + ensemble SR256; d = 3-expert base (E⁹₅₁₁, E³₀, and a middle expert) + ensemble SR256.

Other reported results (no numeric table in the paper):

• FID–CLIP trade-off curves over classifier-free guidance weights 0 to 10 in steps of 0.5 for the 4-expert ensemble vs baseline on COCO. The ensemble is better across the curve §5.1 fig. 7.
• Same comparison for 2-expert vs baseline, shown as images only §5.1 fig. 9.
• Speed/size curves: the paper claims the same inference cost for the ensemble, with values only in a figure §5.1 fig. 13.
• CLIP text vs T5 vs both: curves on COCO and Visual Genome. On COCO, CLIP and T5 alone perform similarly, and the combination is best. On Visual Genome, T5 alone beats CLIP alone §5.2 fig. 8 fig. 18. Mean caption length: 10.62 words (COCO) and 61.92 (Visual Genome) §5.2.
• Qualitative comparisons with DALL·E 2 and Stable Diffusion on multiple entities, text rendering and long captions fig. 10 fig. 11 fig. 12.
• Style transfer with CLIP image embeddings §5.3 fig. 16. Style variants fig. 14. Variations fig. 15.
• Paint-with-words samples §5.4 fig. 17.

Paper vs code

No official code.

Searched: the arXiv abstract and the v5 LaTeX source (no repository link; the only URL is the project page); the project page https://deepimagination.cc/eDiff-I/ (the fetch was not permitted in this run, so it was not read); GitHub repository search for "eDiff-I" and "ediff nvidia". The GitHub results were third-party reimplementations only: cloneofsimo/paint-with-words-sd (MIT) and Mikubill/sd-paint-with-words (no license listed). They are not official and are not used here. The code, code_commit and code_license fields are therefore left empty.

Pitfalls

• Text and table disagree. §5.1 gives 7.04 as the best FID, while Table 1 gives 6.95 for config d §5.1 tab. 1. Use the table value and flag the text value.
• No code, no weights, partly proprietary data. Training uses about one billion pairs, some proprietary, filtered with a CLIP alignment threshold and an aesthetic scorer. The threshold values are not stated §5 Datasets. An implementer cannot reproduce the data or the exact model.
• Unstated sampling settings for Table 1. The classifier-free guidance weight used for the FID numbers is not stated. Only the sweep for the curves is given §5.1. The σ_min and σ_max values are also not stated §3.
• Expert selection at inference is not spelled out. The paper says each expert covers an interval of noise levels and that intervals are equal-area splits of p(σ), so the boundaries are log-normal quantiles. The exact switching rule in the sampler is not given in the text. The pseudocode above marks this as implied.
• Base iteration budget does not add up from the text. Table 5 lists tree-branch iterations, and the base model total is 1.9M in Table 7. The paper does not say how the tree schedule maps onto that total.
• "3kutta" is unexplained. It appears in the sampler row of Table 7 and no text describes it tab. 7.
• High CLIP-image dropout. A 0.9 dropout rate on the CLIP image embedding is very high, so style conditioning is active only about 10% of the time during training. The paper does not explain the choice tab. 7.
• Competitor numbers. The paper does not say whether competitor FIDs were re-measured or copied from their papers tab. 1. Parameter counts include the text encoder, and the authors say so, but the numbers for other models may not use the same rule.
• Ablations use curves. CLIP score is reported only as curves, with no numeric table §5.1 fig. 7 fig. 8. Speed and parameter claims are also figure-only fig. 13.
• Paint-with-words details missing. The user weight w′ has no default. It is not stated whether max(QKᵀ) is taken per head, per sample or over the whole matrix. The method applies to all cross-attention layers §4.3.
• Timing of expert switch. The "prompt switching" evidence (last 7% vs first 40%) comes from individual examples, not a metric §4 fig. 4.
• Figure and table numbers come from LaTeX source order. They were not checked against the compiled PDF. If a number looks off, cite the section instead.
• Version. This card reflects arXiv v5 (updated 2023-03-14). The licence recorded by fetch is CC BY 4.0.
• Hidden costs. Base training used 256 A100 GPUs and each SR model used 128 §5 Optimization. Gradient checkpointing is enabled tab. 7. The ensemble avoids extra cost per sampling step but adds training cost from many stages.

Lineage

• Builds on: the EDM denoiser and sampler formulation (Karras et al.) §3; the Dhariwal and Nichol U-Net, modified here §4.2 app. A; the efficient U-Net block structure from Imagen (Saharia et al.) §4.2 app. A; two-stage denoising by Luhman et al., cited as related §4.1; DEIS sampler (Zhang et al.) tab. 7.
• Compared against: GLIDE, DALL·E 2, Make-A-Scene, Stable Diffusion, Imagen, Parti tab. 1.