Oscillatory State-Space Models

We thank the reviewer for appreciating the merits of our paper and their suggestions for improvement. Below, we address the concerns raised by the reviewer.

• ''I think section 3.2 can be written more accessible.'' We appreciate the reviewer’s feedback and have revised parts of Section 3.2 accordingly. Specifically, we have clarified abstract concepts, including the specific function spaces, by providing definitions in more straightforward terms. Moreover, we have summarized the main statement of Theorem 3.3.

• ''I believe Figure 1 is very important but can be made more explanatory.'' We have added further explanations, e.g., labeling $\bf y(t)$ and $\bf u(t)$ in the figure. Moreover, we have added a new section in the appendix (Section A) that explains the full architecture in detail, including the nonlinear layers. We hope that this together with the changes in the figure made it more explanatory.

• ''To the best of my understanding, the model cannot produce chaotic dynamics. What if the task in hand requires this?'' This is a very interesting question. Indeed, due to the inherent linear nature of state-space models one might expect that learning chaotic time-series using these models might denote a challenge. However, first experiments in this context suggest that state-space models actually perform well on chaotic dynamical systems (Hu et al., 2024), significantly outperforming their chaotic RNN counterparts such as LSTMs and GRUs. We further note that our universality result still holds. However, the Lipschitz constant in this case would need to blow up. That being said, this should not matter locally.

Paper: Zheyuan Hu, Nazanin Ahmadi Daryakenari, Qianli Shen, Kenji Kawaguchi, George Em Karniadakis, "State-space models are accurate and efficient neural operators for dynamical systems". 2024.

We thank the reviewer for appreciating the merits of our paper and their suggestions for improvement. Below, we address the concerns raised by the reviewer.

• ''Limited Analysis of Model Interpretability'' We sincerely thank the reviewer for their insightful comment. We fully agree that exploring model interpretability is both important and intriguing. In response, we have added a new experiment in the revised version of our paper (see Section C.2 in the appendix) to analyze the physics underlying our model. Specifically, we predict simple harmonic motion for varying initial positions and velocities over an extended period. The results demonstrate that LinOSS-IMEX maintains a constant error, whereas the error for LinOSS-IM increases over time. This is due to the fact that LinOSS-IM introduces dissipative behavior and in the asymptotic limit would converge towards a steady-sate for any state matrix values. This is in contrast to the conservative nature of LinOSS-IMEX. We thank the reviewer once again for their valuable suggestion, which has significantly enhanced the interpretability aspect of our work.

• ''No analysis of how the model captures different timescales'' We thank the reviewer for the comment. We would like to point out that most, if not all, our datasets are based on 'real-world' data and thus likely contain information arranged according to multiple scales. The strong performance of LinOSS on these dataset indicates that LinOSS is able to capture multiple scales in the underlying data. One highlight is the weather prediction task, which is known to be multiscale. Again, LinOSS outperforms any other competing method considered in this task, suggesting strong performance when applied to multiscale data. From a theoretical perspective, our universality result also holds for multiscale dynamical systems. Thus, we can show that LinOSS can approximate any multiscale system up to any desired accuracy.

• ''No ablation studies on the impact of different initialization schemes'' We appreciate the reviewer's suggestion and agree with their observation. We have now added a new experiment (Section C.1), testing several different initializations of the state matrix. Additionally, we have provided a detailed discussion of the results in Section 4.4. The findings support our claim that uniform random initialization is a robust choice in this context. Thank you again for highlighting this important aspect.

• ''The implementation details are unclear.'' We fully agree with the reviewer and have now added run times, GPU memory usage, and number of parameters for all models on all six datasets from Section 4.1 in Table 5 in the Appendix. While LinOSS seems to be well within the efficiency range of other state-space models, we would like to highlight that LinOSS is the fastest among all other considered models on two out of six datasets, as well as the second fastest on again two out of six datasets.

• ''Could the model be extended to incorporate coupled oscillations while maintaining stability guarantees?'' This is a very interesting question. Indeed, coupling within one LinOSS layer can be incorporated leveraging the diagonalization trick outlined in the first paragraph of Section 3.1. However, one would need to rely on a fast diagonalization of the underlying state matrix. That being said, the oscillators are already coupled in LinOSS, however, delayed by different layers. Since the output of the previous layer denotes the input of the current layer and due to the fact that the matrix $\bf B$ in equation (1) is dense, this introduces a notion of coupling.

• ''What is the impact of the time step parameter $\Delta t$ on model performance and stability?'' We have added a proposition on the stability of LinOSS-IMEX in Section E.1. There, we have rigorously worked out the impact of $\Delta t$ on the stability of LinOSS. Moreover, motivated by the reviewer's suggestion, we have added a new experiment where we analyze the impact of $\Delta t$ on the performance of LinOSS (Section C.3 and Section 4.4). From our obtained results, we can conclude that while the choice of $\Delta t$ does influence performance, the variations are not substantial.

• ''How does the choice between IM and IMEX variants affect training dynamics and convergence?'' This is a very interesting question. Based on our new experiments analyzing the underlying physics of our proposed LinOSS models, we argue that LinOSS-IM converges faster and has better generalization compared to LinOSS-IMEX on data that includes conserved quantities. However, whenever data exhibits dissipative behavior, or forgetting is crucial, the roles are swapped.

We sincerely hope that we have addressed the concerns of the reviewer satisfactorily in the revised version and would kindly ask the reviewer to update their score accordingly.

