Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
Authors: Yi He
Yiming Yang
Xiaoyuan Cheng
Hai Wang
Xiao Xue
Boli Chen
Yukun Hu
Venue: ICML 2025 poster (2025)
Keywords: Dynamical system
chaos
transformer
auto-regression
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: ChaosMeetsAttention is proposed as a scalable, ergodicity-preserving transformer framework for generating long-term, high-dimensional ergodic chaotic dynamics. We introduce new benchmarks for machine learning on chaos.

Abstract: Generating long-term trajectories of dissipative chaotic systems autoregressively is a highly challenging task. The inherent positive Lyapunov exponents amplify prediction errors over time.
Many chaotic systems possess a crucial property — ergodicity on their attractors, which makes long-term prediction possible. 
State-of-the-art methods address ergodicity by preserving statistical properties using optimal transport techniques. However, these methods face scalability challenges due to the curse of dimensionality when matching distributions. To overcome this bottleneck, we propose a scalable transformer-based framework capable of stably generating long-term high-dimensional and high-resolution chaotic dynamics while preserving ergodicity. Our method is grounded in a physical perspective, revisiting the Von Neumann mean ergodic theorem to ensure the preservation of long-term statistics in the $\mathcal{L}^2$ space. We introduce novel modifications to the attention mechanism, making the transformer architecture well-suited for learning large-scale chaotic systems. Compared to operator-based and transformer-based methods, our model achieves better performances across five metrics, from short-term prediction accuracy to long-term statistics. In addition to our methodological contributions, we introduce new chaotic system benchmarks: a machine learning dataset of 140$k$ snapshots of turbulent channel flow and a processed high-dimensional Kolmogorov Flow dataset, along with various evaluation metrics for both short- and long-term performances. Both are well-suited for machine learning research on chaotic systems.

