DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity
Authors: Haowei Zhu
Ji Liu
Ziqiong Liu
Dong Li
Jun-Hai Yong
Bin Wang
Emad Barsoum
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion
Acceleration
Sparsity
Primary area: generative models
TL;DR: We propose DiffSparse, a differentiable approach to optimize layer-wise token sparsity in diffusion models sampling process.

Abstract: Diffusion models demonstrate outstanding performance in image generation, but their multi-step inference mechanism requires immense computational cost. Previous works accelerate inference by leveraging layer or token cache techniques to reduce computational cost. However, these methods fail to achieve superior acceleration performance in few-step diffusion transformer models due to inefficient feature caching strategies, manually designed sparsity allocation, and the practice of retaining complete forward computations in several steps in these token cache methods.
To tackle these challenges, we propose a differentiable layer-wise sparsity optimization framework for diffusion transformer models, leveraging token caching to reduce token computation costs and enhance acceleration. Our method optimizes layer-wise sparsity allocation in an end-to-end manner through a learnable network combined with a dynamic programming solver. Additionally, our proposed two-stage training strategy eliminates the need for full-step processing in existing methods, further improving efficiency.
We conducted extensive experiments on a range of diffusion-transformer models, including DiT-XL/2, PixArt-$\alpha$, FLUX, and Wan2.1. Across these architectures, our method consistently improves efficiency without degrading sample quality. For example, on PixArt-$\alpha$ with 20 sampling steps, we reduce computational cost by 54% while achieving generation metrics that surpass those of the original model, substantially outperforming prior approaches. These results demonstrate that our method delivers large efficiency gains while often improving generation quality.
.

