Mamba Integrated with Physics Principles Masters Long-term Chaotic System Forecasting

Long-term forecasting of chaotic systems remains a fundamental challenge due to the intrinsic sensitivity to initial conditions and the complex geometry of strange attractors. The paper states: "Conventional approaches, such as reservoir computing, typically require training data that incorporates long-term continuous dynamical behavior to comprehensively capture system dynamics. While advanced deep sequence models can capture transient dynamics within the training data, they often struggle to maintain predictive stability and dynamical coherence over extended horizons."

The paper addresses the specific problem of forecasting the long-term dynamics of chaotic systems given short-term historical observations on their states, which the authors note remains a largely underexplored scientific problem in existing research. Three fundamental challenges are identified: (i) Extracting complex dynamical information from short observation windows, (ii) Maintaining long-term predictive stability and dynamical coherence, and (iii) Reproducing key statistical properties and attractor geometry.

The authors propose PhyxMamba, a framework that integrates a Mamba-based state-space model with physics-informed principles to forecast long-term behavior of chaotic systems given short-term historical observations on their state evolution. The framework consists of three main components:

The paper states: Rather than relying on neural network-based autoencoders for representation learning, we propose a physics-informed pipeline to construct high-dimensional representations for chaotic systems. Using Takens' embedding theorem, the authors construct high-dimensional representations through time-delay embedding: given a univariate trajectory x ∈ RT from a chaotic system's attractor with T observed steps, we define two hyperparameters: the embedding dimension m and the time delay τ, producing representations zt = (xt−(m−1)τ, xt−(m−2)τ, · · ·, xt) ∈ Rm. The parameters m and τ are chosen using CC methods. The combined representations from different variables are defined as Z ∈ RV ×T ×m. The representations are then partitioned into patches and mapped to latent embeddings via linear mapping.

The paper states: we employ a generative training strategy, enabling the model to predict the next patch autoregressively. The prediction for patch Pi is formulated as P̂i = f (P0, P1, · · ·, Pi−1). The authors adopt a residual stacking architecture designed to decompose the underlying system dynamics into distinct components, with each Mamba layer learning to model different aspects of the dynamics. For the l-th Mamba layer: hi(l) = Mamba(l) (r0:i(l−1)), Êi(l) = Decoder(l) (hi(l)), ri(l) = ri(l−1) − Êi(l).

The paper also introduces a multi-patch prediction (MPP) objective, encouraging the model to capture the global dynamics of the underlying system. This uses M auxiliary modules to predict M subsequent patches, with each module containing a dedicated Mamba layer, and a fusion projection layer ψ(·). The training objective combines next-patch prediction loss and multi-patch prediction losses.

The paper states: To mitigate this issue, we introduce an additional student-forcing training phase, wherein the model autoregressively generates W patches based on its historical predictions. The authors employ Maximum Mean Discrepancy (MMD) to quantify the distance between empirical distributions phist, ppred, and pgt formed by these trajectories. The regularization term is defined as: Lreg = MMD2(phist, ppred) + λc MMD2(pgt, ppred), where the first term enforces statistical consistency between the historical dynamics and the model's predictions, while the second aligns the predicted state distribution with that of the ground truth. The kernel κ is implemented as a mixture of rational quadratic kernels.

Datasets: The method is evaluated on three standard simulated chaotic systems: the 3-dimensional Lorenz63, Rössler, and 5-dimensional Lorenz96, plus real-world datasets: the 5-dimensional Electrocardiogram (ECG) and 64-dimensional Electroencephalogram (EEG). Data is resampled according to their Lyapunov time (TL), where each Lyapunov time corresponds to 30 time steps for most systems, and 60 time steps for EEG.

Evaluation Metrics: The paper uses 1-step Error, symmetric mean absolute percentage error (sMAPE), and valid prediction time (VPT) for point-wise accuracy, and the correlation dimension error (Dfrac) and KL divergence between attractors (Dstsp) for statistical fidelity.

