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

Diffusion models [38, 13] have demonstrated superior performance over prior SOTA generative adversarial networks (GANs) in image synthesis tasks [7]. Early diffusion models primarily utilized U-Net architectures. Subsequent work introduced several improvements, such as advanced sampling [16, 39, 21] and classifier-free guidance [12]. Although effective, these models suffered from high generation latency due to processing directly in pixel space, thus limiting their practical applications. The introduction of Latent Diffusion Models (LDMs) [32] marked a significant advancement by encoding pixel space into a more compact latent space through training a Variational Auto-Encoder (VAE). This reduced the computational cost of the diffusion process, paving the way for widely used models like Stable Diffusion Models (SDMs) [28]. More recently, researchers have explored Transformer [41] architectures for diffusion, leading to the development of DiTs [27, 1], which employ a pure Transformer backbone and exhibit improved scalability. Our DiffRatio-MoD proposes a novel dynamic DiT inference framework with differentiable compression ratios and is compatible with all recent DiT models.

DiTs [27] are resource-intensive due to the transformer architecture, with the attention module exhibiting quadratic complexity relative to the number of tokens. Previous work has mainly focused on optimizing DiTs’ deployment efficiency along three dimensions: token, layer, and timestep. For tokens, researchers have introduced techniques like token merging [2] to merge similar tokens, token pruning [43] or image resolution downsampling [37] to remove redundant tokens, and LazyDiffusion [26], which is specialized for the inpainting task and bypasses generating background tokens. For layers, methods such as layer [17] and channel [9] pruning, as well as intermediate feature caching [52, 23, 20], have been proposed to skip redundant computations. For timesteps, strategies include distillation to reduce the required number of timesteps, which has been explored for UNets [33, 21, 54, 15, 53, 34] although there is no reason to believe these techniques cannot apply for Transformers, and asymmetric sampling, which has been applied to Transformer architectures and allocates more samples to undersampled stages and fewer to stages that have already converged [44, 29]. Additionally, to accelerate diffusion T2I models, more specialized techniques have been introduced [3, 4]. In contrast, our proposed DiffRatio-MoD is a learnable and unified dynamic DiT inference framework with differentiable compression ratios across layers and timesteps, exploring the compounded effects of compression across all three axes. However, we do not explore few step distillation (e.g. [53]) in this paper, since it is an orthogonal acceleration.

Model compression [6] offers a static approach to improving inference efficiency, while dynamic inference [47, 56, 50, 46, 31] enables adaptive compression based on inputs, layers, or other conditions. For example, early exiting methods [40, 14, 22] predict the optimal point for early termination within intermediate layers, allowing the model to exit before completing all computations. Dynamic layer-skipping methods [47, 46, 50] selectively execute subsets of layers for each input, often utilizing a gating network to make decisions on the fly. At a finer granularity, researchers have also explored channel skipping [24, 9] and mixture-of-depths (MoD) approaches [31], which select specific subsets of layers for individual tokens rather than processing the entire input uniformly. In contrast, our DiffRatio-MoD is the first to introduce a unified dynamic DiT inference framework that optimizes along three axes: token, layer, and timestep. It also enables differentiable compression ratios that are jointly fine-tuned with the network, enhancing adaptability and efficiency.

Motivated by the need for unified and dynamic compression during DiT inference, DiffRatio-MoD introduces a token-level routing scheme to dynamically learn the importance of each token on the fly. As illustrated in Fig. 1 (a), similar to previous mixture-of-depths (MoD) work [31] for NLP tasks, each DiT layer incorporates a lightweight router using a single linear layer to predict the importance of each token based on the input image/noise and text embedding. This allows us to bypass computations for less important tokens in each layer and directly concatenate their cached activations to the final layer outputs. Consequently, each token is processed by only a selective subset of layers. Visualization of these routers’ predictions reveals that different layers or timesteps favor varying compression ratios—for instance, some layers prioritize generating objects, while others focus on backgrounds—highlighting the need for adaptable compression across layers and timesteps. To achieve such dynamic compression, DiffRatio-MoD incorporates a differentiable compression ratio scheme, as shown in Fig. 1 (b). This scheme includes a learnable scalar parameter that represents a continuous compression ratio, and predefined discrete ratio bins as proxy ratios. The scalar queries the bins to identify lower and upper bin ratios, creating two separate paths with distinct compression ratios. The final output is a linear combination of these two paths, weighted by the proximity of the learned ratio to each bin. We apply a mean-square error (MSE) loss to ensure that the average learned ratio across layers or timesteps converges to the target ratio. By doing so, DiffRatio-MoD learns the adaptive compression ratios in a differentiable manner, resulting in efficient and dynamic mixture-of-depths DiTs.