We thank the reviewer for appreciating the merits of our paper and for the insightful feedback. Below, we address the concerns raised by the reviewer.

• ''novelty of this method'' We argue that the novelty of our work lies in the state-space formulation of oscillatory systems, introducing a fresh perspective to this domain. Specifically, we employ a system of linear oscillators, in contrast to the nonlinear systems explored by Rusch and Mishra (2021). Additionally, achieving an efficient state-space modeling approach required multiple innovations, enabling us to leverage fast associative parallel scans for these systems effectively. Moreover, the linear formulation of these systems allowed for explicit theoretical statements about these models, which is not available for previous nonlinear systems.

• ''Since the oscillators are independent, can the proposed architecture model transient synchronization / desynchronization?'' This is a very interesting question. While the oscillators in each LinOSS layer are not connected, a notion of coupling can be introduced via stacking several layers, where the output of the previous layer is the input of the current layer and the matrix $\bf B$ modulating these inputs is dense, i.e., allows for mixing. Thus effects such as synchronization/desynchronization are possible. Another way to approach this would be to introduce sparse coupling within each layer and simply diagonalizing the resulting state matrix to use fast associative parallel scans. We plan to investigate this in more detail together with other important notions from computational neuroscience in future research.

• ''If $\bf M$ has real eigenvalues with algebraic multiplicity $>1$, would that make the system unstable?'' An algebraic multiplicity bigger than 1 does not necessarily make the system unstable. One example for this is our LinOSS-IMEX model, which has algebraic multiplicity equal to the state dimension $m$. However, we show in Proposition E.1 that this system is stable.

We sincerely hope that we have addressed the concerns of the reviewer satisfactorily in the revised version and would kindly ask the reviewer to update their score accordingly.

We start by thanking the reviewer for appreciating the merits of our paper and the welcome suggestions to improve it. Below, we address the very valid concerns raised by the reviewer and thank the reviewer in advance for their patience in reading our detailed reply.

• Connection between theoretical findings and empirical results: We apologize for the ambiguous statement in our first version of the paper regarding forgetting in LinOSS-IM vs LinOSS-IMEX. We would like to emphasize that the main message of our paper is that we use forced harmonic oscillators as a building block to construct a new state-space model architecture. As the reviewer correctly pointed out, the 'proper' way to dicretize this ODE system is via symplectic time integration method, e.g., implicit-explicit Euler as in the case of LinOSS-IMEX. We argue, however, that there is another viable choice to discretize this system using purely implicit methods. These methods, in contrast to symplectic methods, introduce dissipation to the system. However, we further argue that this can actually be an advantage when learning long-range sequences, as it introduces 'forgetting' in the sense of the equation in lines 260-262. This makes the model more flexible while at the same time can still propagate information over very long timescales, i.e., the absolute eigenspectrum of the implicit LinOSS can be further controlled via the state matrix (equation line 289) and thus can be made arbitrarily close to 1. This flexibility in balancing forgetting with propagating information over long timescales makes LinOSS-IM a strong candidate for learning on long-range sequences. We would thus expect that LinOSS-IM outperforms LinOSS-IMEX in practice unless the underlying dataset contains conserved quantities. We have clarified this now in Remark 1 in the revised paper. We further followed the reviewer's excellent suggestion and provide a new synthetic experiment examining the dissipative nature of LinOSS-IM vs the conservative one of LinOSS-IMEX. To this end, we predict simple harmonic motion for different initial positions and velocities for a long time. We show that LinOSS-IMEX keeps the error constant, while the error of LinOSS-IM grows over time. This is due to the fact that LinOSS-IM introduces dissipative behavior and in the asymptotic limit would converge towards a steady-sate for any state matrix values. This is in contrast to the conservative nature of LinOSS-IMEX. We thank the reviewer again for suggesting this experiment. We have added the experiment to the paper in Section C.2 and a discussion about it in Section 4.4. Finally, we would like to point out that we have implemented an explicit version of LinOSS using the explicit Euler method. However, having tested this architecture on several datasets it continues to return NaN values after some time, since the hidden states of this system explode. Of course, one could take the limit of $\Delta t\rightarrow0$ to stabilize it, but this would be in contrast to the idea of efficiency in state-space models.

• "What are the runtimes of the different methods in the experiments? Run time when not using the equation for matrix inversion? Run time when not using the parallel scan?": We have now added run times, GPU memory usage, and number of parameters for all models on all six datasets from Section 4.1 in Table 5 in the Appendix. While LinOSS seems to be well within the efficiency range of other state-space models, we would like to highlight that LinOSS is the fastest of all other considered models on two out of six datasets, as well as the second fastest on again two out of six datasets. We further note that relying on numerical methods to invert the state matrix would make associative parallel scans obsolete, as the run time is proportional to $log_2(N)$ up to some constant that depends on the structure of the state matrix. Moreover, not using a parallel scan would result in a simple linear RNN cell stacked over several layers. This is known to be exceptionally slow.

• "Why is A diagonal": The majority of recent state-space models choose A to be diagonal. Indeed, this allows for a quick closed-form matrix inversion in our case. Moreover, the resulting matrix $\bf M$ (eq 4 and eq 6) is a 2 x 2 block matrix, where each block is again a diagonal matrix. This leads to a very fast computation of the ODE solution using associative scans. Note that associative parallel scans are significantly slower for dense state matrices. Finally, it was noted in Orvieto et al, 2023, (see Table 2 of this paper) that diagonal state matrices perform better in practice compared to dense matrices.