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

arXiv: https://arxiv.org/abs/2602.05605 · HTML: https://mlsanity.com/paper/2602.05605/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 42 of its 59 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Scaling Autoregressive Models for Content-Rich Text-to-Image Generation](https://mlsanity.com/paper/2206.10789/) (2022)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [WorldSimBench: Towards Video Generation Models as World Simulators](https://mlsanity.com/paper/2410.18072/) (2024)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [DoRA: Weight-Decomposed Low-Rank Adaptation](https://mlsanity.com/paper/2402.09353/) (2024)
- [PixArt-Sigma: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation](https://mlsanity.com/paper/2403.04692/) (2024)
- [Playground v2.5: Three Insights towards Enhancing Aesthetic Quality in Text-to-Image Generation](https://mlsanity.com/paper/2402.17245/) (2024)
- [Token Merging for Fast Stable Diffusion](https://mlsanity.com/paper/2303.17604/) (2023)
- [Accelerating Diffusion Transformers with Token-wise Feature Caching](https://mlsanity.com/paper/2410.05317/) (2024)
- [Token Fusion: Bridging the Gap between Token Pruning and Token Merging](https://mlsanity.com/paper/2312.01026/) (2023)
- [DiTFastAttn: Attention Compression for Diffusion Transformer Models](https://mlsanity.com/paper/2406.08552/) (2024)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [DiffRate : Differentiable Compression Rate for Efficient Vision Transformers](https://mlsanity.com/paper/2305.17997/) (2023)

## Built on by (most cited first, 2 of 2)

- [Beyond Surrogate Gradients: Fully Differentiable Token Pruning for Vision-Language Models](https://mlsanity.com/paper/2605.28051/) (2026)
- [DiffPrune: differentiable information throttling for token pruning in vision-language models](https://mlsanity.com/paper/2608.01985/) (2026)

## Abstract (arXiv)

Diffusion Transformers (DiTs) are costly at high resolution because self-attention scales quadratically with token sequence length. Existing pruning methods do not jointly provide end-to-end learnability, low training overhead, and deterministic token counts for predictable token-dependent computation. We propose Shiva-DiT, based on Residual-Based Differentiable Top-k Selection. Its forward pass executes hard top-k selection, while a residual-aware straight-through estimator propagates gradients to both token scores and the budget k without evaluating a second backbone path. A Context-Aware Router and Adaptive Ratio Policy learn layer- and timestep-dependent retention schedules under a target average budget. Experiments on SD3-Medium, Flux.1-dev, and PixArt-Sigma show consistent reductions in FLOPs and measured latency. On SD3-Medium, Shiva-DiT provides four fidelity-latency operating points and reaches a 1.54x wall-clock speedup with competitive fidelity.