针对高效扩散变换器的自适应层与时间步压缩比

Layer- and Timestep-Adaptive Differentiable Token Compression Ratios for Efficient Diffusion Transformers

摘要 Abstract

扩散变换器(Diffusion Transformers, DiTs)在图像生成质量方面达到了最先进的水平(SOTA),但其高延迟和内存效率低下问题使其难以部署在资源受限的设备上。主要的效率瓶颈在于现有的DiTs对图像的所有区域应用了相同的计算量。然而,并非所有的图像标记具有相同的重要性,某些局部区域需要更多的计算,例如物体。为了解决这一问题,我们提出了DiffCR,这是一种具有可微压缩比的动态DiT推理框架,能够自动学习为每个图像标记动态分配层间和时间步间的计算,从而实现高效的DiTs。具体而言,DiffCR集成了三个特性:(1)标记级路由方案,其中每个DiT层包含一个路由器,该路由器与模型权重联合微调以预测标记重要性分数。这样,不重要的标记可以跳过整个层的计算;(2)层级可微比例机制,不同的DiT层从零初始化自动学习变化的压缩比,在冗余层中实现较大的压缩比,而其他层则保持较少压缩甚至完全不压缩;(3)时间步级可微比例机制,每个去噪时间步学习其自身的压缩比,结果表明噪声较大的时间步具有较高的压缩比,随着图像变得清晰,压缩比逐渐降低。在文本到图像和图像补全任务上的大量实验表明,DiffCR有效捕捉了标记、层和时间步轴上的动态性,在生成质量和效率之间取得了优于先前工作的权衡。项目网站可在https://www.haoranyou.com/diffcr获取。

Diffusion Transformers (DiTs) have achieved state-of-the-art (SOTA) image generation quality but suffer from high latency and memory inefficiency, making them difficult to deploy on resource-constrained devices. One major efficiency bottleneck is that existing DiTs apply equal computation across all regions of an image. However, not all image tokens are equally important, and certain localized areas require more computation, such as objects. To address this, we propose DiffCR, a dynamic DiT inference framework with differentiable compression ratios, which automatically learns to dynamically route computation across layers and timesteps for each image token, resulting in efficient DiTs. Specifically, DiffCR integrates three features: (1) A token-level routing scheme where each DiT layer includes a router that is fine-tuned jointly with model weights to predict token importance scores. In this way, unimportant tokens bypass the entire layer's computation; (2) A layer-wise differentiable ratio mechanism where different DiT layers automatically learn varying compression ratios from a zero initialization, resulting in large compression ratios in redundant layers while others remain less compressed or even uncompressed; (3) A timestep-wise differentiable ratio mechanism where each denoising timestep learns its own compression ratio. The resulting pattern shows higher ratios for noisier timesteps and lower ratios as the image becomes clearer. Extensive experiments on text-to-image and inpainting tasks show that DiffCR effectively captures dynamism across token, layer, and timestep axes, achieving superior trade-offs between generation quality and efficiency compared to prior works. The project website is available at https://www.haoranyou.com/diffcr.

针对高效扩散变换器的自适应层与时间步压缩比 - arXiv