Motivation. We are motivated by the varying computational demands across tokens, where many of them require fewer layers for efficient processing. We start from the same token-level routing scheme as MoD [31]. We remove from MoD two features that were specialized for the acausal NLP task: specifically, we remove the auxiliary loss and auxiliary MLP predictor from Section 3.5 of their paper. To the best of our knowledge, our paper is the first application of MoD to the vision domain, so we next review the routing mechanism, perform some visualizations, and report insights for vision tasks.

Token-level Routing. DiTs process noise and conditional text embeddings as inputs, aiming to denoise and generate images in an end-to-end manner. To predict token importance, we employ a simple yet effective token-level routing scheme from MoD [31]. As illustrated in Fig. 1 (a), each DiT layer incorporates a lightweight router composed of a single linear layer with a sigmoid activation function, predicting each token’s importance on a scale from 0 to 1. After passing through the routers, we select the top-$k$ most important tokens for this layer’s processing, while the activations of other tokens are cached and concatenated with the layer outputs, bypassing the entire layer computation, including both attention and MLPs. To enable gradient flow to the router’s weights during joint fine-tuning with pretrained DiT models, the same as MoD [31], we rescale the top-$k$ token output activations by multiplying them with the router’s predictions. This rescaling ensures that gradients are propagated effectively to the router during backpropagation. The value of $k$ is determined based on compression ratios; we empirically find that 20% or 30% compression offers an optimal trade-off between latency/memory efficiency and minimal degradation in generation quality. Unlike previous token merging techniques [2], where computational savings are not directly proportional to the merging ratios due to the extra overhead, as shown in Fig. 3, the MoD in our approach yields more reductions in actual latency and memory usage due to negligible overhead.

Router Visualization and Insights. To test the router’s effectiveness, we visualize its predictions in Fig. 2. Our observations reveal that: (1) The router effectively captures semantic information, clearly delineating object shapes and achieving an attention-like effect with significantly reduced computational cost; (2) The predicted token importance varies across layers and timesteps. For instance, some layers prioritize object generation, while others emphasize background areas. Additionally, as timesteps progress, the router increasingly captures the semantic contours of objects, underscoring the need for dynamic token importance estimation; (3) The optimal compression ratio differs across layers and timesteps. For example, certain layers designate all tokens as high-importance, showing minimal redundancy, whereas other layers selectively prune object or background tokens with distinct shapes, requiring varying compression ratios. A similar variance is observed across timesteps. In the current MoD, a fixed global compression rate is applied equally to each layer and timestep, rather than adapting to their individual significance. Uniform pruning risks over-pruning critical layers or timesteps while leaving redundant ones less compressed. This motivates us to apply adaptive and dynamic compression ratios across both layers and timesteps.

Motivation. Recognizing that different layers prioritize different objects or background elements and thus benefit from distinct compression ratios, we propose a novel layer-wise differentiable compression ratio mechanism. This approach automatically learns each layer’s compression ratio from a zero initialization in a differentiable manner, adapting to the varying redundancy levels across layers.