Baselines: Two categories are compared: physics-informed dynamical systems models including Koopa, PLRNN, nVAR, PRC, and HoGRC, and general time series forecasting models including NBEATS, NHiTS, CrossFormer, PatchTST, TimesNet, TiDE, iTransformer, DLinear, NSFormer, FEDFormer, and AutoFormer. Zero-shot performance of Timer and Chronos is also investigated.

Training Setup: During training, our model utilizes data spanning one Lyapunov time (1TL) to perform next-patch prediction with teacher forcing. The model then autoregressively predicts the subsequent TL trajectory based on the preceding TL segment in the student forcing stage. For evaluation, models are given observation trajectories of length TL and tasking them with predicting the subsequent trajectories spanning 10TL.

The paper reports: Our model consistently achieves state-of-the-art point-wise forecasting accuracy across all evaluated datasets. Specifically, our model attains a VPT of 9.83TL and 9.43TL on the Rossler simulation and the real-world EEG dataset, respectively, significantly surpassing baselines. The model also excels at preserving the long-term statistical characteristics and geometric integrity of the strange attractors, with Dfrac of 0.049 and Dstsp of 0.034 on the Rossler dataset versus best baseline values of 0.069 and 0.137.

The paper notes: "Strong short-term predictive accuracy does not necessarily guarantee robust long-term forecasting performance. Models exhibiting lower 1-step prediction errors, such as HoGRC, PRC, and NHiTS, often show rapid degradation in their forecasts."

The ablation study reveals: removing physics-informed representation results in the most substantial degradation of the attractor's long-term geometric integrity, while ablation of student-forcing led to the most acute decline in point-wise accuracy. The paper concludes: all architectural designs and training strategies are crucial for achieving the superior performance demonstrated by our proposed method.

Against Noise: "although point-wise accuracy (sMAPE@1) deteriorates notably when σ > 0.1, the long-term consistency metric (Dstsp) remains remarkably stable, demonstrating the model's ability to faithfully preserve the geometric structure and essential statistical invariants of the system's strange attractor."

Against Training Data Ratio: "our approach exhibits significantly higher data efficiency than the best baseline methods. For instance, on the Rossler system, our model requires only 20% of the full training dataset to outperform the best baseline in both point-wise accuracy and long-term geometric consistency."

The paper concludes: "By synergistically integrating a Mamba-based state-space model with physics-informed embedding, our approach effectively achieves point-wise state forecasting accuracy and reproduces key statistical properties of systems. Through extensive experiments on both simulation and real-world chaotic system datasets, we demonstrate the effectiveness in chaotic system forecasting of our model design and superior performance compared with baseline models. The work paves the way for future research in modeling complex dynamical systems under observation-scarce conditions, with potential applications in climate science, neuroscience, epidemiology, and beyond."

Based on the paper, here are the specific improvements I can make to AI systems, along with what the improved systems can do:

Improvement 1: Physics-Informed State Representation for Time Series Models

• Implementation: Add a pre-processing layer that applies Takens’ time-delay embedding to each input variable before feeding it into the model. This creates a higher-dimensional representation (Z ∈ R V×T×m) that captures the underlying attractor manifold. Use the CC method to automatically select embedding dimension (m) and delay (τ).

• What the improved AI system can do: It can learn the global dynamics of a chaotic system from a very short observation window (e.g., 1 Lyapunov time) instead of requiring long, continuous trajectories. This makes it effective for real-world scenarios where only brief historical data is available (e.g., early warning systems for equipment failure, short-term climate anomaly detection).

Improvement 2: Generative Next-Patch Training with Multi-Patch Prediction (MPP)

• Implementation: Replace standard sequence-to-sequence training with a decoder-only, autoregressive next-patch prediction objective. Add M auxiliary prediction heads that forecast M future patches simultaneously during training. The loss function combines next-patch loss (L next) and multi-patch losses (L MPP) with a weighting factor λp.

• What the improved AI system can do: It learns the underlying generative process rather than just a direct input-output mapping. This enables the model to maintain predictive stability and dynamical coherence over long horizons (e.g., 10 Lyapunov times), preventing the error accumulation that plagues standard autoregressive models. It also improves data efficiency, requiring only 20–40% of training data to outperform baselines trained on full datasets.

