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× wall-clock speedup with superior fidelity compared to existing baselines, effectively eliminating ragged tensor overheads.
@article{arxiv.2602.05605,
title = {Shiva-DiT: Residual-Based Differentiable Top-$k$ Selection for Efficient Diffusion Transformers},
author = {Jiaji Zhang and Hailiang Zhao and Guoxuan Zhu and Ruichao Sun and Jiaju Wu and Xinkui Zhao and Hanlin Tang and Weiyi Lu and Kan Liu and Tao Lan and Lin Qu and Shuiguang Deng},
journal= {arXiv preprint arXiv:2602.05605},
year = {2026}
}