Design Choice. Before designing DiffRatio-MoD, we address a key choice: discrete proxy or continuous ratio representation. Previous work [5] uses a discrete proxy with multiple compression ratio candidates and learnable probabilities, but this approach poses three challenges for MoD: (1) it lacks effective initialization, as the final ratio relies on the product of candidates and probabilities, making it difficult to initialize all ratios at zero; (2) MoD requires a differentiable and learnable router, incompatible with discrete proxies that need multiple sets of top-$k$ tokens; and (3) it introduces numerous learnable parameters, complicating training and interpretability, and leading to hard-to-interpret ratio distributions. In contrast, we represent each layer’s compression ratio with a single continuous scalar parameter, reducing the parameter count to 28 for DiT’s 28 layers [3] and directly representing the compression ratio.

Layer-wise DiffRatio-MoD. As illustrated in Fig. 1 (b), we assign each layer a single learnable parameter and introduce discrete MoD compression ratio bins at 10% intervals, ranging from 0% to 100%. During training, the learnable MoD ratio queries the nearest two discrete bins to retrieve the lower and upper bin ratios. For example, a 22% learnable ratio would correspond to 20% as the lower bin and 30% as the upper bin. We then apply a forward pass through the DiT layer (with MoD routers) with each of these bin ratios, producing two output branches. The final output is a weighted linear combination of these branches, where the weights are determined by the proximity of the learnable ratio to each bin—e.g., for 12%, the output would be 80% weighted towards the 10% branch and 20% towards the 20% branch. Although this approach doubles the cost of a forward pass during training, at inference, we simply select the nearest bin as the final compression ratio, eliminating any overhead. To ensure that the ratio converges to our target values, we incorporate an additional MSE loss between the current learned average ratios across all layers in the batch and the target ratio, a hyperparameter.

Ratio Trajectory Analysis. We visualize the training trajectory of compression ratios for all layers during fine-tuning of a Lazy-Diffusion model on an inpainting task in Fig. 4 (a-c). The visualization reveals that: (1) Each layer learns its unique compression ratio, with redundant layers achieving higher compression and critical layers remaining less or entirely uncompressed; (2) The average ratio across layers gradually converges to the target ratio. In this example, with a target of 30%, the final achieved average ratio is approximately 29%, indicating a minor gap. Notably, a trade-off exists between convergence speed and generation quality: a higher MSE loss coefficient for the ratio accelerates convergence but may degrade quality due to overly rapid compression, while a smaller coefficient promotes gradual convergence and maintains quality, albeit with slower training. In practice, we set the coefficient to 0.3 to balance speed and quality effectively; (3) The middle layers exhibit greater redundancy, while the later layers generally have less redundancy and often cannot be compressed. The early layers show variable redundancy levels.

Motivation. In addition to layer-wise ratio variances, we also observe that the model exhibits varying levels of redundancy across timesteps. This motivates us to explore an approach for timestep-wise compression ratios as well.

Timestep-wise Differentiable Ratio. On top of the layer-wise DiffRatio-MoD, we introduce learnable parameters specific to different timestep regions. For the image inpainting task, following the previous SOTA Lazy Diffusion [26], we use 1,000 training timesteps and 100 sampling timesteps, which we evenly divide into 10 regions. We assign 10 learnable parameters per layer accordingly, resulting in a total of 280 learnable parameters. Similarly, for T2I tasks, following PixArt-$\Sigma$ [4] with 20 timesteps, we divide them into 4 regions, assigning 4 parameters per layer, yielding 112 learnable parameters. The same as before, we apply an MSE loss between the averaged learned ratios within the batch and the target ratio to ensure convergence.

Ratio Pattern Analysis. We visualize the learned compression ratio patterns across both timesteps and layers in Fig. 5. For both image inpainting and T2I tasks, we consistently observe that noisy timesteps (corresponding to earlier sampling timesteps or later training timesteps) exhibit higher redundancy and allow for higher compression, whereas timesteps where images become clearer (corresponding to later sampling timesteps or earlier training timesteps) show less redundancy. This learned pattern aligns with previous empirical findings [44], which suggest that high-noise timesteps are associated with convergence regions containing easier samples, allowing for fewer sampling timesteps, while low-noise timesteps involve harder samples and require more frequent sampling.

