A Generalizable Physics-Enhanced State Space Model for Long-Term Dynamics Forecasting in Complex Environments

Response to Reviewer gB8T

Q 4.1: About clarifying the difference between our method and existing hybrid modeling approaches [1-3] that address incomplete physics knowledge.

A 4.1: We have discussed the difference between our model and hybrid modeling approaches as follows. The models in [1–3] are all physics-based NODE methods, which suffer from a shared limitation: difficulty in modeling nonlinear and time-variant systems over long time horizons due to their heavy reliance on initial conditions. This limits their ability to capture long-term sequence dependencies.

PI-VAE [2] is already included as a baseline in Tables 1–4 of our submission, and our method outperforms it. We will include references [1] and [3] in the related work section of the revised version.

[1] Yin et al. Augmenting physical models with deep networks for complex dynamics forecasting. JSTAT 2021

[2] Takeishi et al. Physics-integrated variational autoencoders for robust and interpretable generative modeling. NIPS 2021

[3] Wehenkel et al. Robust hybrid learning with expert augmentation. TMLR 2023

Q 4.2: About (i) comparing with baselines like ODE2VAE; (ii) NODE limitations on initialization sensitivity.

A 4.2: We have included ODE2VAE [4] as a baseline in our experiments. The results are provided in Tables 1-4 [link]. While ODE2VAE introduces uncertainty modeling to alleviate the sensitivity to initial conditions, it still lacks an effective mechanism to dynamically refine predictions based on later observations. As a result, its performance on long-term forecasting remains limited.

Regarding the discussion of initialization in Section 2, we refer specifically to the sensitivity of NODE-based methods to initial conditions $x(t_0)$, not to the initialization of network parameters. This is due to the fact that NODEs solve initial value problems (IVPs), where the trajectory is determined by both the initial state and the learned vector field. In contrast, our Phy-SSM adopts a dynamical VAE framework that refines its predictions over time using the posterior from the preceding time step. This mechanism effectively mitigates errors caused by inaccurate initial conditions.

A motivating example that visually illustrates this improvement is provided in Fig. 2, Appendix A.

[4] Yildiz et al. Ode2vae. NIPS 2019

Q 4.3: About how Phy-SSM handles continuous dynamics and the definition of $\psi(z)$.

A 4.3: The Phy-SSM unit learns continuous dynamics through the parameterized matrices $A_{unk}$ and $B_{unk}$ as showed in Eq. (8), as illustrated in S5 [5].

$\psi(z)$ denotes additional extended state terms, which can include nonlinear functions or constants—such as $\sin{({z})}$, $\cos{({z})}$ and 1. This augmentation enables the representation of certain nonlinear systems in a linear state-space form. Please refer to page 4, right column, lines 192–200 for the detailed explanation.

[5] Smith et al. S5. ICLR 2023

Q 4.4: About defining the knowledge mask when physics is fully unknown.

A 4.4: If physics is unknown, we treat the entire system as unknown and set all mask entries to 1. In this case, our method reduces to a data-driven SSM that learns all dynamics from data without explicit constraints.

Q 4.5: About elaborating on why the L2 norm of the latent states from the prior and posterior distribution is used.

A 4.5: We use L2-norm for its efficiency and strong empirical performance. We have conducted comparisons with alternative metrics(e.g., Chebyshev, cosine). The results are provided in Table 5 [link], where Euclidean distance (L2-norm) performs best in extrapolation tasks.

Q 4.6: About the irregularity details of the three real-world datasets used in the experiments.

A 4.6: Appendix E in our submission provides the details. Drone data is high-frequency and irregularly sampled, recorded at nearly 1010 Hz (minimum: 573.05 Hz, maximum: 1915.86 Hz); COVID-19 contains 10% missing daily records; vehicle dataset includes 5% missing agent observations. We will add clearer descriptions in the revised version.

Q 4.7: About the seemingly marginal contribution of the physics state regularization term and the request for a regularization-only comparison.

A 4.7: We would like to emphasize that the physics state regularization yields approximately a 10% relative improvement in extrapolation performance. It plays a crucial role in guiding the model toward learning more generalized physically representations—particularly under noisy and irregular data conditions.

It is important to note that the regularization term cannot be used independently, as it penalizes the distance between the posterior (output by the sequential encoder) and the prior physics-based latent states predicted by the Phy-SSM unit.

