Accelerate High-Quality Diffusion Models with Inner Loop Feedback

TL;DR

This work tackles the computational cost of diffusion-model inference by introducing Inner Loop Feedback (ILF), a lightweight, trainable module that predicts future backbone features using a selected backbone block's output while keeping the diffusion backbone frozen. By feeding a transformed, rescaled signal from the inner loop to earlier blocks, ILF creates a powerful, yet compact, feedback mechanism that enables high-quality image generation with substantially fewer steps. Training leverages Learnable Feedback Rescaling and Fast Approximate Distillation, requiring only the feedback module to learn and avoiding full-model distillation costs. Experiments on transformer-based backbones like DiT and PixArt demonstrate 1.7x–1.8x speedups with improved Image Reward, CLIP, and FID metrics, validating ILF as a flexible and effective approach for efficient, high-quality diffusion inference.

Abstract

We propose Inner Loop Feedback (ILF), a novel approach to accelerate diffusion models' inference. ILF trains a lightweight module to predict future features in the denoising process by leveraging the outputs from a chosen diffusion backbone block at a given time step. This approach exploits two key intuitions; (1) the outputs of a given block at adjacent time steps are similar, and (2) performing partial computations for a step imposes a lower burden on the model than skipping the step entirely. Our method is highly flexible, since we find that the feedback module itself can simply be a block from the diffusion backbone, with all settings copied. Its influence on the diffusion forward can be tempered with a learnable scaling factor from zero initialization. We train this module using distillation losses; however, unlike some prior work where a full diffusion backbone serves as the student, our model freezes the backbone, training only the feedback module. While many efforts to optimize diffusion models focus on achieving acceptable image quality in extremely few steps (1-4 steps), our emphasis is on matching best case results (typically achieved in 20 steps) while significantly reducing runtime. ILF achieves this balance effectively, demonstrating strong performance for both class-to-image generation with diffusion transformer (DiT) and text-to-image generation with DiT-based PixArt-alpha and PixArt-sigma. The quality of ILF's 1.7x-1.8x speedups are confirmed by FID, CLIP score, CLIP Image Quality Assessment, ImageReward, and qualitative comparisons. Project information is available at https://mgwillia.github.io/ilf.

Figures (16)

• Figure 1: ILF uses a lightweight, learnable feedback module to create a powerful inner loop within a diffusion model. Instead of computing a forward through all backbone blocks, in order, we choose some block, provide its output features as input to the feedback module, then feed those features back to some earlier blocks in the model, modified by a learnable scaling term. The feedback's objective is essentially to predict features corresponding to some future diffusion time step, so the resulting noise prediction is more reliable for the model's current step.
• Figure 2: Change in features across time steps, measured for each block as difference from that block's feature at $t=1000$, normalized by dividing by the maximum difference across both plots. Caching reduces the degree to which the features change over time.
• Figure 3: Change in features across blocks, measured for each block at each time step as difference from the first block's feature at that time step, normalized as in Figure \ref{['fig:motivation_a']}. The trend with caching is similar here as when measuring difference over time.
• Figure 4: We compare typical 20 step diffusion inference to caching for PixArt-alpha, 512x512 images. We cache the middle 14 blocks, recomputing features every other step. Caching, while more efficient, sometimes results in quality degradation -- loss of detail (no faces in leftmost image), less appealing design (middle images), and blurriness (rightmost image, zoom in on eyes, ears, hair, and mouth).
• Figure 5: Caching (top) vs. ILF (bottom). We show how we use a partial diffusion forward pass to compute $\text{f}_{t,b+3}$, which we then use to compute $\text{f}_{t-r,b+3}$. We can use fewer of ILF's heavy steps to ultimately achieve better quality-time trade-offs than caching's cheap steps.