Learning-to-Cache: Accelerating Diffusion Transformer via Layer Caching

最近，扩散变压器在各种任务中展示了前所未有的生成能力。然而，这些令人鼓舞的结果是以缓慢的推理速度为代价的，因为每个去噪步骤都需要对具有大量参数的变压器模型进行推理。在这项研究中，我们做出了一个有趣而令人惊讶的观察：通过引入缓存机制，可以轻松地删除扩散变压器中大部分层的计算，而无需更新模型参数。例如，在U-ViT-H/2的情况下，我们可以删除高达93.68％的缓存步骤计算量（所有步骤为46.84％），而FID下降不到0.01。为了实现这一点，我们引入了一种名为“学习缓存（L2C）”的新方案，该方案可以动态地学习扩散变压器的缓存方法。具体而言，通过利用变压器层的相同结构和扩散的顺序性，我们通过将每个层视为缓存的基本单元，探索了时步之间的冗余计算。为了解决深度模型中指定要缓存和删除的层的指数搜索空间的挑战，我们提出了一种新颖的可微优化目标。然后，我们优化了一个输入不变但时步变化的路由器，最终可以生成静态计算图。实验结果表明，L2C在推理速度相同的情况下，远远优于DDIM和DPM-Solver等采样器以及之前的基于缓存的方法。

本文旨在解决Diffusion Transformer中慢速推理的问题，提出了一种名为Learning-to-Cache (L2C)的方案，通过学习动态地进行缓存，从而减少冗余计算，提高推理速度。

L2C方案通过将每个层视为缓存的基本单元，利用Transformer中层的相同结构和扩散的顺序性，探索了时间步之间的冗余计算，提出了一种可微分的优化目标，最终生成静态计算图，从而大大提高了推理速度。

实验结果表明，L2C方案在推理速度上大大优于DDIM和DPM-Solver等采样器以及之前的基于缓存的方法，同时在FID上只有不到0.01的下降。该论文的实验使用了U-ViT-H/2数据集，并提供了开源代码。

与本文相关的研究包括Diffusion Models、Transformer模型和缓存技术的研究。