💻 computer science

Oscillatory State-Space Models as Inductive Biases for Physics-Informed Neural PDE Solvers

This paper proposes a novel Physics-Informed Neural Network (PINN) architecture that integrates oscillatory state-space dynamics and PDE-aware spectral bases to efficiently solve time-dependent partial differential equations with improved accuracy and reduced memory usage across forward, inverse, and high-dimensional problems.

Original authors: Abhishek Chandra, Taniya Kapoor

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Abhishek Chandra, Taniya Kapoor

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to predict how a wave moves across a pond, or how heat spreads through a metal plate. These are problems described by Partial Differential Equations (PDEs). For decades, scientists have used "Physics-Informed Neural Networks" (PINNs) to solve these. Think of these networks as students who are given the rules of physics (the equations) and asked to guess the answer.

However, these students often struggle with time. They are good at looking at a snapshot, but they get confused about how things evolve smoothly over seconds, minutes, or hours. They tend to "drift" away from the truth or get stuck in loops.

This paper introduces a new type of student called OSSM-PINN. Here is how it works, using simple analogies:

1. The Problem: The "Generic Sequence" Student

Previous methods tried to teach the computer to predict time by treating it like a generic sequence of events, similar to how a language model predicts the next word in a sentence.

The Flaw: If you ask a language model to predict a swinging pendulum, it might guess the pendulum will slowly drift off into space or stop moving entirely. It doesn't "know" that pendulums naturally swing back and forth.
The Result: These models often produce messy, inaccurate predictions, especially for high-speed waves or complex 3D problems, and they require a massive amount of computer memory (like trying to carry a library in your backpack).

2. The Solution: The "Oscillatory" Student

The authors built a new architecture that forces the computer to understand time the way physics does: as a series of natural vibrations.

The Analogy of the Tuning Fork: Imagine a tuning fork. When you strike it, it vibrates at a specific, pure frequency. It doesn't wander off; it oscillates.
How OSSM-PINN Works: Instead of letting the computer guess how time moves, the researchers gave it a "tuning fork" built into its brain. This is called an Oscillatory State-Space Model.
- Time (The Vibration): The computer's internal "memory" is forced to vibrate like a spring or a pendulum. It is mathematically impossible for it to drift away; it must oscillate. This matches how waves, heat, and quantum particles actually behave.
- Space (The Shape): While the time part vibrates, the computer uses a fixed set of "shapes" (like sine waves or specific mathematical curves) to describe the physical space. It's like describing a complex sound not by recording every air molecule, but by listing the specific notes (frequencies) that make up the chord.

3. Why This is a Big Deal

The paper claims this approach solves three major headaches:

Accuracy: Because the computer is "hard-wired" to think in vibrations, it captures the natural rhythm of physics much better. In tests, it was up to 219 times more accurate than previous methods for certain wave problems. It didn't just "guess" the wave; it understood the beat.
Memory Efficiency: Old methods tried to remember every single step of the timeline, which filled up the computer's memory quickly (like trying to write down every frame of a movie). This new method uses a "parallel scan" (like reading a book by jumping to the chapters you need instantly) and a compact "vibration" model. It solved a problem with 100 spatial dimensions (a complexity level that usually crashes computers) without running out of memory.
Versatility: The same "vibrating brain" worked for:
- Forward problems: Predicting how a system evolves.
- Inverse problems: Looking at a messy result and figuring out the hidden rules (like guessing the wind speed based on how a flag flapped).
- Weird shapes: It worked on triangular domains and large areas without needing to redraw the map.

4. The "Cell Swap" Proof

To prove that the "vibration" part was the secret sauce, the researchers did a "cell swap" experiment. They took their new model and replaced the vibrating time-cell with a standard, non-vibrating one.

The Result: The model's performance crashed immediately. This proved that the accuracy wasn't just because the model was bigger or smarter; it was specifically because it was forced to oscillate like a physical system.

5. Limitations (The Fine Print)

The authors are honest about where this method might struggle.