Tasks, Datasets, and Models. Tasks & Datasets. We evaluate the DiffRatio-MoD on two representative image generation tasks using corresponding benchmark datasets: (1) an image inpainting task on an internal dataset of 220 million high-quality images, covering diverse objects and scenes. Masks and text prompts are generated following [26, 51]; and (2) a T2I task on the LAION-5B dataset [36], restricted to image samples with high aesthetic scores, English text, and a minimum text similarity score of 0.24. Models. We integrate our proposed DiffRatio-MoD approach with SOTA models. For the inpainting task, we use Lazy Diffusion (an adapted PixArt-$\alpha$ model with an additional ViT encoder) to generate images at 1024$\times$1024 resolution. For the T2I task, we use PixArt-$\Sigma$ to generate images at 512$\times$512 resolution.

Training and Sampling Setting. For the inpainting task, we fine-tune the model parameters until convergence using the AdamW optimizer [18] with a learning rate of $10^{-4}$ and weight decay of $3\times 10^{-2}$. For sampling, images are generated using IDDPM [25] with 100 timesteps and a CFG factor of 4.5. For the T2I task, we fine-tune the model using a LoRA adapter with a rank of 32 until both training and validation losses converge, and the MSE loss between the current and target compression ratios drops to approximately zero for DiffRatio-MoD models. During training, we calculate the diffusion loss with IDDPM [25] over 1K timesteps. For sampling, we generate images using DPM-solver [19] with 20 timesteps and a CFG factor of 4.5. All training is conducted on a cluster of 8$\times$A100-80GB GPUs.

Baselines and Evaluation Metrics. Baselines. For both the T2I and inpainting tasks, we compare the proposed DiffRatio-MoD against SOTA baselines, including ToMe [2], AT-EDM [43], and our adapted MoD with uniform MoD compression ratio. For the inpainting task, we also compare against RegenerateCrop, which generates a tight square crop around the masked region, similar to popular software frameworks [48, 42], and RegenerateImage, which generates the entire image, as commonly done in the literature [28, 32, 51, 45]. Evaluation Metrics. We assess the generated image quality using FID scores [11], text-image alignment using CLIP scores [10], and efficiency through inference FLOPs, latency, and memory usage, all measured on an A100 GPU. For inpainting and T2I models, we evaluate on 10K images from LAION-400M [35] or LAION-5B [36], excluding training samples, respectively.

Text-to-Image. To assess the effectiveness of our proposed DiffRatio-MoD, we apply our proposed DiffRatio-MoD to the general text-to-image task and compare it with previous token merging [2] and pruning [43] baselines. Specifically, we apply these compression methods on PixArt-$\Sigma$, a SOTA publicly accessible T2I model known for its high-resolution image generation quality and efficiency tradeoffs. As shown in Tab. 1, PixArt-$\Sigma$ with DiffRatio-MoD significantly improves generation quality, achieving 57.83 and 241.11 FID reductions over ToMe [2] and AT-EDM [43], respectively, with comparable or even lower latency ($\downarrow$8.59%$\sim$20.15%) and memory usage ($\downarrow$-2.71%$\sim$0.72%). Also, under similar latency compared to ToMe [2] with 20% compression ratio, DiffRatio-MoD achieves 335.23 FID reductions. Moreover, PixArt-$\Sigma$ with DiffRatio-MoD also achieves comparable image generation quality with uncompressed PixArt-$\Sigma$, while delivering 20.68% and 8.33% latency and memory savings. Note that we compare with fine-tuned PixArt-$\Sigma$ on the LAION datasets for a fair comparison. This set of experiments demonstrates the effectiveness of DiffRatio-MoD for general T2I tasks.

