# Scalable Diffusion Models with State Space Backbone

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

## Key facts

- Cited by 45 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 64 of its 73 references that are in the map.
- First posted on arXiv 8 Feb 2024.

## Placed in

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

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [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)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [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)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://mlsanity.com/paper/1809.11096/) (2018)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Imagen Video: High Definition Video Generation with Diffusion Models](https://mlsanity.com/paper/2210.02303/) (2022)
- [Make-A-Video: Text-to-Video Generation without Text-Video Data](https://mlsanity.com/paper/2209.14792/) (2022)
- [DiffWave: A Versatile Diffusion Model for Audio Synthesis](https://mlsanity.com/paper/2009.09761/) (2020)
- [Cascaded Diffusion Models for High Fidelity Image Generation](https://mlsanity.com/paper/2106.15282/) (2021)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Human Motion Diffusion Model](https://mlsanity.com/paper/2209.14916/) (2022)
- [Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model](https://mlsanity.com/paper/2401.09417/) (2024)

## Built on by (most cited first, 30 of 45)

- [Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models](https://mlsanity.com/paper/2410.11081/) (2024)
- [PlainMamba: Improving Non-Hierarchical Mamba in Visual Recognition](https://mlsanity.com/paper/2403.17695/) (2024)
- [RoboMamba: Efficient Vision-Language-Action Model for Robotic Reasoning and Manipulation](https://mlsanity.com/paper/2406.04339/) (2024)
- [Align Your Flow: Scaling Continuous-Time Flow Map Distillation](https://mlsanity.com/paper/2506.14603/) (2025)
- [ZigMa: A DiT-style Zigzag Mamba Diffusion Model](https://mlsanity.com/paper/2403.13802/) (2024)
- [SkyReels-A2: Compose Anything in Video Diffusion Transformers](https://mlsanity.com/paper/2504.02436/) (2025)
- [DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis](https://mlsanity.com/paper/2405.14224/) (2024)
- [Scaling Diffusion Transformers to 16 Billion Parameters](https://mlsanity.com/paper/2407.11633/) (2024)
- [State Space Model for New-Generation Network Alternative to Transformers: A Survey](https://mlsanity.com/paper/2404.09516/) (2024)
- [DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention](https://mlsanity.com/paper/2405.18428/) (2024)
- [Diffusion Models and Representation Learning: A Survey](https://mlsanity.com/paper/2407.00783/) (2024)
- [A Survey on Visual Mamba](https://mlsanity.com/paper/2404.15956/) (2024)
- [Adversarial Score identity Distillation: Rapidly Surpassing the Teacher in One Step](https://mlsanity.com/paper/2410.14919/) (2024)
- [LinFusion: 1 GPU, 1 Minute, 16K Image](https://mlsanity.com/paper/2409.02097/) (2024)
- [Diffusion-RWKV: Scaling RWKV-Like Architectures for Diffusion Models](https://mlsanity.com/paper/2404.04478/) (2024)
- [Visual Mamba: A Survey and New Outlooks](https://mlsanity.com/paper/2404.18861/) (2024)
- [A Survey of Mamba](https://mlsanity.com/paper/2408.01129/) (2024)
- [SkyReels-A1: Expressive Portrait Animation in Video Diffusion Transformers](https://mlsanity.com/paper/2502.10841/) (2025)
- [SegMAN: Omni-scale Context Modeling with State Space Models and Local Attention for Semantic Segmentation](https://mlsanity.com/paper/2412.11890/) (2024)
- [Dimba: Transformer-Mamba Diffusion Models](https://mlsanity.com/paper/2406.01159/) (2024)
- [Matten: Video Generation with Mamba-Attention](https://mlsanity.com/paper/2405.03025/) (2024)
- [TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training](https://mlsanity.com/paper/2501.04765/) (2025)
- [StateSpaceDiffuser: Bringing Long Context to Diffusion World Models](https://mlsanity.com/paper/2505.22246/) (2025)
- [FLUX that Plays Music](https://mlsanity.com/paper/2409.00587/) (2024)
- [Mamba in Vision: A Comprehensive Survey of Techniques and Applications](https://mlsanity.com/paper/2410.03105/) (2024)
- [QuadMamba: Learning Quadtree-based Selective Scan for Visual State Space Model](https://mlsanity.com/paper/2410.06806/) (2024)
- [SkyReels-Audio: Omni Audio-Conditioned Talking Portraits in Video Diffusion Transformers](https://mlsanity.com/paper/2506.00830/) (2025)
- [ViG: Linear-complexity Visual Sequence Learning with Gated Linear Attention](https://mlsanity.com/paper/2405.18425/) (2024)
- [Ingredients: Blending Custom Photos with Video Diffusion Transformers](https://mlsanity.com/paper/2501.01790/) (2025)
- [Video Diffusion Transformers are In-Context Learners](https://mlsanity.com/paper/2412.10783/) (2024)

## Abstract (arXiv)

This paper presents a new exploration into a category of diffusion models built upon state space architecture. We endeavor to train diffusion models for image data, wherein the traditional U-Net backbone is supplanted by a state space backbone, functioning on raw patches or latent space. Given its notable efficacy in accommodating long-range dependencies, Diffusion State Space Models (DiS) are distinguished by treating all inputs including time, condition, and noisy image patches as tokens. Our assessment of DiS encompasses both unconditional and class-conditional image generation scenarios, revealing that DiS exhibits comparable, if not superior, performance to CNN-based or Transformer-based U-Net architectures of commensurate size. Furthermore, we analyze the scalability of DiS, gauged by the forward pass complexity quantified in Gflops. DiS models with higher Gflops, achieved through augmentation of depth/width or augmentation of input tokens, consistently demonstrate lower FID. In addition to demonstrating commendable scalability characteristics, DiS-H/2 models in latent space achieve performance levels akin to prior diffusion models on class-conditional ImageNet benchmarks at the resolution of 256×256 and 512×512, while significantly reducing the computational burden. The code and models are available at: https://github.com/feizc/DiS.