The "Smoothness" Assumption: This method assumes the solution is made of smooth, repeating waves (like a calm ocean).
The "Shock" Problem: If the physics involves sudden, sharp breaks—like a car crash, a sonic boom, or a shockwave where things change instantly—this "vibrating" approach might need too many notes to describe the jagged edge. The paper suggests future work will need to handle these "jagged" moments better.

Summary

In short, this paper says: "Don't teach a computer to predict time like it's writing a story; teach it to predict time like it's playing a musical instrument." By building the natural rhythm of physics directly into the neural network's architecture, they created a solver that is faster, more accurate, and capable of handling problems that were previously too big to solve.

Technical Summary: Oscillatory State-Space Models as Inductive Biases for Physics-Informed Neural PDE Solvers

1. Problem Statement

Solving time-dependent partial differential equations (PDEs) remains a significant challenge in computational science. While Physics-Informed Neural Networks (PINNs) have emerged as a mesh-free method for simulating PDEs and solving inverse problems, they struggle with time-dependent dynamics. Specifically, standard coordinate-based Multi-Layer Perceptron (MLP) PINNs and recent sequence-model-based approaches (e.g., Transformers, State-Space Models like Mamba) face difficulties capturing high-frequency dynamics, stiffness, long-time integration, and high-dimensional domains.

Recent sequence-model-based PINNs improve upon coordinate MLPs by parameterizing time evolution using general-purpose sequence backbones. However, the authors argue these methods inherit design choices from language and generic time-series modeling that fail to encode the structured modal dynamics inherent to PDE solutions. As visualized in the paper's Figure 1, these models often collapse latent states into smooth, unbounded drifts rather than the oscillatory behaviors dictated by PDE dispersion relations. Furthermore, these approaches suffer from unfavorable memory scaling with sequence length and spatial resolution, limiting their applicability in high-dimensional or large-scale settings.

2. Methodology: OSSM-PINN

The paper proposes Oscillatory State-Space Model Physics-Informed Neural Networks (OSSM-PINNs). This architecture decouples spatial and temporal representations, enforcing a specific inductive bias on the temporal evolution of the solution.

2.1 Modal Factorization

Instead of learning a direct space-time map $(x, t) \mapsto u(x, t)$ , OSSM-PINN approximates the solution via a modal factorization:
$\hat{u}(x, t) = D(x) \sum_{k=1}^{K} c_k(h(t)) \phi_k(x) + g(x)$

Spatial Basis ( $\phi_k$ ): A fixed, analytical basis (e.g., Fourier modes, sine modes, or problem-specific eigenfunctions like Hermite polynomials) is chosen based on the PDE domain and boundary conditions. This allows for closed-form spatial differentiation of any order and exact enforcement of boundary conditions via a boundary factor $D(x)$ and offset $g(x)$ .
Modal Coefficients ( $c_k$ ): The network learns a finite set of time-evolving coefficients rather than the full field.

2.2 Oscillatory Temporal Evolution

The temporal evolution is driven by a Linear Oscillatory State-Space (LinOSS) block, which serves as the inductive bias for the time dynamics.

Latent State: The initial condition is encoded into a latent state $(y_0, z_0)$ by an MLP encoder.
Dynamics: The state evolves according to forced linear oscillators:
$\dot{y} = z, \quad \dot{z} = -A y + B s(t)$
where $A$ is a non-negative diagonal matrix (parameterized as $\text{ReLU}(\hat{A})$ ) ensuring eigenvalues lie on the imaginary axis (undamped oscillation) or inside the unit disk (dissipative dynamics).
Discretization: Two stable discretization schemes are employed:
1. IMEX (Implicit-Explicit): Symplectic for autonomous parts, suitable for conservative dynamics.
2. IM (Implicit): Introduces controlled dissipation, suitable for dissipative dynamics.
Parallelization: The recurrence is associative, allowing the full temporal trajectory to be computed via parallel scan in $O(\log N_t)$ depth, avoiding the $O(N_t)$ cost of sequential integration.

2.3 Training Objective

