Shiva-DiT: Residual-Based Differentiable Top-$k$ Selection for Efficient Diffusion Transformers
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915777803714560 |
|---|---|
| author | Zhang, Jiaji Zhao, Hailiang Zhu, Guoxuan Sun, Ruichao Wu, Jiaju Zhao, Xinkui Tang, Hanlin Lu, Weiyi Liu, Kan Lan, Tao Qu, Lin Deng, Shuiguang |
| author_facet | Zhang, Jiaji Zhao, Hailiang Zhu, Guoxuan Sun, Ruichao Wu, Jiaju Zhao, Xinkui Tang, Hanlin Lu, Weiyi Liu, Kan Lan, Tao Qu, Lin Deng, Shuiguang |
| contents | Diffusion Transformers (DiTs) incur prohibitive computational costs due to the quadratic scaling of self-attention. Existing pruning methods fail to simultaneously satisfy differentiability, efficiency, and the strict static budgets required for hardware overhead. To address this, we propose Shiva-DiT, which effectively reconciles these conflicting requirements via Residual-Based Differentiable Top-$k$ Selection. By leveraging a residual-aware straight-through estimator, our method enforces deterministic token counts for static compilation while preserving end-to-end learnability through residual gradient estimation. Furthermore, we introduce a Context-Aware Router and Adaptive Ratio Policy to autonomously learn an adaptive pruning schedule. Experiments on mainstream models, including SD3.5, demonstrate that Shiva-DiT establishes a new Pareto frontier, achieving a 1.54$\times$ wall-clock speedup with superior fidelity compared to existing baselines, effectively eliminating ragged tensor overheads. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_05605 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Shiva-DiT: Residual-Based Differentiable Top-$k$ Selection for Efficient Diffusion Transformers Zhang, Jiaji Zhao, Hailiang Zhu, Guoxuan Sun, Ruichao Wu, Jiaju Zhao, Xinkui Tang, Hanlin Lu, Weiyi Liu, Kan Lan, Tao Qu, Lin Deng, Shuiguang Machine Learning Artificial Intelligence Computer Vision and Pattern Recognition Diffusion Transformers (DiTs) incur prohibitive computational costs due to the quadratic scaling of self-attention. Existing pruning methods fail to simultaneously satisfy differentiability, efficiency, and the strict static budgets required for hardware overhead. To address this, we propose Shiva-DiT, which effectively reconciles these conflicting requirements via Residual-Based Differentiable Top-$k$ Selection. By leveraging a residual-aware straight-through estimator, our method enforces deterministic token counts for static compilation while preserving end-to-end learnability through residual gradient estimation. Furthermore, we introduce a Context-Aware Router and Adaptive Ratio Policy to autonomously learn an adaptive pruning schedule. Experiments on mainstream models, including SD3.5, demonstrate that Shiva-DiT establishes a new Pareto frontier, achieving a 1.54$\times$ wall-clock speedup with superior fidelity compared to existing baselines, effectively eliminating ragged tensor overheads. |
| title | Shiva-DiT: Residual-Based Differentiable Top-$k$ Selection for Efficient Diffusion Transformers |
| topic | Machine Learning Artificial Intelligence Computer Vision and Pattern Recognition |
| url | https://arxiv.org/abs/2602.05605 |