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
Jul 20, 2026cs.CV

Pixel-Space Diffusion Transformers

Latent diffusion models (LDMs) enable efficient high-resolution image synthesis by denoising in a VAE-compressed latent space. However, fixed visual tokenizers can discard fine textures and structural details, while separate representation and diffusion training creates a mismatch between reconstruction and generation objectives. These limitations have renewed interest in pixel-space diffusion, which models raw pixels directly, removes the VAE bottleneck, and supports end-to-end optimization. This formulation better matches the demands of high-fidelity generation but introduces challenges in high-dimensional modeling, including noise scheduling, loss weighting, token efficiency, and scalable architecture design. Pixel-space modeling also offers a promising basis for unified multimodal systems: raw pixels, text, and task conditions can be represented in a shared token space and jointly processed by a single Transformer, narrowing the gap between visual understanding and generation. This paper reviews Pixel-Space Diffusion Transformers (pDiTs) from the perspectives of model architecture, continuous generative mechanisms, and unified multimodal modeling. We summarize representative methods, identify key technical challenges, and discuss future directions toward high-fidelity, end-to-end vision foundation models that integrate generation and understanding.

Renye Yan, Jikang Cheng, You Wu +8

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

Diffusion TransformersVisual TokenizersTransformer Architectures
AskPDF
May 13, 2026cs.CV

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

Diffusion Transformers (DiTs) deliver remarkable image and video generation quality but incur high computational cost, limiting scalability and on-device deployment. We introduce CoReDiT, a structured token pruning framework for DiTs across vision tasks. CoReDiT uses a linear-time spatial coherence score to estimate local redundancy in the latent token lattice and skips high coherence (redundant) tokens in self-attention. To maintain a dense representation and avoid visual discontinuities, we reconstruct skipped attention outputs via coherence-guided aggregation of spatially neighboring retained tokens. We further introduce a progressive, block-adaptive pruning schedule that increases pruning gradually and allocates larger budgets to blocks and denoising steps with higher redundancy. Across state-of-the-art diffusion backbones including PixArt-α and MagicDrive-V2, CoReDiT achieves up to 55% self-attention FLOPs reduction and inference speedups of 1.33x on cloud GPUs and 1.72x on mobile NPUs, while maintaining high visual quality. Notably, CoReDiT also increases on-device memory head-room, enabling higher-resolution generation.

Zhuojin Li, Hsin-Pai Cheng, Hong Cai +2

Qualcomm AI Research†

Diffusion TransformersVisual Token PruningTransformer ArchitecturesSpatial Coherence
AskPDF
Nov 24, 2025cs.CV

Phase-Aligned RoPE for Mixed-Resolution Diffusion Transformer

Rotary positional embeddings (RoPE) are widely used in diffusion transformers (DiTs) to encode spatial relationships, yet their behavior with mixed-resolution tokens remains underexplored. A natural approach is to rescale token positions from different resolutions into a unified coordinate system before attention, but we show this fails. Our analysis shows that with RoPE, the attention similarity score is a highly structured and periodic function of token distance, so rescaling distances across resolutions moves token pairs to different regions of this periodic function, leading to incorrect attention scores. Motivated by this, we introduce Phase-Aligned Mixed-Resolution Attention (PMA), a training-free mechanism that stabilizes mixed-resolution attention. PMA modifies the RoPE position mapping to enforce a consistent positional scale for every query-key pair, ensuring that relative distances are evaluated under a single reference scale. To further improve local coherence near resolution transitions, we incorporate a lightweight boundary refinement module that softly exchanges features across adjacent scales. Experiments on image and video diffusion models validate our analysis and demonstrate consistent improvements in visual fidelity and computational efficiency.

Haoyu Wu, Jingyi Xu, Qiaomu Miao +2

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

Diffusion TransformersRotary Position EmbeddingRepresentation-Alignment Loss
AskPDF