The model is trained by minimizing a physics-informed loss function comprising:

PDE Residual: The governing equation $F(\hat{u}) = 0$ evaluated at collocation points. Temporal derivatives are computed using high-order (default 6th-order) finite differences on the predicted sequence.
Initial Condition (IC) Loss: Ensures $\hat{u}(x, 0) = u_0(x)$ .
Data Loss (Inverse Problems): Includes sparse observational data for inverse tasks.

3. Key Contributions

Identification of Temporal Inductive Bias: The paper posits that the temporal component is the primary locus for physics-aligned inductive bias in neural PDE solvers. It demonstrates that a state-space recurrence constrained to undamped oscillation is structurally appropriate for PDEs.
Architecture Design (OSSM-PINN): Introduces a novel architecture coupling an oscillatory state-space cell with a fixed spatial basis. This design learns time-evolving modal coefficients rather than unconstrained space-time fields.
Empirical Validation: The method is evaluated across 12 benchmarks, including forward/inverse problems, high-dimensional domains (up to 100D), non-rectangular geometries, and large domains.
Ablation and Analysis: Through cell-swap analysis and component ablation, the authors demonstrate that performance gains stem specifically from the oscillatory eigenstructure of the temporal cell, not merely from increased model capacity or specific basis choices.

4. Experimental Results

The method was evaluated against sequence-model baselines (PINNsFormer, PINNMamba, ML-PINN) and neuro-spectral baselines (NeuSA) under a fixed 24 GiB GPU memory budget.

Accuracy: OSSM-PINN consistently outperformed baselines.
- High-Frequency Convection: Reduced relative mean absolute error (rMAE) by up to 107 $\times$ (IMEX) compared to the best baseline.
- Wave Propagation: Reduced rMAE by 132 $\times$ .
- Euler-Bernoulli Beam: Achieved reductions of 49,000 $\times$ to 98,000 $\times$ on extended benchmarks where baselines failed (Out of Memory or high error).
- Inverse Problems: Successfully recovered unknown coefficients in KdV and Sea Surface Temperature (SST) advection-diffusion problems with relative errors below 0.1%, while all baselines failed.
- High Dimensions: Successfully solved 5D and 100D Schrödinger equations, whereas baselines exceeded memory limits.
Memory Efficiency: The modal factorization and parallel scan approach significantly reduced memory usage, enabling solutions to problems (e.g., coupled 2D flows, 100D Schrödinger) that caused sequence-model baselines to run out of memory (OOM).
Physical Consistency: Spectral analysis confirmed that OSSM-PINN preserves the correct Fourier spectrum and that learned latent frequencies align with analytical PDE dispersion relations, unlike baselines which exhibited spurious high-frequency content or drift.
Optimization Landscape: The loss landscape for OSSM-PINN was found to be smoother (lower Lipschitz constant) compared to baselines, suggesting easier optimization.

5. Significance and Claims

The paper claims that neural PDE solvers should not treat time as a generic sequence. Instead, encoding oscillatory modal structure provides an effective and efficient inductive bias.

Physics-Alignment: By constraining the temporal evolution to oscillatory dynamics, the architecture naturally recovers characteristic frequencies and phase shifts from the PDE residual, rather than learning them as generic sequence patterns.
Efficiency: The decoupling of space and time, combined with parallelizable state-space rollouts and analytical spatial derivatives, addresses the scalability bottlenecks of current sequence-model PINNs.
Generality: The same backbone architecture applies across diverse problem classes (forward, inverse, high-dimensional, non-rectangular) by only changing the spatial basis and residual definition, without architectural re-engineering.

6. Limitations and Future Work

The authors acknowledge that the method assumes the target solution is well-represented by a finite set of spatial modes with structured temporal evolution. This assumption may be less effective for:

Shocks and moving discontinuities.
Intermittent localized structures.
Strongly non-modal dynamics.

Future work is suggested to extend the framework with adaptive, localized, multi-resolution, or domain-decomposed bases, as well as adaptive residual sampling and automated basis selection.