Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction

We genuinely appreciate the reviewer’s time and effort in enhancing our paper.

Link for results and references

Backup link

• Concern 1 & Question 5:

The fair comparison in terms of model expressivity (e.g. by reporting parameter count).

We appreciate the reviewer for raising concerns similar to those of reviewers vgH9 and ARTA. We report the number of parameters in the provided link. Furthermore, we also compare the computational cost in terms of grid scalability and performer attention, which may also be of interest to the reviewer..

• Concern 2 & Question 4:

Lacks evaluation on small chaotic dynamical systems based on ODEs.

Embedding ergodicity into learning methods for chaotic ODEs has been a subject of research and validation in recent literature [1-3], which may be of interest to the reviewer. While the literature did not specifically focus on large-scale chaotic dynamical systems, it is one of the primary objectives of our work.

Establish for very long autoregressive rollouts, e.g. by comparing Lyapunov spectra, on simpler cases.

As per the reviewer’s request, we implemented MNO [1] and our method on the Lorenz 63 system (L63). We collected 200k rollout steps, along with 300k true velocity data, on the invariant measure of the L63 system. We calculated the Lyapunov exponent for both our predictions and the actual data (1 timestep = 0.05s). Based on this foundation, we present the performance in terms of Lyapunov exponent and Lyapunov time in the table, and provide additional visualisation results demonstrating the generalizability of our approach on L63 via the provided link.

• Concern 3:

Clarification of Lorenz 63 dimension

‘1D examples’ refers to systems with a one-dimensional spatial domain or systems that are low-dimensional in terms of their structure. We have revised the text to better distinguish between spatial and phase space dimensionality.

typos

Thanks for bringing the typos to our attention. We have corrected them in the revised manuscript.

Reference

We appreciate the reviewer’s suggestion of the reference, which includes efficient reconstruction approaches through dimension reduction. We have incorporated it into the revised manuscript.

results consistency & error bar for table 1&2

The updated tables 1 and 2 are included in the provided link. Long-term statistics are robust across multiple runs, where the standard deviations are negligible ($\leq$ 1e-4). Hence, we only report the mean values on the (percentage) advantages. For short-term prediction accuracy, we report the standard deviation in the updated Table 1&2. The results are consistent with the experiment part in the manuscript. We will add the standard deviation to the revised manuscript as suggested.

• Question 1:

How is $G$ implemented? How expensive is the conjugation of the unitary operator?

In practice, $G$ is a square matrix in real space. Therefore, the conjugacy relation is equivalent to the transpose $G^T$. The relation between $G$ and $G^T$ preserve the unitarity using our proposed unitary loss in Equation (9) during training. Using the regularized term with Frobenius-norm, the computational complexity is $O(d^3 + 2d^2)$, challenging for large-scale chaos systems. And, we use the Hutchinson trace estimation technique, which effectively accelerates the computation by random projections. This reduces the complexity to $O(kd^2)$ where k is the number of random samples from the $d$-dimensional unit sphere. For further details, please refer to Section 3.3 and Algorithm 1 of the manuscript.

• Question 2:

Autocorrelation pooled?

Yes, the autocorrelation is calculated using the equation provided in Appendix E.2, specifically lines 1101-1117. This equation is closely related to the mixing rate. The autocorrelation is then averaged over all spatial points. We appreciate the reviewer’s suggestion to delve deeper into the subject matter.

• Question 3:

Fig. 3: How long are “1000 steps”?

Roll out 1000-step counts for 20 seconds for Figure 3. We have provided the dataset details in Appendix G.

What is the Lyapunov time of the dataset at hand (if such a quantity is easily accessible for this type of PDE data)?

The Lyapunov time is not directly accessible for the datasets we have. Given the high dimensionality of the data and the time constraints, we will provide a rigorous report of this statistic in the revised manuscript. We would like to know if the reviewer would prefer us to present the Lyapunov times as [mean, maximum] values.

Please let us know if we have addressed the concerns and increased your confidence in our work.

We sincerely thank reviewer ARTA for carefully reviewing our manuscript, providing valuable feedback and reconzing the strengths of our work. We'd like to address your concerns in the initial review and answer you questions as follows:

Essential Link for results and references to this Rebuttal

Backup link

• Concern 1 & Question 2 (also mentioned by reviewer vgH9):

However, the paper only enforces this unitarity approximately via a soft loss term (using Hutchinson’s trace), rather than proving exact unitarity.

Yes, we use Hutchinson’s trace estimation to approximate unitarity. This approach is motivated by two key reasons:

