Shiva-DiT: 효율적인 디퓨전 트랜스포머를 위한 잔차 기반 미분 가능한 Top-k 선택 방법

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

디퓨전 트랜스포머(DiT)는 자체 주의(self-attention) 메커니즘의 2차 함수적 스케일링으로 인해 상당한 계산 비용을 발생시킵니다. 기존의 가지치기 방법은 미분 가능성, 효율성, 그리고 하드웨어 오버헤드를 위한 엄격한 정적 예산 요구 사항을 동시에 만족시키지 못합니다. 이러한 문제를 해결하기 위해, 우리는 잔차 기반 미분 가능한 Top-k 선택 방법을 활용하여 이러한 상충되는 요구 사항을 효과적으로 조화시키는 Shiva-DiT를 제안합니다. 우리의 방법은 잔차를 고려한 스트레이트-쓰루 추정기를 활용하여 정적 컴파일을 위한 결정적인 토큰 수를 강제하면서, 잔차 그래디언트 추정을 통해 엔드 투 엔드 학습 가능성을 유지합니다. 또한, 컨텍스트 인식 라우터와 적응형 비율 정책을 도입하여 자율적으로 적응형 가지치기 일정을 학습합니다. SD3.5를 포함한 주요 모델에 대한 실험 결과, Shiva-DiT는 기존의 기준 성능보다 1.54배 빠른 처리 속도를 달성하고, 더 높은 정확도를 제공하며, 불규칙한 텐서 오버헤드를 효과적으로 제거하여 새로운 Pareto 최적점을 제시합니다.

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.