Diffusion transformer caching optimized by solver-aware target selection
Rethinking What to Cache in Few-Step Diffusion Transformers: Solver-Aware Target Selection
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Generating images with diffusion transformers takes many steps and lots of computing time. The authors found that reusing intermediate data from previous steps (caching) can speed this up, but picking which data to reuse matters more after speeding up the process with fewer steps. They created a method called AutoTarget that tests some runs without caching to find the best data to reuse for each setup. This reduces computing without hurting image quality much, making the generation faster and more efficient.
What this means in practice
- •For machine learning engineers: Speed up image and video generation in diffusion transformer models by optimizing when and what data to cache during sampling.
- •For graphics software developers: Build faster image and video synthesis tools that maintain quality by incorporating solver-aware cache target selection.
Authors
Shuo Yang, Lihao Fang, Yi Zhang, Haixiang Wang, Xincheng Ye, Shufan Chen, Jipeng Guo, Youqing Wang
Abstract
Diffusion Transformers (DiTs) can generate high-quality images and videos, but generating each sample requires multiple costly DiT forward passes. Two common ways to accelerate DiT sampling are step distillation, which reduces the number of sampling steps, and caching, which skips some DiT evaluations by reusing a tensor computed at an earlier step. Most caching methods decide in advance which tensor to reuse. After distillation, adjacent sampling steps are farther apart. Reusing a tensor across this larger gap introduces more error, so choosing what to cache becomes especially important. We therefore introduce AutoTarget, a method that chooses the cached tensor for a given model, solver, and reuse schedule. AutoTarget uses a small set of runs without cache reuse to measure the error caused by reusing each candidate tensor, then selects the candidate with the lowest error. We also analyze how an error at one reuse step affects the final sample. For Euler sampling, we identify cache targets that produce the same trajectory and show why a stored solver update may not. Experiments on distilled image and video DiTs show that the best cache target changes with the model, image resolution, and solver. AutoTarget reduces DiT evaluations and retained cache storage. Generation quality remains close to the corresponding uncached run. On the tested PixArt-LCM and FLUX.1-schnell settings, its calibration ranking matches the ranking from held-out cached runs. To help others reproduce the method, we provide its core implementation on GitHub at https://github.com/wali1024-offical/AutoTarget.