Title: SparseDiT: Token Sparsification for Efficient Diffusion Transformer
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion Transformer

Sparse Tokens

Efficient DiT
Abstract: Diffusion Transformers (DiT) are renowned for their impressive generative performance; however, they are significantly constrained by considerable computational costs due to the quadratic complexity in self-attention and the extensive sampling steps required. While advancements have been made in expediting the sampling process, the underlying architectural inefficiencies within DiT remain underexplored. We introduce SparseDiT, a novel framework that implements token sparsification across spatial and temporal dimensions to enhance computational efficiency while preserving generative quality. Spatially, SparseDiT employs a tri-segment architecture that allocates token density based on feature requirements at each layer: Poolingformer in the bottom layers for efficient global feature extraction, Sparse-Dense Token Modules (SDTM) in the middle layers to balance global context with local detail, and dense tokens in the top layers to refine high-frequency details. Temporally, SparseDiT dynamically modulates token density across denoising stages, progressively increasing token count as finer details emerge in later timesteps. This synergy between SparseDiT’s spatially adaptive architecture and its temporal pruning strategy enables a unified framework that balances efficiency and fidelity throughout the generation process. Our experiments demonstrate SparseDiT’s effectiveness, achieving a 55\% reduction in FLOPs and a 175\% improvement in inference speed on DiT-XL with similar FID score on 512$\times$512 ImageNet, a 56\% reduction in FLOPs across video generation datasets, and a 69\% improvement in inference speed on PixArt-$\alpha$ on text-to-image generation task with a 0.24 FID score decrease. SparseDiT provides a scalable solution for high-quality diffusion-based generation compatible with sampling optimization techniques.  Code is available at https://github.com/changsn/SparseDiT.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes SparseDiT, a new architecture that avoids redundant computations in diffusion transformers, leading to increased throughput at minimal or no quality cost. Compared to the normal DiT, the authors introduce two components. First, poolingformer, which simply removes the computation from queries and keys, and simply average pools and repeats the values of the first layer's attention to save the quadratic attention computation (in the first few transformer blocks). And second, sparse-dense token modules (SDTM), which cross-attend to a pooled version of the dense tokens, perform attention with sparse tokens, and then upsample the tokens back to dense tokens. The authors propose a timestep-wise pruning strategy for the number of tokens in the SDTM modules to achieve a balance between performance and speed. The authors show the superiority of their method on class-conditional and text-conditional image generation, as well as video generation.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nThe paper is well-organized and well-written, with a clear common thread. A key strength lies in the comprehensive experimental evaluation, which spans class-conditional and text-conditional image generation as well as video synthesis, demonstrating the broad applicability. The proposed modules are simple, yet effective in increasing the throughput of diffusion transformers. The spatially-adaptive design based on their empirical analysis of attention is interesting. Additionally, the timestep-aware pruning strategy is both intuitive and effective. Overall, I think it is a solid and simple method to improve efficiency.\n\n**Weaknesses**\n\n- Lack of baselines: While the paper provides many internal comparisons, it lacks comparisons to other strong baseline methods. In the supplementary material, the authors compare to DyDiT and ToMeSD, however, there exist other test-time caching methods, some of which are mentioned in the related works section, but not empirically compared to (i.e., [1, 2, 3]). I would consider direct performance comparisons to be crucial since some of these methods avoid computing tokens entirely and do not require retraining the base model at all. The authors also do not compare their method to attention compression methods (e.g., [4]).\n- Potential mischaracterization of flow-based methods: The paper describes flow-based approaches (e.g., Rectified Flow) as methods that “straighten transport trajectories” (l.99), which may oversimplify their role. In practice, Rectified Flow provides a more general framework that modifies aspects like the interpolant and training target [5], rather than explicitly straightening trajectories. While trajectory straightening can be achieved through techniques like ReFlow [2], which trains on image-noise pairs from a pre-trained flow model, this is not a property of Rectified Flow per se (albeit found in some papers empirically, e.g. [5, 7]). This conceptual ambiguity also makes the comparison between DDIM and RFlow in Table 5 somewhat unclear. Since DDIM is closely related to Euler discretization in Rectified Flow, a more representative comparison of samplers would involve solvers such as DPM [8] or DPM++ [9].\n- Lack of analysis protocol: The proposed tri-segment design and pruning strategy requires prior analysis on a trained model. Hence, the method can be considered as a post-training strategy. However, the analysis on which the authors ground their method doesn't establish a clear and systematic protocol for selecting key hyperparameters (e.g. number of SDTM modules). While the justification for using poolingformer layers is reasonable, it remains unclear how the authors infer the distinction between local and global attention solely from the normalized attention variances across layers. A more systematic protocol and detailed analysis would definitely strengthen the paper.\n- Missing isolated analysis of SDTM vs DiT block: The SDTM module introduces additional parameters and FLOPs through the two cross-attention operations and the two linear mixing matrices $W_1, W_2 \\in \\mathbb{R}^{C \\times C}$. I would highly appreciate a direct comparison of throughput, FLOPs, and parameter counts between a standard DiT block and SDTM. This would clarify whether the efficiency gains are truly due to token sparsity.\n- Confound metrics from finetuning: The improved FID scores reported e.g., in table 1, could also stem from continued training of the DiT rather from the proposed architectural changes. A comparison to the DiT-XL and PixArt-$\\alpha$ baselines with continued training would help isolate the effect of the proposed method.\n\n**General Questions and Clarifications**\n- Usually DiT models contain high norm tokens. Did you observe similar patterns in your analysis?\n- What kind of token upsampling do you use for your method? How do you handle different sizes for floating point numbers of $r$?\n- l266ff: You mention that SparseSiT gives a more precise structure. Is this consistent behaviour or are these just artefacts from the initial noise? For example, sometimes one seed might work better for your method than for the baseline. Can you provide any support for this claim or an intuitive explanation?\n\n**Minors**\n- writing (L96ff and L193ff)\n- improved readability of the settings for DiT-B and DiT-XL in experimental section (e.g. in form of a table)\n- label of table 5 and table 7 seem to be switched\n- lack of other metrics for the T2I setting (e.g. CLIP score, etc)\n\n\n**References**\n- [1] Wimbauer, F., Wu, B., Schoenfeld, E., Dai, X., Hou, J., He, Z., ... & Wang, J. (2024). Cache me if you can: Accelerating diffusion models through block caching. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 6211-6220).\n- [2] Ma, X., Fang, G., & Wang, X. (2024). Deepcache: Accelerating diffusion models for free. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 15762-15772).\n- [3] Lou, J., Luo, W., Liu, Y., Li, B., Ding, X., Hu, W., ... & Ma, C. (2024). Token Caching for Diffusion Transformer Acceleration. arXiv preprint arXiv:2409.18523.\n- [4] Yuan, Z., Zhang, H., Pu, L., Ning, X., Zhang, L., Zhao, T., ... & Wang, Y. (2024). Ditfastattn: Attention compression for diffusion transformer models. Advances in Neural Information Processing Systems, 37, 1196-1219.\n- [5] Ma, N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden-Eijnden, E., & Xie, S. (2024, September). Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision (pp. 23-40). Cham: Springer Nature Switzerland.\n- [6] Liu, X., Gong, C., & Liu, Q. (2023, January). Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. In The Eleventh International Conference on Learning Representations (ICLR).\n- [7] Schusterbauer, J., Gui, M., Ma, P., Stracke, N., Baumann, S. A., Hu, V. T., & Ommer, B. (2024, September). FMBoost: Boosting Latent Diffusion with Flow Matching. In European Conference on Computer Vision (pp. 338-355). Cham: Springer Nature Switzerland.\n- [8] Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. (2022). Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35, 5775-5787.\n- [9] Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. (2025). Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. Machine Intelligence Research, 1-22.\n\nQUESTIONS:\nI believe the proposed method has clear potential. However, to consider increasing my final rating, I would appreciate clarification on the following key points (which were already mentioned above):\n\n1. Clarification on baselines and direct empirical comparison to other test-time caching or attention compression methods.\n2. Direct comparison of FLOPs, parameters, and throughput for a DiT block versus SDTM, to isolate the efficiency gains due to token sparsity.\n3. Isolate the gain from the proposed architectural components compared to continued training.\n\nLIMITATIONS:\nThe limitations section mainly consists of a single sentence that describes the main limitation of requiring manual analysis of the structure design. This neither discusses the limitation nor provides any explanation or solution to this problem. Hence, I would consider this as insufficient.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces SparseDiT, a novel framework designed to tackle the significant computational inefficiency of Diffusion Transformers. The core contribution lies in a unified spatio-temporal token sparsification strategy, motivated by the different attentions shown in different layers and denoising steps. Spatially, it utilizes a tri-segment architecture: a lightweight Poolingformer captures global features in early layers, an innovative Sparse-Dense Token Module (SDTM) alternates between sparse and dense tokens to efficiently process mid-level features, and standard Transformers refine details in final layers. Temporally, it employs a dynamic pruning strategy, progressively increasing the number of active tokens throughout the denoising process to match the evolving demand for detail, thereby optimizing computation across both space and time.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper derives its core contribution from a careful empirical and conceptual analysis of attention distributions across Transformer layers in DiT, revealing meaningful structural patterns that inform the proposed sparsification mechanism. This insight leads to a well-designed spatial token compression strategy, which is both intuitive and effective, and clearly contributes to improving computational efficiency. The implementation is elegant and shows strong alignment between observation and design. However, the proposed temporal sparsification—adjusting the pruning rate based on diffusion timesteps—feels more arbitrary and lacks sufficient theoretical or empirical grounding. It appears less principled compared to the layer-wise design. A more critical issue lies in the experimental evaluation: the main text only includes with/without SparseDiT comparisons, while the appendix contains limited and coarse comparisons with ToMeSD and DyDiT. Many related methods mentioned in the Related Work section—such as other token pruning, merging, or flow strategies—are not included in the empirical evaluation at all. This makes it difficult to assess the actual advantage and novelty of SparseDiT in the broader context of efficient models.\n\nQUESTIONS:\n1.\tTemporal sparsification design: Could you provide more justification—either theoretical or empirical—for the choice of T/4 as the divide point of the denoising stage? Why was the specific linear interpolation chosen, and have you considered or tested alternative scheduling strategies?\n2.\tLimited baseline comparisons: The main paper includes only ablation-type comparisons (with/without SparseDiT), which makes it hard to gauge the method’s performance relative to prior work. Could you include more detailed quantitative and qualitative comparisons with relevant baselines such as ToMeSD, DyDiT, and other token sparsification or merging techniques on more complex  datasets?\n3.\tRelated work inclusion: Several optimization strategies mentioned in the Related Work section (e.g., adaptive token routing, dynamic masking, or hybrid dense-sparse attention methods) are not covered in the experiments. What is the rationale for excluding these, and how do you expect SparseDiT would compare to them?\n\nLIMITATIONS:\nA key limitation of this work also stated in limitation part is its heavily manual design choices across both spatial and temporal sparsification modules. Many core architectural decisions—such as the pruning schedule over diffusion steps, the sparse-dense transition structure,—are introduced without theoretical grounding or extensive empirical search, making them appear heuristic rather than principled. This raises concerns about the generality and robustness of the design. Another notable limitation is the insufficient engagement with prior work. While some related methods are briefly mentioned in the Related Work section, the paper does not clearly acknowledge how existing techniques (e.g., token merging from ToMeSD) have inspired or relate to the proposed method. A deeper and more explicit comparison, both conceptual and empirical, would significantly strengthen the work’s positioning and clarify its contribution boundaries.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors aim to improve the DiT method by pruning its architecture. They propose a new architecture called SparseDiT, which achieves performance comparable to the original DiT while using significantly fewer FLOPs after fine-tuning pretrained DiT models. The main approach is to: (1) apply average pooling to extract sparse tokens, (2) alternate self-attention on the sparse tokens with cross-attention between the sparse and dense tokens, and (3) apply a few dense attention layers at the end. In a fine-tuning setting, the authors show that their proposed method can achieve performance comparable to DiT while significantly reducing FLOPs.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1. In Figures 1 and 3, the authors show that the early-layer attention maps in DiT are highly uniform. This is an insightful observation that provides a clear motivation for using average pooling.\n\n2. The authors validate the proposed SparseDiT across various benchmarks, including image and video generation, demonstrating improved efficiency in all cases. The experiments are solid.\n\n3. The performance gains appear consistent across different model sizes, showing strong potential for scalability.\n\nWeaknesses\n\n1. My major concern is with the experimental setup. The authors only report performance after fine-tuning a pre-trained model and do not include results from training SparseDiT from scratch. The claims would be much stronger if the proposed architecture also performs comparably to the original DiT when pre-trained. Otherwise, the method should be compared to knowledge distillation techniques, which also create smaller, more efficient models from a pre-trained teacher.\n\n2. The architectural design is non-uniform and involves many hyperparameters, including:\n\n(1) The number of layers for the PoolingFormer, the sparse-dense modules, and the final dense block.\n\n(2) The specific pruning ratio schedule used across different timesteps (t).\n\nI am concerned that without a principled way to determine these hyperparameters, the method's broader applicability may be limited.\n\nQUESTIONS:\n1. As raised in the weaknesses section, how does the model perform when trained from scratch? Is its effectiveness limited to a fine-tuning context, or can it achieve strong performance without relying on weights from a pre-trained dense model?\n\n2. Could the authors provide more insight into the selection of the model's hyperparameters?\n\n(a) Is there a principled method for determining the number of layers for each module (e.g., PoolingFormer, SDTM)? How sensitive is the model's performance to this specific configuration?\n\n(b) Similarly, what is the rationale behind the design of the pruning schedule across different timesteps (t)? How sensitive is the final performance to changes in this schedule?\n\n3. The model's non-uniform architecture raises questions about potential optimization challenges. Conventional transformers benefit from well-studied optimization strategies due to their uniform layer structure. Does the heterogeneity of SparseDiT introduce any new difficulties for training stability? For instance, when training from scratch, do the different modules require distinct learning rates or optimizers to converge effectively?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces SparseDiT, a novel framework for enhancing the computational efficiency of Diffusion Transformers (DiT) through token sparsification across spatial and temporal dimensions. The authors employ a tri-segment architecture (Poolingformer, Sparse-Dense Token Modules, and dense transformers) to dynamically adjust token density, achieving significant reductions in FLOPs (up to 55%) and improvements in inference speed (up to 175%) while maintaining generative quality, as demonstrated on tasks like image and video generation.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n1. Clear motivation. The token sparsification strategy is well-motivated by empirical observations of attention patterns in DiT layers.\n\n2. The paper demonstrates impressive performance gains across multiple benchmarks (e.g., ImageNet, video datasets, and text-to-image generation) with minimal degradation in output quality (e.g., FID).\n\n3. Overall, this paper is easy to follow and is also written well.\n\n\n## Weaknesses\n\n1. My main concern is the practical value of this technique. Specifically, today's text to image/video generation models are pretrained on large-scale datasets with many parameters. Thus, finetuning makes it rather challenging as it may change the underlying distribution of their initial generation space. For this reason, training-free methods seem to be more likely preferred by the community.\n\n2. Lack deeper analysis of sparse/dense tokens. For example, in L301-302, the paper claims that \"Sparse tokens capture global structures, while dense tokens capture detailed information.\", however, it would be more convincing if the authors can provide some visualizations or comparisons, e.g. more sparse tokens help to capture global or more dense tokens help to generate detailed information.\n\nQUESTIONS:\nPlease see the weakness.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}