Image Inpainting. We further extend the DiffRatio-MoD to the inpainting task. Specifically, we apply it on top of the SOTA Lazy Diffusion (LD) [26], which uses a DiT decoder to generate only the masked areas rather than the entire image, leveraging a separate ViT encoder to capture the global context of the input masked images. We compare our DiffRatio-MoD approach against two types of baselines: (1) RegenerateImage and RegenerateCrop, and (2) LD with previous token merging [2] or pruning [43] techniques. As shown in Tab. 2, our DiffRatio-MoD consistently outperforms all baselines in terms of accuracy-efficiency tradeoffs. For example, LD with DiffRatio-MoD achieves FID reductions of 47.35 and 189.93 compared to LD with ToME [2] or AT-EDM [43], while achieving similar or up to 23.61% and 13.63% higher latency and memory savings. Also, under similar memory usage compared to ToMe [2] with 30% compression ratio, LD with DiffRatio-MoD achieve 265.94 FID reduction while delivering up to 21.54% latency savings. Moreover, compared to RegenerateImage, our method achieves 73.51%/60.26% FLOPs and latency savings when inpainting 256² mask sizes within 1024² images. Notably, like Lazy Diffusion, our method’s complexity scales with mask size, while RegenerateImage generates based on full image resolution, making it less efficient for smaller mask sizes. In comparison to RegenerateCrop, our method achieves significantly higher image generation quality ($+$41.01 FID) while also delivering 22.96% memory savings. Note that all memory measurements are taken with a batch size of 128. This set of experiments validates the effectiveness of our DiffRatio-MoD when applied to image inpainting tasks.

We conduct ablation studies on DiffRatio-MoD, analyzing the contributions of the three enablers described in Sec. 3. As shown in Tabs. 2 and 1, we report the performance of LD or PixArt-$\Sigma$ with MoD (Sec. 3.2), DiffRatio-MoD-L (Sec. 3.3), DiffRatio-MoD-LT (Sec. 3.4) for inpainting and T2I tasks, respectively. The results consistently demonstrate that all components of our DiffRatio-MoD contribute to the final performance. Specifically, MoD alone achieves average FID reductions of 44.08 and 207.01 compared to ToMe [2] and AT-EDM [43], with comparable or even lower latency and memory usage. DiffRatio-MoD-L and DiffRatio-MoD-LT further enhance generation quality, achieving additional FID reductions of 4.81/4.92 for the inpainting task and 10.5/12.1 for the T2I task.

Also, a key benefit of our DiffRatio-MoD is that during fine-tuning, the average compression ratios across all layers gradually converge to the target ratio, producing a series of “by-product” models with a range of compression ratios. As shown in Fig. 7, we visualize the model trajectory with corresponding FID scores and compression ratios for both inpainting and T2I tasks. The observations indicate that, for T2I, the FID gradually increases with compression ratio, achieving the desired results. In contrast, for inpainting, the FID gradually decreases as the compression ratio increases. This difference arises because, compared to T2I, inpainting tasks and LD models are more sensitive to pruning and require longer fine-tuning to boost the generation quality. This is also reflected in Tabs. 2 and 1, where FID increases post-pruning for inpainting, while it even decreases for T2I.

Visual Examples. We select challenging input prompts to evaluate the qualitative results of our proposed DiffRatio-MoD. As shown in Fig. 6, the examples demonstrate that DiffRatio-MoD achieves comparable or even superior generation quality compared to the RegenerateCrop baseline and even uncompressed LD or PixArt-$\Sigma$ for inpainting and T2I tasks, respectively. Note that ToMe and AT-EDM are omitted here due to their poor generation quality when applied to DiTs, even at a mere 10% compression ratio.

Human Preference Scores. We use a computer vision model to estimate likely human preferences and assess the models’ ability to generate high-quality, contextually relevant images. Specifically, we generated 2K samples for the T2I task and used HPSv2 [49] to evaluate human preferences for images generated by different methods. As shown in Tab. 4, for T2I, we apply all compression methods to PixArt-$\Sigma$[4]. DiffRatio-MoD achieves a higher human preference score of 4.685/0.847 compared to previous compression methods, ToMe [2] and vanilla MoD [31], respectively.