Published as a conference paper at ICLR 2026
DIFESPARSE: ACCELERATING DIFFUSION TRANSFORM- ERS WITH LEARNED TOKEN SPARSITY
Haowei Zhu Ji Liu', Ziqiong Liu', Dong Li', Junhai Yong', Bin Wang Emad Barsoum I Advanced Micro Devices, Inc. "Tsinghua University BNRist (haoweiz, liuji, ziqioliu, d. 1.li, ebarsoum ammd. com, yongjheetsinguaa edu. cn,
edu.cn cn
ABSTRACT
Diffusion models demonstrate outstanding performance in image generation, but their multi-step inference mechanism requires immense computational cost. Pre- vious works accelerate inference by leveraging layer or token cache techniques to reduce computational cost. However, these methods fail to achieve superior acceleration performance in few-step diffusion transformer models due to inef- ficient feature caching strategies, manually designed sparsity allocation, and the practice of retaining complete forward computations several steps in these token cache methods. To tackle these challenges, we propose a differentiable layer-wise sparsity optimization framework for diffusion transformer models, leveraging token caching to reduce token computation costs and enhance acceleration. Our method optimizes layer-wise wise sparsity allocation in an end-to-end manner through a learn- able network combined with dynamic programming solver. Additionally, our proposed two-stage training strategy eliminates the need for full-step processing in existing methods, further improving efficiency. We conducted extensive experi- ments on a range of transsoorrre models, including DiT-XL/2, PixArt-o, FLUX, and 1. Across these architectures, our method consistently improves efficiency without degrading sample quality. For example, on PixArt-a with 20 sampling steps, reduce computational cost by 54% while achieving generation metrics that surpass those of the original model, substantially outperforming prior approaches. These results demonstrate that our method delivers large efficiency gains while often improving generation quality.
INTRODUCTION
In recent years, diffusion models have made remarkable progress in the field of image generation.
Among them, the Stable Diffusion series (Rombach et al. 2022 Podell et al.) 2023 Tian et al. 2024
Esser et al. 2024) has achieved significant success in high-quality image generation (Zhu et al. 2024
2025a). This advancement is largely attributed to the effectiveness of diffusion probabilistic models
(DPM) Ho et al. 2020) and the powerful U-Net (Ronneberger et al. 2015) architecture, which allows high resolution synthesis with detail preservation. Additionally, some recent works (Peebles & Xie 2023b Chen et al. 2024b Tian et al. 2024) have explored the integration diffusion models with Transformer-basoa architectures, demonstrating outstanding performance. In particular, scaling laws have been leveraged to expand the model size of Transformers (Vaswani et al. (007)). further enhancing precision and generative quality, these large-scale models leverage improved expressivity and enhanced generalization, pushing the boundaries of generative artificial intelligence.
However, despite these advances, the substantial computational cost associated with diffusion models presents a significant challenge for real-world deployment. The inference of such large models requires extensive computational resources, which can hinder practical applications. Addressing this issue requires innovations in model acceleration techniques to enable broader accessibility and usability of diffusion-based generative models. Existing methods of diffusion model acceleration typically focus on sampler optimization (Song et al., 2020 Lu et al.] (022a), model pruning (Fang al. 2023b Zhang et al. 2024
Fang et al. 2023a), distillation (Yin et al. 2024b Luo et al., 2023
"Corresponding author.

---

Published as a conference paper at ICLR 2026
Salimans & Ho 2022), and feature caching (Selvaraju et al. 2024 Ma et al. 2024 Liu al. 2025a b).
Feature caching methods leverages temporal redundancy to reuse intermediate features, achieving significant speedups. They become popular in the field of diffusion model accelcration due to non- training diffusion model and easy integrating into the original inference pipeline. Previous methods cache and reuse coarse-grained, layer-level features, whereas token cache methods (Zou et al., 2025 2024; Zhang et al. 2025) reuse token-level features, achieving better acceleration performance. However, these approaches require manual sparsity allocations and hand-crafted schedules that preserve several full forward passes during denoising, which limits the acceleration potential of token-level feature caching.
To address these challenges. we propose Diffspaase, a learnable framework for optimizing layer-wise sparsity allocation in diffusion transformer models. Our approach dynamically determines the optimal sparsity configuration across all layers and inference steps, ensuring that the overall pruning rate is met in an end-to-end manner through a model-driven process. Moreover, parsS eliminates the need for complete forward computations in predefined steps required by existing methods, further enhancing efficiency.
Specially, our approach formulates the token cache optimization as dynamic programming- based sparsity allocation problem. We innovatively design a learnable sparsity cost predictor, which predicts cost matrix that quantifies the sparsity costs associated with target sparsity rates for all layers across every denoising step. Then we propose a dynamic programming approach to determine the optimal sparsity configurations for all layers cost while satisfying the required sparsity rate. Finally, we introduce a token selector that dynamically selects a specific proportion of tokens for reuse, leveraging the learned sparsity ratio to accelerate inference. To optimize the learnable sparsity cost predictor, we utilize a perceptual distillation loss that minimizes the degradation in generation quality. Furthermore, we introduce a two-stage training strategy that eliminates the need for complete forward computations in predefined steps required by existing methods while also improving accuracy. We have conducted extensive experiments on various transformer-bart -based baselines, and the pruning results outperform other SOTA pruning methods by a large margin. For example, pruning 54% of tokens yields an FID of 27.79, our method substantially better than the oftte-of-the-ar methods ToCa (28.35) and TaylorSeer (29.08), while achieving a higher speedup (1.91 x) on PixArt-a. These results underscore the practical effectiveness of our method. Our contributions are summarized as follows:
the relevant denoising steps, minimizing the overall sparsity
We propose DiffSparse, a differentiable approach to optimize layer-wise token sparsity in diffusion models sampling process. By integrating sparsity cost predictor, dynamic programming solver, and adaptive token selector, automates sparsity allocation and token reuse without manual heuristics.
We introduce two-stage training strategy that eliminates the need for predefined complete forward computations in several steps required by existing methods, fully unlocked the acceleration potential of token-level feature caching.
Extensive experiments on diverse foundation models prove that our method surpasses existing SOTA methods by a large margin, setting new efficiency-accuracy benchmarks.
2
RELATED WORK
Diffusion Transformer Models. The integration of transformers into diffusion models has signif- icantly advanced generative modeling. improving scalability and performance. Diffusion models, which generate data by iteratively denoising from noise distribution. Traditionally, diffusion models relied on CNNs, but recent studies demonstrate the effectiveness of transformers (Peebles & Xie 2023b Chen et al.) 2024b Tian et al.) 2024 Brooks et al. 2024; Chen et al., 2024a Wu et al. 2025). Diffusion Transformer (DiT) (Peebles & Xie 2023b replaces the U-Net backbone with a transformer, leveraging longerange dependencies and efficient scaling to achieve superior image generation. PixArt (Chen et al., 2024b) builds on this by introducing a hierarchical transformer architecture and novel noise schedule, excelling in high-resolution and text-to-iman to-image synthesis. Although diffusion transformer models have achieved great success, the substantial computational overhead from the iterative denoising process makes them inefficient for industrial deployment.
2

---

Published as a conference paper at ICLR 2026
Acceleration of Diffusion Models. Diffusion acceleration is critical research area focused on reducing computational costs and improving inference efficiency while preserving high-quality generation. Recent advancements can be categorized into sampler optimization (Song al. [202T
Lu al. 2022a b), model pruning (Fang et al. 2023b Zhang et al. (024), distillation (Salimans & Yin et al. 2024a), and feature caching (Li al. 2023 Ma al., 2024 Zhu al. 2025b).
Ho 2022
Sampler optimization reduces the number of denoising steps during inference using deterministic
or adaptive strategies to approximate the denoising process efficiently. Model pruning removes
redundant parameters and achieving speedups with structured pruning (Fang et al. 2023a). Other
strategies, such as Rectified Flow (Liu et al. 2022) and knowledge distillation (Yin al. 2024a)
accelerates inference by matching model outputs in fewer steps without quality loss.
Feature caching is particularly effective for DiT architectures. Methods such as FORA (Selvaraju et al. 2024) and A-DiT Chen et al. (202cc) reuse attention and MLP representations, while FitFttttttn ( (uaan et al.
A-DiTCChen
2024) further reduces redundancies in self-attention. Dynamic strategies like TeaCache (Liu et al. 2025a) estimate tinestep-dependent differences, and TaylorSeer (Liu et al. 2025b) introduced a "cache-then-forecass paradigm that predicts and updates cached features, though its advantage is most evident with long-range caching. SpeCa (Liu et al. (025c) further enhance the performance with speculative sampling. Complementary to these are token cache methods (Zou et al., 2025 2024 Zhang et al., 2025 You al., 2025), which apply fine-grained, error- error-guided token-wise caching to dynamically update features, achieving substantial acceleration without compromising quality. More discussion with existing methods are presented in Appendix.
In this paper, we introduce DiftSSaarse feature-caching approach for accelerating diffusion trans- former models. These models typically require only few dozen sampling steps and have seen growing adoption in industry. Unlike prior works (Zou et al. 2025 2024). DiftSparse employs a token-level cache within an end-to-end learning framework that casts model acceleration under fixed compression ratio as a layer-wise sparsity optimization problem across timesteps, eliminating the need for manually tuned sparsity or acceleration parameters. To address inefficiencies exist- ing approaches, which depend on predefined full-step computation schedules, we also propose two-stage training protocol that adaptively allocates computation where it is most needed.
METHOD
In this section, we start with a brief introduction to the diffusion transformer model and the token cache strategy. We then present the challenges of the existing token eaching approaches. Finally, we present our DiffSparse approach, which builds upon the token cache strategy for acceleration and optimizes the layer-wise token sparsity of diffusion transformer model in a learnable manner, enhancing accuracy while maintaining the sparsity requirement.
3.1 PRELIMINARY
Diffusion Models. Diffusion models are a class of generative models that construct a Markov chain of latent variables by progressively adding Gaussian noise to data samples and then reversing this process to synthesize new samples. Given an initial data sample zo, the forward diffusion process transforms the data through series of steps:
(( xx-1) = N/ 145 V 2e 8.1),
where / is the time step, { denotes a predefined variance schedule. After T steps, the data is nearly transformed into an isotropic Gaussian distribution, i.e., g(xT) R (0,
The reverse process is parameterized by a noise prediction network, which aims to recover the original data by iteratively removing the added noise, and modeled as:
Po (x xx)) = N' (10,,, 1), (o)r))))
where μe and Eg learned functions. Because the network is applied at each timestep in the multi-step denoising process, the repeated evaluations of the noise prediction network dominate the computational cost, accounting for the majority of the model's 's Moating-poin operations (FI.Oss).
3

---

Published as a conference paper at ICLR 2026
Diffusion Transformer. The Diffusion Transformer (Chen et al. 2024b) is novel architecture that synergizes the iterative refinement capabilities of diffusion processes with the representational power of transformers. In this framework, the input is represented as a set of tokens X E N where N denotes the number of tokens and D their dimensionality. The network architecture is composed of stacked blocks, each integrating three key components: a self-attention (SA) layer, a cross-attention (CA) layer, and a multi-layer perceptron (MLP) layer. The self-attention mechanism enables the model to capture long-range dependencies among tokens. In parallel, the cross-attention module facilitates the incorporation of conditioning information, enhancing the model's s ability to generate contextually relevant outputs. The subsequent MLP further refines these token representations through non-linear transformations.
significant advantage of the Diffusion Transformer lies in its ability to iteratively refine token representations during the denoising process, leading to improved sample quality. This layered approach allows the model effectively balance global context and local details, thereby offering enhanced performance in complex generative tasks.
Froken-Wise Feature Caching Approach. Prior work (Ma et al., 2024) has demonstrated that features adjacent timesteps exhibit high similarity, leading to significant redundancy. To exploit this redundancy for computational efficiency, previous approaches (Ma et al. [0244] Wimbauer et al., 2024) have introduced caching mechanisms that reuse features accelerate processing. The token- wise feature caching approach (Zou et al. 2025) operates at a finer granularity by caching features at the individual token level, enabling more effective exploitation of the redundancy.
Token-wise feature caching mechanism begins by computing and storing the intermediate token features X = (.o, .i..... ................................ 2 -1) from each self-attention, cross-attention, and MLP layer into cache C' at the initial timestep t. In subsequent timesteps, a predefined cache ratio R. determines the proportion of tokens reused from the cache C for each layer each timestep. The R selected tokens based on token importance rank, denoted as ICache. will bypass re-computation by reusing their cached values, while the remaining tokens ICompute = ] \ [Cache are recomputed. For given layer f. the computation for each token ă; is formulated as:
F(a:) 7 ((x)) ) + (1 - C+)
where Y/2 O for i; E [Cache and % 1 for át; E Coopptee To mitigate error accumulation from reused features, the cache dynamically updated for tokens in /Compute via:
C(a,) < F(i:).
This token-wine wise feature caching approach effectively reduces redundant computations by leveraging the high similarity of features across adjacent timesteps, thus significantly accelerating the inference process while maintaining robust feature representations.
Challenges in Existing Token Caching Approaches. While token caching methods (Zou et al. 2025) have shown great promise in speeding up diffusion transformers, key limitations remain. First, they require manually setting a reuse sparsity rate for each layer at every timestep. resulting in large. hard-to-tune parameter space. This manual process hampers performance and scalability. A learnable adaptive sparsity strategy could unlock further gains. Second, current methods still depend on a full-step design (several steps without caching) to maintain generation quality. However, this compromises the efficiency of basen operations. Replacing this with dynamic caching tailored to diffusion transformers can better balance quality and speed. In this paper, we propose an intelligent framework that jointly learns optimal sparsity across layers and removes the reliance on full-step computation, significantly improving both performance and flexibility.
3.2 DIFFSPARSE APPROACH
To automate per-layer sparsity selection and remove the reliance on full-step designs, we propose DiffiSparse, an efficient token caching framework for diffusion transformers. DifiSparse learns layer-wise sparsity end-to-end by combining a learnable sparsity cost predictor with a dynamic programming solver to find optimal sparsity configurations across layers and denoising steps. It also adopts a two-stage training scheme that gradually replaces full computation steps with cache-based ones, improving efficiency without sacrificing performance. As illustrated in Figure]] DiffSparse
4

---

Published as a conference paper at ICLR 2026
Full Steps
Sampling
DIT
Xo
LPIPS Loss
Token Cache
x CT
Blook2
Block
Blook
DIT
X
Compute
Reuse Frozen
Token
Token Selector
Predicted Masks
Tumable
Masks
STE
Sparsity Cost
Dynamic
0
0000 ..000))
Programming Solver
+0
(11,111 0.00)
Predictor
(1111, ... 1,1)
Cost Matrix
f
Layer Sparsity
CarslichteMMall
[5]
R
x
Figure DiffSparse uses a learnable sparsity-cost predictor and dynamic programming to learn per-layer sparsity under target ratio R. We generate binary masks from the chosen sparsity maps and candidate masks. token selector reuses features from previous diffusion steps to skip unimportant tokens and speed sampling. To enable gradient flow through the binary masks, apply Straight- Through Estimation (STE) and train our model using full-step sampling targets with LPIPS loss.
comprises three components: token selector, sparsity cost predictor, and dynamic programming solver. The cost predictor estimates a cost matrix representing the sparsity cost for various predefined rates across all layers and denoising steps (excluding the first). The dynamic solver then identifies the optimal sparsity pattern under a global sparsity constraint R. Based on this, the token selector determines which tokens reuse and which recompute at each layer. Training is guided by perceptual distillation loss, integrated into a two-stage training pipeline for effective learning.
Token Selector. We employ Token Selector that assigns each token è; an importance score used decide which tokens freshly computed and which remain cached. The score is a composite, layer-wise quantity of the form:
S(i)) = B A s (5.))
where each sy(4;) is scalar signal capturing a different criterion (for example, self-attention influ- ence, cross-attention focus, cache-reuse frequency, etc.), and (N weighting hyperparameters that balance these criteria. The operator B(-) optional and denotes a spatial bonus operation that promotes spatially uniform coverage of selected tokens (implemented, e.g., by boosting tokens that are local maxima within a & neighborhood). Other choices for B are possible (e.g. smooth kernels or distance-basee adjustments).
Given the per-token scores (ā;; in a layer with N tokens, we sort tokens by descending score and select the top K tokens according to predefined sparsity ratio R. We emphasize that our contribution orthogonal to any particular token-ranking heuristic; the choice of scoring components (e.g. self-attention influence, cross-altention terms, spatial bonus) is optional and can be replaced by alternative ranking methods. Detailed descriptions and comparisons of specific token-ranking strategies provided in the Appendix A.5.1 Empirically, our allocation scheme yields consistent gains across different token-rankin feen-ranning methods (see Table
Learnable Sparsity Cost Predictor. We propose a learnable sparsity cost predictor to adaptively determine layer-wise sparsity in diffusion transformers (DiTs) while balancing inference efficiency and computational cost. Given a DiT with L layers operating over T denoising timesteps, our goal to generate binary mask M E (0,1)* 1) for each layer / and timestep that selects K2 tokens
5

---

Published as a conference paper at ICLR 2026
for full computation and reuses features for the remaining N - Kr, tokens. This is formalized as constrained optimization over a candidate sparsity set S, where |SS] denotes the number of predefined sparsity configurations. For layer containing N tokens, let S denote the set of sparsity rates, each a value between O and I, at which we retain a corresponding fraction of tokens. For instance, if 0.70, 0.75, 1.0), which
N = 256 and we choose step size of 32 tokens, corresponds to retaining {0, 64, 128, 192, 256} tokens, respectively. Our objective is to learn the relative cost of applying different sparsity rates across layers. Our experimental results (Table 4 demonstrate that the learned sparsity predictor generalizes across resolutions, so that sparsity allocation trained at low resolution remains effective at higher resolutions.
obtain S {0, 0.25, 0.50,
We implement the sparsity cost predictor using (T' x L) x |S] learnable parameters, where T is the number of timesteps, is the number of layers, and |S|| the size of the candidate sparsity set. The predictor outputs a normalized cost matrix C' E R ( L)x where each entry Cic.n. quantifies the cost of applying sparsity configuration E S to layer at timestep t. We minimize the cumulative cost while ensuring the total sparsity meets a predefined overall pruning rate R. The sorted token set X E RNND enables efficient mask selection by prioritizing tokens with high scores.
Importantly, The cost predictor's size depends only on T, T., and [S], not on token Lekennse length N. Empirically, we found that simply increasing [S| beyond moderate size yields diminishing or negative returns ( and experiments show the learned cost predictor transfers across resolutions (Table 4 demonstrating scalability to high resolutions and robustness to token-length variation.
Dynamic Programming Solver. To determine the optimal sparsity configuration while satisfying a global sparsity constraint, we employ dynamic programming approach minimize the overall cost across layers. Formally, we define the state function:
i 8 =r,
F(G,)) = min Σ
s.t. Σ
(6)
(ech i=1
1=1
where F(i, r) represents the minimum achievable cost when assigning sparsity levels to the first i layers under a total sparsity constraint r. The recursive formulation is given by:
F(i,r)
min sES,
(F( - s) + Ci,a).
Here, the transition considers all possible sparsity levels s that can be allocated to layer ensuring that the total sparsity constraint is maintained. The algorithm iteratively computes F((,r) for i T and 0, .., R, followed by backtracking step to reconstruct the optimal sparsity allocation, where R R T. This approach operates with a time complexity of O(LL -T) [S)), making it computationally feasible for practical deep learning scenarios. To reduce the number of redundant state computations and lower overall complexity, we implement pre-pr pre-pruning strategies. For example, when target sparsity ratio R = 43%, |S| = 5, T = 20, and L = 28, it requires about 4 hours (including DP optimization and fine-tuning) of total training time. The DP solver runs in approximately Rx0 seconds for the configurations reported, but it is not executed at inference time. At inference, the model only uses the precomputed masks. Since the direct conversion of the predicted cost matrix C' to a discrete mask M is non-differentiable, we utilize the Brraigh-TThrouee Estimator (STE) (Jng et al. 2016) to approximate the gradients of the discrete mask with respect to the cost predictions. This approach facilitates end-to-end optimization of the sparsity cost predictor.
Training Loss. To guide the optimization of the pruned Diffusion Transformer, we employ the Learned Perceptual Image Patch Similarity (L.PISS) loss (Zhang et al. 2018) as perceptual dis- tillation loss. In our framework, the original model prior to token pruning serves as the teacher network, while the pruned model is treated as the student network. Both models generate outputs via a multi-step sampling process inherent to diffusion models.
Let 2o and 2á denote the multi-step sampling outputs from the teacher and student networks, respectively. The LPIPS loss is then defined as:
(8) which measures the perceptual similarity between the outputs. During training, gradients are back- propagated solely through the student network, as the teacher network's 's parameters detached (i.e.,
LiPWS = LPIPS 26)),

---

Published as a conference paper at ICLR 2026
its gradients are not computed). This setup ensures that the student model is effectively distilled to mimic the perceptual characteristics of the teacher model, thereby achieving acceleration through token pruning while preserving output quality.
Two-Stage Training Strategy. matrices for full-stee step positions and layer sparsity components. In the first stage, we follow (Selvaraju et al. 2024 Zou et al, 2025) to preset y full-step positions and independently optimize the step cost matrix Cy E RT'² encoding temporal sparsity decisions and the layer sparsity cost matrix C E ((LxT)x [S] governing token retention per layer. We first solve Cj via dynamic programming identify [Tyl optimal full-step positions with minimal cumulative cost. For these selected steps, we warm-start layer sparsity optimization by subtracting ō from the predicted costs:
We propose a two-stage training framework to optimize the cost
c(tce) < c(tile) - E I E [L,.... L) = N.
(9)
This strategy preserves inter-layer cost ranking while leveraging full-step error correction capabilities.
In the second stage, integrate step and layer costs by modifying layer sparsity entries using Equation | The unified cost matrix is then fine-tuned to systematically redistribute FLOPs across sampling steps. Unlike existing methods (Selvaraju et al. 2024 Zou et al.) 2025) that rigidly enforce full steps for noise correction, our approach dynamically optimizes sparsity patterns through differentiable cost interaction. The pseudocode is provided in the supplementary materials.
4
EXPERIMENTS
4.1 EXPERIMENT SETTINGS
Model Configurations. We conduct experiments on four widely used DiT- models across various generation tasks: (1) PixArt-a with 20 DPM Solver++ (Lu et al. (022b) steps and FLUX.1- schnell (Labs) 2024) with steps for text-to-image generation and (2) DiT-XL/2 with 50 DDIM (Song et al. (021) steps for class-conditional image generation. (3) 1an2.1-1.3B (Wan et al. 2025 with 25 now-matching sampling steps for text-to-video generation. We define the candidate set as the range from 0 to with interval of 0,25, yielding |S| = 5 token sparsity candidates. More details of the implementation are provided in the supplementary material.
Training. For PixArt-a (Chen et al., 2024b) and FLUUX.1-ccheell (Labs 2024), we train the learnable sparsity-cost predictor on 10,000 captions randomly sampled from the COCO (Lin et al. 2014) train dataset. For
(Peebles & Xie 2023a), we use 10,000 ImageNet (Deng et al. 2009) train category indices, and for I we sample 10,000 captions from WebVid-1OM (Bain al. 2021) for training. During training we information, which do not overlap with the evaluation set.
image data, only captions class-conditioning
We leverage the layer sparsity configuration in the token-cache- based model (Zou al., 2025) to initialize our sparse router training. All the models are trained with AdamW optimizer. The sparsity cost predictor is trained in two stages. For the first stage, the layer sparsity cost component is optimized for 1 epoch with learning rate of = 1.0, while the step cost component is trained separately using 1) = 0.01 to capture temporal patterns across denoising steps. For the second stage, we integrate the step cost into the layer-wish costs with 8 10 and then fine-tuned for epoch with 1 = 0.1 to optimize layer sparsity allocation. Training requires approximately 4-10 hours on AMD MI250 GPUs with 80GB memory per experiment.
Evaluation. For rext-to-image generation, we evaluate on the COCO dataset (Lin et al. 2014) using 30,000 samples at 256 x 256 resolution and PartiPrompts Yu et al. (022) with 1,632 samples. Image quality is quantified by FID-30k (Heusel et al., 2017). which compares generated images against originals, while text-image alignment is measured by two complementary metrics: CLIP. Score (computed with CLIP-ViT-Large- 14 (Hessel et al., 202I)) and Image Reward (Xu et al. 2023), metric shown to more accurately reflect human preferences. For class-conditional image generation, 50,000 images at 256 x 256 resolution are generated from 1,000 ImageNet (Deng et al.) 2009) classes and evaluated using the FID-50k metric. We evaluate text-to-video generation using the VBench framework on 950 prompts, generating 4,750 videos at 256 x 256 resolution, each lasting seconds frames per second, and assess them across 16 metrics.
7

