提出可微分的动态选顶方法，让扩散模型更快更稳。
Shiva-DiT: Residual-Based Differentiable Top-$k$ Selection for Efficient Diffusion Transformers
- 用残差引导的可微选顶机制，实现高效且可训练的注意力剪枝。
- 在SD3.5上实现1.54倍的推理加速，同时保持高质量生成效果。
- 适合需要低延迟、高效率部署的生成式AI系统开发者。
扩散变换器（DiTs）因自注意力的二次计算开销导致高昂的计算成本。现有剪枝方法难以同时满足可微性、效率和硬件资源的严格静态预算要求。为此，本文提出Shiva-DiT，通过基于残差的可微分Top-k选择，有效调和这些矛盾需求。利用残差感知的直通估计器，该方法在保证静态编译确定性令牌数量的同时，通过残差梯度估计维持端到端可学习性。此外，引入上下文感知路由模块与自适应比例策略，使模型能自主学习动态剪枝策略。在主流模型如SD3.5上的实验表明，Shiva-DiT建立了新的帕累托前沿，在保持更优生成保真度的同时实现1.54倍的墙钟速度提升，有效消除了不规则张量带来的开销。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。