Improvement 3: Two-Stage Training with Student Forcing and Distribution Matching

• Implementation: After the initial teacher-forcing phase, add a second training stage where the model autoregressively generates W patches based on its own predictions (student forcing). During this stage, add a regularization term (L reg) that uses Maximum Mean Discrepancy (MMD) with a mixture of rational quadratic kernels to minimize the distance between the predicted state distribution and both the historical and ground-truth distributions.

• What the improved AI system can do: It becomes robust to its own prediction errors during inference, avoiding the train-test mismatch that causes performance degradation. The MMD regularization ensures the model reproduces the invariant measure and fractal geometry of the strange attractor, preventing collapse into trivial states or spurious attractors. This is critical for applications requiring statistical fidelity, such as generating realistic synthetic data for stress-testing financial models or simulating disease spread for public health planning.

Improvement 4: Residual Stacking Architecture for Decomposed Dynamics

• Implementation: Use a stack of Mamba layers where each layer l learns the residual between the input and the cumulative prediction of previous layers. Each layer outputs a component (Ê i(l)) that represents a distinct aspect of the dynamics, and the final prediction is the sum of all components.

• What the improved AI system can do: It can decompose complex, composite system dynamics into interpretable components, each modeled by a separate layer. This improves forecasting accuracy by allowing the model to capture both fast and slow timescales separately, and it provides a level of interpretability that is useful for scientific discovery (e.g., identifying which components drive short-term vs. long-term behavior).

Improvement 5: Adaptive Observation Length at Inference

• Implementation: Design the model to accept variable-length input sequences at inference time (as long as the length is a multiple of the patch size). This is enabled by the generative next-patch training paradigm.

• What the improved AI system can do: It can generate accurate long-term forecasts even when the available historical context is shorter than the training window. The paper shows that the model maintains stable long-term geometric fidelity (Dstsp) even with only 0.2TL of input, and in some cases, shorter inputs yield better short-term accuracy than longer ones. This flexibility is crucial for real-time applications where data may be incomplete or delayed.

Summary of Capabilities of the Improved AI System:

• Long-horizon forecasting: Predicts chaotic system behavior for up to 10 Lyapunov times with high point-wise accuracy (e.g., VPT of 9.83TL on Rossler, 9.43TL on EEG).

• Statistical fidelity: Reproduces the attractor’s correlation dimension and KL divergence with errors as low as 0.049 and 0.034, respectively, on the Rossler system.

• Data efficiency: Achieves superior performance with only 20–40% of the training data required by state-of-the-art baselines.

• Noise robustness: Maintains stable long-term statistical properties even under significant observational noise (σ up to 0.5).

• Flexibility: Can operate with variable-length historical inputs and is computationally efficient (e.g., 3.65M parameters, 184ms inference time for 10TL forecasting on Lorenz63).

Long-term forecasting of chaotic systems remains a fundamental challenge due to the intrinsic sensitivity to initial conditions and the complex geometry of strange attractors. Conventional approaches, such as reservoir computing, typically require training data that incorporates long-term continuous dynamical behavior to comprehensively capture system dynamics. While advanced deep sequence models can capture transient dynamics within the training data, they often struggle to maintain predictive stability and dynamical coherence over extended horizons. Here, we propose PhyxMamba, a framework that integrates a Mamba-based state-space model with physics-informed principles to forecast long-term behavior of chaotic systems given short-term historical observations on their state evolution. We first reconstruct the attractor manifold with time-delay embeddings to extract global dynamical features. After that, we introduce a generative training scheme that enables Mamba to replicate the physical process. It is further augmented by multi-patch prediction and attractor geometry regularization for physical constraints, enhancing predictive accuracy and preserving key statistical properties of systems. Extensive experiments on simulated and real-world chaotic systems demonstrate that PhyxMamba delivers superior forecasting accuracy and faithfully captures essential statistics from short-term historical observations.