---

Published as a conference paper at ICLR 2026
4.2
MAIN RESULTS
Results on Text-to-lmage Generation. We compare pirtsS with existing methods under identical sparsity budgets. FORA, DeepCache (CVPR'24) and TaylorSeer (ICCV'25) evaluated with cache interval = 2, while DiCache, ToCa (ICLR'25) and DuCa are tested using their respective optimal configurations. Table shows that DiftSparse delivers both faster inference and improved generation quality compared with existing methods. At roughly 1.74x speed-up. existing methods suffer degraded image quality, while DiffSparse achieves a strong FID of 26.91 (vs. TaylorSeer's 29. 39.08 and ToCa's 28.35). This corresponds to relative +5.1% improvement in FID of parss over ToCa. Pushing further, parts attains 1,91 acceleration while producing an FID that surpasses the original (full) model. This improvement stems from learned sparsity schedule that accelerates convergence of the generated image distribution and improves visual fidelity, while preserving semantic aligmment with the conditioning signal. We provide additional text-to-image comparisons Appendix A.6 and also present more qualitative visual comparison in Appendix A.7
Table Results of text-to-image generation on MS-COC02017 with PixArt-a and 20 DPM++ steps.
Method
MACs (T)4
Speedupt
FID-30k4
CLIPT
PixArt-a (Chen
2024b)
2.86
1.00x
28.20
0.163
50% steps
1.43
1.74x
37.57
0.158
FORA (N = 2)
Selvaraju et al.
2024
1.43
1.64x
29.67
0.164
DeepCache (N 2) (M
2024
1,48
1.61 X
29.61
0.163
DiCache (Bu et al.
2025
1,63
1.77
28.19
0.164
ToCa (Zou al. 2025
1.64
1.75x
28.35
0.164
DuCa iZou el al. 2024
1.78x
27.98
0.164
TaylorSeer (Liu ct
2025b
1.57
1
29.08
0.163
DiftSparse [R 43%)
1.64
1.74x
26.91
0.164
DiffSparse (R 54%)
1.30
1.91
27.79
0.164
Table Results class-conditional generation with DiT-XL/2 and 50 DDIM steps on ImageNet.
Method
MACs (T)
Speedupt
FID.
sFID
Precision
Recall t
DDIM-50 steps
11.44
1.00x
2.26
4.29
0.80
0.60
DDIM-40 steps
9.14
1.24x
2.39
4.28
0.80
0.59
DDIM-25 steps
5.73
1.96x
3.01
4.60
0.79
0.58
DDIM-20 steps
4.58
2.42x
3.48
4.64
0.79
0.56
FORA
4.13
2.12x
3.88
6.74
0.79
0.56
ToCa
4.97
2.09
3.05
4.70
0.79
0.57
DuCa
4.94
2.10x
3.04
4.70
0.79
0.57
DiffSparse
4.97
2.07
2.81
4.61
0.80
0.59
compares faster sampler DDIM with fewer steps, FORA, ToCa, DuCa and DiftSparse. Our method achieves better speed- accuracyy balance by reallocating computation the most important layers. At the same acceleration ratio, Dift'Sparse improves the FID from 3.05 to 2.81, outperforming ToCa by 8% at 2.07x acceleration. demonstrating its ability to preserve detail and improve image fidelity in diffusion model acceleration.
Results on Class-Conditional Image Generation. Table 2
Table 3: Comparison in text-to-video generation for 1- with 20 sampling steps on VBench.
Table 4;
Comparison PixArt-a using 20 sampling steps at 512x512 resolution.
Method
MACs (T)
Speedup
VBench f
Method
MAC's (T)
FID |
CLIP
Wan 2.1 1.3B
43.866
1.00x
43.82
50% steps
21.933
1.86x
43.14
PixArt-a:
10.851
21.95
0.164
DuCa (R = 54%)
20.332
1.69x
43.56
50% steps
5.426
25.05
0.163
DuCa (R = 59%)
18.124
1.68x
43.30
ToCa
5.993
23.02
0.165
DiffSparse
18.124
2.05x
43.83
DiffSparse
5.986
22.42
0.165
Results Text-to-Vide Generation. Table presents comparison between DiffSparse and
DuCa (Zou et al. 2024) on Wan2.1-1.3B (Wan et al. (025) using 20 sampling steps. The methods
comprehensively evaluated across 16 aspects defined in VBench (Huang al. 2024). We adopt
8

---

Published as a conference paper at ICLR 2026
DuCa's bosm- token ranking compatible with FlashAttention (Dao et al. 2022 for faster inference. DiftSparse achieves the highest VBench score while minimizing computational cost and inference time. At the same compression ratio, it delivers greater speedup by skipping partial layers with zero sparsity, and its adaptive, layer-wise sparsity allocation preserves model quality.
4.3 ABLATION STUDIES
Comparison of Two Stage Training. In this work, we adopt a two-stage training strategy. The first stage independently trains cost matrices for full-step and layer sparsity. In the second stage, the learned full-step cost merged into the layer sparsity optimization, and the layer sparsity is subsequently fine-tuned. This design enables the model to initially leverage the full--teep correct errors and to learn layer sparsity cost values, followed by a gradual reduction of the full-step influence. Results proved that two-stage approach achieves better performance, with FID of 26.91 compared to 27.40 from the single-stage baseline.
Comparison of Important Scores. Table compares three importance scores: attention (Equation 10 cosine similarity and the & norm. The cosine similarity is computed between the current input token and cached tokens. The lg norm is the norm value input tokens. The attention- based score attains the best FID, followed by the similarity measure, which captures token redundancy effectively, Norm-based scoring introduces noise and performs worst, confirming that accurate importance estimation is critical for optimal token selection.
Table 5: Ablation study on token importance metrics.
Table 6: Ablation study on distillation loss func- tions.
Method
Base.
w/ DiftSparse
Method
FID 4
CLIP t
Norm
29.05
28.89 (-0.16)
L2
27.68
0.164
Similarity
29.00
28.07 (-0.93)
SSIM LPIPS
27,46
0.164
Attention
28.35
26.91 (-1.44)
26.91
0.164
Table 7: Ablation study of sparse interval.
Table 8: Ablation of warm-start strength 6.
Interval
[S]
FID + CLIP f
FID
CLIP 1
0.1
11
27.96
0.163
0.125
9
27.91
0.163
0
27.40
0.163
0.25
5
26.91
0.164
27.01
0.164
0.5
3
27.54
0.164
10
26.91
0.164
1.0
2
28.22
0.162
20
26.95
0.164
Comparison of Training Losses. We compare L2, SSIM (Wang al. 2004), and LPIPS losses in Table LPIPS outperforms the others, yielding the best FID. L2 loss penalizes pixel-wise squared errors and often produces overly smooth images that lack fine details. SSIM enforces local structural similarity but may over-penalize perceptually good images that differ spatially from the original. By measuring distances in learned perceptual feature space, LPIPS avoids these pitfalls and better preserves image quality during training.
We distribute sparsity uniformly across layers by token count and evaluate different granularity settings in Table 7 granularity of 0.5 yields minimal within-layer variation, which hinders convergence, while 0.5 limits the range of sparsity choices. The optimal granularity is 0.25, producing sparsity rates [0, 0.25, 0.50, 0.75, 1.0] (corresponding to candidate token counts of [0, 64, 128, 192, 256] for sequence length of 256) and delivering the best performance.
Comparison of Sparse Intervals.
Generalization on Higher Resolution Models. As the token sequence length increases with image resolution, peak memory usage during training grows substantially, even though the size and computational cost of our cost matrix remain unchanged. This makes direct training at very high resolutions impractical. To address this, we investigate whether a sparsity predictor trained at lower resolution can be transferred to higher resolution without retraining. As shown in Table the sparsity predictor learned at 256 x 256 resolution achieves a lower FID than ToCa on 512 x 512 images while maintaining comparable CLIP-Score to the original PixArt model. These results demonstrate that our method generalizes effectively to higher resolutions, enabling model acceleration with limited memory and training cost.
9

---

Published as a conference paper at ICLR 2026
Compared with GA Search. We compared DiftSp against traditional search methods such as random search and genetic algorithms and found that in the vast sparsity space they underperform. After 1,000 iterations 500 images, these methods yield FID scores of 28.34 and 27.94, respectively, compared with 26.91 for DiftSparse. Moreover, they require about 16 hours, whereas parse completes training in roughly hours. These results show that differentiable learning framework discovers more effective layer-wise sparsity allocations and delivers superior acceleration.
Comparison of Warm-Start Constant 5. Algorithm uses a warm-start constant = 10 for the two-stage optimization. Intuitively, injects the Stage-1 prior (the timesteps selected to remain full-step) into Stage 2 by lowering the cost of the "full"" candidate at those timesteps. In effect, larger more strongly encourages preserving full computation at the Stage-I selected steps. To quantify this effect we evaluated 0 {0, 5, 10, 20). Table 8 reports the results on PixArt-a with T = 20. A moderate warm-start (Ō = 10) recovers most of the benefit, while å = 0 (no warm-start) removes the Stage- 1 prior and yields noticeably worse performance.
4.4 QUALITATIVE ANALYSIS
Visualization of Generated Images. We provide detailed visual comparisons among our proposed method, ToCa, and the original PixArt-a across various sparsity ratios. in Figure 2 reveals that DiftSparse consistently maintains high fidelity, even under aggressive pruning conditions. Moreover, DiftSparse effectively preserves the semantic content of the text prompt, ensuring that the generated images remain closely aligned with the original descriptions. In contrast, baseline methods exhibit noticeable degradation in both visual quality and text-image alignment higher pruning ratios, further highlighting the strength and efficiency of DiftSparse.
Original
ToCa
Tayloreer 1.8
DiftSparse 1.74
DiffSparse 1.91
Two people standing on the beach with kite.
A guy standing in the grass is ready to throw something.
A goat with horns is standing in a grassy field.
Figure 2: Comparison of our method with the baseline (PixArt-a with DPM-Solver++ using 20 steps) and existing methods under different acceleration rates.
5 CONCLUSION
We introduce learnable token sparsity allocation framework to accelerate diffusion transformers. By formulating sparsity allocation as a dynamic programming problem and employing a two stage training strategy, method substantially reduces computational cost while preserving generative quality. Extensive experiments across various foundation models and datasets demonstrate improved acceleration ratios without compromising image quality of our method.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Max Bain, Arsha Nagrani, Gül Varol, and Andrew Zisserman. Frozen in time: A joint video and image encoder for end-to-end retrieval. IEEE International Conference on Computer Vision, 2021.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. URL httttp:://ppeaa.... cot.ttaiaccooreeerrcccic
1d- simulators
Jiazi Bu, Pengyang Ling, Yujie Zhou, Yibin Wang, Yuhang Zang. Dahua Lin, and Jiaqi Wang. Dicache: Let diffusion model determine its cache. arXiv preprint arXiv:2508.1733 2025.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, al. Pixart-o: Fast training of diffusion transformer for 2023.
photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.0042
Junsong Chen, Yue Wu, Simian Luo, Enze Xie, Sayak Paul, Ping Luo, Hang Zhao, and Zhenguo Li. Pixart-0: Fast and controllable image generation with latent consistency models. 2024a.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-ine -o--gee--mgeeeeee-e synthesis. In International Conference on Learning Representations, 2024b.
Pengtao Chen, Mingzhu Shen, Peng Ye, Jianjian Cao, Chongjun Tu, Christos-Savee -Savvas Bouganis, Yiren Zhao, and Tao Chen. 5-dit: A training-free acceleration method tailored for diffusion transformers. arXiv preprint arXiv:2406.01125, 2024c.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Aur Rudra, and Christopher R'c. Flashattention: Fast and memory-efficient exact attention with io-awareness. ArXiv, abs/2205, 2022. URL
httpss//aapi. wttp::/sssmmttiiiccchheee httpi//a.....gg/oorssss::::.:... :001:0pPP..1:29111111111111
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. leee, 2009.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024.
Gongfan Fang, Xinyin Ma, and Xinchao Wang. Structural pruning for diffusion models. In Advances in Neural Information Processing Systems, 2023a.
Gongfan Fang, Xinyin Ma, and Xinchao Wang. Structural pruning for diffusion models. arXiv preprint ar arXiv:2305. 10924, 2023b.
Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference- free evaluation metric for image captioning. arXiv preprint arXiv:2104.082 000018888 2021.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-663 2020.
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P. Kingma, Ben Poole, Mohammad Norouzi, David J. Fleet, and Tim Salimans. Imagen video: High definition video generation with diffusion models, 2022. URL https:://////....t.. http:////xxvvooo org/abs// 1221 0. 02303
11

---

Published as a conference paper at ICLR 2026
Ziqi Huang, Yinan He, Jiashuo Yu, Fan Zhang, Chenyang Si, Yuming Jiang, Yuanhan Zhang, Tianxing Wu, Qingyang Jin, Nattapol Chanpaisit, et al. Vbench; Comprehensive benchmark suite for video generative models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 21 2180-22188, 2024.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizatioon with gumbel-softmax. arXiv preprint 1r1:::1611.0/ 144, 2016.
Black Forest Labs. Flux.
2024.
Senmao Li, Taihang Hu, Fahad Shahbaz Khan. Linxuan Li, Shiqi Yang, Yaxing Wang, Ming-Ming Cheng, and Jian Yang. Faster diffusion: Rethinking the role of unet encoder in diffusion models. 2023.
arXiv preprint
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision- ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Purt V 13. pp. 740-755. Springer, 2014.
Feng Liu, Shiwei Zhang, Xiaofeng Wang, Yujie Wei, Haonan Qiu, Yuzhong Zhao, Yingya Zhang, Qixiang Ye, and Fang Wan. Timestep embedding tells: It's time to cache for video diffusion model. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 7353-7363, 2025a.
Jiacheng Liu, Chang Zou, Yuanhuiyi Lyu, Junjie Chen, and Linfeng Zhang, From reusing to forecasting: Accelerating diffusion models with taylorseers. arXiv preprint MXXX::5533306223 2025b.
Jiacheng Liu, Chang Zou, Yuanhuiyi Lyu, Fei Ren, Shaobo Wang, Kaixin Li, and Linfeng Zhang. Speca: Accelerating diffusion transformers with speculative feature caching. arXiv preprint arXiv:2509.1 11628, 2025c.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 030000 2022.
Ziming Liu, Yifan Yang, Chengruidong Zhang, Yiqi Zhang, Lili Qiu, Yang You, and Yuging Yang. Region-adaptiie sampling for diffusion transformers. arXiv preprint arXiv:2502. 10389, 2025d.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu, Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint 27X0092211.01095. 2022b.
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthesizing high-resolution images with few-step inference. arXiv preprint arX1v:2310.00337. 2023.
Siwei Lyu. Deepfake detection: Current challenges and next steps. pp. 1-6, 2020.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. Deepeache: Accelerating diffusion models for free. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15762-15772, 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp.
2023a.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195 4195-4205, 2023b.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolutee image synthesis. arXiv preprint arXin:2307.011952 2023.
12

---

Published as a conference paper at ICLR 2026
Junxiang Qiu, Shuo Wang, Jinda Lu, Lin Liu, Houcheng Jiang, Xingyu Zhu, and Yanbin Hao. Accelerating diffusion transformer via error-optimized cache. arXiv preprint arXiv:2501.1924. 2025.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp
Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assine assisted Cievventioo-MICCOE 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 231. Springer, 2015.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems, 35:36479-36494, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005// 2022.
Pratheba Selvaraju, Tianyu Ding, Tianyi Chen, Ilya Zharkov, and Luming Liang. Fora: Fost-forward caching in diffusion transformer acceleration. arXiv preprint arXiv:2407.01425, 2024.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2021.
Wenzhang Sun, Qirui Hou, Donglin Di, Jiahui Yang, Yongjia Ma, and Jianxun Cui. Uniep: A unified caching and pruning framework for efficient video generation. arXiv preprint arXiv:2502.0433 043993, 2025.
Yuchuan Tian, Zhijun Tu, Hanting Chen, Jie Hu, Chao Xu, and Yunhe Wang. U-dits: Downsample tokens in u-shaped diffusion transformers. arXiv preprint arXiv:2405.02730 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Team Wan, Ang Wang, Baole Ai, Bin Wen, Chaojic Mao, Chen-Wel Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang. et al. Wan: Open and advanced large-scale video generative models. arXiv preprint arXiv:2503.203 2025.
Zhou Wang, Alan Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from eтror visibility to structural similarity. IEEE transactions image processing, 13(4):600-612, 2004.
Felix Wimbauer, Bichen Wu, Edgar Schoenfeld, Xiaoliang Dai, Ji Hou, Zijian He, Artsiom Sanakoyeu, Peizhao Zhang, Sam Tsai, Jonas Kohler, et al. Cache me if you can: Accelerating diffusion models through block caching. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 6220, 2024.
Wei Wu, Xi Guo, Weixuan Tang, Tingxuan Huang, Chiyu Wang, and Chenjing Ding, Drivescape: High-resolution driving video generation by multi-view feature fusion. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 17187-17196, 2025.
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: learning and evaluating human preferences for rext-to-image generation. In Proceedings of the 37th International Conference on Neural Information Processing Systems, pp. 15903-15935. 2023.
13

---

Published as a conference paper at ICLR 2026
Tianwei Yin, Michaël Gharbi, Taesung Park. Richard Zhang. Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. In NeurIPS, 2024a.
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6613-6636666 6623. 2024b.
Haoran You, Connelly Barnes, Yuqian Zhou, Yan Kang, Zhenbang Du Wei Zhou, Lingzhi Zhang, Yotam Nitzan, Xiaoyang Liu, Zhe Lin, et al. Layer-and timestep-adaptive differentiable token compression ratios for efficient diffusion transformers. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 18072-18082, 2025.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al. Scaling autoregressive models for content- rich text-to-image generation. arXiv preprin! 220.:208660888 2(3):5, 2022.
Zhihang Yuan, Hanling Zhang, Lu Pu, Xuefei Ning. Linfeng Zhang, Tianchen Zhao, Shengen Yan, Guohao Dai, and Yu Wang. Ditfastatm:: Attention compression for diffusion transformer models. Advances in Neural Information Processing Systems, 17:1196-1219, 2024.
Dingkun Zhang, Sijia Li, Chen Chen, Qingsong Xie, and Haonan Lu. Laptop-diff: Layer pruning and normalized distillation for compressing diffusion models. arXiv preprint arXiv:2404. 1098, 2024.
Evelyn Zhang, Jiayi Tang, Xuefei Ning, and Linfeng Zhang. Training-free and hardware-friendly acceleration for diffusion models via similarity- token pruning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 9878-9886, 2025.
Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as perceptual metric. 2018. URL
1801103922
Wangbo Zhao, Yizeng Han, Jiasheng Tang, Kai Wang, Yibing Song, Gao Huang, Fan Wang, and Yang You. Dynamic diffusion transformer. arXiv preprint arXiv:2410.03456 2024.
Haowei Zhu, Ling Yang, Jun-Hai Yong, Hongzhi Yin, Jiawei Jiang, Meng Xiao, Wentao Zhang, and Bin Wang. Distribution-awwit aware data expansion with diffusion models. Advances in Neural Information Processing Systems, 37:102768-1001 102795, 2024.
Haowei Zhu, Tianxiang Pan, Rui Qin, Jun-Hai Yong, and Bin Wang. Recon: Region-controllable data augmentation with rectification and alignment for object detection. arXiv preprint arXiv:2510.1578.. 2025a.
Haowei Zhu, Dehua Tang, Ji Liu, Mingjie Lu, Jintu Zheng, Jinzhang Peng. Dong Li, Yu Wang. Fan Jiang, Lu Tian, et al. Dip-go: A diffusion pruner via few-step gradient optimization. Advances in Neural Information Processing Systems, 37:92581-92604 2025b.
Chang Zou, Evelyn Zhang, Runlin Guo, Haohang Xu, Conghui He, Xuming Hu, and Linfeng Zhang. Accelerating diffusion transformers with dual feature caching. arXiv preprint 1rX1v:2412.18911 2024.
Chang Zou, Xuyang Liu, Ting Liu, Siteng Huang, and Linfeng Zhang. Accelerating diffusion trans- formers with token-wise feature caching. In The Thirteenth International Conference on Learning Representations, 2025,
14

---

Published as a conference paper at ICLR 2026
A
APPENDIX
A.1 ETHICAL STATEMENT
Generative models have shown impressive capabilities in content creation (Chen al. al. 2022). but their high inference costs hinder rapid deployment. Our method offers efficient acceleration strategy for diffusion models, achieving near-lossless speedup without retraining and maintaining compatibility with various architectures. This generalizability makes it well-suited for fast deployment on mobile and edge devices.
2023 Rombach
However, generative models pretrained large-scale internet data may reflect inherent social biases and stereotypes. There is also potential for misuse, such as in DeepFake (Lyu) 2020) creation, which can cause serious societal harm. As the cost of generation decreases, the risk of irresponsible use increases. Therefore, it's essential to establish regulations, foster a well-govern community, and provide clear usage guidelines ensure the responsible application of generative technologies.
A.2 REPRODUCIBILITY STATEMENT
To support reproducibility, we provide detailed pseudocode for the proposed method (Appendix A.5.3). full training and evaluation protocols including all hyperparameters and optimizer settings (Section 4.1 and Appendix A.5), dataset descriptions and preprocessing steps, and the computing environment for experiments (Section 4.1) Where applicable, we report evaluation metrics and include instructions sufficient to reproduce the experimental pipelines described in the main text and appendices. We plan to release anonymized code, trained model checkpoints, and exact run scripts upon paper acceptance to facilitate full replication.
A.3 THE USE OF LARGE LANGUAGE MODELS (LLMs)
We did not rely on LLMs for research ideation, experiment design, data analysis, or the generation of technical content. Any use of LLMs was limited to minor, general-purpot editorial assistance (proofreading, grammar, phrasing, and formatting suggestions); all such suggestions were reviewed and revised by the authors. The paper's 's conceptual contributions, algorithms, experiments, results, and conclusions were produced solely by the authors, and no LLM is credited contributor.
A.4
MORE DISCUSSION WITH EXISTING WORKS
A.4.1 MORE RELATED WORKS
Diffusion Transformer Models. Recent work has improved the efficiency and scalability of transformer-based diffusion models. Hybrid CNN-tranforrmer architectures (Saharia et al., 2022) combine local inductive biases with global attention, and transformer-based video generation (Ho) ctt al. 2022) demonstrates strong temporal modeling. These results establish transformers as versatile backbone for diffusion, motivating efforts on optimization, faster inference, and stronger conditional generation. Nevertheless, the iterative denoising loop still incurs substantial computational overhead that limits industrial deployment.
Acceleration of Diffusion Models. Several recent methods target inference cost directly: EOC leverages prior knowledge to improve caching (Qiu et al. 2025), while designs such as UniCP and RAS further boost efficiency (Sun al.) 2025 Liu et al. 2025d). DyDiT (Zhao et al. 2024) accelerates inference by skipping unimportant tokens and slimming per-layer width, whereas Dif'Sparse reduces compute by reusing cached features. The two strategies are complementary and be combined for larger speedups. DiffSparse computes token importance with a training free compositional intostttionallaattntiion score and lears compact layer-wise predictor, leading to much faster convergence (on the order of 10 iterations versus DyDiT's ~ 2 x 10" fine-tuning steps). Moreover, by optimizing a global T Ttep objective with a pwwaniiiiporramminggg. solver, DiftSparse coordinates sparsity across timesteps and layers and is validated across multiple architectures and generation tasks.
In our main configurations, the differentiable optimization requires xs 4 hours of training versus 16 hours for a genetic-algorithn
Comparison with Search-basee and Training- Graining-based based Methods.
15

---

Published as a conference paper at ICLR 2026
search baseline, DiftSS attains better FID while using less optimization time. The learned sparsity predictor is compact (size (T x L) x (S] and often transfers from 256 x 256 training to 512 x 512 evaluation, reducing the need for retraining higher resolutions. By contrast, distillation- pipelines can demand orders of magnitude more compute: reported distillation efforts (DMD2) involve (110--104) GPU-hours (e.g., SD11.5: 1,664 GPU-hr; SDXL.: 8,192 GPU-hr), far exceeding the cost of our method and frequently relying on private data. Importantly, pift'Sparse also improves some distilled models: for example, on 4-step distilled model 1.81 x speedup with no measurable quality drop (see Table 9h.
we observe
A.5
MORE IMPLEMENTATION DETAILS
A.5.1 TOKEN SELECTOR
We rank tokens using a composite importance score that integrates four criteria: self-attention influence, cross-attention influence, cache reuse frequency, and uniform spatial distribution. This composite score, which has demonstrated effective prior work (Zou et al. 2025), is defined for each token à as follows:
S( = B 81 (f:) + A2 82(;)) + Àa sa()))
(10)
where 81 (dx) = Σÿ= Osj quantifies the self-attention contribution of token i,, with 0ij being the (ith))th element of the normalized self-attention matrix. A higher value indicates that the token exerts significant influence on others, meaning error in its representation may easily propagate. The term (2,) = Σ}= oú log(o.;) represents the entropy of the cross-attention weights Oıj, measuring how the control signal influences token 2,, with lower entropy indicating more focused guidance. Additionally, 83 (ă.) ni denotes the number of times token 2; has been reused from the cache since its last computation, where a higher n,, suggests possible accumulated errors, thus necessitating a fresh computation. The spatial bonus function Bl-) promotes uniform spatial distribution of the selected tokens by adding a bonus value X to the score of é; if it has the highest composite score within its l neighborhood. For each layer, tokens are ranked in descending order based on S(xx), and the top K tokens are selected for computation and cache updates according to a predefined sparsity ratio R.
We adopt the hyperparameter settings recommended by ToCa Zou et al. 2025) and DuCa Zou al. (2024) (which were shown be optimal for that setup) and therefore do not include ablation experiments for these parameters. since tuning them is not central to our contribution. Specifically, for PixArt-a, we set Å₁ 0.0, A2 1.0, A3 0.25/3, A 0.4, and & 4. For DřT, we use A1 1.0, A2 = 0.0, AS 0.25/3. A = 0.6, and k = 2. For FL.UX.1-schnek we set A 0.0, A2 = 1.0, A3 = 0.25/3, \4 0.4, and 4.
Besides, for 1an2.1, we select tokens with smaller norms in their value matrix substitutes for those with high attention map scores, following a strategy shown to effective in DuCa Zou et al. (0024). Notably, our method does not introduce new token-selector Instead, it can be applied to existing token-selection methods and uses a differentiable sparsity-cost matrix to assign the model an optimal sparsity level. As shown in Table 5 across various token- -ieeniimpottance metries approach consistently yields substantial gains.
A.5.2 LAYER SPARSITY COST PREDICTOR
We define a sparsity router for each layer. For PixArt-a and 1 model, each transformer block consists of a self-attention layer, a cross-attention layer, and an MLP layer, with each layer being assigned an individual sparsity value. In contrast, the DiT model does not include a cross-attention layer, thus the corresponding predictor for cross-attention layer is removed. Additionally, the FLUX model contains an image MLP layer, a text MLP layer, and standard MLP layer, each which is assigned its own sparsity value.
A.5.3 TWO-STAGE TRAINING
We present the pseudocode for our two-stage training algorithm in Algorithm illustrating the training details of our sparsity cost predictor.
16

