Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer neural operators for PDEs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Replaces self-attention in a Diffusion Transformer (DiT) with a learnable convection-diffusion-reaction PDE block solved in the Fourier domain (PDE-SSM), giving near-linear spatial mixing while matching DiT image generation quality.

Problem

Self-attention in diffusion transformers has quadratic cost in the number of spatial tokens and lacks a strong spatial/locality inductive bias, limiting efficient high-resolution generative modeling.

Approach

Generalizes 1D state space models to N-D space by replacing the ODE with a learnable diffusion-convection-reaction PDE whose Green's function is computed in closed form in the Fourier domain, yielding O(N log N) global spatial mixing as a drop-in replacement for attention in DiT.

Key results

• On CIFAR-10, PDE-SSM-DiT reaches FID 4.18 (50K) versus 4.25 for DiT and 4.19 for U-Net, with 34.2M parameters.
• On LSUN-Churches and Oxford-Flowers, FID drops from 55.04 to 40.71 and from 61.50 to 49.56 relative to DiT.
• On ImageNet256 (DiT-S/2 scale), PDE-SSM-DiT gets FID 49.37 versus 48.44 for DiT, i.e. slightly worse.

Surprising findings

• On ImageNet256 the PDE-SSM variant is slightly worse than DiT in FID (49.37 vs 48.44) despite better IS, which the paper glosses over as comparable.

Limitations

• Spectral translation-equivariant mixing lacks content-dependent token-to-token interactions that attention provides.
• Reported comparisons use attention-tuned DiT hyperparameters and small models; FID differences are small and mostly single-run.
• Mixed results: ImageNet256 FID is slightly worse than DiT.

Experimental setup

Datasets: CIFAR-10 ImageNet64 ImageNet256 CelebA-HQ LSUN-Churches Oxford-Flowers

Benchmarks: CIFAR-10 ImageNet64 ImageNet256 CelebA-HQ LSUN-Churches Oxford-Flowers

Metrics: FID MSE sFID IS precision recall wall-clock time

Baselines: DiT U-Net DDSM iDDPM SLCGAN

Builds on: S4 Mamba DiT flow matching Vision Mamba

Closest papers by meaning

• DiffusionPDE: Generative PDE-Solving under Partial Observation (NeurIPS 2024 · similarity 0.88)
• LaVin-DiT: Large Vision Diffusion Transformer (CVPR 2025 · similarity 0.87)
• PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations (ICML 2025 · similarity 0.87)
• Scaling Diffusion Transformers Efficiently via $\mu$P (NeurIPS 2025 · similarity 0.86)
• Plug-and-Play Diffusion Distillation (CVPR 2024 · similarity 0.86)
• Efficient Generative Transformer Operators for Million-Point PDEs (NeurIPS 2026 · similarity 0.86)