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

• URL: http://arxiv.org/abs/2603.13663v1
• Date: Sat, 14 Mar 2026 00:16:10 GMT
• Title: PDE-SSM: A Spectral State Space Approach to Spatial Mixing in Diffusion Transformers
• Abstract summary: We propose PDE-SSM, a spatial state-space block that replaces attention with a learnable convection-diffusion-reaction partial differential equation.<n>Our results show that, analogous to 1D settings where SSMs attention, multi-dimensional PDE operators provide an efficient, inductive-bias-rich foundation for next-generation vision models.
• Score: 17.569349375085206
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Unsupervised Adaptation of PDE Foundation Models [17.002625091063937]
  We propose an unsupervised PDE-based finetuning framework that eliminates the need for ground-truth solutions.<n>Our method achieves performance comparable to supervised LoRA finetuning without requiring any ground-truth solutions.
  arXiv  Detail & Related papers  (2026-08-07T10:01:50Z)
• PDE-Transformer: A Continuous Dynamical Systems Approach to Sequence Modeling [4.1812935375151925]
  We propose PDE-Transformer, a sequence modeling paradigm that casts the forward pass of a Transformer as the numerical discretization of a continuous reaction-diffusion system.<n>In our framework, token embeddings evolve under a partial differential equation whose nonlocal integral term models self-attention.<n>We design an Adaptive PDE Diffusion Layer that enforces local smoothness in feature space with linear time complexity.
  arXiv  Detail & Related papers  (2025-09-27T08:58:47Z)
• Physics-Informed Distillation of Diffusion Models for PDE-Constrained Generation [19.734778762515468]
  diffusion models have gained increasing attention in the modeling of physical systems, particularly those governed by partial differential equations (PDEs)<n>We propose a simple yet effective post-hoc distillation approach, where PDE constraints are not injected directly into the diffusion process, but instead enforced during a post-hoc distillation stage.
  arXiv  Detail & Related papers  (2025-05-28T14:17:58Z)
• Mechanistic PDE Networks for Discovery of Governing Equations [52.492158106791365]
  We present Mechanistic PDE Networks, a model for discovery of partial differential equations from data.<n>The represented PDEs are then solved and decoded for specific tasks.<n>We develop a native, GPU-capable, parallel, sparse, and differentiable multigrid solver specialized for linear partial differential equations.
  arXiv  Detail & Related papers  (2025-02-25T17:21:44Z)
• Advancing Generalization in PINNs through Latent-Space Representations [71.86401914779019]
  Physics-informed neural networks (PINNs) have made significant strides in modeling dynamical systems governed by partial differential equations (PDEs)<n>We propose PIDO, a novel physics-informed neural PDE solver designed to generalize effectively across diverse PDE configurations.<n>We validate PIDO on a range of benchmarks, including 1D combined equations and 2D Navier-Stokes equations.
  arXiv  Detail & Related papers  (2024-11-28T13:16:20Z)
• Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers [53.79279286773326]
  We present Unisolver, a novel Transformer model trained on diverse data and conditioned on diverse PDEs.<n>Unisolver achieves consistent state-of-the-art on three challenging large-scale benchmarks, showing impressive performance and generalizability.
  arXiv  Detail & Related papers  (2024-05-27T15:34:35Z)
• AdjointDEIS: Efficient Gradients for Diffusion Models [2.0795007613453445]
  We show that continuous adjoint equations for diffusion SDEs actually simplify to a simple ODE.<n>We also demonstrate the effectiveness of AdjointDEIS for guided generation with an adversarial attack in the form of the face morphing problem.
  arXiv  Detail & Related papers  (2024-05-23T19:51:33Z)
• The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit [38.89510345229949]
  We study the covariance matrix of a modified Softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. To achieve a well-defined limit, the Transformer's attention mechanism is modified by centering the Softmax output at identity. We show, through simulations, that the differential equation (SDE) indexed by the depth-to-width ratio provides a surprisingly good description of the corresponding finite-size model.
  arXiv  Detail & Related papers  (2023-06-30T16:10:36Z)
• PDE+: Enhancing Generalization via PDE with Adaptive Distributional Diffusion [66.95761172711073]
  generalization of neural networks is a central challenge in machine learning. We propose to enhance it directly through the underlying function of neural networks, rather than focusing on adjusting input data. We put this theoretical framework into practice as $textbfPDE+$ ($textbfPDE$ with $textbfA$daptive $textbfD$istributional $textbfD$iffusion)
  arXiv  Detail & Related papers  (2023-05-25T08:23:26Z)
• Solving High-Dimensional PDEs with Latent Spectral Models [74.1011309005488]
  We present Latent Spectral Models (LSM) toward an efficient and precise solver for high-dimensional PDEs. Inspired by classical spectral methods in numerical analysis, we design a neural spectral block to solve PDEs in the latent space. LSM achieves consistent state-of-the-art and yields a relative gain of 11.5% averaged on seven benchmarks.
  arXiv  Detail & Related papers  (2023-01-30T04:58:40Z)
• Learning to Accelerate Partial Differential Equations via Latent Global Evolution [64.72624347511498]
  Latent Evolution of PDEs (LE-PDE) is a simple, fast and scalable method to accelerate the simulation and inverse optimization of PDEs. We introduce new learning objectives to effectively learn such latent dynamics to ensure long-term stability. We demonstrate up to 128x reduction in the dimensions to update, and up to 15x improvement in speed, while achieving competitive accuracy.
  arXiv  Detail & Related papers  (2022-06-15T17:31:24Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.