---

Published as a conference paper at ICLR 2026
Algorithm 1 Two-Stage Training Strategy for Cost Matrix Optimization
Input:
Step cost matrix Cf E I2Tx2; Layer sparsity cost matrix Cj E ( Total steps T; Number of layers L; Candidate set S: Desired full-step count [Tyl: Mutation constant ó = 10.
Stage 1: Initialization and Warm-Starting
Solve Cy via dynamic programming to obtain optimal full-step set Ty. Integrate C5 into the tuned C₁ to form the unified cost matrix:
for each / T do
for 1 = 1 to L do
Update cost: C C(LLN)) & c(1..M) 8.
end for end for
Stage 2: Unified Cost Optimization
Fine-tune the integrated C₁ using differentiable cost interactions to systematically redistribute FLOPs across sampling steps.
Output: Optimized layer cost matrix Ct.
A.5.4 SEARCH-BASED APPROACHES
In this paper, we compare our method with search-based approaches. For the GA algorithm, we start by initializing population of 50 (T L) layer-sparsity vectors that satisfy the sparsity requirements. Each candidate evaluated using FID value, which serves as the fitness score. subsequent iterations, we select the best- performiffoogggg individuals as parents for crossover operations and introduce mutations with a probability of 0.01 to maintain population diversity. This iterative process continues until the individual with the highest fitness score is identified. Besides, the random search algorithm generates population candidates that meet the sparsity requirements in a completely random manner. Their fitness is also evaluated using the FID value, and the optimal solution is updated iteratively until the maximum number iterations reached.
A.5.5 ABOUT THE RETRAINING REQUIREMENT
you change the model architecture significantly (e.g.. different number of layers or a different block structure), retraining or at least fine-tuning is required when the temporal or architectural axes change (7 L)) because its parameters are tied to (T, L, [S]). but not usually when only token length (image resolution) increases. Given the modest one-time training cost (4-10 GPU-hours in our experiments) and the measurable quality, speed improvements, we believe the overhead is justified for deployed models where inference cost matters.
A.6 MORE EXPERIMENTS
A.6.1 COMPARISON ON DISTILLED MODEL.
Feature caching leverages redundancy across timesteps but provides little benefit for distilled diffusion models with only one or two steps. Nevertheless, we still evaluate DiffSparse FLUX. 1-schnell (Labs 2024) with steps at 256x256 resolution on the PartiPrompts (Yu et al. 2022) dataset. As Table shows, Diff'Sparse attains the same acceleration rate ToCa but yields a higher Image Reward, confirming its effectiveness reallocating computation to the most critical layers and delivering lossless speedup.
Table Comparison text-to- -ot-tto-mmae generation for FLUX, FLUX.I-schnell -chhnell on PartiPrompts.
Method
MACs (T)
Image Reward*
FLUX. 1 1-schnell
13.247
1.064
75% Steps
9,936
1.063
ToCa
7.313
1.063
DiffSparse
7.316
1.184
17

