Shiva-DiT: Residual-Based Differentiable Top-$k$ Selection for Efficient Diffusion Transformers

TL;DR

Shiva-DiT introduces a residual-based differentiable top-$k$ selection method for diffusion transformers, enabling efficient, static-budget pruning with improved speed and fidelity.

Contribution

It proposes a novel residual-aware straight-through estimator and adaptive pruning schedule for static, differentiable token selection in diffusion transformers.

Findings

Achieves 1.54× speedup over baseline models.

Establishes a new Pareto frontier with better fidelity.

Eliminates ragged tensor overheads.

Abstract

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.