arXiv:2412.12444cs.LGcs.AI2024-12AAAI被引 48

通过懒惰计算跳过冗余步骤,加速扩散Transformer推理。

LazyDiT: Lazy Learning for the Acceleration of Diffusion Transformers

  • 引入懒惰学习层,动态跳过可复用的计算步骤。
  • 在多种分辨率下比DDIM更快,延迟相近但性能更优。
  • 适合移动端部署,兼顾速度与生成质量。

扩散Transformer已成为众多生成任务的主流模型,展现出优异性能。然而其推理速度慢,因每一步去噪都需要完整运行参数量巨大的Transformer。本文发现:在连续去噪步骤间,输出相似性较高,且可通过输入线性近似;因此无需每步都执行完整计算。为此,我们提出LazyDiT——一种利用缓存前步结果跳过冗余计算的懒惰学习框架。通过在模型中嵌入懒惰学习层,实现动态跳步。实验表明,LazyDiT在多个扩散Transformer模型、不同分辨率下均优于DDIM采样器。此外,在移动设备上实现时,其性能超越DDIM,而延迟相近。

原文摘要 · Abstract (English)

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

扩散模型加速推理懒惰计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。