---

Published as a conference paper at ICLR 2026
A.6.2 ABLATIONS THE ATTENTION BASSDD SCORE
Table 5 compares three importance metries (attention-basee) score, cosine similarity, norm) and shows the attention-based score performs best overall. In Table we further remove one component time (-$1, $1, 82. $3. -) to answer how much each term contributes independently.
Table 10: Ablations of the attention-base score text-to-image generation for PixArt-a.
Variant
FID
CLIP
DiftSparse
26.91
0.164
81 (self-attention influence)
27.11
0.164
82 (cross-attention focus)
27.48
0.163
S3 (cache-ruuse frequency)
27.23
0.164
B (spatial bonus)
27.05
0.164
A.7 QUALITATIVE ANALYSIS
A.7.1 VISUALIZATION OF LAYER SPARSITY
Figure B] shows the learned layer wise sparsity allocation for PixArt-a at 256x256 resolution with 20 sampling steps under 1.74 speedup, with the first step omitted because always fully computed in cache-based acceleration methods. In the self attention layers, sparsity is higher (that is, the layers are more cacheable) in early time steps and shallow layers, while the cross attention layers sparsity is lower in later time steps and deeper layers, suggesting that textual semantics are most important in the initial layers. The MLP layers receive more computation in early steps and shallow layers, with reduced sparsity in deep layers at early steps and in shallow layers at later steps. In addition, [ggure]ddmoostratcc that our method can redistribute the computation across all steps, reducing the dependence on fully computed steps. It shows that additional resources allocated to MLP layer, this might be attributed to its ability of correct errors introduced by caching.
-
(a) Self Attention layer.
(b) Cross Attention layer.
(c) MLP layer.
Figure 3: Visualization of predicted layer sparsity of PixArt-a with 20 steps. In the figure, the x-axis denotes different network layers, the y-axis denotes sampling time steps, and the color gradient from blue to yellow indicates increasing sparsity.
A.7.2 MORE VISUALIZATION OF GENERATED IMAGES
Figures and 5 present further visualizations, including additional comparison samples and higher- resolution results. They confirm that approach delivers markedly higher acceleration ratios compared to the baseline, while preserving performance quality.
18

