cs.LGMar 14, 2026

PDE-SSM: A Spectral State Space Approach to Spatial Mixing in Diffusion Transformers

Authors: Eshed Gal, Moshe Eliasof, Eldad Haber

Organizations: University of British Columbia, Vancouver, BC Canada · University of Cambridge, Cambridge, United Kingdom

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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(Nlog⁡N)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.

Figures & tables

Figure 1 : Visualizing the PDE-SSM Convolutional Kernels. By sampling the learnable parameters ξ=(Bγ,ζ)\xi=(\mathcal{B}_{\gamma},\zeta) , our PDE-SSM can represent a diverse family of convolutional kernels. The examples show kernels that are (from left to right): localized, directionally blurred (anisotropic diffusion), shifted (convection), and a combination of effects. This flexibility allows our PDE-SSM model to learn a rich basis for spatial feature mixing, including non-local connections.
Figure 2
Figure 2 : CIFAR-10 Images: (a) real images; (b) DiT; (c) PDE-SSM-DiT. Visual quality is comparable, in congruence with Table 1 .
Figure 3 : LSUN-Churches generations: (a) real images; (b) DiT; (c) PDE-SSM-DiT.
Method#Params (M)MSE ↓\downarrowFID (50K) ↓\downarrow
DiT29.63.82×10−23.82\times 10^{-2}4.25
U-Net32.63.78×10−23.78\times 10^{-2}4.19
PDE-SSM-DiT (Ours)34.23.76×10−2\mathbf{3.76\times 10^{-2}}4.18
Table 1 : PDE-SSM-DiT matches attention-based DiT and U-Net performance on CIFAR-10, indicating quality parity at equal settings , isolating the effect of the block swap ( RQ1, RQ2 ).
MethodParams (M)FID ↓\downarrow
DiT3322.9
U-Net3222.2
PDE-SSM-DiT (Ours)3122.1
Table 2 : Unconditional Generation on ImageNet- 64×6464{\times}64 .
Figure 4 : Oxford-Flowers generations: (a) real images; (b) DiT; (c) PDE-SSM-DiT.
MethodCelebALSUNOxford
HQ 6464ChurchesFlowers
DiT9.1755.0461.50
U-Net11.0651.5657.20
PDE-SSM-DiT (Ours)7.0140.7149.56
Table 3 : FID (2048 features) on CelebA-HQ 6464 , LSUN-Churches and Oxford-Flowers.
ArchitectureParametersFIDMSE
DiT32.99 M21.088.02e-2
PDE-SSM-DiT(Ours)26.48 M18.365.37e-2
Table 4 : Training results on CelebA HQ 256
ImagePatchParams (M)Time (s)
 SizeAttnPDE-SSMAttnPDE-SSM
3229.0112.761.09×10−11.09{\times}10^{-1}2.14×10−22.14{\times}10^{-2}
3248.9812.763.34×10−23.34{\times}10^{-2}2.40×10−22.40{\times}10^{-2}
3289.1312.761.60×10−21.60{\times}10^{-2}2.40×10−22.40{\times}10^{-2}
6429.3112.784.66×10−14.66{\times}10^{-1}5.58×10−25.58{\times}10^{-2}
6449.0612.789.92×10−29.92{\times}10^{-2}5.77×10−25.77{\times}10^{-2}
Table 5 : DiT vs. PDE-SSM-DiT training step runtimes (in seconds). Parameters in millions (M).
Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

 CIFAR-10ImageNetCelebA-HQLSUN-ChurchesOxford-Flowers
Max iter400,000400,000400,000400,000400,000
Learning rate5×10−45\times 10^{-4}5×10−45\times 10^{-4}5×10−45\times 10^{-4}5×10−45\times 10^{-4}5×10−45\times 10^{-4}
Hidden size384384384512512
Batch size6464643232
Image size3264646464
Layers1212121212
Appendix
Table 6 : Experimental settings and hyperparameters for different experiments.
Figure 5 : SSM maps visualization: (a) real simple image; (b) simple image SSM map; (c) real church image; (d) church SSM map.
ModelPatch SizeTime (s)
DiT-XL80.1288
DiT-XL40.9633
DiT-XL210.4166
PDE-SSM-DiT-XL(Ours)80.1509
PDE-SSM-DiT-XL(Ours)40.6315
PDE-SSM-DiT-XL(Ours)22.3376
Appendix
Table 7 : Iteration time vs. patch size for DiT-XL and PDE-SSM-DiT-XL models.
ModelVariantTime (s)
DiTXL10.4162
DiTL5.9323
DiTB2.1223
DiTS0.9863
PDE-SSM-DiT(Ours)XL2.3435
PDE-SSM-DiT(Ours)L1.6020
Appendix
Table 8 : Iteration time for DiT and PDE-SSM-DiT models.
ModelFID ↓\downarrowsFID ↓\downarrowIS ↑\uparrowPrecision ↑\uparrowRecall ↑\uparrow
DiT48.449.0430.780.450.53
PDE-SSM-DiT-S-2 (Ours)49.3711.2134.090.470.48
Appendix
Table 9 : Imagenet256 results.
Figure 6 : ImageNet 6464 training. (a) The achieved FID score is consistent with the internal FID score. (b) All methods converge at a similar rate and to an FID score that is similar.
Figure 7 : Sample celebrity faces generated by PDE-SSM-DiT (left) and DiT (right), trained on CelebA-HQ 256.

Explore similar work

CardsList

  1. Pixel-Space Diffusion Transformers

    Jul 20, 2026Renye Yan, Jikang Cheng, You Wu +8Diffusion TransformersVisual Tokenizers

    Peking University, China · Nanjing University, China · Stanford University, USA +3

  2. CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers

    May 13, 2026Zhuojin Li, Hsin-Pai Cheng, Hong Cai +2Diffusion TransformersVisual Token Pruning

    Qualcomm AI Research†

  3. Phase-Aligned RoPE for Mixed-Resolution Diffusion Transformer

    Nov 24, 2025Haoyu Wu, Jingyi Xu, Qiaomu Miao +2Diffusion TransformersRotary Position Embedding

    Stony Brook University, NY, USA · UNC-Charlotte, NC, USA