Response to Reviewer xWcK

Q3.1: About handling uncertainty in domain knowledge within our model.

A3.1: We do not explicitly handle uncertainty in domain knowledge in this work. However, such uncertainty can be modeled within the unknown dynamics. A possible approach is to apply conformal prediction [1] as a post-processing step to estimate uncertainty in the deep SSM outputs. We will explore this in future work.

[1] Kamile et al. Conformal time-series forecasting. NIPS 2021

Q3.2 About the lack of ablations comparing L2 with other regularization metrics.

A3.2: Following your suggestion, we have added comparisons with alternative metrics, including Chebyshev distance, and cosine distance. The results are presented in Table 5 [link], where our method achieves the best performance in extrapolation tasks. We will include these results in the revision.

Chebyshev distance focuses on worst-case deviations, while cosine distance captures directional similarity, which may not fully penalize deviations in magnitude. Since our goal is to measure the overall discrepancy between two physical state trajectories, the L2 norm is not only empirically effective but also conceptually the most appropriate metric.

Q3.3: About why MLP outperforms standard SSMs in ablation study, explain performance gains in extrapolation, and the placement of ablation study in the paper.

A3.3: We would like to clarify that we do not replace the entire Phy-SSM unit with MLPs in the ablation study. Instead, we only replace the SSM layer of approximating unknown terms inside the Phy-SSM unit (highlighted in the blue rectangle in Fig. 1) with an MLP. we will include a clearer and more detailed explanation in the revised version.

The observed improvements—particularly in extrapolation—are due to the SSM layer’s ability to model long-range dependencies, which MLPs cannot effectively capture. This allows the Phy-SSM to learn more generalizable unknown physical dynamics, especially over extended horizons.

Finally, per your suggestion, we will move the ablation study from the appendix into the main body of the paper.

Q3.4: About typos and structural suggestions.

A3.4: We will correct all listed typos and adjust the section order as suggested in the revised version.

Q3.5: About the computational cost of Phy-SSM compared to standard SSMs and MLPs (Appendix H).

A3.5: As clarified in A3.3, we do not replace the entire Phy-SSM unit with MLPs. To address your concern, we report the computational costs of the standard SSM, our Phy-SSM module, and a NODE-based baseline (GOKU). The results are provided in Table 6 [link].

Phy-SSM is slightly slower than the standard SSM due to the additional structure and regularization components introduced for physics integration. However, it is faster than NODE-based models, as it preserves the parallel-scan acceleration capability inherent in SSMs.

Q3.6: About the initialization for the Phy-SSM module and the S5.

A3.6: In the Phy-SSM module, the known dynamics $A_{knw}$ is initialized using known physical parameters. The unknown dynamics $A_{unk}$ is initialized based on the output of the deep SSM layer. For S5, we follow the default HiPPO initialization [2].

[2] Gu et al. "Hippo: Recurrent memory with optimal polynomial projections." NIPS 2020

Response to Reviewer ooWi

Q2.1: About validating physical integration beyond the ablation drone experiment, especially for nonlinear and multi-variable systems.

A2.1: We evaluate our method on three real-world nonlinear, multi-variable systems: COVID-19, drone, and vehicle dynamics as presented in Tables 1–4 of our submission. Compared to data-driven models like S5, our physics-integrated approach consistently outperforms baselines, validating its effectiveness under diverse dynamical systems.

Q2.2: About the disjoint assumption in Proposition 1 and the generalizability of the knowledge mask definition to overlapping dynamics.

A2.2: The uniqueness result in Proposition 1 still holds when a term includes both known and unknown parts (e.g., a known multiplier applied to an unknown function). This is because we can treat the entire term as "unknown". For example, in the COVID-19 model (Eq. 28), the term $-\frac{I}{N} \cdot (*)$ is handled by learning the unknown part with a deep SSM and reintroducing the known factor $-\frac{I}{N}$ in post-processing.

The same principle applies to the knowledge mask design. For disjoint dynamics, we assign 1 to unknown terms and 0 to known ones. For overlapping terms, the mask marks the full expression as unknown, with known components injected after training. We will introduce this guideline of mask in our revision.

Q2.3: About using L2 regularization and other principled alternatives like Lagrangian multipliers.

