Title: PDE-SSM: A Spectral State Space Approach to Spatial Mixing in Diffusion Transformers
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: PDE state-space models, flow-matching diffusion transformers

Abstract: The success of vision transformers—especially for generative modeling—is limited by the quadratic cost and weak spatial inductive bias of self-attention. We propose PDE-SSM, a spatial state-space block that replaces attention with a learnable convection–diffusion–reaction partial differential equation. This operator encodes a strong spatial prior by modeling information flow via physically grounded dynamics rather than all-to-all token interactions. Solving the PDE in the Fourier domain yields global coupling with near-linear complexity of $O(N \log N)$, delivering a principled and scalable alternative to attention. We integrate PDE-SSM into a flow-matching generative model to obtain the PDE-based Diffusion Transformer PDE-SSM-DiT. Empirically, PDE-SSM-DiT matches or exceeds the performance of state-of-the-art Diffusion Transformers while substantially reducing compute. Our results show that, analogous to 1D settings where SSMs supplant attention, multi-dimensional PDE operators provide an efficient, inductive-bias-rich foundation for next-generation vision models.

TL;DR: We present a PDE inspired state-space mechanism for Diffusion Transformers that lifts 1D SSMs to 2D via learnable convection-diffusion-reaction dynamics in Fourier space, delivering near O(N log N) mixing and DiT-level quality with lower compute.