---

Published as a conference paper at ICLR 2026
Odiginal
DiffSparse 1.74 x
DirfSparse L91 x
Original
DiffSparse 1.74
Diff'Spase 1.91
A girl looking
beautiful view the Rockies.
large piece meat surrounded by vegetables.
a close up shirtless man wearing lie.
dog wearing baseball hat over eyes.
Figure Comparison our method with the baseline (PixArt-a with DPM-Solver++ using 20 steps) under different acceleration rates.
Original
ToCa
DifiSparee
Original
TbCa
DiffSporee
person siiting bed with dog his lap.
woman holding surfboard walking into ocesn loward dog.
who appears be hercing sheep clostng two big fance doors.
Sheep are running across green field grass.
sitling green grass ground.
Figure 5: Comparison between our DiflSparse, and ToCa with the baseline (PixArt-a with DPM- Solver+ using 20 steps under 512x512 resolution).
19


Reviews:
[{"review_id": "aQi7R0RqQo", "reviewer": "Reviewer_ZHnY", "summary": "DiffSparse is a novel framework for accelerating Diffusion Transformer models through learnable token-level sparsity optimization. Unlike existing token caching methods that require manual sparsity configuration and depend on predefined full-step computations, DiffSparse automatically learns optimal layer-wise sparsity allocation in an end-to-end manner while eliminating the need for complete forward passes at certain steps.", "questions": "1. Does the method require retraining for each new model, new resolution, or  new sampling step count? Is this training overhead worthwhile?\n2.While Table 5 demonstrates that the Attention strategy outperforms other  strategies overall, it fails to analyze the individual contributions of \nthe four components (self-attention influence, cross-attention focus, cache-reuse frequency, etc)  within the Attention method .", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The method comprises three key components: (1) a Token Selector that computes importance scores based on self-attention influence, cross-attention entropy, cache reuse frequency, and spatial distribution to determine which tokens should be recomputed versus reused from cache; (2) a Learnable Sparsity Cost Predictor with trainable parameters that outputs a normalized cost matrix quantifying the quality impact of different sparsity rates for each layer at each timestep; and (3) a Dynamic Programming Solver that finds the optimal sparsity configuration across all layers and timesteps by minimizing cumulative cost while satisfying a global sparsity constraint , with gradients flowing through discrete decisions via Straight-Through Estimation.", "weaknesses": "1. Symbol M is overloaded with two conflicting meanings within the same section.  First, M (introduced in Section 3.2, \"Learnable Sparsity Cost Predictor\") represents a binary mask matrix for token selection at each layer and timestep,  with dimension equal to the number of tokens N. Second, M appears as the  summation upper bound in Equation 5, representing the count of importance  criteria combined in the token ranking score. These two usages denote entirely  different mathematical objects—a vector versus a scalar integer—creating  potential confusion for readers.\n2.Compared to \"training-free\" competing methods (such as ToCa), this approach  requires additional training but does not demonstrate a clear advantage in  acceleration performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZSYsp8FZBf", "reviewer": "Reviewer_9MrV", "summary": "This paper proposes DiffSparse, a learned token-sparsity allocation framework for accelerating diffusion transformer models.It addresses inefficiencies in existing token caching methods by automating layer-wise sparsity allocation using a differentiable cost predictor and dynamic programming solver.Experiments demonstrate moderate acceleration with limited quality degradation across several diffusion transformer backbones.", "questions": "1.How sensitive is performance to δ and |S|? The limited ablations do not paint a full picture of robustness.\n\n2.Could you include comparisons with state-of-the-art training-free caching approaches such as DiCache and DeepCache.\n\n3.Since the Token Selector computes attention-based ranking signals and dynamically applies binary masks at each step, could the authors quantify the corresponding inference-time overhead and its impact on real wall-clock latency?\n\n4.Given the training cost (4–10 GPU hours with full teacher inference), is the amortized gain meaningful vs. training-free caching or distillation approaches?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The motivation behind the proposed method is well articulated. \n\n* Validates across diverse architectures (DiT, PixArt, FLUX, Wan2.1) and tasks (text-to-image/video, class-conditional), demonstrating robustness.", "weaknesses": "* The training process is relatively complex, requiring a bespoke two-stage strategy and 4-10 hours of training on 8 high-end GPUs.\n\n* Results on higher resolution (512×512) are based on one configuration and show only marginal benefit,Table 4 shows the proposed method does not outperform all baselines in FID.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PKi01SWktD", "reviewer": "Reviewer_c1BY", "summary": "This paper introduces DiffSparse, a framework to accelerate diffusion transformer models by optimizing layer-wise token sparsity during the sampling process. It employs a learnable sparsity cost predictor and dynamic programming to allocate computation efficiently, eliminating the need for manual sparsity tuning and full-step computations. The proposed two-stage training strategy enhances acceleration while preserving image generation quality. Extensive experiments on various models demonstrate significant speedups (up to 1.91×) with improved or comparable performance to state-of-the-art methods.", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The proposed differentiable layer-wis sparsity optimization is novel and effective.", "weaknesses": "1. What is the training cost for FLUX and Wan?\n\n2. The paper would benefit from reorganization. For instance, the related work section is overly long and could be trimmed and it would be better to include visualizations directly in the main paper to enhance clarity and accessibility.\n\n3. The discussion should also include other token-based acceleration methods, such as DyDiT [1].\n\n[1] Dynamic Diffusion Transformer, 2025 ICLR", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eVQfffbOWj", "reviewer": "Reviewer_wjJ7", "summary": "This paper focuses on the issue of high computational cost of diffusion models, and proposes the **DiffSparse** framework, which achieves acceleration through optimization via learnable layer-wise sparsity. This method eliminates the need for manual adjustment of sparsity parameters and realizes automated sparsity allocation through end-to-end learning, making it more efficient than traditional feature caching methods that rely on manual sparsity allocation. DiffSparse performs excellently on a variety of diffusion models and possesses cross-resolution generalization capability. Meanwhile, compared with search methods such as genetic algorithms, it also features higher training efficiency.", "questions": "### **1. Similarity of Dynamic Sparsity Allocation in Other Feature Caching-Based Diffusion Acceleration Works**\n Are there other papers that use Feature caching to accelerate diffusion models and have similar dynamic sparsity allocation ideas?\n\n### **2. Feature Recomputation and Utilization Logic After Sparsity Determination**\nAlthough the method in this paper aims to achieve dynamic sparsity allocation, it lacks mention of how features are recomputed and utilized after sparsity is determined. For example, does this paper directly adopt previous work, or it make modifications based on previous methods?\n\n### **3. Input, Cost Definition, and Calculation Logic of Learnable Sparsity Cost Predictor**\nWhat is the input of the Learnable Sparsity Cost Predictor? What is \"Cost\" defined as, and how is the Cost of a specific layer of the model at a certain timestep calculated?\n\n### **4. Rationale for Using Formula (9) to Integrate Step and Layer Costs**\n Formula (9) is intended for warm-start. Why can it be used to integrate step and layer costs? I cannot see the logic connection for now.\n\n### **5. Lack of Performance Comparison with Other Non-Feature-Caching-based Acceleration Methods**\nAlthough this paper focuses on accelerating diffusion models using Feature Caching, the experiments lack performance comparisons with non-Feature caching acceleration methods, such as Model pruning, Rectified Flow, and knowledge distillation.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "### 1. Innovative Method Design & Problem Formulation\nRedefines diffusion transformer acceleration as a **differentiable layer-wise sparsity allocation problem**, abandoning the manual sparsity schedules of prior token cache methods (e.g., ToCa). It combines a learnable sparsity cost predictor, dynamic programming solver, and adaptive token selector into a unified framework. Additionally, it eliminates predefined full-step computation via a two-stage training strategy to unlock token caching’s full acceleration potential, and achieves cross-resolution generalization where low-resolution-trained sparsity predictors work for high-resolution models without retraining.\n\n### 2.  Robust Experimental Validation\nIt conducts comprehensive experiments, validating on 4 models (DiT-XL/2, PixArt-α, FLUX, Wan2.1) across image/video generation—for example, PixArt-α achieves 54% cost reduction (1.91× speedup) with better FID than ToCa. Thorough ablation studies confirm the optimality of two-stage training, LPIPS loss, and 0.25 sparsity interval. It also outperforms search methods like genetic algorithms with 4× faster training.", "weaknesses": "### 1. Insufficient Analysis of Computational Overhead\n- **Unreported training/inference overhead of DiffSparse components**: The paper claims the dynamic programming solver has low complexity $O((L \\cdot T)^2 \\cdot |S|)$ but does not report the **actual runtime overhead** of the sparsity cost predictor or the dynamic programming solver during inference. \n- **Lack of memory consumption breakdown**: While noting cross-resolution generalization reduces memory usage, the paper does not provide quantitative data (e.g., peak memory for 256×256 vs. 512×512 training with DiffSparse) or compare it to baselines.\n\n### 2. Limited Generalization Validation\n- **No testing on ultra-low-step models**: The paper excludes distillation methods but does not test DiffSparse on **1–2 step distilled diffusion models** (e.g., SDXL-Turbo). \n\n### 3. Incomplete Ablation of Key Components\n- **No ablation of dynamic programming solver necessity**: The paper claims the dynamic programming solver optimizes global sparsity, but it does not compare against a \"no solver\" baseline (e.g., using the cost predictor’s output directly without global constraint optimization). This fails to prove the solver’s contribution to performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)