Fast Sampling Through The Reuse Of Attention Maps In Diffusion Models

TL;DR

This work tackles the latency of diffusion-model sampling by introducing a training-free method that reuses attention maps across sampling steps, leveraging the ODE formulation of reverse diffusion. It defines memory-based reuse with a binary schedule $oldsymbol{C0}$ and develops two strategies, HURRY (late reuse) and PHAST (greedy-perturbed search), to locate near-optimal reuse patterns, supported by a Lyapunov-exponent-like analysis. Empirically, attention-map reuse yields higher PSNR than equivalent-step reductions and remains competitive on FID and CLIP across multiple models and datasets, albeit at the cost of increased memory to cache maps. The approach offers a practical, model-agnostic path to speed up high-quality image synthesis without retraining or distillation, with potential refinements in memory management and extensions to reuse other internal tensors.

Abstract

Text-to-image diffusion models have demonstrated unprecedented capabilities for flexible and realistic image synthesis. Nevertheless, these models rely on a time-consuming sampling procedure, which has motivated attempts to reduce their latency. When improving efficiency, researchers often use the original diffusion model to train an additional network designed specifically for fast image generation. In contrast, our approach seeks to reduce latency directly, without any retraining, fine-tuning, or knowledge distillation. In particular, we find the repeated calculation of attention maps to be costly yet redundant, and instead suggest reusing them during sampling. Our specific reuse strategies are based on ODE theory, which implies that the later a map is reused, the smaller the distortion in the final image. We empirically compare our reuse strategies with few-step sampling procedures of comparable latency, finding that reuse generates images that are closer to those produced by the original high-latency diffusion model.

Figures (6)

• Figure 1: This figure compares a step-reduced sampler with our best reuse strategy of (approximately) the same latency. The reuse strategy clearly outperforms the step-reduced sampler at producing realistic images that match the original 20-step sampler. Red circles have been added to the images on the bottom row to help readers identify some of the key differences between the 20-step and 13-step DDIM sampler.
• Figure 2: Normalised L1-distance between the Self- (left) and Cross- (right) attention maps $A(s)$ and $A(s-1)$ for an unperturbed flow. This is generated from 200 random ImageNet prompts. The shaded region includes one standard deviation.
• Figure 3: The L1-distance between a sample produced by a normal DPM and a DPM where the attention map is perturbed at sampling step s. Specifically, the pre-softmax attention map is perturbed (in proportion to its norm) at step s. The results are scaled into the range [0,1] and then averaged over 200 random ImageNet prompts; the shaded region includes one standard deviation. The orange curve is of the form $k_1e^{-k_2s}$, tuned to approximate the empirical results between steps 1 and 18 (inclusive).
• Figure 4: A comparison of PSNR (averaged over 200 ImageNet labels as prompts) between our reuse strategies and several alternative heuristic reuse strategies, for a 20-step DDIM sampler with 10 reuse steps. For qualitative analysis, a sample for the prompt 'Sealyham terrier Sealyham' is provided for each strategy.
• Figure 5: This figure compares the PSNR (averaged over 200 ImageNet labels as prompts) offor PHAST, HURRY, and the (base) DDIM sampler for comparable latencies. The PSNR is taken over 200 ImageNet labels as prompts.