DiffCR

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

Haoran You^1,2, Connelly Barnes², Yuqian Zhou², Yan Kang², Zhenbang Du¹, Wei Zhou¹,
Lingzhi Zhang², Yotam Nitzan², Xiaoyang Liu², Zhe Lin², Eli Shechtman², Sohrab Amirghodsi², Yingyan (Celine) Lin¹,

¹Georgia Institute of Technology, ²Adobe Research

Paper Code

Our DiffCR adaptively learns to route computation across layers and timesteps for image tokens in a differentiable manner.

Abstract

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 key 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 DiT models. Specifically, DiffCR integrates three features: (1) A token-level routing scheme where each DiT layer includes a router that is jointly fine-tuned 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 both 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.

Overview of the proposed DiffCR framework: (a) Token-level routing scheme and (b) Differentiable compression ratios.

Visualization of MoD Routers' Prediction

Visualization of the router's predictions: (a) For inpainting tasks, where inputs are masked images with text prompts, we follow the previous SOTA method Lazy-Diffusion to generate only the masked area rather than the entire image; (b) For text-to-image (T2I) tasks, where inputs are noise and text prompts, we follow PixArt-Σ for generation. Each visualization includes the router's prediction map with values ranging from 0 to 1. The generated image at each corresponding timestep is shown on the left, while the router's prediction maps across various layers and timesteps are displayed on the right.

Visualization of Layer-Wise Differentiable Ratios

Visualization of the compression ratio trajectory during fine-tuning: (a) For inpainting tasks, and (b) For text-to-image (T2I) tasks. The visualization includes: (1) Trajectories for each of the 28 layers in DiT models; (2) Average ratio trajectory across all layers; and (3) The final learned ratio distribution across 28 layers.

Visualization of Timestep-Wise Differentiable Ratios

Visualization of the learned ratio patterns across both timesteps and layers for the (a) inpainting task and (b) T2I task.

DiffCR over SOTA Baselines

Comparisons of our DiffCR with previous uncompressed models and SOTA compression methods: (a) T2I tasks, where DiffCR is applied to PixArt-Σ, and (b) Inpainting tasks, where DiffCR is applied to Lazy-Diffusion models.

Qualitative Visual Examples

Visual comparisons of our DiffCR with previous uncompressed models and SOTA compression methods: (a) Inpainting tasks, where DiffCR is applied to Lazy-Diffusion, and (b) T2I tasks, where DiffCR is applied to PixArt-Σ.