Feature caching across timesteps

Overview

Diffusion sampling runs the denoiser once per step, so inference cost grows with the step count 2312.00858 §1. Feature caching exploits the fact that the coarse, high-level features of a U-Net change little between adjacent steps. Some steps run the full network and store an intermediate feature; the steps in between reuse it and recompute only the shallow part of the network 2312.00858 §3.2 2312.00858 §3.3.

The invariant: the sampler's update rule is unchanged, and the weights are not retrained. Only the noise prediction is approximated on the cached steps.

Variants

The only variant with a card so far is the one below. Other caching methods would be added here once their papers have cards.

Uniform and non-uniform 1:N caching (DeepCache)

A branch index m picks one skip connection. A full pass at step t caches the feature that enters up block m from the main branch. Each later step in the block runs only down blocks 1 to m and reuses the cached feature 2312.00858 §3.3 2312.00858 eq. 4 2312.00858 eq. 5. The uniform schedule puts full passes every N steps, so there are ceil(T / N) of them 2312.00858 eq. 6. The non-uniform schedule spaces the full passes more densely where the authors observed feature similarity to drop, using a power p and a centre c 2312.00858 §3.3 2312.00858 eq. 7. Card: DeepCache: Accelerating Diffusion Models for Free.

Reported settings: N between 2 and 20 on the LDM, CIFAR-10 and LSUN models, and N=5 as the Stable Diffusion setting in the main table 2312.00858 tab. 1 2312.00858 tab. 2 2312.00858 tab. 3. The sampler is DDIM or PLMS, and the authors state other fast samplers also work 2312.00858 §4.1 2312.00858 app. A. The method was not tested with other samplers in the paper.

When to use which

Implementation sketch

The core is a loop over blocks of N steps. The first step of each block runs the full network and keeps the cached feature. The others run the shallow part and reuse it. The schedule function decides where the full passes fall.

# Full passes at the start of each N-step block, uniform 1:N [§3.3] [eq. 6]
# Input:  T total steps; interval N
# Output: sorted list of steps that run the full network
def full_steps_uniform(T, N):
    return list(range(0, T, N))

# One block of N steps: a full pass, then N - 1 partial passes [alg. 1]
# Input:  latent x, step t, interval N, skip branch m
# Output: the N updated latents
def cached_block(x, t, N, m):
    # Full pass: run every block and keep the feature entering up block m [eq. 4]
    cached, x_next = full_forward_and_cache(x, t, m)
    latents = [x_next]
    # Partial passes: run down blocks 1..m on the newest latent [eq. 5]
    for n in range(2, N + 1):
        # Recompute the shallow path only
        eps = partial_forward(latents[-1], t - n + 1, m, cached)
        # Apply the sampler's own update with this step's timestep
        latents.append(sampler_step(latents[-1], eps, t - n + 1))
    return latents

Where variants differ: the skip branch m, the schedule that produces the full pass steps, and the sampler. The paper's own code does not match its non-uniform LDM setup in every respect; see the card below.

Contradictions

• Stable Diffusion headline. The abstract's 2.3x speed-up with a 0.05 CLIP drop matches the non-uniform N=6 row on PartiPrompts 2312.00858 app. F 2312.00858 tab. 13. The main table's N=5 row gives 2.15x with the same CLIP score of 29.46 2312.00858 tab. 3. The paper does not list the N=6 row in the main table.
• LDM FID drop. The abstract gives a drop of 0.22 against the regenerated baseline of 3.37 2312.00858 tab. 1. The text gives the change as "from 3.39 to 3.59". The 3.39 is the uniform N=2 result, not the baseline.
• Non-uniform LDM with N=20. The main table gives FID 7.11 2312.00858 tab. 1. The appendix comparison gives 7.36 2312.00858 app. C 2312.00858 tab. 11. The selected-value table gives centre 120 and power 1.4, which the sweep gives as FID 7.13 2312.00858 app. B 2312.00858 tab. 8. No single configuration reproduces the main value.
• Small mismatches. PLMS-50 on COCO shows CLIP 30.30 in the main table and 30.24 in the appendix 2312.00858 tab. 3 2312.00858 app. F. The sFID value 50.6 for DDIM-91 in the sampling-steps comparison is far from the other rows and is probably a typo; the paper does not confirm it 2312.00858 tab. 4.
• Paper and public code. The released LDM code fixes the centre at 0, but the paper's LDM table uses centres of 110 and 120 2312.00858 app. B. The full list of code differences is on the card: DeepCache card.

Papers

• DeepCache: Accelerating Diffusion Models for Free: training-free 1:N feature caching for U-Net diffusion samplers, with uniform and non-uniform schedules.