LazyDiT: Lazy Learning for the Acceleration of Diffusion Transformers

TL;DR

LazyDiT introduces a lazy learning framework for diffusion transformers that reduces redundant computations during inference, significantly speeding up the process while maintaining high performance, even on mobile devices.

Contribution

It proposes a novel lazy learning approach that reuses previous computations in diffusion transformers, enabling faster inference without sacrificing accuracy.

Findings

LazyDiT outperforms DDIM across multiple models and resolutions.

The method achieves better performance than DDIM on mobile devices with similar latency.

Experimental results validate the efficiency and effectiveness of lazy computation reuse.

Abstract

Diffusion Transformers have emerged as the preeminent models for a wide array of generative tasks, demonstrating superior performance and efficacy across various applications. The promising results come at the cost of slow inference, as each denoising step requires running the whole transformer model with a large amount of parameters. In this paper, we show that performing the full computation of the model at each diffusion step is unnecessary, as some computations can be skipped by lazily reusing the results of previous steps. Furthermore, we show that the lower bound of similarity between outputs at consecutive steps is notably high, and this similarity can be linearly approximated using the inputs. To verify our demonstrations, we propose the \textbf{LazyDiT}, a lazy learning framework that efficiently leverages cached results from earlier steps to skip redundant computations. Specifically, we incorporate lazy learning layers into the model, effectively trained to maximize laziness, enabling dynamic skipping of redundant computations. Experimental results show that LazyDiT outperforms the DDIM sampler across multiple diffusion transformer models at various resolutions. Furthermore, we implement our method on mobile devices, achieving better performance than DDIM with similar latency. Code: https://github.com/shawnricecake/lazydit