Chaos Meets Attention:
Transformers for Large-Scale Dynamical Prediction
Yi He" Yiming Yang" 12
Xiaoyuan Cheng
Hai Wang? Xiao Xue
Boli Chen Yukun Hu'
Abstract
1. Introduction
Generating long-term trajectories of dissipative chaotic systems autoregressively is a highly chal- lenging task. The inherent positive Lyapunov exponents amplify prediction errors over time. Many chaotic systems possess a crucial property ergodicity on their attractors, which makes long-term prediction possible. State-of-the-arr methods address ergodicity by preserving sta- tistical properties using optimal transport tech- niques. However, these methods face scalabil- ity challenges due to the curse of dimensionality when matching distributions. bottleneck, we propose a scalable transformer- based framework capable of stably generating long-tern high-dimensional and high-resolution chaotic dynamics while preserving ergodicity. Our method is grounded in a physical perspec- tive, revisiting the Von Neumann mean ergodic thcorem to ensure the preservation of long-term statistics in the L space. We introduce novel modifications to the attention mechanism, mak- ing the transformer architecture well-suited for learning sarge chaotic systems. Compared methods,
Predicting and modeling the behavior of chaotic systems is crucial for various applications. including weather forecast- ing (Lorenz & Haman, 1996), climate modeling (Trevisan & Palatella, 201 I), and understanding turbulent flows (Ot- tino et al., 1990). In recent years, the use of autoregressive models to predict the long-term behavior of chaotic systems has presented a significant challenge to the ML/DL commu- nity (Gilpin, 2021; Mikhaeil et al., 2022). This challenge stems from the presence of positive Lyapunov exponents, hallmark of chaotic systems. These exponents amplify small perturbations in the initial state, leading to exponential accumulation of errors over time (Holden, 2014). Conse- quently, accurately predicting the long-term behavior of chaotic systems remains a formidable task. To address these challenges, two primary methods and one emerging trend have been reviewed for learning chaotic systems: a) se- quence models that excel in learning temporal patterns; b) operator- based models, which learn operators in function spaces without knowing the prior differential equations; and c) the transformer- models, with recent advancements in synthesizing sequence and operator-based methods.
overcome this
Sequence models. Sequence models have shown notable success in short-term predictions by minimizing temporal mean squared errors (MSEs). Prominent methods include recurrent neural networks (RNNs) (Madondo & Gibbons, 2018; Dudukcu et al., 3023), long short-term memory net- works (LSTMs) (Sangiorgio & Dercole, 2020; Langeroudi et al., 2022), and reservoir computing methods (Pathak et al., 2018; Yan al., 2024). These approaches have been applied to classic 1-D examples, such as the Lorenz 63, Lorenz 96, and Kuramoto-Sivashinkk equations. However, due to the inherent instability of chaotic systems, these methods often experience exponential error accumulation over time.
to operator- bassd and transformer- based our model achieves better performances across five metrics, from short-term prediction accuracy to long-term statistics. In addition our method- ological contributions,
introduce new chaotic system benchmarks: machine learning dataset of 140k snapshots of turbulent channel flow and processed high-dimonn nigh-dimensional Kolmogorov Flow dataset, along with various evaluation metrics for both short- and long-term performances. Both are well-suited for machine learning research on chaotic systems.
To address this issue, recent advances leverage ergodicity, a key property of many chaotic systems on strange attractors (Eckmann & Ruelle, 1985; Young, 2002), stabilize the method performance in long-term predictions by aligning predicted distributions with ground-truth distributions. Two notable approaches utilize optimal transport theory (Jiang et al., 2024; Schiff et al., 2024) to improve long-term pre- diction accuracy, i.. incorporating a regularized transport term to penalize discrepancies between the predicted and
Eqqual contribution 'Dynamic Systems, University College London, United Kingdom Sttttsstical Science, University Col- lege London, United Kingdom 'Center for Computational Science, University College London, United Kingdom "Electronic and Elec- trical Engincering, University College London, United Kingdom. Correspondence to: Yukun Hu <yukun, ucl.ac.uk>.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
true distributions. However, matching distributions often encounters the curse of dimensionality (Kloeckner, 2020; Zhang et al., 2023; Vögler, 2023), particularly when chaotic systems are high-dimensional and characterized by mul- timodal probability distributions (Sengupta, 2003; Zelik, 2022). High- nigh--dimensioal chaotic systems, such as turbulent flows, typically manifest in phase spaces with dimensions exceeding O (10") and exhibit non- Gaussian statistical prop- erties (Biferale et al., 2006). The limitation of distribution- matching methods lies in their computational intractability for such high-dimensional spaces, compounded by the dif- ficulty of accurately capturing the intricate statistical struc- tures of non- Non-Gaussian distributions (Korotin et al., 2021).
dictions for dissipative chaotic systems by enforeing the forward invariance via the constraint. Yet no loss function in MNO theoretically guarantees the ergodicity for long- term prediction, leaving potential discrepancies in statistical properties unresolved.
Transformer-basedd models. Leveraging the flexibility and universality of transformers, researchers started to in- tegrate sequence models and operator theory with trans- former architectures to model complex dynamical systems (Li et al., 2022b; Guibas et al., 2021; Hao al., 2023; Kissas et al., 2022; Zhang & Gilpin, 2024). For example, the Fourier Transformer has been proposed as an extension of the Fourier Neural Operator to solve a variety of PDEs (Li et al., 2022b; Guibas et al., 2021; Hao et al., 2023; Chen et al., 2024). To reduce the computational cost of standard scaled-dot product attention, several pioneering works have eliminated the Softmax operation and utilized matrix asso- ciativity to achieve linear complexity attention, which has shown promise for large-scale PDE modeling (Cao, 2021; Li et al., 2024). Despite the significant advancements in transformer- basedd methods for dynamical systems, most efforts have primarily concentrated on general PDEs. This leaves notable research gap: the lack of well-designed transformer architecture tailored specifically for large-scale chaotic systems.
Operator-based models. Rooted in operator theory, pperator-based approaches offer a promising alternative for modeling dynamic systems (Marchenko, 2012). Un- like sequence models, these methods do not require prior knowledge of the underlying systems. Instead, they lever- age function approximation theory to identify an operator that captures the system's is evolution within a predefined function space (Smale& & Zhou, 2007). Two representative methods are transfer operator (Jorgensen, 2001) and Koop- man operator (adjoint form of transfer operator) (Brunton et al., 2017; Cheng et al., 2023; Brunton et al., 2021; Mezić, 2021; Cheng et al., 2025). The transfer operator models the evolution of probability density functions, while the Koopman operator focuses on the evolution of observable functions. Classical transfer and Koopman operators lever- age kernel methods within the framework of reproducing kernel Hilbert space or Banach space (Klus et al., 2020;
To address this gap, we first propose a tailored transformer- based model for predicting chaotic systems. Our key contri- butions can be summarized as follows:
Ikeda et al., 2022; Hou et al., 2023; Yang et al., 2025). These approaches face computational bottlenecks due to the prohibitive cost of matrix inversion (Bousquet & Her- rmann, 2002) for nigh-dimenssonal chaotic systems. Conse- quently, applying kernel-basee techniques to such systems remains a significant challenge. In recent years, deep leam- ing techniques have been integrated with operator theory and have yielded significant advancements in solving par- tial differential equations (PDEs). Methods such as Deep Operator Network (DeepONet) (Lu et al., 2021) and Fourier Neural Operator (FNO) (Li et al., 2020) have shown their effectiveness as approximators for initial value problems in PDEs. Many variants of these two methods have been developed to solve large-scaled PDEs, such as U-shaped operator
Designed for modeling chaotic systems and aligned with their intrinsic properties, we introduce A3M mod- ifications into factorized attention (Figure 1(b)) to re- design the scaled-dot product attention, and enhance it with random Fourier positional encoding for improved representation of complex chaos patterns.
Building on the Goopman-Neuman ergodic theorem (Neumann, 1932; Mezić, 2021), we propose novel loss function with a unitary constraint on the forward operator (Figure 1(c)), which is scalable to capture long-term statistics of large-scale chaotic systems and ensure stability without directly matching probability distributions.
FNO (Rahman et al., 2022), multiwavelet- (Gupta al., 2021). However, chaotic systems are highly sensitive initial values (Kolesov & Rozov, 2009), purely predicting the long-term behavior based on learned opera- tors cannot guarantee long-term stability. Hence, more problem-specific framework is necessary to model chaotic systems. Markov Neural Operator (MNO) (Li et al., 2022a) as a closely related framework, introduced hard constraints upon the architecture of FNO to enhance long-term pre-
We introduce datasets of two high-dimensional chaotic systems: 1) turbulent channel flow with 140k snap- shots in the 3D simulation; 2) the Kolmogorov Flow simulation of 185k 2D vorticity states. Both datasets with details in Appendix D and G have been carefully processed to ensure consistency and usability, mak- ing them well-suited for early-stage machine learning research on ergodic chaotic systems.
2

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
We benchmark our algorithm with the state-of-the-art operator-based and transformer-basee methods two chal- lenging chaotic systems. The results consistently demon- strate superior short-term prediction accuracy and long-term statistical consistency in high resolution. Our core frame- work is illustrated in Figure I.
where & the feature map as Φ :M- R. with ø E F L and function space F on M. The operator describes the evolution of learnable features. The time evolution of these features can be obtained by iteratively applying g in L space. By applying a decoding function Φi, the state 2k+1 can be reconstructed. The choice of L space is moti- vated by the well-developed theory of ergodicity, stemming from the contributions of Von Neumann (Neumann, 1932).
2. Preliminary
Theorem 2.2 (Von Neumann Ergodic Theorem (Neumann, T) be MPT. If E L2, then
Notation. (X, denotes the measure space, with set X, Borel r-algebra B and measure p.. The forward map on the state space is denoted as T'. The Lebesgue space with the p-norm is denoted as C*(X, pi), abbreviated as L in this paper. In particular, the Lebesgue space L is equipped with an inner product structure as (). Feature functions G, ib belong the L space. C denotes the transpose of real operator, and D denotes the conjugate transpose of complex operator. The symbol z} indicates the i-th indexed state at the j-th time step, and &a represents the coordinate of the i-th index.
1932)). Let (X, B, B, Nim N Ek=o g = ã where Φ is invariant. T is ergodic, then ē odu.
This theorem demonstrates that for an ergodic measure- preserving transformation (MPT) T, the associated operator G is unitary and preserves the norms in the L space. This implies that the time averages of the L functions converge to their space averages, ensuring the long-term predictability of chaotic systems. By combining with Equation 2, the operator G preserves the ergodic property during long-term predictions. Our objectives are twofold: (1) to achieve accurate short-term predictions and (2) to capture statistical properties by preserving ergodicity in the _ space.
Problem formulation. The objective is to forecast the behavior of ergodic chaotic systems, described as
2k+1 = T(*k),
2 E M,
(1)
3. Method
where T"; : :M + M is nonlinear forward map on com- pact set .M C x
The state zy represents physical quantities defined on a uniform grid with dimen- sions S1, S2, Sas, with the total number of grid points N S x 5 x x SM. Each component of z bounded and represents physical quantity at the kth grid point.
In this section, we present our customized transformer ar- chitecture designed for modeling ergodic chaotic systems. This method is structured into three key components: 1) introducing an attention mechanism to identify spatial cor- relations based on random Fourier features, which naturally align with chaotic system properties as discussed in section 3. 2) efficient processing of multi-dimensional tensors with capturing extreme values by factorizing attention and the proposed A3M block, detailed in section 3.2), and 3) an efficient constraint for learning the unitary operator in the high-dimensional feature space to preserve ergodicity using Hutchinson's 's stochastic trace estimation with details in section 3.3).
Definition 2.1 Mleasure-PPeservvin Transformation and Ergodic (Cornfeld et al., 2012)). Let (X, B, T)) define a measure-preserving transformation (MPT), where for every E E B, the measure satisfies p(T-1E) = μ(E), An MPT called ergodic if, for any invariant set F, either yxE) = or p(X E') = 0. In this case, j is referred as an ergodic measure.
Intuitively, ergodicity ensures spatial statistics are compati- ble with temporal statistics. This implies that the distribution of system trajectories follows an invariant measure that re- mains consistent over time. Many chaotic systems exhibit this property on their attractors, such as Lorenz systems (Shi et al., 2020), Kuramoto- -Siassii-- equation (Yang, 2006) and turbulent fluids (Galanti & Tsinober, 2004; Hairer & Mattingly, 2006). By preserving ergodicity, the learned model can capture the invariant statistical behaviors of the attractors.
3.1. Attention Mechanism with Random Fourier
Give three sets of feature functions (or vectors), namely the queries { keys [K and values {b') J== the self-attention mechanism (Vaswani, 2017; Han et al., 2022) calculates a weighted average of values 6(2') N Σ h(g', k'))",, where o', g',k',e" v² E Rdd The attention weights h((,-) are defined as the scaled-dot product with Softmax function in the original transformer (Vaswani,
exp( (t0" k///vZ)
2017): h(g', kd) Σn exp (g" VV)/Va) It is interesting
Learning the forward map of chaotic systems while pre- serving ergodicity is challenging. To address this, our model adopts an operator to forward in feature space:
to interpret the weighted average o(x') E L2 as a learnable feature of state 2 as shown in Equation (2).
((z++1) = (2x)
The queries/keys/valles are typically obtained through learn-
3

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
Our Transformer
Axial Mean-Max-Min (A3M) Attention block
Unitary Operator
Decode
Unitary Operator
out
(b-1 1) Axial Attention for pooling type
Unitary Constraint Gg
Forward Dynamics g
Learnable Projection
A3M Attention (Min)
+
Alleatiun
Aste
A3M Attention (Max)
Attantion
Asic
by
A3M Attention (Mean)
Altention
Axial Mean
Attention
Pooling
(b-2) 3M pooling types
Pooled Axial States
RFF Positional Encoding
3M
Mica Pooling
Max pooling
Min Pooling
Pooling
Linear Projection
Type
Real Part
(a)
(b)
Figure 1. The framework of ChaosMeettAutention Figure ((a) presents the main structure of our transformer, highlighting the modified
positional encoding, attention mechanism, and unitary operator. Figure 1(b) illustrates the efficient Axial Mean-Max-Ming (A3M) attention, designed to capture both statistical moments and extreme values in the physical field. Finally, Figure I(c) aligns with the Von Neumann mcan ergodic theorem described in 2.2, demonstrating how the unitary operator preserves statistical properties.
able linear projections:
to how mixing rates in physical systems determine the scale of spatial correlations.
g¹ 2' 4.
='
0 = =
(3)
where z' E Re represent the i-th input physical quantity with dimensions din in the domain M,, and {Wg. We, W.} are learnable projection matrices. Standard self-attention mechanisms do not explicitly encode spatial correlations. They compute similarity in the learned feature space and apply probability- weeibilittywwiightedd averaging via the Softmax op- eration. This formulation lacks an explicit spatial distance representation. To overcome this limitation, we introduce distance-based Gaussian kernel to capture spatial corre- lations based on relative distances. This kernel function decays smoothly with increasing distance, aligning with the spatial mixing behaviors observed in chaotic systems.
RFF Positional Encoding. Based on Bochner's theorem (Bochner al., 1959), random Fourier features (RFFs) can be used to approximate stationary kernel. For sim- plicity, we demonstrate the applications of RFFs in a 1- dimensional grid domain, which can be easily extended to higher-dimenssoonl grid domains. We use a random Fourier map 0 to featurize input grids, which project input coordi- nate index the surface of a sphere with set of sinusoids: ((i) 1 cos BEx), 1 sin(2n BG))), where & is the vm
coordinate of i-th index, B E R with B, _N(,, for i = 1,..., m. From trigonometrie identities, the inner product of (()) and 0(j) is formulated as:
To efficiently compute this kernel, we employ random Fourier features (Rahimi& & Recht, 2007; Tancik al., 2020) for the approximation of Gaussian kernel using randomly sampled Fourier basis functions. This distance-based formu- lation shares an interesting connection with the topological mixing property of chaotic systems (Xiong & Yang, 1991), In chaotic systems, topological mixing describes how any localized region of the state space eventually spreads and overlaps with other regions, allowing local perturbations to propagate throughout the space. Similarly, the Gaussian kernel-ber based attention mechanism enables each point to in- teract with its surroundings with a strength that smoothly decays with distance, providing a natural way to model how information or influence propagates through space in physical fields. The kernel's bandwidth controls the charac- teristic length scale of these spatial interactions, analogous
((i)T(()) =0(i j)
lim m-+00
cos(2 - E2) (4)
m j=1
mexp(-mor
I|E E 11))
The last line of Equation (4) is from the characteristic func- tion of Gaussian distribution (Shiryaev, 2016) with details in Appendix B. By constructing the attention mechanism with RFF, the corresponding queries and keys become g(g") q (i) and ((k)) = k-((j). Let Q = ig(a"),... . }}, K [g( ).... (I.NN)) and V = [o', Tol,...,,0 N. The fea-

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
ture x((²) is expressed as:
where E denotes the i-th coordinate axis of grids. Here, g' (q¹) and g'(k)) for are pooling from multi-dimensional tensors Q, E R xxxxx Xd along all dimensions except the j-th di- mension, e.g.
1,..., M and i 1, 1,... Sj
o(z*)
vV²
N
Σ (9((") ((k("))((kk))) g(k")))
N
(5)
Σ
Qi
Σ
0(i - () ,00000000
I 5.
s=1
N
with mean aggregation. Different from (Li et al., 2024). we propose separate attention heads in the transformer block to apply axial attention with mean, max, and min poolings, collectively named the A3M Attention blocks (as shown in Figure 1(b)). This design captures both statistical moments and extreme values in the field (Yeung et al., 2015; Faraz- mand & Sapsis, 2017; Sapsis, 2021), which are critical for chaotic systems like turbulent flows, where both statistical properties and local extremas significantly influence dynam- ics. By recursively calculating the quadrature of the integral kernel along each axis from to M, the number of required matrix multiplications is reduced to S + S2 + ... + SM achieving considerable improvement in computational efficiency. Figure ((b) illustrates the proposed attention mechanism.
exp( (xsp(-oo'10 - - ko') Σ
Here, ((=") represents a quadrature of the integral kernel, capturing the weighted average of values. The spatial cor- relation between the i-th index and other indices 8 decays according to the Gaussian kernel, leading to stronger in- fluence from nearby points, aligning with the topological mixing property (see section 3.1). The parameter o in Equa- tion (5) controls the decay rate of these spatial correlations: a smaller a leads to slower decay and potential underfitting, while larger a causes rapid decay, increasing the risk of overfitting. An ablation study section 4.3 examines the impact of different a values on model performance.
3.2. Factorizing as Multi-Dimenssonal Tensor
3.3. Embedding Ergodicity on Transformer
The computation complexity of the attention mechanism increases quadratically with the total domain grid points N. Unlike natural language processing (NLP) sequences, the domain M has a uniform grid structure S₁ x S2 x x SM. Directly applying attention as in Equation (5) requires Si x S2 x .... x Sa times matrix multiplications, which becomes prohibitively expensive as the grid resolution in- creases. To address this, we leverage a tensor factorization approach from (I.i et al., 2024), which reduces the num- ber of matrix multiplications from S x S2 S1 + S2 + + SM- The basic idea is to factorize the atten- tion operation along each axis of the grid domain, resulting in linear increase in computational complexity with grid resolution rather than a quadratic or cubic increase.
By connecting chaos to operator theory and Von Neu- mann's statement in Theorem 2.2, it provides a framework for preserving the statistical behavior of chaotic systems through the operator G. Specifically, the eigenvalue (or spectrum) of the operator 9... lies on the complex unit circle Eigen(G) C {r E C Ir| = 1} (Avigad, 2009).
To identify the unitary operator in the ergodic state, we leverage the intrinsic property that G and its conjugate trans- pose G' satisfy G-g = Ia. This property confirms that ç is unitary, preserving both norms and inner products, and establishes well-defined backward dynamic. Specifically, the backward operator g-1 is equivalent to G". The conju- gate transpose relationship G'g = Gg = Ia ensures that G is invertible, with its inversion fully consistent with unitarity, as shown in Figure 1(c).
x SM to
Multi-dimensional tensor product. The data in the grid domain can be represented as a tensor A E R S1 x291x8xx-xx3 The product with a matrix W E RR x on the m-th mode in tensor of shape S x S2 .... x J x SM, is defined as:
To incorporate the theorem in practice, we introduce a regu- larized term to penalize it, and the loss function is expressed as:
Sm A Witm-
(A mm W) W...e.........
arg min Ezwu [1l)0(=)) d(2x+1)ll2 + Munnttoy (9)].
G
According to the principles of multi-dimensional tensor, the attention mechanism with positional encoding along each axis can be expressed as:
(8)
where the first term represents the forward loss, and Cunitory (9) is the unitary regularization with coefficient X E (0, 1]. To enhance the scalability of the framework, we propose a novel approach by incorporating Hutchinson's 's stochastic trace estimator.
81
0(2') - - ... Σ ( ... ........MMMMM....MMMM......
5

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
To express the unitary operator on the torus, dis- cussed in (Das, 2023), it can be represented as the form U(1) x U(L) x x U(), where (11 = {eex((i)
4. Numerical Experiments
In this section, we comprehensively evaluate the perfor- mance of our model on two high-resolution chaotic sys- tems. Kolmogorov flow has external periodic forcing and structured large-scale dynamics, while turbulent channel flow involves flow-driven turbulence with more complex, anisotropic behavior. The experiments are repeated with three random seeds, and the mean value is reported in the results tables and figures in section 4.1, 4.2 and Appendix F. We also include details of computational cost to demon- strate fair comparison with baselines and the efficiency our methods on increasing state resolutions in Table 7 and 8.
d/2
0 E (m]} is the circle group. It is well-known that S2() is isomorphic to (1). Since the parametriza- tion of neural networks is typically in the real space in- stead of complex space, the real representation of prod- uct circle group (11) x U(1) ... (1) corresponding x (2) can be embedded into
to SO(2) x SO(2)
4/2
SO(d), where the parametrized dimension d is an even number. Here, (O(d) = {A E Rdxd AAAT = A Ta, (A) = 1} denotes the d-dimensional special orthog- onal group. In such a situation, Lunitary (9) can be defined as:
Baselines. We compare our transformer with four com- petitive baselines covering different learning mechanisms in dynamical systems, including: (a) Markov neural oper- ator (MNO) (Li et al., 2022a), an FNO-based model de- signed capture long-term dissipative chaotic behavior in sequential data by incorporating tailored constraint; (b) UNO (Rahman et al., 2022), U-shaped FNO architecture that incorporates skip connections; (c) Multiwavelet- Operator (MWT) (Gupta et al., 2021) a neural operator using wavelet transformations to effectively model dynam- ical behaviors; and (d) Factformer (Li et al., 2024), an attention- based model for learning PDEs, designed to fac- torize attention computations along different spatial axes to achieve scalability. With the specialized design, MNO has demonstrated wate-of-the-are performance in capturing the long-term statistics of chaotic systems.
Lunitary (9) E Bn--Uni((($
|
1 TTTTGU
1
(9)
Here, is random unit vector sampled uniformly from the unit (d - 1)-sphere Sd-1 C _dd i.e., ]|u|l2 = 1. The regularized term Lunitory (g) constrains 9 within (d).. The sample number of the unit random vector k, and the upper error bound of the stochastic trace estimator is scaling as (+72)) (Meyer al., 2021). The number of sample & is (Tk
(0)(1/5) where e measures the upper error proportional to 2² bound with at least 1 - 8 probability. We impose a soft constraint on our loss functiion making our learned physics align with Von Neumann's 's mean ergodic theorem in 2.2.
Benchmarks. We choose two turbulent fluid dynamics as popular benchmarks for learning states from chaotic sys- tems: Kolmogorov Flow (KF256): 2D shear flow char- acterized by sinusoidal velocity fields in one direction and external forcing in the perpendicular direction. In our ex- periments, we utilize a 256x256 resolution vorticity field to study the short- and long-term performance (He, 2025). Tur- bulent Channel Flow (TCF); 3D canonical wall-bounded flow where turbulence develops between two parallel planes due a pressure gradient, characterized by turbulent veloe- ity fields, enhanced mixing, and distinct near-wall and core regions. We extract the 3-channel 2D velocity field from the 3D simulation and evaluate our approach with baselines on the resolution 192 x 192 over all channels (He et al., 2025). We generated and prepared these datasets with details in Appendix and G.
Compared to the penalization in (Cheng et al., 2025), our proposed loss function demonstrates higher efficiency when applied to large-scale chaotic systems than using Frobenius- norm regularized term (Golub & Van Loan, 2013). The computational complexity of the Frobenius-nowe norm regular- ized term is Odd + 2d²), whereas the complexity of the stochastic trace estimator is O(kd²). Given that d is typ- ically large for approximating evolution in L space, our method provides a more computationally efficient approach.
Algorithm Lunitary with Hutchinson's 's Stochastic Trace Estimator
Require: Operator & Rdxd batch size B 1: Initialize Lunitory
2: for b 1 to B do
3: Sample (b()_ Unif()d-1))
4: (0h) = (b) " (
We evaluate the models in terms of short- and long-term performances. For short-term performance, we use the relative norm L following (Li et al., 2022a). Due to the high dimensionality of our experiments, achiev- ing robust estimations of long-term statistics is challenging. To ensure a valid comparison, we assess long-term perfor-
Evaluations.
5: Lunitory & Lunitary + (g) -1//B
6: end for
7: return Cunitary
6

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
Ours
MNO
UNO
MWT
FactFormer
1.0
Tree
Thus
Tree
True MINI resuto
Tree PectForerr
Our rosilss
MMO rosu/o
UNO rocils
corrlation
0.8
0.0
0.6
0.
0.6
0.6
0.4
0.2-
02-
0.0
0.0
0.0
0.0
0.0
10
16
20
10
10
20
lag steps
Time lag steps
lag stops
lag stepps
lag stops
Figure 2. Time correlation of our long-term predictions of TCF and the baselines.
True TKE
MNO Err
UNO Err
MWT Err
FactFormer Err
Ours Err
0.4
0.2
0.0
Figure 3. TKE error long-term predictions of TCF, The left-most snapshot indicates the TKE of ground truth trajectories; and the rest snapshots present the absolute percentage error of TKE in 1000-step rollout predictions from baselines and our approach.
mance using three distinct metrics: (I) the mean energy ab- solute percentage error (ME-APE), which quantifies the ab- solute deviation of predicted energy in each frequency mode from the ground truth; (2) weighted variant (ME-LRw), which assigns higher weights to modes with greater energy, emphasizing low-frequency components following (Schiff al., 2024); and (3) the absolute difference in estimated mixing rates (AA) between true and generated trajectories, key indicator of chaotic systems that reflects how quickly the system loses memory of initial conditions. The mixing rate is estimated from time correlation function, which we also compare in the KF256 and TCF to analyze the detailed temporal decay. Meanwhile, we also compare the abso- lute difference of the estimated Turbulent Kinetic Energy (TKE), which describes the turbulent energy distribution of the underlying dynamical system; and Kullback-Leibler di- vergence (KLD) principle components, which measures the statistical distance between the estimated distributions from the generated data and the true data. More details on the evaluation metrics are provided in Appendix C and F.
and 20.0% in matching the energy spectrum according to ME-APE, ME-LRw, and the mixing rate AA, respectively. These results highlight the effectiveness of our method on both short- and long-forecasting horizons.
Table 1. Short-term and long-term performance of baselines and our transformer prediction on KF256.
Short-term (Rel-L/")
Long-term
Modules
r= 5
7
25
ME-APE
ME-I.Rw
AA
KL.D
UNO
0.92
1.32
0.22
0.36
0.11
0.38
MWT
0.95
1.32
0.26
0.39
0.17
0.40
0.37
Ours
0,86
1.27
0.10
0.15
0.08
0.29
Advaniage (%)
6.52
1.55
23.1
21.0
20.0
21.6
Detailed results of time correlation to derive the mixing rate and the TKE absolute errors are provided in Appendix F.1, where Figure demonstrates our model captures the corre- lation decay of the underlying system, and Figure 6 further validates the advantage of applying A3M in capturing ex- treme values and the unitary operator to maintain invariant statistics.
4.1. Kolmogorov Flow
Table demonstrates the better performance of our trans- former model across both short- and long-term metrics tasks on the KF256 dataset. In the short-term evaluation, our model achieves the lowest relative I error, with a 6.52% and 1.55% improvement over the best baseline for E {5, 25} steps rollout. For long-term statisties, our approach consistently outperforms existing methods in all metrics, achieving notable improvements of 23.1% 21.0%,
4.2. Turbulent Channel Flow
From the results in Table 2, method achieves consis- tently better performance for both short- and long-term pre- dictions. For short-term metrics, we achieve the lowest relative 2 -errors, with 25.0% and 7.14% improvement over baselines at autoregressive forward steps = 5 and = 25, respectively, demonstrating improved accuracy in

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
short-term dynamics. For metrics evaluating energy spec- trum accuracy (ME-APE, ME-LRw) and the mixing rate error ( Δ our model outperforms all baselines by 30.8%, 25.0%, and 27.3%, respectively. These gains highlight our model's ability to maintain precise statistical and physical consistency over long extended horizons.
Table 3. Short-term and long-term evaluation on KF128 of key modules in our transformer.
Short-term (Rel-TA)
Long-term
Modules
5
25
ME-APE
ME-LRw
AA
Base
1.02
1.43
0.17
0.23
0.11
+ A3M Att.
0.98
1.38
0.15
0.21
0.10
Unitary Op.
0.92
1.30
0.14
0.19
0.08
Table 2. Short-term and long-term prediction performance of base-
Advantage (%)
6.12
5,80
6.67
9,52
20.0
lines and transformer predictions for TCF.
Short-term (Rel-L2)
Long-term
sures the preservation of long-term statistics in the L2 space.
Modules
==5
25
ME-APE
ME-LRw
AA
KLD
This property is reflected in the improved long-term met- rics presented in Table 3. Additionally, the A3M pooling effectively captures local extreme values, contributing to improvements in both short- and long-term performance.
MNO
0.08
0.24
0.84
1,95
0.11
3,73
UNO
0,97
0.19
0.86
2.07
0.12
4.89
MWT
0.04
0.14
0.21
0.19
0.11
2.70
FactFormer
0.09
0.17
0.13
0.12
0.23
3.34
Ours
0.13
0.09
0.09
0.08
1.97
Advantage (%)
25.0
7.14
30.8
25.0
27.3
27.0
Secondly, we investigate the effect of RFF positional encod- ing on our transformer model by examining the influence of the hyperparameter o, as described in (Tancik et al., 2020). Specifically, we explore values of a in the range [1, 4, 8, 16, 32). The kernel bandwidth a determines the characteristic length scale of spatial interactions, akin to how mixing rates in physical systems influence the scale of spatial correlations. When the bandwidth is large, spa- tial interactions are highly localized, making each point nearly self-determined and ignoring spatial correlation from other positions, Conversely, a small bandwidth results in almost uniform interactions, representing a state of maximal mixing. From the evaluation results shown in Table we observed that for 2 8, the model more effectively cap- tured high-frequency spectrums, counting more localized interactions. In contrast, for a E [1,4). the model exhibited better performance in capturing low-frequency spectrums. This phenomenon aligns with thcoretical insights discussed in (Tancik et al., 2020). From the attention maps in Figure 8, the results also indicate that a small produces overly smoothed attention that leads to underfitting: while large values of 32 produce noisy attention and excessively large a may lead to overfitting.
Moreover, Figure demonstrates a more consistent mixing rate with the ground truth, characterized by a monotonic decrease in correlation. In contrast, models such as MNO and UNO, which are based on the integer Fourier spectrum, exhibit more periodic correlation pattern. This may con- tract with the mixing property of chaotic systems. We also compare the accuracy of TKE in Figure 3. The top and bottom boundaries of all other models exhibit significant errors, whereas our model demonstrates more stable per- formance across the entire domain, which is attributed to the mnitary-consrraie operator, preserving the invariant statistics.
4.3. Ablation Study
We revisit the Kolmogorov Flow system in a lower resolu- tion of 128 128 KF128, setting training and evaluation random seed as O, to conduct extensive ablation studies to validate the design of our approach efficiently.
Firstly, we evaluate the effectiveness of the unitary opera- tor and the three-channel pooling method as key modules of our transformer. We compare three model versions as follows: (1) the basic version, which is composed of one Mean-pooling dimension reduction module, four factorized attention blocks, and one forward operator learned by simple neural network; (2) an introduction of a Min- and a Max-pooling module to the basic version; (3) further unitary-constrained version introduced with Lunitary on the operator. Table includes the metries evaluated for each configuration, providing insight into how these modules affect short- and long-term performance.
Table 4. Short-term and long- long-term evaluation on KF128 of Kernel Bandwidth positional encoding.
Short-term (Rel-T.2)
Long-term
Bandwidth
t T= 5
f == 25
ME-APE
ME-LRw
AA
1
0.98
1.33
0.22
0.27
0.17
4
0.94
1.33
0.19
0.21
0.08
8
0.92
1,30
0.14
0.19
0.08
16
0.92
1.30
0.15
0.19
0.07
a
0.93
1.30
0.14
0.20
0.11
The introduction of a unitary-constraiin operator and the Min- and Max-pooling module consistently enhances the modeling performance of ergodic chaotic systems across both short- and long-term metrics. According to the Von Neumann mean ergodic theorem, the unitary operator en-
5. Conclusion
The autoregressive generation of long-term trajectories for
large-scale chaotic systems remains significant challenge
8

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
in machine learning. We propose novel transformer-based framework specifically designed for large-scale ergodic chaotic systems. Our approach enhances the transformer architecture by redesigning the attention blocks with A3M pooling in factorized attention to learn the topology mix- ing property and capture the extreme behavior, leading to improved long-term trajectory generation. Moreover, we introduce a novel loss function inspired by the classical von Neumann ergodic theorem, ensuring the preservation of long-term statistical properties in chaotic systems. Exper- imental results show that the proposed model achieves supe- rior performance over both operator-based and transformer- based methods and excels in both short-term accuracy and the preservation of long-term statistical properties, achiev- ing relative 23. % gain among short- and long-term metrics on the high-resolution KF256 dataset and 30.8% on the challenging 3-channel high-resolution TCF dataset.
Academy of Sciences USA, 2021.
Avigad, J. The metamathematics of ergodic theory. Annals of Pure and Applied Logic, 157(2-3):64-75 2009.
Biferale, L.,, Boffetta, G., Celani, A., Lanotte, A., and Toschi, F. Lagrangian statistics in fully developed turbulence. Journal of Turbulence, (7):N6, 2006.
Bochner, S. et al. Lectures on Fourier integrals, volume 42.
Princeton University Press, 1959.
Bousquet, O. and Herrmann, D. On the complexity of learn- ing the kernel matrix. Advances in neural information processing systems, 15, 2002.
Brunton, S. L.., Brunton, B. W., Proctor, J. L.,, Kaiser, E., and Kutz, J. N. Chaos as an intermittently forced linear system. Nature communications, 8(1):19, 2017.
Brunton, S. L., Budišić, M., Kaiser, E., and Kutz, J. N. Modern koopman theory for dynamical systems. arXiv preprint arXiv:2102.12086 2021.
Limitations and future work. Our approach currently focuses on ergodic chaotic systems and does not generalise well non-ergodic system, where long-term statistical prop- erties not well-defined. Future research could explore and extend the framework to accommodate such cases. Ad- ditionally, our method relies on uniform grid structures, which may limit its applicability to more complex chaotic systems. Future work could explore mesh-free techniques, such as graph-based transformers, to model chaotic physical fields without relying on grid-based data structures.
Cao, S. Choose transformer: Fourier galerkin. Ad- vances in neural information processing systems, 34: 24924-24940, 2021.
Chen, T., Zhou, H., Li, Y., Wang, H., Gao, C., Shi, R., Zhang, S., and Li, J. Building flexible machine learning models for scientific computing at scale. arXiv preprint arXiv:2402. 16014, 2024.
Acknowledgments
Cheng, X., Yang, Y., Jiang, W., and Hu, Y. Keec: Embed control an equivariant geometry. arXiv preprint 2rXiv:2312.00555 2023.
The authors thank the reviewers and the program com- mittee of ICML 2025 for their insightful feedback and constructive suggestions that have helped improve this work. We also acknowledge the support from the Univer- sity College London (Dean's Prize, Chadwick Scholarship), the Engineering and Physical Sciences Research Council projects (EP/W007762//1, the United Kingdom Atomie En- ergy Authority (NEPTUNE 2057701-TNN-33, and the Royal Academy of Engineering (IF-2455-1--A116)
Cheng, X., He, Y., Yang, Y., Xue, X., Cheng, S., Guiles, D., and Hu, Y. Learning chaos in a linear way. In The Thir- teenth International Conference Learning Represen- tations, 2025. URL https: https://openreriew /ppnnevvie forum?ide iu-uunddd-ihhiiiiiy
net/
Cieslak, T., Oleszkiewicz. K., Preisner, M., and Szumaúska, M. Kinetic energy represented in terms of moments of vorticity and applications. Journal of Mathematical Fluid Mechanics, 21(4):53, 2019.
Impact Statement
This paper aims advance machine learning applications in studying the long-term properties of large-scale chaotic sys- tems. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Corneeld, P., Fomin, S. V., and Sinai, Y. G. Ergodic theory, volume 245. Springer Science & Business Media, 2012.
Das, S. Lie group valued koopman eigenfunctions. Nonlin- earity, 36(5):2149, 2023.
Dresdner, G., Kochkov, D., Norgaard, P., Zepeda-Núñez, L.,
References
Smith, J. A., Brenner, M. P., and Hoyer, S. Learning to
Alieva, A., Kochkov, D., Smith, J. A., Brenner, M., Wang, Q., and Hoyer, S. Machine learning accelerated com- putational fluid dynamics. Proceedings of the National
correct spectral methods for simulating turbulent flows.
2022. doi: 10-48550/ARXIV22007.000 00056.. URL https:

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
Dudukcu, H. V., Taskiran, M., Taskiran, Z. G. C., and Yildirim, T. Temporal convolutional networks with rnn approach for chaotic time series prediction. Applied soft computing, 133:10994$$ 2023.
Large-Scale Dynamical Prediction. 5 2025. doi:
10.552200//91112222.22
URL https://rdr.
Eckmann, J.-P. and Ruelle, D. Ergodic theory of chaos and strange attractors. Reviews of modern physics, 57(3):617, 1985.
29118212.
Holden, A. V. Chaos, volume 461. Princeton University Press, 2014.
Farazmand, M. and Sapsis, P. A variational approach to probing extreme events in turbulent dynamical systems. Science advances, :e 701533, (3155,333 2017.
Hou, B., Sanjari, S., Dahlin, N., Bose, S., and Vaidya, U. Sparse learning of dynamical systems in rkhs: An operator-theoretie approach. In International Conference Machine Learning, pp. 13325-13352. PMLR, 2023.
Galanti, B. and Tsinober, A. Is turbulence ergodic? Physics Letters A, 330(3-4):173-18,, 2004.
Gilpin, W. Chaos as an interpretable benchmark for fore- casting and data-driven modelling. In Thirty-fifth Confer- Neural Information Processing Systems Datasets and Benchmarks Truck (Round 2), 2021. URL https: / //openreview net /ww/teeeuueee forum? en
Hou, S., Sterling, J., Chen, S., and Doolen, G. A lattice Boltzmann subgrid model for high Reynolds number flows. Pattern formation and lattice gas automata, pp. 151-166, 1995.
Ikeda, M., Ishikawa, I., and Schlosser, C. Koopman and perron perron-frobenius frobenius operators on reproducing kernel banach spaces. Chaos: An Interdisciplinary Journal of Nonlinear Science, 32(12), 2022.
Golub, G. and Van Loan, C. F. Matrix computations. JHU press, 2013.
Guibas, J., Mardani, M., Li, Z., Tao, A., Anandkumar, A., and Catanzaro, B. Efficient token mixing for transformers via adaptive fourier neural operators. In International Conference on Learning Representations, 2021.
JHTDB. Johns hopkins turbulence databases, 2023. URL phaa jhu edu/.
Jiang, R., Lu, P. Y., Orlova, E., and Willett, R. Training neural operators preserve invariant measures of chaotic attractors. Advances in Neural Information Processing Systems, 36, 2024.
Gupta, G., Xiao, X., and Bogdan, P. Muld operator learning for differential equations. Advances in neural information processing systems, 34:24048-240M 2021.
Jørgensen, P. Ruelle Operators: Functions which Are Harmonic with Respect Transfer Operator: Func- tions which Harmonic with Respect to Transfer Operator, volume 720. American Mathematical Soc.. 2001.
Hairer, M. and Mattingly, J. C. Ergodicity of the 2d navier- stokes equations with degenerate stochastic forcing, An- nals of Mathematics, pp. 1032, 2006.
Han, K., Wang. Y., Chen, H., Chen, X., Guo, J., Liu, Z., Tang, Y., Xiao, A., Xu, C., Xu, Y., al. A survey vision transformer. IEEE transactions pattern analysis and machine intelligence, 45(1):87-110, 2022.
Kissas, G., Seidman, J. H., Guilhoto, L. F., Preciado, V. M., Pappas, G. J., and Perdikaris, Learning operators with coupled attention. Journal of Machine Learning Research, 23(215): 2022.
Hao, Z., Wang, Z., Su, H., Ying, C., Dong, Y., Liu, S., Cheng. Z.. Song. J., and Zhu, J. Gnot: A general neu- ral operator transformer for operator learning. In Inter- national Conference on Machine Learning, pp. 12556- 12569. PMLR, 2023.
Kloeckner, B. R. Empirical measures: regularity a counter-curse to dimensionality. ESAIM: Probubility and Statistics, 24:408-4444 -43400844444000 2020.
Klus, S., Schuster, L. and Muandet, K. Eigendecomposi- tions of transfer operators in reproducing kernel hilbert spaces. Journal of Nonlinear Science, 30:283-315. 2020.
He, Y. Benchmark dataset kl256 for chaos meets atten- tion: Transformers for large-scale dynamical prediction, may 2025. URL https: /tt.////iii//// org/10. 5281/ zenodo. 148C1580.
Kolesov, A. Y. and Rozov, N. K. On the definition of chaos'. Russian Mathemutical Surveys, 64(4):701, 2009.
He, Y. Xue, X., Hu, Y., Yang, Y., Cheng, X., and Wang, H. Benchmark dataset Turbulent Channel Flow for Chaos Meets Attention: Transformers for
Korotin, A., Li, L., Genevay, A., Solomon, M., Filip- pov, A., and Burnaev, E. Do neural optimal transport
10

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
solvers work? continuous wasserstein-2 benchmark. Advances in neural information processing systems, 34: 14593-14605, 2021.
Mikhaeil, J., Monfared, Z., and Durstewitz, D. On the difficulty of learning chaotic dynamics with rnns. Ad- vances in Neural Information Processing Systems, 35: 11297-11312, 2022.
Langeroudi, M. K., Yamaghani, M. R., and Khodaparast, S. Fd-Istm: A fuzzy Istm model for chaotic time-series prediction. IEEE Intelligent Systems, 37(4):70-78, 2022.
Neumann, J. v. Proof of the quasi-ergodic hypothesis. Pro- ceedings of the National Academy of Sciences, 18(1): 70-82, 1932.
Latt, J., Chopard, B., Malaspinas, O., Deville, M., and Mich- ler, A. Straight velocity boundaries in the lattice boltz- mann method. Physical Review E Statistical, Nonlinear, and Soft Matter Physics, 77(5):056703, 2008.
NVIDIA. 2023.
Two-equation turbulent channel flow,
URL httt::::/doocs.nidia..oooo deee_laarriinng//mdiiiiaoooouoae
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:2010.00 08895, 2020.
wewpllent_chanee.hheee
Ottino, J. M. et al. Mixing, chaotic advection, and turbu- lence. Annual Review of Fluid Mechanics, 22(1):207-254, 1990.
Li, Z.,, Liu-Schiaffin,, M., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandku- mar, A. Learning chaotic dynamics in dissipative systems. Advances in Neural Information Processing Systems, 35: 16768-16781, 2022a.
Pathak, J., Hunt, B., Girvan, M., Lu, Z., and Ott, E. Model- free prediction of large spatiotemporally chaotic systems from data: A reservoir computing approach. Physical review letters, 120(2):024102, 2018.
Li, Z., Meidani, K., and Farimani, A. B. Transformer for partial differential equations' operator learning. arXiv preprint XrXXv:2205.10 13671, 2022b.
Rahimi, A. and Recht, B. Random features for large-scale kernel machines. Advances in neural information pro- cessing systems, 20, 2007.
Rahman, M. A., Ross, Z. E., and Azizzadeneshelii K. U-no: U-shaped neural operators. arXiv preprint D/Xiv:224.111222 2022.
Li, Z., Shu, D., and Barati Farimani, A. Scalable trans- former for pde surrogate modeling. Advances in Neural Information Processing Systems, 36, 2024.
Sangiorgio, M. and Dercole, F. Robustness of Istm neural networks for multi-step forecasting of chaotic time series. Chaos, Solitons & Fractals, 139: 10044, 2020.
Lorenz, E. N. and Haman, K. The essence of chaos. Pure and Applied Geophysics, 147(3):598-599, 1996.
Lu, L., Jin, P., Pang. G.. Zhang, Z., and Karniadakis, E. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218-23) 2021.
Sapsis, T. P. Statistics of extreme events in fluid flows and waves. Annual Review of Fluid Mechanics, 53 2021.
Schiff, Y., Wan, Z. Y., Parker, J. B., Hoyer, S., Kuleshov, V., Sha, F,, and Zepeda-Núñez, L. Dyslim: Dynamics stable learning by invariant measure for chaotic systems. arXiv preprint arXiv:2402.04467. 2024.
Madondo, M. and Gibbons, T. Learning and modeling chaos using Istm recurrent neural networks. In Proceedings of the Midwest Instruction and Computing Symposium, Duluth, Minnesota, pp. 6-7, 2018.
Sengupta, A. Toward a theory of chaos. International Journal of Bifurcation and Chaos, 13(11)3314-3223 2003.
Marchenko, V. A. Nonlinear equations and operator alge- bras, volume 17. Springer Science Business Media, 2012.
Shi, Y., Tian, X., and Wang, X. The space of er- godic measures for lorenz attractors. arXiv preprint arXiv:2006.08/93 2020.
Meyer, R. A., Musco, C., Musco, C., and Woodruff, D. Hutch++: Optimal stochastic trace estimation. In Sympo- sium on Simplicity in Algorithms (SOSA), pp. 142-155.
SIAM, 2021.
Shiryaev, A. N. Probability-l, volume 95. Springer, 2016.
Mezić, Koopman operator, geometry, and learning of dynamical systems. Not. Am. Math. Soc, 6887) 1105, 2021.
Smagorinsky. General circulation experiments with the primitive equations: I. the basic experiment. Monthly Weather Review, 91(3):99-164 1963.

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
Smale, S. and Zhou, D.-X. Learning theory estimates via integral operators and their approximations. Constructive approximation, 26(2):153-177 2007.
for lattice boltzmann simulation of high reynolds ber turbulent flows. Communications Physics, 7(1):338, 2024.
Sommerficld, A. Ein beitrag zur hydrody erk- larung der turbulenten flussigeiiseweegung. Atti Congr. Int. Math. 4th, 1908.
Yan, M., Huang, C.,, Bienstman, P., Tino, P., Lin, W., and Sun, J. Emerging opportunities and challenges for the future of reservoir computing. Nature Communications, 15(1):2056, 2024.
Succi, S. The Lattice Boltzmann Equation for Fluid Dynam- ics and Beyond. Oxford University Press, 2001.
Yang, D. Random attractors for the stochastic kuramoto- sivashinsky equation. Stochastic analysis and applica- tions, 1286)1128511303 2006.
Tancik, M., Srinivasan, P., Mildenhall, B., Fridovich-Keil. S., Raghavan, N., Singhal, U., Ramamoorthi, R., Barron, J., and Ng. R. Fourier features let networks learn high fre- quency functions in low dimensional domains. Advances in neural information processing systems, 33:7537-7547, 2020.
Yang, Y., Cheng, X., Giles, D., Cheng, S., He, Y., Xue, X., Chen, B., and Hu, Y. Tensor-var: Variational data assimilation in tensor product feature space. 2025.
Yeung. P., Zhai, X., and Sreenivasan, K. R. Extreme events in computational turbulence. Proceedings of the National Academy of Sciences, 112(41):1263331122
Temam, R. anfinite-dimensinn dynamical systems in me- chanics and physics, volume 68. Springer Science & Business Media, 2012.
2015.
Young, L.-S. What are srb measures, and which dynamical systems have them? Journal of statistical physics, 108: 733-754, 2002.
Trevisan, A. and Palatella, L.. Chaos and weather forecasting: the role of the unstable subspace in predictability and state estimation problems. International Journal of Bifurcation and Chaos, 21(12):3389-344 2011.
Zelik, S. Attractors. then and now. arXiv preprint arXiv:2208. 12101, 2022.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Zhang, J., Ma, P., Zhong, W., and Meng, C. Projection- based techniques for high-dimensional optimal transport problems. Wiley Interdisciplinary Reviews: Computa- tional Statistics, 15(2)=e187,, 2023.
Virtanen, P., Gommers, R., Oliphant, T. E.,, Haberland, M., Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, J., er al. Seipy 1.0: fundamental algorithms for scientific computing in python. Nature methods, 17(3):261 272, 2020.
Zhang, Y. and Gilpin, W. Zero-shot forecasting of chaotic systems. arXiv preprint arXiv:2409.15771 2024.
Vögler, D. The curse of dimension in multi-marginad optimal transport. PhD thesis, Universität München, 2023.
Wan, Z. Y., Baptista, R., Boral, A., Chen, Y.-F., Anderson, J., Sha, F., and Zepeda-Nunez, L. Debias coarsely, sample conditionally: Statistical downscaling through optimal transport and probabilistic diffusion models. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL https: / openreview n
Xiong, J. and Yang, Z. Chaos caused by topologically mixing map. Technical report, International Centre for Theoretical Physics, 1991.
Xue, X., Yao, H.-D., and Davidson, Synthetic turbulence generator for lattice Boltzmann method the interface be- tween RANS and LES. Physics Fluids, 34(5):055118, 2022.
Xue, X., Wang, S., Yao, H.-D., Davidson, L., and Coveney, P. V. Physics informed dataenriven near-wall wall modelling
12

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
A. Table of Notations
Notations
Meaning
(* p) (X, B, ())
function spaces defined using a natural generalization the p-norm
measure space
P T
Lebesgue measure
nonlinear forward map
0
parameters of neural networks
feature functions in L space
S
Number of grid points in i-th coordinate
N
I (=1
M
S
Total Number of grid points
C
Number of physical quantities
d
feature dimensions in L space
Number of samples from d dimensional unit sphere
m
Number of RFF samples
Random unit vector sampled from d d-dimensionall unit sphere
13

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
B. Convergence of Random Fourier Features
Let X N(,, o²) be a Gaussian random variable. The characteristic function defined as:
xxx = [((
V2TTO |_00
00
etz e
(a-1)
20
dx.
(10)
Complete the square in the exponent:
(x - 11)
x² 2 + u²
il.
i1.
20
20
1
(x² 24 + 112 =2ttaxx)
20²
(11)
1
(x=
22( + itoo2 Jo: + 122)
20²
20 ((x (1 + itoo))² - (H + ito + 122)
Substituting back:
00
(u-
(t)
dx
V2To
00
2
(1+ttt
c itp
e
2
dr
12n
-00
integral Gaussian density
The integral equals V since it's the integral of Gaussian density (after appropriate substitution). Therefore:
(t) etg 2
This completes the proof of the characteristic function for a Gaussian distribution with parameters μ and a². In our case, t is replaced by the relative Euclidean norm in the grid domain M and is zero in random Fourier sampling. Thus, when samples m is sufficiently large, 8(i - j) in Equation (4) converges to
O j) Ele 2Tb(s, =s1]] To [1:,- 611
(12)
14

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
C. Ergodicity and Mixing
Mixing and ergodicity are correlated concepts the mathematical analysis dynamical systems (Cornfeld et al., 2012), particularly within the framework of measure theory. Consider meesure-presernee dynamical system (X, B, ju, T), where T: X + X is transformation that preserves the measure μ. Ergodicity is defined by the property that for any measurable set A E B, i (A) A then ((A) is either 0 or 1. This implies that time averages equal space averages for integrable functions, formally expressed as
N Σ x( [x/ dp for p-almost every x E X.
lim N-
N
n=0
Mixing a stronger condition where, for any two measurable sets A, B B, the measure of their intersection under iteration satisfies
(M ((T"(A)) - (A))((B) as 72 20.
This signifies that the system asymptotically "forgets" its initial state, leading to statistical independence of A and B over time. The mixing rate quantifies the speed at which this convergence occurs, often characterized by the decay of temporal correlations. Mathematically. if there exists a function ((n) such that
(T(T"(A)n -
((n),
(13)
and ((n) decays to zero at a certain rate (e.g., exponentially ((n) x exp(-An)), this rate ((n) is referred to as the mixing rate. faster mixing rate implies more rapid approach equilibrium, which is crucial for applications ranging from statistical mechanies the analysis of randomized algorithms. Together, ergodicity ensures the thorough exploration of the state space, while mixing and its associated rate provide quantitative measure of how efficiently the system achieves statistical uniformity.
15

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
D. Experiment settings.
In this section, we provide the details numerical experiments and baseline implementations on the high-resolution Kolmogorov flow and turbulent channel flow.
D.I. Kolmogorov flow
Governing Equations The Kolmogorov flow system is a classic model for studying fluid instabilities and turbulence in wwo-dimensional incompressible flows (Temam, 2012). It is described by nonlinear, incompressible Wavier-Stoker equation driven by a sinusoidal forcing term.
iu a1
+ (u Vju+ Vp Au +f in domain 2T] x [0,
Re
(14)
V u 0 in domain j0, 2m]² x [0, [],
0 represents the external forcing the y-direction. where u(x, t) is the velocity field, p is the pressure, and f sin kyy)
Re is the Reynolds number calculated by Re uL/v, where is the kinematic viscosity, L is the characteristic length and is the speed (Sommerfield, 1908). The equivalent vorticity formulation for = 8x +
Ju,
0u is
By
Dw 0t
1 W in domain [0, 2m x (0, Re
+ n V
(15)
We use the spectral method to generate vorticity states with the pseudo-spectral solver in the jax-cfd toolbox (Dresdner al., 2022). The initial conditions follow those initializing methods in (Alieva et al., 2021). We set the maximum speed as the 256 256 grid with temporal resolution of
Umar 9.9 Uz Byy 7 for each direction), le 3, and 1.7e 3.
Datasets The datasets of vorticity states of the Kolmogorov flow system consist of 150 training trajectories, 40 validation trajectories, and 30 testing trajectories in total. Each trajectory contains 500 frames for 10 seconds with unique initial state. KF256 dataset uses the full resolution (256 x 256) of the generated states, and KF128 dataset downsamples the state resolution by half (128 x 128). Both the ablation studies and the baseline models were trained and evaluated on the corresponding sets.
D.2. Turbulent channel flow
This simulation employs the 3DQ19 lattice model for the lattice Boltzmann method (LBM) (Succi, 2001), which features dimensional and 19 discretized velocity directions. Via Chapman-Enskog Enskog analysis, LBM can recover the Navier-Stokes
equations and beyond (Succi, 2001). The lattice cell is specified by its position at time t, and is characterized by discretized set of speeds c, where E {0, 1, Q 1} with Q = 19. The evolution equation for the distribution functions can be written as:
+ + At) = f(x, t) Ω [f( t) f" (x,+))
(16)
where Ω denotes the Bhatnagar- Saattnaggr-Grros--ee (BGK) collision kernel (Succi, 2001). The collision kernel relaxes the distribution function towards the local Maxwellian distribution function f.²:
1 (x, = wip (x, 1 C; (x, C. (x, [u (x, t) u (xx,1) t)]
(17)
20h
2
where Wt denotes the direction based weights (w; 1/3 for i=0, 1/18 for the six nearest neighbours and 1/36 for the remaining directions), p (x, t) is the cell macroscopic density, u (x, t) is the cell macroscopic velocity. In Equation (16), At symbolizes the lattice Boltzmann time step, which is set to unity. The LBM kinematic viscosity v is defined as
T 2- Ar.
(18)
with c, representing the speed of sound, and ct equating to 1/3 in Lattice Boltzmann Units (LBU). Now, we apply the Smagorinsky subgrid-scale turbulent model within the LBM framework, the effective viscosity Var (Smagorinsky, 1963; Hou et al., 1995) is modeled as the sum of the molecular viscosity, Vo. and the turbulent viscosity, 14:
16

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
Veft Mo + U4.
4/2 Camag A |||
(19)
where | is the filtered strain rate tensor, Csmag is the Smagorinsky constant, A represents the filter size. We apply the total viscosity Pt into Equation (18) to add the turbulent viscosity. Eventually, macro-scale quantities such as density and momentum are derived from the moments of the distribution function fi t), the discrete velocities c,:
Q-1
((x, - Σ [ix,
(20)
==0
u(x, t) t(x,t)
1
Σ f(xxt)c.
(21)
Dataset Contribution The datasets derived from the macroscopic velocity field computed using lattice Boltzmann numerical solver. A 2D plane (192 x 192) was extracted at the cid-crross-section of the flow direction in a 3D turbulent channel flow. Each pixel in the plane contains velocity vector field information u, including Uay Bly, and Mz- The datasets comprise 240 training trajectories, 24 validation trajectories, and 24 test trajectories, with each trajectory containing 595 frames. The simulations were executed on 4,096 CPUs over 144 hours, resulting in a total computational cost of approximately 300,000 CPU hours. More details of the dataset are available in Appendix G
D.3. Baseline implementations
In this subsection, we present the details of the baseline implementations. For all the baselines, the same experimental settings and model configurations are used for both the KF256 and TCF experiments.
For the U-shaped Fourier Neural Operator (UNO) (Rahman et al., 2022), we adopt the original model configuration described in Rahman et al. (2022) for the Kolmogorov flow with resolution of 256 x 256. The UNO architecture consists of four Fourier blocks, each with 24 frequencies per channel, a width of 64, and a hidden dimension of 96. includes skip connections and is connected via a channel-wise MLP at the end of each block. The official implementation is utilized from https: /githtb....//gg////
maann.
For Markov Neural operator (MNO) (Li et al., 2022a), we used the official implementation from https: / huub com/
/tattoooomaarooooo/ neural www.nneeuuuurreeeeeeeeee www.roorrr........ ree mmain. For the important choice of dissipativity hyperparameters in MNO, we follow their settings in their numerical experiments on Kolmogorov flow with dissipativity regularization coefficient a = 0.01, scaling down factor À = 0.5, and attractor radius as 156.25 x where the square root of total number of grid points N. The MNO model consists of four 2d Fourier layers with 20 frequencies per channel and width 64.
Apart from Fourier transform- based baselines, the Multiwavelet-basee Operator (MWT) (Gupta et al., 2021) uses wavelet transformations to capture localized, multi-scale dynamics effectively. The MWT model is composed of four 2D multiwavelet blocks, each utilizing 32 Legendre- wavelets. The network features a hidden dimension of 128, applies ReLU activations between the multiwavelet blocks, and generates outputs through a channel-wise MLP. The official implementation can be found from https: githtb...hu..go///g////gg....h.t www hatmm//aurrv771553 htt:///wo/// -wwtopeeattrrrr
For Factformer (Li et al., 2024), we use the official implementation provided at https: /tt:::///gtttbbb. giggthtthh..... com/ Barattiaab/FactFoooo www.aactrrreeeereeeeeeeeeeeeeeee htt///ree//mmmmm Factformer adopts a scalable architecture combining self-attention mechanisms with low-rank approximations to reduce computational complexities. The factformer consists of four factorized attention blocks. Each attention block has 16' heads with dimension 64. The hidden dimension set 256 with two-layered channel-wise MLP the forward block. To align with other methods, we set the temporal dimension as 1 and follow the auto-regression paradigm for training and evaluation as described in (Li et al., 2024).
All models are optimized using the Adam optimizer with the L.2 norm between predictions and ground truth, except for MNO. For MNO, the loss function consists of two terms: (1) the first-order Sobolev norm between predictions and ground
increased number of heads from 8 to 16 to accommoaate the higher resolution and align with our implementations.
17

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
truth, and (2) dissipative loss term (refer (Li et al., 2022a) for details). The learning rate (LR) and scheduler are kept the specified in their original implementations. All models trained for 50 epochs using their default batch size settings.
D.4. Our transformer implementation details
For the benchmarks, the major hyperparameter configurations of our transformers listed in Table 5.
Table 5. Major hyperparameter configurations of our transformers.
Hyperparameters
Kolmogorov Flow
Turbulent
(KF128)
(KF256)
Channel Flow
Number of blocks
4
4
Attention heads
16
Kernel bandwidth
8
8
Latent dimension
256
256
256
Input function
2D Conv
2D Conv
2D Conv
Output function
MLP
MLP
MLP
The Latent dimension determines the number of channels in the latent state representation. The Number of blocks specifies the factorized attention blocks incorporated into the encoder of our transformer, influencing its ability to process input data, The Atlention heads defines the number of heads used in the multi-head attention mechanism, enabling the model to capture diverse patterns effectively. The kernel bandwidth, represented by the sigma parameter of the Gaussian distribution, governs the sampling of the random Fourier spectrum to approximate the kernel. Finally, the operator dimension specifies the latent dimensionality of the bidirectional unitary operator. We use Adam as the optimizer train the models. The learning rate initiates from 1e 4 and decays by half for every 10 of 50 epochs, following the scales and settings of related chaos works (Li et al., 2022a; 2024). At the end of our model, we employ a three-layer CNN residual block with circular padding to enforce periodic boundary conditions and project back to physical variables. The choice of latent dimension and attention heads follows the setting of the corresponding lowest relative error in the ablation study of (Li et al., 2024).
D.5. Computational resources for machine learning experiments
Experiments on KF256 and TCF were submitted as resource-restricted jobs to shared compute cluster. Ablation studies on KF128 were conducted within a group workstation. Computational resources for these machine learning studies used in each chaos system experiment are listed in Table 6.
Table Computational resources by experiment.
Experiment
Hardware
KF128
GeForce RTX4090 GPUs, 24 GB
KF256
A 100 GPU, 40GB
TCF
1 A100 GPU, 40GB
18

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
E. Evaluation Criteria
In this section, we include further details about the criteria used in Section 4.
E.1. Short-term accuracy evaluation
Relative 1.2 norm. The stepwise accuracy the model's predictions for physical states, such as velocity and vorticity, can be evaluated using the relative norm L2, as described in (Li et al., 2022a; 2024). The relative error is defined as
NV
2
21112
Relative L (e-rrro((k)) =
(22)
N
1
ll
where of and ziz are the prediction and ground-truth states at time step i of jth test trajectory.
E.2. Long-term statisties evaluation
Matching energy spectrum Evaluation of the matching of the energy spectrum is important to quantitatively assess the model performance (Wan et al., 2023). The energy spectrum is calculated as
2
E(k)= Σ 12
Σ
Σ u(xa,) exp (-j2mk Dijj
(23)
where E(k) represents the energy spectrum wavenumber k, and ū(k) denotes the Fourier transform of the velocity field.
Mean energy absolute percentage error (ME-APE)
No
Latuee
ME-APE
x 100
(24)
N
where 1t Exred and tx represent the predicted and true energy values, respectively. We evaluate this metric to uniformly evaluate the relative energy spectrum absolute error.
Mean energy log ratio (ME-LR)
Pepred (k)
ME-LR Σ wk log
(25)
Emn (k)
where wx denotes the weight assigned to each wavenumber. This metric measures the logarithmic error between the predicted and true energy values. We further define wik Σ. Ein (1) to capture high energy modes as ME-LRw.
(
Turbulent Kinetic Energy. We evaluated the turbulent kinetic energy (TKE) of long-term rollout predictions of baselines and our approach. The TKE represents the mean kinetic energy per unit mass associated with turbulent velocity fluctuations, which quantifies the intensity of the turbulence and the energy distribution across scales. It is defined as:
TKE
d
where = - 2.)² and represent the spatial domain dimensionality. The TKE derived from vorticity field can approximated from (Cieslak et al., 2019). To evaluate the models, we compute the grid-wise absolute difference between the ground truth TKE and the estimated TKE from the long rollouts of learned models as:
ErrortkE ||K TKE
19

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
Kullback-Leibler Divergence. We evaluate the statistical distance between the predicted and ground truth distributions of the state field using the Kullback-Leibler (KI.) divergence. The KI. divergence quantifies the information loss when the predicted distribution Q is used to approximate the true distribution P. It is defined as:
(L (P || Q) = Epllog
where we estimate P and Q via gauss i an_kde of principal components in samples of the ground truth and predicted fields, respectively. In our setup, the principal components and samples are drawn over the full spatial grid and temporal range of the rollout. A lower KL divergence indicates better agreement in the distributional shape and concentration of vorticity statistics.
The mixing rate of chaotic system quantifies how quickly its state variables lose the memory of initial conditions and become statistically independent. It measures the decay of correlations and the rate which distributions approach equilibrium. For system with state the autocorrelation function is given as
Mixing rate.
C(t) - (=)=
The system is said to mix at exponential rate if Ctt) e where À is the mixing rate (see more explanation in Equation (13) at Appendix C). Empirically, we estimated the autocorrelation function up to time-lag K from N true and generated system trajectories such as
T- K-t
N
T
Č(I) N(T - K
(zk - (2(44 2), where E
(26)
NT
k=
t=0
where z/ is the tth snapshot in ith (true/generated) trajectory. To estimate the mixing rate, we use nonlinear least squares optimization to fit the normalized autocorrelation function C(1) C(t)/((0) an exponential decay model of the form -Axr, where À is the fitting parameter. It minimizes the sum of squared errors between the empirical and model values over time lags. This approach allows for robust estimation of the mixing rate. Finally, we make the absolute difference between the estimated values from ground truth and generated rollouts to form metric.
"I practice, we the curve_fis function from the SciPy library (Virtanen al., 2020), which provides efficient estimation of the parameters by minimizing the residual sum squares.
20

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
F. Experimental results visualization
In this section, we include more visualized experimental results supplementary to the section 4.
F.1. Kolmogorov flow
k=6
k=11
k=16
k=21
k=26
k=
True
MNO
UNO
MWT
FactFormer
Ours
Absolute Error
Figure 4. Absolute of short-term KF256 predictions of baselines and
Ours
MNO
UNO
MWT
FactFormer
0.4
True
20ar MN0 recils
Tree
Troe
True
Denforrve
correlation
0.2
03
0.:
02
03
0.0
0.0
10
15
143
1.5
30
301
Timee lagl
Time
Time
Timne
Time
Figure 5. Time correlation of our long-term predictions KF256 and the basclines.
21

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
True TKE
MNO Err
UNO Err
MWT Err
FartFormer Err
Ours Err
0.4
0.0
Figure 6. TKE emor of long-term predictions of KF256. The left-most snapshot indicates the TKE of ground truth trajectories; and the rest snapshots present the absolute percentage error of TKE in 1000-step rollout predictions from baselines and our approach.
F.2. Turbulent channel flow
For this 3D dataset, the visualization includes short-term prediction states and long-term statistics of the 2D cross-section states the forward direction.
k=1
k=6
k=11
k=16 16
k=21
k=26
k=31
True
MNO
OO
UNO
MWT
FactFormer
Ours
0.0
0.1
0.3
0,4
0.5
Absolute Error
Figure 7. Absolute error of short-term TCF predictions of baselines and ours.
22

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
F.3. Further details on ablation study
In this section, various attention maps of the ablation study section 4.3 on the kernel bandwidth of RFF positional encoding are included in Figure 8.
Depth = 1
Depth = 4
X-axis
Y-axis
X-axis
Y-axis
125
125
100
100-
75
75-
75-
75
50
50
50
25-
25
25
50
100
0
100
0
50
100
50
100
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
1.0
-1.0
-0.5
1.0
125
125
4
100
100
11
75
75-
75
75-
30
50
25
25
0
50
100
0
50
100
0
50
1000
-0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
11.0
0..5
0.0
0.5
1.0
125
125
125
125
8
100
100
100
100
|I
75-
75
75
30
50
25
25
25
100
50
100
100
50
100
-0.5
0.01
0.5
1.0
-0.5
0.0
0.3
1.0
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
125
125
125
0 1
16
100
100
100
75
75-
75
50
50
50
25
25
25
50
100
50
100
0
50
200
50
100
-1.0
-0.5
0.0
0.5
1.0
0.5
1.0
0.0
0.0
125
125
125
32
100
100
100
75
75
75
75
0 =
50
50
50
50
25
25
25
25
100
0-
50
50
100
1..0
-0.5
0.0
0.5
1.0
1.00
0.5
0.0
0.5
1.0
1.0
0.5
0.0
0.5
1.0
-0.5
0.0
0.5
Figure 8. Impact of applying random Fourier positional encoding on the attention maps with respect to the kernel bandwidth o (1,4, 8, 16, 32}. From the attention maps of the first and the last block of our implementation, we observed that small E (1,4) produced overly smoothed attention maps, while the large value a = 32 produced noisy attention maps.
23

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
E.4. Computational costs
Table Computational of method scaling grid resolution. We report the model performance of computational runtime, memory usage and floating-point operations sccond (FLOPs) per forward pass, across a range of grid resolutions and demonstrated significant time and memory efficiency on increasing resolutions consistently. The results are gathered on the datasets given model configurations in Table on A100 GPU listed in 6 with 7, 325, 663 parameters.
Grid Resolution
Runtime
Memory Usage
FLOPS per Forward Pass
32x32
24 ms
174 MB
2.5 GB
x4x64
27 ms
537 MB
10.5 GB
x28x 128
ms
1986 MB
50.2 GB
256
58
7684 MB
268 GB
Table Computational costs baselines and ours on TCF predictions. We report model performance computational runtime,
memory usage and Moating-point operations per second (FLOPs) forward pass. The results are gathered given model configurations Table A 100 GPU listed in 6. Attention mechanism overall computational heavy than neural operator methods, tensor factorization and axial attention make the computation tractable for large scale chaos states.
Models
Parameter Count
Runtime
Memory Usage
FLOPs per Forward Pass
MNO
6,467,425
31 ms
377 MB
3,45 GB
UNO
17,438,305
12 ms
769 MB
6.88 GB
MWT
5,089,153
50 ms
313 MB
9.52 GB
FactFormer
6,083,009
53 ms
6889 MB
239 GB
Ours
7,325,665
58 ms
7684 MB
268 GB
24

---

Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction
G. Turbulent channel flow as a benchmark for AI4Science
We employ turbulent channel flow as a benchmark to assess the generality of our approach for turbulent flows. Unlike isotropic turbulence, turbulent channel flow is bounded by two parallel planes with no-slip boundary conditions. The Lattice Boltzmann Method (LBM), an alternative to directly solving the Navier-Stoker equations, is inherently parallelization- friendly, making it particularly advantageous for tigh-performnce computing applications. In recent decades, LBM has proven its versatility, enabling applications that range from micro-nano fluidics to macroscopic turbulent flows (Succi, 2001).
In this study, the computational domain for the turbulent channel flow simulation is defined with dimensions Lg x Ly x L 1024 x 192 x 192, where z. y. and 2 denote the streamwise, vertical, and spanwise directions, respectively. The friction Reynolds number is set to Re,- = 180 which is equivalent to Re = 3250. Periodic boundary conditions are applied the streamwise and spanwise directions, while the vertical direction governed by no-slip boundary conditions (Latt et al., 2008). This configuration distinguishes our dataset from existing studies that primarily rely on 2D turbulent cases, as our dataset is based on fully
simulations.
A key motivation for utilizing 3D simulations, even when analyzing 2D cross-sectional results, lies in their ability to capture realistic flow phenomena that are inherently three-dimensional. These include features such as secondary flows, coherent vortical structures, and fully developed turbulence. Such complexities are absent in 2D turbulent channel flow datasets (NVIDIA, 2023), which neglect variations along the third dimension. As result, 2D datasets fail to adequately represent the intricacies of real- realdworld turbulent flows, limiting their utility in validating models or exploring phenomena where three-dimensional effects are critical. Furthermore, existing high-resolution 3D turbulent channel flow datasets often derived from direct numerical simulations (DNS), which, while highly accurate. come with significant drawbacks. These datasets can require up to 100 terabytes of storage (JHTDB, 2023), resulting in an unnecessarily high memory cost that makes them impractical for early-stage machine learning research. In comparison, our datasets are designed to be three orders of magnitude smaller, approximately 100 GB, striking a balance between resolution and efficiency. This makes them see for heatiee--eaannnn--peeeeee tasks, enabling broader accessibility and usability for researchers.
The fully developed 3D turbulent channel flow simulation follows the configuration outlined in reference (Xue al., 2022). The simulation begins from initial zero-velocity field, with square block of size 20 x 20 100 grid points positioned at = 192. A volumetric force is applied uniformly across the domain to drive the flow. The simulation is run for 50 domain- domain-through times to establish initial flow characteristics. After this initial phase, the block is removed, and the simulation is continued for an additional 50 Nomain-through times, allowing the flow to fully develop into a turbulent state. Data sampling begins after this stage, focusing on the 2D cross- at x = 512. Data is collected over a further 100 turnover times, ensuring that the samples represent fully developed turbulence. The simulation timestep is set to A 0.02 s, and it operates in dimensionless units (800 timesteps), as is typical for LBM simulations. The detailed transformation from dimensionless units to physical units can be found in reference (Xue al., 2024). To ensure robustness and generality, we conducted three independent 3D turbulent channel flow simulations under these conditions, generating comprehensive dataset for analysis.
y
z
(a) 3D-Turbulent Channel Flow for illustration
(b) The cross section in direction
Figure 9, Turbulent Channel Flow dataset for illustration
25


Reviews:
[{"review_id": "nzd8PLPo9O", "reviewer": "Reviewer_vgH9", "summary": "The paper introduces a transformer-based framework for predicting large-scale chaotic systems. The authors tackle a key challenge in dynamical system forecasting -- the amplification of prediction errors due to positive Lyapunov exponents -- by using ergodicity. Their approach includes:\n\n- A modified attention mechanism (A3M Attention) that captures statistical moments and extreme values in chaotic systems.\n- A loss function inspired by the von Neumann mean ergodic theorem, aimed at enhancing long-term statistical stability.\n- A large-scale dataset featuring 140k snapshots of turbulent channel flow to benchmark prediction accuracy.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer's constructive feedback and recognition of our work. We address the concerns and questions as follows:\n\n[Link to visualizations and references](https://anonymous.4open.science/r/ChaosMeetsAttention/README.md)\n\n[Backup](https://filebin.net/37p4dxup0t320143)\n\n* Concern 1 & part of Question 1:\n\n    `Whether the proposed loss function strictly enforces ergodicity in practice, any proof, and preserves in long-term predictions.`\n\nWe 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).\n\n* Concern 2:\n\n    `KL Divergence and other metrics`\n\nWe 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.\n\n`Autocorrelation Decay Rate & (Max) Lyapunov Exponent Deviation`\n\nWe 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).\n\n* Concern 3 & Questions 2, 3, 6:\n  \n  `Computational cost (runtime, memory, FLOPs) and scaling with grid size`\n\nWe 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.\n\n* Concern 4:\n\n  `Relevance to PINN and physics hybrid transformer`\n\nWe 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.\n\n* Concern 5 & Question 4:\n\n `Tuning strategy for kernel bandwidth`\n\nDue 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.\n\n* Question 1:\n\n`Ergodic loss vs traditional methods`\n\nWe appreciate the reviewer’s insightful questions regarding the ergodic loss.\n\n(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.\n\n(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].\n\nOur 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.\n\n* Question 5:\n\n  `Alternative pooling methods`\n\nWe implemented the mean-pooling method for comparison with our full A3M method. Results are reported in the link.", "claims_and_evidence": "The paper presents compelling claims regarding the effectiveness of its transformer model in maintaining long-term statistical properties while enhancing short-term prediction accuracy. While the claims are well-supported by empirical results and the von Neumann ergodic theorem is leveraged to justify the approach, it remains unclear whether the proposed loss function strictly enforces ergodicity in practice; thus, additional theoretical validation, such as a formal proof or convergence analysis, would further strengthen the argument.", "methods_and_evaluation_criteria": "While the evaluation criteria used in the paper (relative $L^2$ norm for short-term accuracy, ME-APE, ME-LRw, and $\\nabla \\lambda$ for long-term statistical consistency) are reasonable, additional and more rigorous evaluation metrics could provide a more comprehensive assessment of the model's performance. Could the authors consider integrating some of these additional metrics to provide a more holistic evaluation of their method? For instance: \n\n- Kullback-Leibler (KL) Divergence or Wasserstein Distance, or/and the Hellinger distance: Please see Sect. 4.2  of this paper for more details https://proceedings.mlr.press/v202/hess23a/hess23a.pdf\n\n- Autocorrelation Decay Rate: Measuring how the model preserves the autocorrelation structure over time can indicate its ability to maintain memory of system dynamics.\n\n- Computational Cost-to-Accuracy Ratio: A more explicit comparison of runtime, memory usage, and FLOPs per training iteration would provide a clearer picture of the model’s scalability.\n\nAnd maybe this one: \n\n- (Max) Lyapunov Exponent Deviation: Evaluating how well the predicted trajectories preserve the Lyapunov exponents (or max LE) of the true system would provide a more direct measure of long-term dynamical consistency.", "theoretical_claims": "The paper provides a well-motivated theoretical foundation using the von Neumann mean ergodic theorem, but key aspects require further clarification. The claim that the operator $G$ is unitary and preserves norms in $L^2$ space is critical, and verifying whether these properties hold under the proposed modifications is essential for the method's validity. The regularization term in the loss function aims to enforce unitarity, but assessing whether it sufficiently guarantees ergodicity in long-term predictions is crucial. Also, while the authors argue that their approach captures invariant statistical behaviors, a more rigorous justification -- particularly for high-dimensional chaotic systems -- \n would strengthen confidence in their claims regarding stability and accuracy.", "experimental_designs_or_analyses": "The experimental design is robust, utilizing diverse datasets that represent various chaotic systems and comparing against multiple state-of-the-art baselines. However, a more detailed computational complexity analysis comparing memory consumption and training time across baselines would strengthen the overall evaluation of the model's efficiency.", "relation_to_broader_scientific_literature": "The paper is well-situated within the literature on chaotic system prediction, referencing key works on operator-based learning, transformers for PDEs, and ergodicity-based learning approaches. It builds upon advancements in Fourier Neural Operators (FNOs) and Markov Neural Operators (MNOs) while introducing a novel transformer-based framework specifically designed to enhance prediction accuracy while preserving ergodicity, a crucial property highlighted by Eckmann & Ruelle (1985) and Young (2002). Moreover, the introduction of a new chaotic system benchmark dataset enriches the field by providing a standardized resource for evaluating machine learning methods in chaotic dynamics. These contributions collectively advance the intersection of machine learning and dynamical systems theory, offering both methodological improvements and practical benchmarks for future research.", "essential_references_not_discussed": "I also recommend considering the following papers on learning chaotic dynamics with RNNs and combining deep learning techniques with Koopman operator theory:\n1. https://proceedings.mlr.press/v202/hess23a/hess23a.pdf\n2. https://arxiv.org/pdf/2410.23467\n\nAlso, the paper does not sufficiently discuss PINNs and other hybrid methods that integrate physics-based constraints into ML models. Works such as:\n1. Karniadakis et al., \"Physics-Informed Neural Networks for PDEs\" (Journal of Computational Physics, 2021)\n2. Raissi et al., \"Physics-Informed Machine Learning for Dynamical Systems\" (PNAS, 2019).", "other_strengths_and_weaknesses": "**Strengths**\n\nOriginality: The introduction of A3M Attention and ergodicity-based loss function presents a novel approach to chaotic system modeling.\n\nSignificance: The paper tackles an important problem in ML-driven dynamical system forecasting with well-motivated solutions.\n\n**Weaknesses**\n- The paper mainly focuses on comparing transformer-based and operator-based models but does not discuss how physics-informed neural networks (PINNs) or hybrid methods (e.g., physics-constrained transformers) might perform on chaotic systems.\n- Choosing the kernel bandwidth ($\\sigma$) in RFF positional encoding may significantly affect performance. It would be helpful to have a more in-depth discussion on strategies for tuning hyperparameters.\n- Even though the paper mentions improved efficiency, a deeper look into complexity - like memory usage and training time comparisons - \n would provide more clarity.", "other_comments_or_suggestions": "It would be helpful to: \n- Elaborate on the strategies used for hyperparameter tuning of random Fourier features.\n- Provide a comprehensive comparison of runtime and memory usage with baseline methods.\n\nI am happy to increase my score if the authors can address my main concerns.", "questions_for_authors": "1. How does the von Neumann ergodic loss stack up against traditional loss functions when it comes to stability and long-term statistical preservation? Also, how does the ergodicity-based loss function compare to traditional distribution-matching methods in terms of computational efficiency?\n2. How does the A3M attention mechanism perform compared to other modifications of attention (e.g., Linformer, Performer, or Reformer) in chaotic dynamics?\n3. Regarding the computational efficiency of your approach, how does training time scale with increasing grid resolution? \n4. Given that $\\sigma$ in RFF positional encoding may significantly affect performance, how sensitive is the model to variations in $\\sigma$ across different datasets? How sensitive is the model to hyperparameter tuning, particularly for the kernel bandwidth σ\\sigma in the positional encoding?\n5. What if we swapped out the A3M pooling approach for a different aggregation method? Would using just max-pooling or mean-pooling give us competitive results?\n6. Can you provide scalability benchmarks comparing the runtime and memory usage of your method against baselines?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0OVjbHPY2y", "reviewer": "Reviewer_FGZz", "summary": "The paper investigates the problem of predicting the evolution of ergodic chaotic systems with transformers. To that end, the paper introduces a set of modifications to the traditional transformer architecture that overcome crucial bottlenecks in terms of scalability. Moreover, the paper introduces a novel regularization term that is grounded in physical perspective and helps the model preserve long-term statistics. The paper evaluates its proposed method on two turbulent fluid dynamics systems (Kolmogorov Flow and Channel Flow) and compares the performance to a multitude of baselines, showing superior performance across all metrics. Finally, the paper releases its data as a novel chaotic system benchmark.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n- **Ablation on distance-based Gaussian kernel**\n\nWe 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.\n\nRegarding 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.\n\n  | Method           | Runtime (per batch, size=3) |\n  | ---------------- | --------------------------------|\n  | Base               | 0.0261 ± 0.00016                 |\n  | + A3M Att.      | 0.0244 ± 0.0008                   |\n  | + Unitary Op.  | 0.0995 ± 0.0077                   |\n\n- **Suggested References**\n\nWe appreciate the reviewer’s recommendations. After careful consideration, we’ve integrated the relevant references into the introduction section of our revised manuscript.\n\n- **Benchmark Dataset**\n\nWe 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.\n\n- **Typos**\n\nWe 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.\n\nWe 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](https://anonymous.4open.science/r/ChaosMeetsAttention/README.md), which may be of interest to the reviewer.", "claims_and_evidence": "The paper is exceptionally well written. All claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The paper evaluates its proposed method on two turbulent fluid dynamics, Kolmogorov Flow and Channel Flow, both commonly used to learn states from chaotic systems. Moreover, the paper considers a variety of competitive baselines, both operator- and attention-based. As a result, the experimental evaluation is convincing and shows strong results.", "theoretical_claims": "I am unfamiliar with the mathematical machinery required to evaluate the correctness of the theoretical claims.", "experimental_designs_or_analyses": "The paper presents the first evaluation of transformer-based methods on (ergodic) chaotic systems and justifies this choice with strong empirical results. The paper proposes a set of modifications to the standard transformer architecture: (i) axial mean-max-min (A3M0 attention, (ii) 2D positional encodings based Random Fourier Features (RFF), and (iii) a loss regularization term based on a unitary constraint. Most of these design choices are either empirically validated (for i) or theoretically motivated (for iii).\n\nHowever, it would be nice to also show an ablation for the distance-based Gaussian kernel, i.e., A3M without RFF (Table 3 only ablates the full A3M). Moreover, the paper showing the running times in Table 3 would further strengthen the paper's claims.", "relation_to_broader_scientific_literature": "The paper adequately discussed the related work.", "essential_references_not_discussed": "While not directly related in terms of the type of chaotic system evaluated, the paper would benefit from citing Gilping (2021, 2023), who evaluates transformers on more than 130 chaotic dynamical systems, and Zhang et al. (2025), who evaluate transformers on elementary cellular automata.\n\n\n**References**\n\nChaos as an interpretable benchmark for forecasting and data-driven modelling.\nWilliam Gilpin.\nNeurIPS 2021.\n\nModel scale versus domain knowledge in statistical forecasting of chaotic systems.\nWilliam Gilpin.\narXiv:2303.08011\n\nIntelligence at the Edge of Chaos.\nShiyang Zhang, Aakash Patel, Syed Rizvi, Nianchen Liu, Sizhuang He, Amin Karbasi, Emanuele Zappala, David van Dijk.\nICLR 2025", "other_strengths_and_weaknesses": "The benchmark dataset will perhaps be the paper's biggest contribution in the long run and should, therefore, be featured more prominently (i.e., in the main paper).", "other_comments_or_suggestions": "List of typos and mistakes:\n* L177 (left): \"an attention mechanism to identif**ies**\" -- replace with \"identify\"\n* L178 (left): \"based on **random Fourier using** Random Fourier features\" — replace duplicates\n* L445 (left) \"none **of** which\"", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NBBXNt92TK", "reviewer": "Reviewer_ARTA", "summary": "The paper introduces a transformer-based model for predicting long-term trajectories in high-dimensional chaotic systems. It modifies standard attention mechanisms using Axial Mean-Max-Min (A3M) attention with random Fourier features to capture spatial correlations. It uses a unitary-constrained loss (based on the Von Neumann ergodic theorem) to preserve long-term statistical properties. The approach is scaled efficiently with tensor factorization and shows good performance on chaotic systems.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n[Essential Link for results and references to this Rebuttal](https://anonymous.4open.science/r/ChaosMeetsAttention/README.md)\n\n[Backup link](https://filebin.net/37p4dxup0t320143)\n\n* Concern 1 & Question 2 (also mentioned by reviewer vgH9):\n\n`However, the paper only enforces this unitarity approximately via a soft loss term (using Hutchinson’s trace), rather than proving exact unitarity.`\n\nYes, we use Hutchinson’s trace estimation to approximate unitarity. This approach is motivated by two key reasons:  \n\n1. Directly applying the unitary loss involves eigen decomposition and computing the Frobenius norm, which is computationally infeasible for high-dimensional latent states.  \n2. Hutchinson’s trace estimation is flexible to implement as a randomized projection and is optimized adaptively during training.  \n\nEmpirically, 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.  \n\n\n* Concern 2 & Question 3:\n\n  `References and computational trade-offs to other linear-time attention methods like Performer.`\n\n  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.\n\n  | Models                 | Parameter count | Runtime    | Memory usage | FLOPs per forward pass |\n  | ---------------------- | --------------- | ---------- | ------------ | ---------------------- |\n  | MNO                    | 6467425         | 31ms       | 377MB        | 3.45GB                 |\n  | UNO                    | 17438305        | 12ms       | 769MB        | 6.88GB                 |\n  | MWT                    | 5089153         | 50ms       | 313MB        | 9.52GB                 |\n  | FactFormer             | 6083009         | 53ms       | 6889MB       | 239GB                  |\n  | Ours                   | 7325665         | 58ms       | 7684MB       | 268GB                  |\n  | **Unitary+Performer** | **7717931**    | **111ms** | **2938MB**  | **271GB**             |\n\n* Concern 3 & Question 1:\n\n`Sensitive to the kernel bandwidth? How to tune it on different datasets.`\n\nWe 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.\n\nThe 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.\n\n* Concern 4: \n\n`Expanding the discussion on applicability to non-ergodic systems would be valuable.`\n\nWe 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.  \n\nFinally, 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.", "claims_and_evidence": "The paper’s claims are supported by experiments on Kolmogorov flow and turbulent channel flow. Performance improvements, scalability via factorized attention, and ergodicity preservation through a unitary loss are supported by quantitative results. The lack of generalization to non-ergodic systems is acknowledged.", "methods_and_evaluation_criteria": "Yes. The methods (modified transformer architecture with A3M attention and a unitary loss to preserve ergodicity) make sense for modeling chaotic systems. The evaluation criteria and datasets (Kolmogorov flow and turbulent channel flow) are appropriate for capturing short-term accuracy and long-term statistical behavior.", "theoretical_claims": "I read the derivation using the Von Neumann ergodic theorem, which the paper appears to apply correctly to argue that the operator G should be unitary in L2. However, the paper only enforces this unitarity approximately via a soft loss term (using Hutchinson’s trace), rather than proving exact unitarity.", "experimental_designs_or_analyses": "The experimental design seems sound. The paper evaluates its model on two datasets (Kolmogorov flow and turbulent channel flow), using several metrics and three random seeds (though it would be nice to see the standard deviation in addition to the mean). Ablation studies validate the usefulness of key components. More benchmarks would strengthen the paper.", "relation_to_broader_scientific_literature": "The paper extends transformer models to chaotic systems by combinding neural operator learning (e.g., Fourier Neural Operators) with ergodic theory (Von Neumann’s theorem). It builds on prior work by addressing scalability in operator-based methods and introduces efficient tensor factorization with the A3M attention mechanism.", "essential_references_not_discussed": "One work that might be worth mentioning for context is the Performer (Choromanski et al., 2020), which appears to use random features to achieve linear-time attention.", "other_strengths_and_weaknesses": "Strengths:\n- Combines Transformer architectures with ergodic theory, using A3M attention and a unitary loss to address long-term chaotic dynamics.  \n- Demonstrates improvements on benchmarks (Kolmogorov flow, turbulent channel flow) and introduces a new turbulent channel flow dataset.  \n- Includes evaluations and ablation studies that support the key contributions.\n\nWeaknesses:\n- The method does not generalize to non-ergodic settings.  \n- Unitarity is only enforced approximately with soft regularization rather than exactly.  \n- Performance seems sensitive to the kernel bandwidth.", "other_comments_or_suggestions": "The paper is generally well-written, but a few suggestions could improve it further:\n- Consider including more details on hyperparameter tuning and reporting variability (e.g., standard deviations).\n- Expanding the discussion on applicability to non-ergodic systems would be valuable.\n- A clearer explanation of scalability limits and potential computational trade-offs might help.", "questions_for_authors": "1. Could you elaborate on how to tune the kernel bandwidth in the RFF positional encoding and how its optimal value varries across different datasets?\n2. Can you provide more insight into the impact of enforcing unitarity only approximately via soft regularization?\n3. Could you clarify the computational trade-offs of your tensor factorization approach compared to other linear-time attention methods (e.g., Performer)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u9O0Nx56ab", "reviewer": "Reviewer_BWmX", "summary": "The paper addresses the challenging task of accurately forecasting high-dimensional chaotic systems using a transformer-based approach. By leveraging ergodicity and modifying attention mechanisms, the proposed framework effectively handles high-dimensional chaotic dynamics while preserving long-term statistical properties. The model introduces “A3M” attention blocks to capture extreme values and uses a novel loss function inspired by the Von Neumann ergodic theorem to maintain long-term consistency by enforcing unitarity of the dynamics forward map. Experiments on turbulent flows demonstrate superior performance compared to existing methods both in terms of short-term predictions as well as measures of long-term statistics.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We genuinely appreciate the reviewer’s time and effort in enhancing our paper. \n\n[Link for results and references](https://anonymous.4open.science/r/ChaosMeetsAttention/README.md)\n\n[Backup link](https://filebin.net/37p4dxup0t320143)\n\n* Concern 1 & Question 5:\n\n`The fair comparison in terms of model expressivity (e.g. by reporting parameter count).`\n\nWe 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..\n\n* Concern 2 & Question 4:\n\n`Lacks evaluation on small chaotic dynamical systems based on ODEs.`\n\nEmbedding 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.\n\n`Establish for very long autoregressive rollouts, e.g. by comparing Lyapunov spectra, on simpler cases.`\n\nAs 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.\n\n* Concern 3:\n\n`Clarification of Lorenz 63 dimension`\n\n‘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.\n\n`typos` \n\nThanks for bringing the typos to our attention. We have corrected them in the revised manuscript.\n\n`Reference` \n\nWe appreciate the reviewer’s suggestion of the reference, which includes efficient reconstruction approaches through dimension reduction. We have incorporated it into the revised manuscript.\n\n`results consistency & error bar for table 1&2` \n\nThe updated tables 1 and 2 are included in the provided link.\nLong-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.\n\n* Question 1:\n\n`How is $G$ implemented? How expensive is the conjugation of the unitary operator?`\n\nIn 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.\n\n* Question 2:\n\n`Autocorrelation pooled?`\n\nYes, 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.\n\n* Question 3:\n\n`Fig. 3: How long are “1000 steps”?`\n\nRoll out 1000-step counts for 20 seconds for Figure 3. We have provided the dataset details in Appendix G.\n\n`What is the Lyapunov time of the dataset at hand (if such a quantity is easily accessible for this type of PDE data)?` \n\nThe 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.\n\nPlease let us know if we have addressed the concerns and increased your confidence in our work.", "claims_and_evidence": "Most of the claims are backed by experimental evidence, however, the manuscript lacks some clear information on consistency of the results as well as how fair the comparison is in terms of model expressivity (e.g. by reporting parameter count).", "methods_and_evaluation_criteria": "The proposed method is only benchmarked on PDE-type data, but lacks evaluation on simpler baselines, such as small chaotic dynamical systems based on ODEs (e.g. Lorenz63 or Rössler system). In these cases, established measures for long-term statistics for *very long* autoregressive rollouts exist, e.g. by comparing Lyapunov spectra. It would be interesting how the model fares in these simpler settings.", "theoretical_claims": "I did not check the correctness of any proofs or theoretical claims.", "experimental_designs_or_analyses": "The authors perform experiments for section 4 for three different random seeds and report the mean in e.g. Table 1 and 2. However, to get a measure of consistency of the results, the manuscript should also report some form of error bars via either standard deviation or standard error of the mean. The same goes for Autocorrelation plots in Fig. 2.", "relation_to_broader_scientific_literature": "The paper presents a novel transformer-based approach specifically designed to improve prediction of high-dimensional and complex chaotic systems with a focus on capturing long-term statistics. The work extends existing approaches by enhancing the traditional transformer architecture with axial mean-max-min (A3M) attention blocks and a unitary operator constraint framework. These innovations enable the model to effectively capture both local extreme values and statistical invariants, thereby improving prediction accuracy over both short-term and long-term horizons. The method is novel in that it introduces a new attention block tailored to chaotic systems prediction as well as combining existing approaches in the operator learning field to enable new state-of-the-art forecasting capabilities.", "essential_references_not_discussed": "The paper misses literature for a specific class of methods to mitigate chaos-induced problems when modeling chaotic DS. Specifically, “Generalized Teacher Forcing” as introduced in Hess et al. (ICML, 2023) does not need to explicitly match distributions of invariant measures to successfully model chaotic dynamics and can be applied to various sequence models, including operator based methods.", "other_strengths_and_weaknesses": "**Strengths**: The paper clearly motivates each and every compartment of the architecture, combines and advances existing methods in the field and validates everything through comparison methods and ablation studies.", "other_comments_or_suggestions": "- “These approaches have been applied to classic 1-D examples, such as the Lorenz 63, Lorenz 96, and Kuramoto-Sivashinsky equations.” (p. 1): This may cause confusion as e.g. the Lorenz63 system is a 3D system (three dynamical variables). The manuscript should clearly define what is meant with 1D in this context.\n\n**Typos and similar**:\n- “to identifies” (p. 4, l. 177)\n- “ergodic measure-preserving transformation (MPT)” (p. 3. ll 161-162): MPT has already been previously defined.\n- “The number of sample k” (p. 6, l. 289)\n- Kolmogorov is written as “Kolmogrov” on several occasions in the text (e.g. p. 6 l. 321, heading of section 4.1)", "questions_for_authors": "1. How exactly is $\\mathcal{G}$ implemented? Which form does it take in practice? How expensive is the inversion/complex conjugation of the unitary operator $\\mathcal{G}$?\n2. how is the Autocorrelation (Fig. 2) computed? Is it pooled over the entire grid? \n3. Fig. 3: How long are “1000 steps”? What is the Lyapunov time of the dataset at hand (if such a quantity is easily accessible for this type of PDE data)?\n4. Can the architecture be benchmarked on simpler, low-d dynamical systems in forms of ODEs? This should facilitate longer rollouts ($>> 1000$ steps) to check whether the model also captures long-term statistics in these simple cases.\n5. Did the authors account for equal expressivity in terms of number of parameters for comparison in e.g. Tables 1 and 2? Can those numbers be reported somewhere?\n\n*I am happy to increase my score if the authors address my concerns*.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This submission received four expert reviews, with all reviewers acknowledging the authors' rebuttal. The panel agreed that this is a high-quality paper tackling an important ML problem with novel ideas, excellent presentation, very systematic methodology and exemplary testing/validation of every component of the proposed methodology. The reviewers also praised the originality of the proposed modification to the transformer architecture, and the used benchmark dataset which needs to receive more attention in the main part of the paper as it is likely to attract more attention in the future.

The author-reviewer discussion led to a very constructive exchange with the authors presenting high-quality additional material addressing many of the reviewers initial concerns (or requests for clarifications). I strongly encourage the authors to revise the paper along these lines, i.e., to include the new insights/results that are present in the author-reviewer discussion. This will substantial benefit the paper and strengthen its contribution and visibility in the ML community.
Decision: Accept (poster)