A2.3: We have done additional experiments using other metrics, as shown in Table 5 [link]. We can see that L2 norm performs best. While Lagrangian multipliers are theoretically sound, they are difficult to apply in DNN models due to non-convexity and large searching space. We leave this for future work.

Q2.4: About recent physics-guided transformers (e.g., PDE-Refiner) or Neural ODEs with uncertainty baselines, and the rationale for choosing SSMs.

A2.4: We have included Contiformer [1], a state-of-the-art physics-based transformer, as a baseline. PDE-Refiner [2] is not suitable for irregularly sampled dynamics due to its reliance on fixed-step one-step MSE loss. Per your suggestion, we also added ODE2VAE [3], a Neural ODE with uncertainty-aware priors. Our method still achieves the best performance (see results in this [link] Tables 1–4).

We choose SSMs as the core module of our approach because (i) they capture long-term dependencies via HiPPO memory, and (ii) their continuous-time formulation naturally handles irregular data.

[1] Chen et al. Contiformer. NeurIPS 2023

[2] Lippe et al. Pde-refiner. NeurIPS 2023

[3] Yildiz et al. Ode2vae. NeurIPS 2019

Q2.5: About generalization of COVID-19 data and lack of large-scale benchmarks.

A2.5: In the COVID-19 experiment, we train the model on data from six countries and evaluate on two unseen countries (Ireland and Spain), demonstrating generalization across regions.

Beyond that, we have already included the large-scale autonomous driving (AD) dataset nuScenes [4], which contains 1000 driving scenes from Boston and Singapore—two cities known for dense traffic and challenging driving conditions. This dataset is widely used in AD research [5–7]. Experiment shows the generality of our method in learning vehicle dynamics from real-world data.

[4] Caesar et al. Nuscenes. CVPR 2020

[5] Girgis et al. Autobot. ICLR 2022

[6] Ren et al. Safety-aware motion prediction with unseen vehicles for autonomous driving. ICCV 2021

[7] Zhang et al. G2LTraj. IJCAI 2024

Q2.6: About the simplification of vehicle dynamics and the lack of actuator delay modeling.

A2.6: Following prior works [8,9], we incorporate partially known vehicle dynamics into our Phy-SSM model. Experimental results show that our method performs the best. Nevertheless, we agree with the reviewer that accounting for actuator delay is a valuable direction, and we plan to extend our method to address this scenario in future work.

[8] Rajamani, Vehicle dynamics and control. 2011

[9] Mao et al. Phy-Taylor. TNNLS 2023

Q2.7: About the formality of the uniqueness proof and the lack of broader theoretical discussions (e.g., stability).

A2.7: We use a contradiction-based argument for proving uniqueness, which is a standard formal method. We will refine the proof in the revised version for clarity and rigor.

As for stability analysis, it is beyond our current scope, our focus is on learning the dynamics model rather than control design. we recognize its importance and leave it in future work.

Response to Reviewer VtD1

Q1.1: About understanding the method’s complexity and presentation clarity.

A1.1: To help you better understand our work, we outline the problem, motivation, and key contributions below.

This work addresses the problem of long-term dynamical forecasting in complex environments where data are noisy and irregularly sampled. Our motivation is that SSMs can effectively capture long-range dependencies in sequential data and model continuous dynamical systems, while the incorporation of physics knowledge improves generalization ability.

Our key contributions can be summarized as follows:

1. We propose Phy-SSM, a novel approach that integrates partially known physics into state-space models to improve generalization for long-term forecasting in complex environments.
2. To enhance long-term prediction accuracy, we introduce a physics state regularization term that constrains latent states to align with system dynamics.
3. We provide a theoretical analysis demonstrating the uniqueness of solutions in our framework.
4. Extensive experiments on three real-world applications, including ve hicle motion prediction, drone state prediction, and COVID-19 epidemiology forecasting, demonstrate the superior performance of Phy-SSM over the baselines in both long-term interpolation and extrapolation tasks.

Additionally, based on feedback from the other reviewers, we have:

1. Included ODE2VAE as a new baseline to enable a more comprehensive comparison with uncertainty-aware NODE-based models, as shown in Tables 1–4 [link].
2. Added detailed ablation studies evaluating our choice of the L2 norm for regularization against alternative distance metrics (e.g., cosine distance, Chebyshev distance), as shown in Table 5 [link].