1. Directly applying the unitary loss involves eigen decomposition and computing the Frobenius norm, which is computationally infeasible for high-dimensional latent states.
2. Hutchinson’s trace estimation is flexible to implement as a randomized projection and is optimized adaptively during training.

Empirically, we demonstrate that the unitarity of the learned operators is well approximated. This is visualized in the eigenvalues plot of the learned operators from our ablation models (Base model and +Unitary) in the provided link.

• Concern 2 & Question 3:

  References and computational trade-offs to other linear-time attention methods like Performer.

  We thank the reviewer for the suggested references and have properly cited them in the final version. We also conducted experiments with Performer attention in our setup (Unitary+Performer) on the TCF benchmark during the rebuttal time. Visualization results are available in the link. The following table compares the runtime, memory consumption, and FLOPs per forward pass of baselines, Performer attention and our method. The results are evaluated on a cuda device with _CudaDeviceProperties(name='NVIDIA A100-SXM4-40GB', major=8, minor=0, multi_processor_count=108, L2_cache_size=40MB). Although attention mechanism is overall computational heavy than neural operator methods, tensor factorization and axial attention manage to make the computation tractable for large scale chaos states. With A3M attention and unitary constraint, our model effectively achieves better performance on the benchmarks. Additionally, we observe that Performer attention trades off memory usage on runtime from the empirical results.

  Models	Parameter count	Runtime	Memory usage	FLOPs per forward pass
  MNO	6467425	31ms	377MB	3.45GB
  UNO	17438305	12ms	769MB	6.88GB
  MWT	5089153	50ms	313MB	9.52GB
  FactFormer	6083009	53ms	6889MB	239GB
  Ours	7325665	58ms	7684MB	268GB
  Unitary+Performer	7717931	111ms	2938MB	271GB

• Concern 3 & Question 1:

Sensitive to the kernel bandwidth? How to tune it on different datasets.

We would like to thank the reviewer for raising the concern regarding kernel bandwidth. Through ablation experiments in Section 4.3, we found that selecting the bandwidth within a reasonable interval consistently yields stable results, indicating that tuning the bandwidth is straightforward and efficient, only requiring a coarse cross-validation.

The strategy for tuning the kernel bandwidth depends on the mesh size and viscosity (i.e., Reynolds number) of the system, which help identify a suitable bandwidth interval. If the data characteristics are unfamiliar, we recommend starting with a wider interval and then performing a grid search within that range. We appreciate the reviewer’s suggestion and will include the tuning strategy in the final version of the paper.

• Concern 4:

Expanding the discussion on applicability to non-ergodic systems would be valuable.

We thank the reviewer for this suggestion. Extending our methods to non-ergodic systems is a significant open problem. While our current work focuses on ergodic chaotic systems, we agree that exploring non-ergodic systems is an interesting direction for future research.

Finally, we thank reviewer ARTA once again for the time and effort invested in reviewing our paper. We believe the changes made in response to the reviewer’s comments have significantly improved our manuscript. We look forward to your further feedback.

We sincerely appreciate reviewer FGZz for taking the time to thoroughly review our manuscript, offering valuable feedback, and acknowledging the strengths of our work, particularly in the areas of efficient transformers for large-scale chaos systems, physics-inspired regularisation terms, and the contribution of our dataset to the community. Below, we respond to each of the reviewer's comments and suggestions.

• Ablation on distance-based Gaussian kernel

We thank the reviewer's valuable suggestion. It should be clarified that we actually included an ablation study of Gaussian kernels in Section 4.3 (Table 4) and Appendix D (Figure 4). In particular, we observed that extremely large sigma values like 32 lead to a rapid decay of spatial correlations, which is effectively equivalent to the case without a kernel. We will clarify this in the revised version.

Regarding computational efficiency, we evaluated the runtime per batch for the ablation models presented in Table 3. The results are evaluated on a cuda device with _CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', major=8, minor=9, total_memory=24GB, multi_processor_count=128, L2_cache_size=72MB) with a batch size of 3 over 100 runs.

• Suggested References

We appreciate the reviewer’s recommendations. After careful consideration, we’ve integrated the relevant references into the introduction section of our revised manuscript.

• Benchmark Dataset

We appreciate the reviewer’s acknowledgement of the significance of our proposed benchmark dataset. We will highlight this contribution even more in the final version of our manuscript.

• Typos

We appreciate the reviewer's meticulous review and for bringing the typos to our attention. All the identified typos have been meticulously corrected in the revised manuscript.

We sincerely thank the reviewer once again for their insightful comments and suggestions. We believe that the revisions made in response have significantly improved the quality of our manuscript. We eagerly anticipate any further feedback. In the meantime, we have prepared new visualisations and the code related to trying performer attention via the link here, which may be of interest to the reviewer.

We greatly appreciate the reviewer's constructive feedback and recognition of our work. We address the concerns and questions as follows:

Link to visualizations and references

Backup

• Concern 1 & part of Question 1:

  Whether the proposed loss function strictly enforces ergodicity in practice, any proof, and preserves in long-term predictions.

We thank the reviewer for this insightful question. The ergodicity guarantee through the unitarity constraint is supported by Von Neumann’s theorem (cited in Lines 156-160). Imposing hard unitarity constraints requires complex parameterisation and intensive matrix decomposition during training, which are computationally intensive. Instead, we use Hutchison trace estimation to efficiently approximate the constraint through random projections, enabling its application in large-scale chaotic dynamics. Empirical results show that the learned operators maintain unitary, as visualized in the eigenvalues plot (see link). To verify long-term stability, we use the metrics suggested by the reviewer (next section).

• Concern 2:

  KL Divergence and other metrics

We thank the reviewer for the reference. We evaluated the KL divergence (KLD) of model predictions on both datasets (KF256 and TCF). Updated Tables 1 and 2 are presented in Section B: Updates to Reviewer BWmX (see link). Our method achieves the lowest KLD, demonstrating its effectiveness in preserving long-term statistics and capturing the underlying invariant distribution.

Autocorrelation Decay Rate & (Max) Lyapunov Exponent Deviation

We appreciate the suggestion of these metrics. The autocorrelation decay rate, which relates to the mixing rate, is presented in Appendix E, with visualizations in Figures 3 and 6. Estimating the Lyapunov exponent for high dimensional PDE system was challenging. Instead, we provide the estimated Lyapunov exponent deviation for a lower-dimensional case using our method (see link).

• Concern 3 & Questions 2, 3, 6:

  Computational cost (runtime, memory, FLOPs) and scaling with grid size

We collected relevant results and organized them into two informative tables (see link). Analytical insights on attention mechanisms are shared in Response to Reviewer ARTA: Concern 2. Our method demonstrates linear increase of computational cost in grid size, indicating the strong scalability and computational efficiency of A3M attention.

• Concern 4:

  Relevance to PINN and physics hybrid transformer

We have added a discussion on this aspect in the manuscript. PINN methods require explicit knowledge of differential equations, which are assumed to be unknown in our setting. In contrast, our approach only assumes ergodicity, without relying on such prior knowledge, and focuses on embedding physical properties within transformer architectures. Our method provides a distinct perspective from PINN in incorporating physical knowledge.

• Concern 5 & Question 4:

Tuning strategy for kernel bandwidth

Due to space limits, we kindly refer the reviewer to Response to Reviewer ARTA: Concern 3. The ablation study on kernel bandwidth (Table 4, Appendix C) shows that selecting a reasonable interval yields stable results.

• Question 1:

Ergodic loss vs traditional methods

We appreciate the reviewer’s insightful questions regarding the ergodic loss.

(1) The ergodic loss is designed to enforce unitarity in the forward operator, thereby maintaining energy throughout forecasting. This is crucial for stability and long-term statistical consistency in chaotic systems. Traditional loss functions like MSE focus on short-term accuracy, which can lead to model drift and instability in long-rollout scenarios.

(2) Compared to traditional distribution-matching methods [5,6], our ergodic loss demonstrates superior computational efficiency and scalability. Distribution-matching methods like MMD [1,3] involve kernel operations that scale poorly with sample size, typically requiring large batches for high-resolution data and leading to significant memory consumption. Moreover, MMD with kernels (e.g., rational quadratic, RBF) requires careful bandwidth tuning, as results are sensitive to bandwidth choices [2].

Our ergodic loss enforce unitary dynamics using matrix operations without such hyperparameter tuning. Using stochastic Hutchinson trace estimation, we reduce the computational complexity to $\mathcal{O}(kd^2)$, where $d$ is the latent dimension and $k$ is the number of random draws from the unit sphere [4]. Typically, $d \ll D$ (the system dimension), and the convergence rate scales as $\mathcal{O}(1/\sqrt{k})$. In practice, we found $k \approx 1000$ sufficient for $d = 256$, balancing accuracy and computational cost.

• Question 5:

  Alternative pooling methods

We implemented the mean-pooling method for comparison with our full A3M method. Results are reported in the link.