Training neural operators to preserve invariant measures of chaotic attractors
Authors: Ruoxi Jiang
Peter Y. Lu
Elena Orlova
Rebecca Willett
Venue: NeurIPS 2023 poster (2023)
Keywords: Neural operators
contrastive learning
optimal transport
chaotic attractors
invariant measures

Abstract: Chaotic systems make long-horizon forecasts difficult because small perturbations in initial conditions cause trajectories to diverge at an exponential rate. In this setting, neural operators trained to minimize squared error losses, while capable of accurate short-term forecasts, often fail to reproduce statistical or structural properties of the dynamics over longer time horizons and can yield degenerate results. In this paper, we propose an alternative framework designed to preserve invariant measures of chaotic attractors that characterize the time-invariant statistical properties of the dynamics. Specifically, in the multi-environment setting (where each sample trajectory is governed by slightly different dynamics),  we consider two novel approaches to training with noisy data. First, we propose a loss based on the optimal transport distance between the observed dynamics and the neural operator outputs. This approach requires expert knowledge of the underlying physics to determine what statistical features should be included in the optimal transport loss. Second, we show that a  contrastive learning framework, which does not require any specialized prior knowledge, can preserve statistical properties of the dynamics nearly as well as the optimal transport approach. On a variety of chaotic systems, our method is shown empirically to preserve invariant measures of chaotic attractors.

Training neural operators to preserve invariant measures of chaotic attractors
Ruoxi Jiang"
Peter Y. Lu* Department of Physics University of Chicago Chicago, II. 60637 lup@uchicago.
Elena Orlova
Department of Computer Science
Department of Computer Science
University of Chicago Chicago, II. 60637
University of Chicago
Chicago, II. 60637 eorl eorlova@uchicago.com
roxie62@uchicago.e edu
Rebecca Willett
Department of Statistics and Computer Science University of Chicago Chicago, IL 60637
Caglitt@uchicaoo edu
Abstract
Chaotic systems make long-horizon forecasts difficult because small perturbations in initial conditions cause trajectories to diverge at an exponential rate. In this setting, neural operators trained minimize squared error losses, while capable of accurate short-term forecasts, often fail to reproduce statistical or structural properties the dynamics over longer time horizons and can yield degenerate results. In this paper, propose alternative framework designed to preserve invariant measures chaotic attractors that characterize the time-invarian statis- tical properties of the dynamics. Specifically, in the mulli-enniiooment setting (where each sample trajectory is governed by slightly different dynamics), we consider two novel approaches to training with noisy data. First, we propose loss based on the optimal transport distance between the observed dynamics and the neural operator outputs. This approach requires expert knowledge of the underlying physics to determine what statistical features should be included in the optimal transport loss. Second, show that contrastive learning framework, which does not require any specialized prior knowledge, can preserve statistical properties of the dynamies nearly as well the optimal transport approach. On variety of chaotic systems, our method is shown empirically to preserve invariant measures of chaotic attractors.
Introduction
Training fast and accurate surrogate models emulate complex dynamical systems is key many scientific applications of machine learning, including climate modeling [1], fluid dynamics [2, 3]. plasma physics [4], and molecular dynamics [5, 6). Fast emulators powerful tools that can be used for forecasting and data assimilation [1, 4, 7]. sampling to quantify uncertainty and compute statistical properties [1, 5, 6]. identifying latent dynamical parameters [8, 9], and solving a wide range of inverse problems [10]. Specifically, neural operator architectures [2, 11, 12] have been shown to be promising physics-informed surrogate models for emulating spatiotemporal dynamics, such as dynamics governed by partial differential equations (PDEs).
Equaal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

One key feature of many of the dynamical systems in these applications is chaos, which is charac- terized by a high sensitivity to initial conditions and results in a theoretical limit on the accuracy of forecasts. Chaotic dynamics ensure that no matter how similarly initialized, any two distinct trajectories will diverge at an exponential rate while remaining confined a chaotic attractor [13]. At the same time, chaos is fundamental to many critical physical processes, such as turbulence in fluid flows [14] as well as mixing and ergodicity [13]- prropertiss that underpin the fundamental assumptions of statistical mechanics [15].
Chaos not only presents a barrier to accurate forecasts but also makes it challenging to train emulators, such as neural operators, using the traditional approach of rolling-out multiple time steps and fitting the root mean squared error (RMSE) of the prediction, as demonstrated in Fig. 1. This is because RMSE training relies on encouraging the emulator to produce more and more accurate forecasts, ideally over long time horizon, which is severely limited by the chaotic dynamies. This problem is exacerbated by measurement noise, which, in combination with the high sensitivity to initial conditions, further degrades the theoretical limit on forecasting. Unfortunately, this is precisely the setting for many real- weal-world scientific and engineering applications. While accurate long-term forecasts are impossible in this setting, it is still possible to replicate the statistical properties of chaotic dynamics.
Specifically, we can define natural invariant measure on chaotic attractor that characterizes the time-invariant statistical properties of the dynamics on the attractor [13]. By training neural operator to preserve this invariant measure-ov or equivalently, preserve time-invariant statistics- we can ensure that the neural operator is properly emulating the chaotic dynamics even though it is not able to perform accurate long-term forecasts. In this paper, we introduce two new training paradigms to address this challenge. The first uses an optimal transport- bansp objective to ensure the invariant measure is preserved. While effective, this approach requires expert knowledge of invariant measures to determine appropriate training losses. The second paradigm uses a contrastive feature loss that naturally preserves invariant measures without this prior expert knowledge. These new losses, which are both intended to preserve the invariant statisties that govern the long-term behavior of the dynamics, are used in combination with the standard RMSE loss evaluated over a short time horizon, which ensures any remaining predictability in the short-term dynamics is correctly captured.
We operate in the multi-enniirnment setting, in which parameters governing the system evolution may be different for each train and test sample. This setting is more challenging than the more typical single-enviromment setting because requires emulators to generalize over broader range domain. Most practical use cases for emulators-whert the computational costs of generating training data and training an emulator are outweighed by the computational gains at deployment-are in the multiennvirnment setting.
1.1 Contributions
This paper makes the following key contributions. Firstly, we identify, frame, and empirically illustrate the problem of training standard neural operators on chaotic dynamics using only RMSE- namely, that the sensitivity to initial conditions in combination with noise means that any predictive model will quickly and exponentially diverge from the true trajectory in terms of RMSE. As such, RMSE is poor signal for training a neural operator. We instead suggest training neural operators to preserve the invariant measures of chaotic attractors and the corresponding which are robust to the combination of noise and chaos. Secondly, we propose a direct optimal transport-basee approach to train neural operators to preserve the distribution of a chosen set of summary statistics. Specifically, we use a Sinkhomn divergence loss based on the Wasserstein distance to match the distribution of the summary statistics between the model predictions and the data. Next, we propose general-purpose contrastive learning approach to learn a set of invariant statistics directly from the data without expert knowledge. Then, we construct a loss function to train neural operators to preserve these learned invariant statistics. Finally, we empirically test both of these approaches and show that the trained neural operators capture the true invariant statistics and therefore the underlying invariant measures of the chaotic attractors-much more accurately than baseline neural operators trained using only RMSE, resulting in more stable and physically relevant long-term predictions.
statistics,
2

---

1.2 Related work
Neural operators. The goal of neural operator, proposed in the context of dynamical modeling, is to approximate the semigroup relationship between the input and output function space [2, 11, 12, 16- 191. is particularly effective in handling complex systems governed by partial differential equations (PDEs), where the neural operator is designed to operate on an entire function or signal. The architecture designs vary significantly in recent works, including the Fourier neural operator (FNO) which uses Fourier space convolution [2], the deep operator network (DeepONet) Lu et al. [11], which consists of two subnetworks modeling the input sensors and output locations, and other modern designs like transformers [12] and graph neural networks [19].
Multi-environment learning. Recent developments in multi-environment learning often involve adjusting the backbone of the neural operator. For instance, Dyad [9] employs an encoder to extract umme-invarian hidden features from dynamical systems, utilizing& a supervisory regression loss relative to the context values. These hidden features are then supplied as additional inputs to the neural operator for decoding in new environment, In contrast, Context-informed Dynamics Adaptation (CoDA, Kirchmeyer et al. (20]) decodes the environment context via a hypernetwork, integrating the parameters of the hypernetwork as auxiliary input to the neural operator. Like our approach, Dyad learns Lime-invarian features from dynamical systems. However, their methods require a time series of measurements instead of just initial conditions. In addition, both CoDA and Dyad alter the backbone of the neural operator, making it challenging to adapt to new neural operator architectures. In contrast, our method uses the learned fime-invariant features to determine loss function that can be applied to any neural operator architecture without changes to either the backbone or the inputs.
Long-term predictions. Lu al. [21] try recover long-term statistics by learning a discrete-time stochastic reduced-order system. From the perspective of training objectives, Li et al. [16] propose using Sobolev norm and dissipative regularization facilitate learning of long-term invariant measures. The Sobolev norm, which depends on higher-order derivatives of both the true process and the output of the neural operator, can capture high-frequency information and is superior to using only RMSE. However, its effectiveness diminishes in noisy and chaotic scenarios where it struggles to capture the correct statisties. The dissipative loss term, regularizes the movement of the neural operator with respect to the input, and attempts to make the emulator stay on the attractor. However, in the case of cunliieenviromment setup for chaotic systems, our goal to learn a neural operator capable of modeling long-tern term distributions with regard to different contexts. Using the dissipative loss as a regularization will cause the neural operator to be insensitive to the context and fail model different attractors.
Prior [22] and concurrent work [23] training recurrent neural networks (RNNs) have suggested that using teacher forcing methods with a squared error loss on short time sequences can produce high-quality emulators for chaotic time series. Instead, we focus on training neural operators on fully observed high-dimensional deterministic chaos and find that using only an RMSE loss generally fails to perform well on noisy chaotic dynamics. Other concurrent work [24] has suggested using dynamical invariants such as the Lyapunov spectrum and the fractal dimension to regularize training for reservoir computing emulators. We show that emulators trained using our approaches also correctly capture the Lyapunov spectrum and fractal dimension of the chaotic attractors without explicitly including these invariants, which are difficult to estimate in high dimensions, in our loss.
Wasserstein distance and optimal transport. Optimal transport theory and the Wasserstein metric have become powerful theoretical and computational tools. for a variety of applications, including generative computer vision models [25], geometric machine learning and data analysis [26], particle physics [27], as well as identify conservation laws [28] and fitting parameterized models for low- dimensional dynamical systems [29]. In particular, Yang et al. [29] uses the Wasserstein distance as an optimization objective to directly fit a parameterized model to data from a low low-dimensional attractor. They compute the Wasserstein distance by turning the nonlinear dynamics in the original state space into linear PDE in the space of distributions and then solving a consstrird optimization problem. This has some nice theoretical properties but generally scales poorly to high-dimensional dynamical systems. Similarly, concurrent work [30] suggests modeling dynamics by fitting a Fokker- Planck PDE for the probability density of the state using the Wasserstein distance. However, since this also requires estimating the full probability distribution of the state on a mesh grid, is again challenging to scale this method to high-dimensiooal systems. In contrast, our proposed optimal transport approach focuses on fitting deep learning- based neural operators to high-dimensional chaotic
3

---

Haogun wno of wow.rrtaaaaar....
-10
64
05
66
Wasomamber
(a) Error metrics vs, noise scale r. (b) Predictions from emulators.
(c) Energy spectrum.
Figure 1: The impact of noise invariant statistics vs. RMSE and Sobolev norm. (a) We show the impact of noise on various error metrics using ground truth simulations of the chaotic Luramoto-Sivashins
(KS) system with increasingly noisy initial conditions (u +1) well as with added measurement noise Uug + +n) + 7/. Here, (a(+) refers to the ground truth solution the differential equation (1) for the KS system given an initial condition, and r ~ (O r² I). where o² is the temporal variance of the trajectory (uo) and r is the noise scale. Relative RMSE and Sobolev norm [16], which focus short-term forecasts, deteriorate rapidly with noise Tj, whereas the invariant statistics have a much more gradual response to noise, indicating robustness. (b) The emulator trained with only RMSE degenerates at times into striped patches, while ours is much more statistically consistent with the ground truth. (c) Again, the emulator trained with only RMSE performs the worst in terms of capturing the expected energy spectrum over long-term prediction.
dynamical systems by first choosing set of physically relevant summary statistics and then using the computationally efficient Sinkhorn algorithm [31] to compute our optimal transport loss- an entropy-regulatee approximation for the (squared) Wasserstein distance on the distribution the summary statistics.
Feature loss and contrastive learning. Contrastive learning [32-34] charges encoder with generating features that are unaffected by transformations by promoting the similarity of features between different transformations of the same image and by maximizing the feature distance between distinct images. The recent advancements in contrastive learning mainly result from using extensive data augmentation
motivate the encoder learn semantic representations. In addition to its success in image recognition, contrastive learning has proven to be effective in modeling fine-grained structures [35] and in scientific applications [10]. Training data generators or solving inverse problems generally requires output to be coherent across multiple structural measurements, a requirement that pixel-wise MSE typically fails to meet. For this reason, Johnson al. [36] computes MSE on deep features of the classification models as a structural distance metric to train generative models. Tian al. [37] propose the encoder of contrastive learning as an unsupervised alternative for calculating feature loss.
2 Problem Formulation
Consider dynamical system with state space u u governed by
du G(u, Φ),
(1)
where is the governing function, and o E 4 is a set of parameters that specify the dynamics. We will denote trajectories governed by these dynamics (1) as u(t), particular points the trajectory as
u, U, and sequence of K consecutive time points the trajectory as U (u, }==/
I
Our experiments use equally spaced time points with a time step At. We refer to settings where Ø varies as multt-enviromment settings.
Our goal to learn an emulator a multi-environment setting that approximates the dynamics (1)
with data-driven neural operator ÿe ld x &b + L that makes discrete predictions
ú := 9o (ū, Ø),
(2)

---

where f are the parameters of the neural operator 9g. In the multi-environment setting, we will have data from variety of environments n E 21, N). and thus our training data { ( J consists of trajectories U! u:L and the corresponding environment parameters p(n).
Predicted sequence notation. We will denote sequence of autonomously predicted states (n) (using the neural operator 9e) with an initial condition ui, as
1 1.744. == (u,, _o ( g(n)), 90 o Go (m) (())), ................................
go (u (nt)
We will often use a concatenated sequence (with total length + 1) of these autonomous prediction sequences (each with length h + 1), which we will denote
(m) B.lIK
((w) 1-1+h 144411420-1 @ .... (n)
(m)
(4)
where is concatenation.
Chaotic dynamical systems and invariant measures of chaotic attractors. We focus on chaotic dynamical systems that have one more chaotic attractors, which exhibit a sensitive dependence on initial conditions characterized by a positive Lyapunov exponent [13]. For each chaotic attractor A, can construct natural invariant measure (also known a physical measure)
1
8 (t) dt, I'
= lim T-300 T
(5)
where Su(t) is the Dirac measure centered on a trajectory uA((() that is in the basin of attraction of A [13]. Note that, because A is an attractor, any trajectory (t(t) () in the basin of attraction of A will give the same invariant measure M. i.c. the dynamics are ergodic on A. Therefore, any time-invariant statistical property SA of the dynamics on A can be written as
f s(ua(()) dt
= Ep = = s(u) dua())
lim T T
(6)
for some function s(u) and a trajectory (t) in the basin of attraction of A. Conversely, for any s(u). (6) gives a time-invariant property S of the dynamics.
In this work, we assume each sampled trajectory the data from a chaotic attractor and therefore, has inve-invariant statistical properties characterized by the natural invariant measure of the attractor.
Noisy measurements. Because of the sensitivity to initial conditions. accurate long-term forecasts (in terms of RMSE) are not possible for time scales much larger than the Lyapunov time [13]. This is because any amount of noise or error in the measurement or forecast model will eventually result in exponentially diverging trajectories. The noisier the data, the more quickly this becomes a problem (Fig. 1). However, in the presence of measurement noise, invariant statistics of the noisy trajectory
( = (t) + 7
(7)
can still provide a useful prediction target. The invariant statistics will be characterized by a broadened measure (the original measure convolved with the noise distribution pr))
lim T T I'
1 lim 1 T
dur(o dt PAA Pn
Py (u)(
(8)
T-200
and are therefore given by
ŠA Eag(s] = = (u)
d(i((()
(9)
lim T + / (() dt lim T-x T
1
I" s(ua(t))
=iim
(t)
which can be a good approximation for S = A [s] even in noisy conditions and does not suffer from the exponential divergence of RMSE (Fig. l).
5

---

(3) Neural Operator
Training
Optimal Transport
t FIIKK
Bo
Sto) bIAK
Training Loss
distribution
II fiest" K+1 2 10011 1
matching
U
BI For for Sin Slae
U BE*K
sumreary
I+1I*
B fu. E oon (((2)))) B (U
(c) Unsupervised Contrastive Learning
Random
Embedding
(O0) O.L I=
(16
fy shared weights
Repution
frr
Attract
Figure Our proposed approaches for training neural operators. (a) Neural operators are emulators trained to take an initial state and output future states in recurrent fashion. To ensure the neural operator respects the statistical properties of chaotic dynamics when trained on noisy data, propose two additional loss functions for matching relevant long-term statistics. (b) We match the distribution of summary statistics, chosen based on prior knowledge, between the emulator predictions and noisy data using optimal transport loss. (c) In the absence of prior knowledge, take advantage of self-supervised contrastive learning to automatically learn relevant tume-invarian statistics, which can then be used to train neural operators.
3 Proposed Approaches
3.1 Physics-informed optimal transport
We propose optimal transport-based based loss function to match the distributions of a set of summary statistics s(u) = (x() (()) (u)) between the data s(u,)) Jls and the model predictions $y = s(ūr,) ~ Ms. Here, es and are the probability measures of the summary statistics for the true chaotic attractor A and the attractor Á learned by the neural operator, respectively. Specifically, (s') = 6x(u))>' (u) du and (S') (사(ū)- d (ū) dú. By choosing the set of summary statistics using expert domain knowledge, this approach allows us directly guide the neural operator toward preserving important physical properties of the system.
Optimal Seansportt-aassrsterrr distance. We match the distributions of summary statisties using the Wasserstein distance (Js, jaa-aa distance metric between probability measures defined in terms of the solution to optimal transport problem:
=EIl(HoHHa) / c(s, (s,
W ( (zs) = inf
(10)
We use a quadratic cost function c(s, s) ||ss - áI[2 (i.e., W is the 2- W-ssssesteein distance), and the map I[(He, (4) must be valid transport map between ba and (i.c., joint distribution for s with marginals that match 14 and (s)) [38].
In the discrete setting with distributions represented by samples S (s and $ (8;)f=2 the Wasserstein distance is given by
min TET
(11)
2
where the cost matrix Cjj f|ss 1,112. The of valid discrete transport maps II consists of all matrices 7' such that Vi,j. Tú 2 0, o, = 1, and T. 1.
Entropy-regulariized optimal transport-Sinkkne
divergence. Exactly solving the optimal transport problem associated with computing the Wasserstein distance is computationally prohibitive,

---

especially in higher dimensions. A common approximation made to speed up computation is to introduce an entropy regularization term, resulting convex relaxation of the original optimal transport problem:
(T = min T Tell
(12)
2.5
where ((T) - - Tiy log Tij is the entropy of the transport map. This entropy-regularized optimal transport problem can be solved efficiently using the Sinkhorn algorithm [31].
As 7 0, the cttropy-regularize Wasserstein distance W W W reduces to the exact Wasserstein distance (11). For ? > O, we can further correct for an entropic bias to obtain the
Sinkhorn divergence [39, 40]
(T W (' WS,
(S (S S) = W/ (S.S)² = 2
(13)
2
which gives us our optimal transport loss. Combined with relative root mean squared error (RMSE)
u. urlp
AMSE U, Ú) K+1
E
(14)
11 1|
EU,
for short-term prediction consistency 12, 16, 18], our final loss function is
(())
loT (SoT(s/n) (s(n) (n) ((m)
I.t+K (.aT4K) + PRMSE (U
(x) K.TKK + (15)
(n)
E [.....N] N) IE(O....L-KK. K}
((m) where S/i+K == {s(u) u E U1ttk $(m)
U(n)
IhIIK = {s(ū) E Ú
J. and > 0 is
hyperparameter. Note that S(n)
I:h:I+ and (())
I.RMSE:I-KK implicitly depend on weights A.
3.2 Contrastive feature learning
When there is an absence of prior knowledge pertaining the underlying dynamical system, or when the statistical attributes are not easily differentiable, we propose an alternative contrastive learning-baat baaning-based approach to learn the relevant invariant statistics directly from the data. We first use contrastive learning to train an encoder to capture invariant statistics of the dynamics in the multi-environment setting. We then leverage the feature map derived from this encoder to construct a feature loss that preserves the learned invariant statistics during neural operator training.
Contrastive learning. The objective of self-supervised learning is train encoder (U) (with parameters v to compute relevant invariant statistics of the dynamics from sequences U with fixed length K + We do not explicitly train on the environment parameters á but rather use contrastive learning approach that encourages the encoder fo to learn variety of lime-invariant features that are able distinguish between sequences from different trajectories (and therefore different Φ).
general purpose general-purpose
contrastive learning framework using the Noise Contrastive Estimation (InfoNCE) loss has been shown to preserve context-aware information by training to match sets of positive pairs while treating all other combinations negative pairs [32]. The selection of positive pairs is pivotal to the success
of contrastive learning. In our approach, the key premise is that two sequences U(n) U-1KK J
from the same trajectory U(n) both sample the same chaotic attractor, i.e. their statistics should be similar, treat any such pair of sequences positive pairs. Two sequences U(2)
i-1+k: U (m)
H-H+K
from different trajectories are treated as negative pairs. This allows us to formulate the InfoNCE loss as:
InfoNCE(1117)
1 (16)
exp
E DE( LEEL....... [I.....]
log
exp (Fr exp( (U (1 fa (L((m) (.mmmm++))//))
E myfn He (0.0...... (IK
J,JE (OJ.E........e b-K

---

The term in the numerator enforces alignment of the positive pairs which ensures we obtain time- invariant statistics, while the term the numerator encourages uniformity, i.e. maximizing mutual information between the data and the embedding, which ensures we can distinguish between negative pairs from different trajectories [41]. This provides intuition for why our learned encoder fø identifies relevant time-invariant statistics that can distinguish different chaotic attractors.
Contrastive feature loss. To construct our feature loss, we use the cosine distance between series of features of the encoder network f [42]:
fcL (U, = D (
(17)
where f gives the output the 1-th layer of the neural network. The combined loss that we use for training the neural operator is given by
Ala (U/ U (irtiiife) + CRRSS ( 1.T+K: UTbRmsE: ]. (18)
((0)
IE rE K}
where À > 0 is hyperparameter.
Experiments
We evaluate our approach on the ID chaotic Sorramtoo-Svvnshiin (KS) system and finite- dimensional Lorenz 96 system. In all cases, we ensure that the systems under investigation remain in chaotic regimes. We demonstrate the effectiveness of our approach in preserving key statistics in these unpredictable systems, showcasing our ability to handle the complex nature of chaotic systems. The code is available at: https : /tt///gitttttt www.rooieeee///oooooooooooooee wwwoooiieen/euuaal..mmam htwe//eullllloeeee www.rraoroprrattrrr_rooom soarrattrrsoo__cccc.e. cha_s
Experimental setup. Our data consists of noisy observations u(l) with noise ~ r020,) where o² the temporal variance of the trajectory and
scaling factor. Baselines. We primarily consider the baseline as training with RMSE [2]. We have additional baselines in Appendix B. including Gaussian denoising and Sobolev norm loss with dissipative regularization [16]. Backbones. We use the Fourier neural operator (FNO, [18]). Evaluation metrics. We use a variety of statistics-based evaluation metries and other measures that characterize the chaotic attractor. See Appendix C.1 for details.
4.1 Lorenz-96
As is common test model for climate systems, data assimilation, and other geophysical applications [43-45], the Lorenz-96 system is a key tool for studying chaos theory, turbulence, and nonlinear dynamical systems. It is described by the differential equation
du;
di (ui+1 w;-2);;
(19)
Its dynamics exhibit strong energy-conserving non-linearity, and for a large F 2 10, it can exhibit strong chaotic turbulence and symbolizes the inherent unpredictability of the Earth's climate.
Experimental setup. When using optimal transport loss, we assume that expert knowledge is derived from the underlying equation. For Lorenz-96, we define the relevant statistics as s(u) == du:
dt (u441 4;-2) uz-2)u;-1- u,), We generate 2000 training data points with each d(m) randomly sampled from uniform distribution with the range [10.0, 18.0). We vary the noise level r from 0.1 to 0.3 and show consistent improvement in the relevant statistics. Results. The results presented Table 1, and predictions and invariant statistics are shown in Fig. (refer to C.4 for more visualizations).
4.2 Kuramoto-Sivashinkky
Known as a model for spatiotemporal chaos, Kuramoto Sivashinsky (KS) has been widely used to describe various physical phenomena, including fluid flows in pipes, plasma physics, and dynamics of certain chemical reactions [46]. It captures wave steepening via the nonlinear term a models dispersion effects through 8'u, S and manages discontinuities by introducing hyper-viscosity via #
au
a'a.
Ou du
Φ 82u
(20)
dr
dz 0z4²
8

---

Sampled trajectory
Histogram 24
Histogram 2d
Sampled trajectory
Histogram 2d
Histogram 2d
Truch
Sh 3h,100
Phe
0T (ours)
Oh Piz
On Po
(ours)
Bu kar
du dut
da/ol
O'w/O
(a) Lorenz-V9
(b) Kuramoto Sivashinsky
Figure 3: Sampled emulator dynamics and summary statistic distributions. We evaluate our proposed approaches by comparing them baseline model that is trained solety using relative RMSE loss. We conduct this comparison on two dynamical systems: (a) Lorenz-96 and (b) Kuramoto- Sivashinsky (KS). For each system, show visual comparison of the predicted dynamics (left) and two-dimensional histograms relevant statistics (middle and right). We observe that training the neural operator with our proposed optimal transport (OT) or contrastive learning (CL) loss significantly enhances the long-term statistical properties of the emulator, as seen in the raw emulator dynamics and summary statistic distributions. The performance of the CL loss, which no prior knowledge, is comparable to that of the OT loss, which requires an explicit choice of summary statistics.
Experimental setup. For the KS system, define s(u) Ou 0t
[47]. We generate 2000 training data points, with each (fi)) being randomly selected from a uniform distribution within the range of [1.0, 2.6). Results. We report our results over 200 test instances in Table and visualize the predictions and invariant statistics in Fig. 3.
Ou
8'y
dz
5 Discussion and Limitations
We have demonstrated two approaches for training neural operators chaotic dynamics by preserving invariant measures the chaotic attractors. The optimal transport approach uses knowledge of the underlying system to construct physics-informed objective that matches relevant summary statistics, while the contrastive learning approach uses a multi-cnvironment setting to directly learn relevant invariant statistics. In both cases, emulator to reproduce the invariant statistics of chaotic attractors from noisy data when compared with traditional RMSE-only training that focuses short-term forecasting.
significant improvement in the ability of the
We also find, high-noise settings. that both our approaches give similar or lower leading LE errors than an emulator trained on RMSE loss alone, despite the fact that our method is only encouraged to match invariant statistics of the final attractor rather than dynamical quantities. evidence that the fractal dimension of the contrastive learning approach is closer to the true attractor. However, we note that the fractal dimension is difficult reliably estimate for high-dimensional chaotic attractors [48].
also see

---

Training
Histogram Error 4
Energy Spec. Error
Leading LE Error 4
FD Error 4
FRMSE
0.056 (0.051, 0.062)
0.083 (0.078, 0.090)
0.013 (0.006, 0.021)
1,566 (0.797, 2.309)
0.1 LOT ÉRMSE
0.029 (0.027. 0,032)
0.058 (0.052, 0.064)
0.050 (0.040. 0.059)
1.424 (0.646, 2.315)
fcL + FRMSE
0.033 (0.029, 0.037)
0.058 (0.049, 0.065)
0.065 (0.058, 0.073)
1.042 (0.522, 1.685)
FRMSE
0.130 (0.1 0.142)
081 (0.172, 188)
0.170 (0.166, 0. 1919
2.481 (1,428, 3,807)
0.2 foT ÉRMSE
0.039 (0.035, 0,042)
0.086 (0.079, 0.095)
0.016 (0.006. 0.030)
2,403 (1.433, 3,768)
fct. + FRMSE
0.073 (0.066, 0,080)
0.131 (0.117, 0.149)
0.012 (0.006, 0.018)
1,681 (0.656, 2.682)
FRMSE
0.215 (0.204, 0.234)
0.291 (0.280, 0.305)
0,440 (0.425, 0.463)
3.580 (2.333, 4.866)
0.3 fOT FRMSE fCr. + ÉRMSE
0.057 (0.052, 0.064)
0.123 (0.1 16, 1.1
0.084 (0.062, 0. 134)
3.453 (2.457, 4,782)
0.132 (0.111, 0.151)
0.241 (0.208, 0.285)
0.064 (0.045, 0.091)
1.894 (0.942, 3.108)
Table 1: Emulator performance on Lorenz-96 data with varying noise scale 0.1, 0.3. The median (25th, 75th percentile) of the evaluation metrics (Appendix C.1) are computed on 200 Lorenz-96 test instances (each with 1500 time steps) for the neural operator trained with (1) only RMSE loss RMSE: (2) optimal transport (OT) and RMSE loss for + ÉRSE (using prior knowledge to choose summary statistics); and (3) contrastive learning (CL) and RMSE loss fcI. + RRMSE (without prior knowledge). We show significant improvements on the long-term statistical metrics including L histogram error of the chosen statistics S(u) == [du (u++1 - U₁ 24-2));;11 u:): relative error of Fourier energy spectrum; and absolute error of estimated fractal dimension (FD). For high noise, OT and CL training also improve the leading Lyapunov exponent (LE) of the neural operator.
Training
Histogram Error |.
Energy Spec. Error 4
Leading LE Error 1
LRMSR
0,390 (0.325, 0.556)
0.290 (0.225, 0.402)
0.101 (0.069, 0.122)
for + ERMSE
0.172 (0. 146, 0.197)
0.211 (0. 188, 0,250)
0.094 (0.041, 0.127)
fCL ÉRMEE
0.193 148, 0.247)
0.176 130, 0.245)
0.108 (0.068, 0.132)
data with noise scale 0.3. The median (25th, 75th percentile) of the evaluation metrics (Appendix C.1) are computed on 200 Luramoto-Sivashinkt test instances (each with 1000 time steps) for the neural operator trained with (1) only RMSE loss /RMSE; (2) optimal transport (OT) and RMSE loss foT + ÉRMSE (using prior knowledge to choose summary statistics); and (3) contrastive learning (CL. and RMSE loss fCL + &RME (without prior knowledge). We again show significant improvements in the long-term statistical metrics including L1 histogram error of the chosen statisties S(u) { &u Ou 82u relative error of Fourier energy spectrum. The fractal dimension (FD) is highly unstable in high dimensions [48] and could not be estimated for this dataset.
Table 2: Emulator performance
dt dx da. J and
Limitations. Because we rely on invariant measures, our current approach is limited to trajectory data from attractors, i.c we assume that the dynamics have reached an attractor and are not in transient phase. We also cannot handle explicit time dependence, including dimendepeneent forcing or control parameters. For the optimal transport approach, choosing informative summary statistics based on prior knowledge is key to good performance (Appendix B.3). For the contrastive learning approach, the quality of the learned invariant statisties also depends on the diversity of the environments present in the muntiiennvirommett setting, although our additional experiments show that we can still obtain good performance even with minimal environment diversity (Appendix B.4).
Future work. In the future, we may be able to adapt our approaches to allow for mild time dependence by restricting the time range over which we compute statistics and select positive pairs. This would allow us to study slowly varying dynamics as well as sharp discrete transitions such as tipping points. We can also improve the diversity of the data for contrastive learning by designing new augmentations or using the training trajectory of the neural operator to generate more diverse negative pairs. We will investigate generalizing our approaches to other difficult systems, such as stochastic differential equations or stochastic PDEs, and we would like to further study the trade-offs and synergies between focusing on short-term forecasting (RMSE) and capturing long-term behavior (invariant statistics). In addition, we would like to investigate and compare training methods [22-24] across different architectures.
Broader impacts. While better emulators for chaotic dynamics may be used in wide range of applications, we foresee no direct negative societal impacts.
10

---

Acknowledgments and Disclosure of Funding
The authors gratefully acknowledge the support of DOE grant DE-SC0022232, AFOSR grant FA9550- 18-1-0166, and NSF grants DMS- DMS-2023109 and DMS-1925100 In addition, Peter Y. Lu gratefully acknowledges the support of the Eric and Wendy Schmidt in Science Postdoctoral Fellowship, Schmidt Futures program.
References
[1] Jaidcep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadeneshelii Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. Fourcastnet: A global data- driven high-resolution weather model using adaptive fourier neural operators, 2022.
[2] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference Learning Representations, 2021.
URL httpa: hwtp//eiepeeeeeeeee/eee/eeeeeeee wttpriiiwwne//fouurtteee htte//oruumii=ccBP9QQY55.
[3] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 1596-1611, Curran Associates, Inc., 2022.
[4] A. Mathews, M. Francisquez, J. W. Hughes, D. R. Hatch, B. Zhu, and B. Rogers. Uncovering turbulent plasma dynamics via deep learning from partial observations. Phys. Rev. 104: 025205, Aug 2021. doi: 10.1 10..110/hs/seeE... 104.02520... URL bttps://link. www.appaapppp.... org/doi/ 10. 1110/////aReer..
025205.
[5] Oliver T. Unke, Stefan Chmiela, Huziel E. Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T. Schütt,, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine learning force fields. Chemical Reviews, 121(16):10142-1111 Aug 2021. ISSN 0009-2665. doi: w...0//aa..chemre..oo
[6] Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, Feb 2023. ISSN 2041-1723. doi:
10.1038841467-000-66629..
[7] Yuming Chen, Daniel Sanz-Alonso, and Rebecca Willett. Autodifferentiable ensemble kalman filters. SIAM Journal Mathematics of Data Science. 4(2):801-883 2022.
[8] Peter Y. Lu, Samuel Kim, and Marin Soljačié. Extracting interpretable physical parameters from spatiotemporal systems using unsupervised learning. Phys. Rev. X, 10:03 1056, Sep 2020. doi: 10. 10.1103/Physeexx..00 10.031056 URL https:///iik aps. app..org//oi/10// 1103/Physeee....... 10. 031056.
[9] Rui Wang, Robin Walters, and Rose Yu. Meta-learning dynamics forecasting using task inference. Advances in Neural Information Processing Systems, 35:21640-21653, 2022.
[10] Ruoxi Jiang and Rebecca Willett. Embed and emulate: Learning to estimate parameters of dynamical systems with uncertainty quantification. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022. URL httpps://openreviee
[11) Lu Lu, Pengzhan Jin, Guofei Pang. Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218-23) 229, Mar 2021. ISSN 2522-5839. doi: 10. 1038/ s42256-021 021-0030--5
[12] Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning, arXiv preprint arXiv:2205.13677 2022.
11

---

[13] Alfredo Medio and Marji Lines. Nonlinear Dynamics: A Primer. Cambridge University Press, Cambridge, 2001. doi: 10.1017/CB00978511
[14] Peter Davidson. Turbulence: An Introduction for Scientists and Engineers. Oxford University 0000020008000000001...
Press, 06 2015. ISBN 9780198722588. doi: 10.
[15] R. Dorfman. An Introduction to Chaos in Nonequilibrium Statistical Mechanics. Cam- bridge Lecture Notes in Physics. Cambridge University Press, 1999. doi: 10.1017/
1809780511628877
[16] Zongyi Li, Miguel Schiafiiii, Nikola Borislavov Kovachki, Kamyar Arizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Learning chaotic dynamics in dissipative systems. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL attps://openrevier met//forummidd1136:F1111
[17] Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data aug- mentation for neural pde solvers. In International Conference on Machine Learning, pages 2241-2256. PMLR, 2022.
[18] Jayesh K Gupta and Johannes Brandstetter. Towards multi-spattooeemporal-sce generalized pde modeling. arXiv preprint arXiv:2209. 2022.
[19] Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations, 2022. URL httpa:// openreview... opearevieer.et//ooruniiieieeeeee
[20] Matthicu Kirchmeyer, Yuan Yin, Jérémie Donå, Nicolas Baskiotis, Alain Rakotomamoniy. and Patrick Gallinari. Generalizing to new physical systems via context-informed dynamics model. In International Conference on Machine Learning, pages 1283-11301 PMLR, 2022.
[21] Fei Lu, Kevin K Lin. and Alexandre J Chorin. basa-based stochastic model reduction for the kuramoto-
kuramoto-sivashinsky sivashinsky equation. Physica D: Nonlinear Phenomena, 340:46-57, 2017.
[22] Jonas M Mikhaeil, Zahra Monfared, and Daniel Durstewitz. On the difficulty of learning chaotic dynamics with RNNs. In S Koyejo, S Mohamed, A Agarwal, D Belgrave, K Cho, and A Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 11297-11312. Curran Associates, Inc., 2022. URL https: /proceedinoo/cssss .aagggs.neure neurips cc/paper_fitteeeee saoer_ffles/pppr//// 2022/ 2022/file/49000
Paper- Conference .df. pdf.
[23] Florian Hess, Zahra Monfared, Manuel Brenner, and Daniel Durstewitz. Generalized Teacher Forcing for Learning Chaotic Dynamics. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato. and Jonathan Scarlett, editors, Proceedings of the 40th Inter- national Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 13017-13049. PMLR, 2023. URL httpp::/pproceeddiiieee mlr press/ v202/ heas23a. html. html.
[24] Jason A Platt, Stephen G Penny. Timothy A Smith, Tse-Chun Chen, and Henry D Abarbanel. Constraining chaos: Enforcing dynamical invariants in the training of reservoir computers. Chaos, 33(10), 1 October 2023. URL http://dx. doi. org/10. 01683/5..156999
[25] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Doina Precup and Yee Whye Tch, editors, Proceedings of the 34th Interna- tional Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pages 214-223, PMLR, 06-11 Aug 2017. URL https:://prccedaii v70/arjovakyy7a..hoo html.
mlr mll..ress//
[26] Abdelwahed Khamis, Russell Tsuchida, Mohamed Tarek, Vivien Rolland, and Lars Petersson. Earth movers in the big data era: A review of optimal transport in machine learning, 2023.
[27] Patrick T. Komiske, Eric M. Metodiev, and Jesse Thaler. Metric space of collider events. Phys. Rew. Lett. 123:041801, Jul 2019. doi: 10.1 N0.1103/hysReev.ett. 123.04 18.004111. URL. https: //link, aps wwwo//og///o//oooo www.///yssseeeeettttoomm 123. 041801.
12

---

[28] Peter Y. Lu, Rumen Dangovski, and Marin Soljačić. Discovering conservation laws using optimal transport and manifold learning. Nat. Commun., 14(1):4744, 7 August 2023. URL
http://bb /ttp://dx.oo doi. woo.o.o///.... 1038/ 94138/441677000 40444667--33---7777
[29] Yunan Yang, Levon Nurbekyan, Elisa Negrini, Robert Martin, and Mirjeta Pasha. Optimal transport for parameter identification of chaotic dynamics via invariant measures. SIAM Journal on Applied Dynamical Systems, 22(1):269-3 -10, 2023, doi: 10.1 10.1137/1MM1421333
[30] Jonah Servinick-Greenhuuse Robert Martin, and Yunan Yang. Learning dynamics on invariant measures using PDE-constrained optimization. Chaos, 33(6), 1 June 2023. URL http://dx. doi.org/10 10. 0149673.
[31] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, edi- tors, Advances in Neural Information Processing Systems, volume 26, Red Hook, NY, 2013. Curran Associates, Inc. URL ptttp::/pprcceediie wwtp.rtpp.cc/ppprr///////1 #22-dcc77bbbe77113777fff96232223
pdf.
[32] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
[33] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729-9738, 2020.
[34] Zhirong Wu, Yuanjun Xiong, Stella X. Yu, and Dahua Lin. Unsupervised feature learning via non-parametric instance discrimination. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2018.
[35] Xiao Zhang and Michael Maire. Self-supervised visual representation learning from hierarchical grouping. Advances in Neural Information Processing Systems, 33: 13:16579-11635 2020.
[36] Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In Computer Vision ECCon 2016: 14th European Conference, Amsterdam, The Netherlands, October |l 11-14, 2016, Proceedings, Part II 14, pages 694-711. Springer, 2016.
[37] Yunjie Tian, Lingxi Xie, Xiaopeng Zhang, Jiemin Fang, Haohang Xu, Wei Huang, Jianbin Jiao, Qi Tian, and Qixiang Ye. Sawan generation for self-supervised visual representation learning, arXiv preprint arXiv:2111.13163 2021.
[38] Cédric Villani. Optimal Transport: Old and New. Springer Berlin Heidelberg, Berlin, Heidel- berg, 2009, ISBN 978-3-540-71050-9, doi: 10. 00177997-3-54-7711 6..
[39] Jean Feydy, Thibault Séjourné, François- Xavier Vialard, Shun-ichi Amari, Alain Trouve, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 2681-2690, Naha, Okinawa, Japan, 16-18 Apr 2019. PMLR. URL https:/ ptrtpe//prooeeeiiggs mlr press/v89 /tte////////yyeeeeeee html.
[40] Hicham Janati, Marco Cuturi, and Alexandre Gramfort. Debiased Sinkhorn barycenters. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 4692-4701, Online, 13-18 Jul 2020. PMLR. URL https: /atp:////roocdgg.s.. mlr hatt.prrsssii11 j janati20a. html.
[41] Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere, In International Conference on Machine Learning, pages 9929-99299993 9939. PMLR, 2020.
[42] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unrea- sonable effectiveness of deep features as perceptual metric. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 586-595, 2018.
13

---

[43] Dirk Leendert van Kekem. Dynamics of the lorenz-96 model: Bifurcations, symmetries and waves, 2018. https:/reesearch.cu rug dynamics-of of-
theee-ofthheeteeee the- 1ore---enz--aae 96- model
aafurracionss-symetrieess -
and waves.
[44] Kody JH Law, Daniel Sanz-Alonso, Abhishek Shukla, and Andrew M Stuart. Filter accuracy for the Lorenz 96 model: Fixed versus adaptive observation operators. Physica D: Nonlinear Phenomena, 325:1-13, 2016.
[45] Andrew J Majda and John Harlim. Filtering Complex Turbulent Systems. Cambridge University Press, 2012.
[46] James M Hyman and Basil Nicolaenko. The kuramoto-sivashinsky equation: a bridge between pde's and dynamical systems. Physica Nonlinear Phenomena, 18(1-3):113-126, 1986.
[47] Pantelis Vlachas, Georgios Arampatzis, Caroline Uhler, and Petros Koumoutsakos. Multi- scale simulations of complex systems by learning their effective dynamics. Nature Machine Intelligence, 4(4):359- -366, 2022.
[48] H. S. Greenside, A. Wolf, Swift, and T. Pignataro. Impracticality of a box-counting algorithm for calculating the dimensionality of strange attractors. Phys. Rev. A, 25:3453-3456 Jun 1982. doi: 10.1103/PtyyRRe2225.44444 URL attps://links: org. httt//////.....///////eeee.. 3453.
[49] Bowen Zhu, Jingwei Hu, Yifei Lou, and Yunan Yang. Implicit regularization effects of the sobolev norms in image processing. arXiv preprint urXiv:2109.06255 2021.
[50] Chuanfu Wang, Chunlei Fan, and Qun Ding. Constructing discrete chaotic systems with positive Iyapunov exponents. International Journal of Bifurcation and Chaos, 28(07):1850084, 2018.
[51] Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhomn' 's algorithm. In Jennifer Dy and Andreas Krause, editors, Pruceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1367-1376. mlr www.....r./..............e.eret. press/v80// /wdddddeessyyyeeee html.
PMLR, 10-15 Jul 2018. URL
[52] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL httpp://ppearevvoeeneet//// Bkg6RicqY7.
14

---

A Additional Discussion
A.1 Contrastive feature learning vs. physics-informed optimal transport
Our two proposed approaches both aim to train neural operators to preserve the invariant measures of chaotic attractors. Ultimately, both methods have strong performance according to a variety of metrics, but the contrastive learning approach requires no prior physical knowledge and is often faster. Both approaches have significant advantages over the standard approach of minimizing RMSE, which fails to preserve important statistical characteristics of the system.
The primary conceptual difference between our approaches is that the optimal transport approach relies on prior domain knowledge, while the contrastive learning approach learns from the data alone. Specifically, the tlysiie-informed optimal transport approach requires a choice of summary statistics, so it is much more dependent our prior knowledge about the system and its chaotic attractor. Contrastive feature learning does not require a choice of statistics and instead learns useful invariant statistics on its own-although it may still be useful to have known relevant statistics to use as evaluation metrics for hyperparameter tuning.
The neural operators trained using the optimal transport loss perform the best on the L1 histogram distance of the summary statistics S (Table and 2). which unsurprising given that the optimal transport loss specifically matches the distribution of S. If we look at a different evaluation metric, e.g. the energy spectrum error or leading LE, the contrastive learning loss, in some cases, performs better even without prior knowledge. This suggests that contrastive feature learning may be capturing a wider range or different set of invariant statistics. We also see that the quality of the emulator degrades if we choose less informative statistics (Appendix B.3).
A.2 Interactions and trade-offs between short-term prediction and long-term statisties
Comparing results in Tables and 2, we find that emulators trained using our approaches perform significantly better on the long-term evaluation metrics (e.g. I, histogram distance and energy spectrum error) by trading off a bit of performance in terms of short-term RMSE (Tables 11 and 12). This suggests that training an emulator to model chaotic dynamics using purely RMSE can result in overfitting and generally poor model of long-term dynamics. However, we still believe that RMSE is a useful part of the loss, even for long-term evaluation metrics, since it is the only term that is directly enforcing the short-term dynamics. In some cases, we find that the invariant statisties SA of trajectories generated by our trained neural operators are closer to the true statistics S. than the statistics directly computed using the noisy training data S i.e., SSA SA < JSA SAl, despite being trained on the noisy data (Table 3). We attribute this result to the RMSE component of our losses, which provides additional regularization by enforcing the short-term dynamics, combined with the inductive biases of the neural operator architectures, which are designed for modeling PDEs.
Training
Histogram Error S5 S + Histogram Error |5^ S 4
ÉRMSE
0.056 (0.051, 0.062)
0.029 (0.023, 0.036)
0.1 foT + FRMSE
0.029 (0.027, 0.032)
0.029 (0.023, 0.036)
ECL + CRMSE
0.033 (0.029, 0.037)
0.029 (0.023, 0.036)
ÉRMBR
0.130 (0.118. 0.142)
0.101 (0.086, 0.128)
0.2
foT ÉRMSE
0.039 (0.035, 0.042)
0.101 (0.086. 028)
fCL ÉRMSE
0.073 (0.066, 0.080)
0.101 (0.086, 0. 0.128)
FRMSE
0.215 (0.204, 0,234)
0.211 10.190, 0..55)
0.3 for + FRMSE
0.057 (0.052, 0.064)
0.213 (0.190, 0.255)
fCL + FRMSE
0.132 (0.111, 0.151)
0.213 (0.190, 0.255)
Table 3: Histogram error of neural operator predictions vs. noisy training data on Lorenz-96. Averaging over 200 testing instances with varying ((*),, we compared the L₃ histogram error of the predicted dynamics from the neural operator |SA - SAl with the histogram error of the raw noisy training data A - S.l. This shows that, even though the neural operator is trained on the noisy data, the statistics of the predicted dynamics often better than those computed directly from the noisy data. Here, SA is computed from the noiseless ground truth data.
15

---

B Additional Experiments and Evaluation Metrics
We have performed several additional experiments that act as points of comparison, help us better understand the behavior of our methods under variety of conditions, and provide useful insights for future applications of our approaches.
B.1 Denoising with Gaussian blurring
Gaussian blurring, often used as a denoising technique for images, employs a Gaussian distribu- tion to establish a convolution matrix that's applied to the original image. The fundamental idea involves substituting the noisy pixel with weighted average of surrounding pixel values. A key hyperparameter in Gaussian blurring is the standard deviation of the Gaussian distribution. When the standard deviation approaches zero, it fundamentally indicates the absence of any blur. Under such cireumstances, the Gaussian function collapses blur effect. In Table 4, we present the results from applying Gaussian blurring to noisy data during training based solely on RMSE. Despite the effectiveness of the widely adopted denoising approach. our findings indicate that Gaussian blurring may not be ideally suited for our purpose of emulating dynamics. This is primarily because significant invariant statistics might be strongly correlated with certain figuerryyy signals that could be affected by the blurring preprocessing.
single point, leading to the elimination of the
Training
Histogram Error J
Energy Spec. Error,
Leading LE Error |
ERMSE (o, = 0.1)
0.390 (0.326, 0.556)
0.290 (0.226, 0.402)
0.098 (0.069, 0.127)
IRMSE (0g = 0.5)
1.011 (0.788, 1.264)
0.493 (0.379, 0,623)
0.098 (0.041, 0.427)
ARMSE
0.390 (0.325, 0.556)
0.290 (0.225, 0.402)
0.101 (0.069, 0.122)
Table Emulator performance with Gaussian blurring on Kuramoto-Sivashins noise scale r = 0.3. Averaging over 200 testing instances with varying g(*), we show the performance of (1) the application of Gaussian blurring as a preliminary denoising effort with a small standard deviation (o, = 0.1); (2) the application of Gaussian blurring with larger standard deviation (o) = 0.5); and (3) training purely on RMSE without any blurring preprocessing. The results suggest that the application of Gaussian blurring might further degrade the results, as the high-frequency
data with
signals associated with invariant statistics can lost.
B.2 Sobolev norm baseline
We recognize that there are alternative methods that strive to capture high-frequency signals by modifying training objectives. For instance, the Sobolev norm, which combines data and its deriva- tives. has been found to be quite effective in capturing high-frequency signals [16, 49]. However, its effectiveness can be significantly curtailed in a noisy environment, especially when noise is introduced to high-frequency domain, minimizing the Sobolev norm then fails to accurately capture relevant statistics, as shown in Tables 5 and
Training
Histogram Error 4
Energy Spee. Frrr
Leading L.E Error
FD Error J.
RRISE
0.215 (0.204, 0.234)
0.291 (0.280, 0,305)
0.440 (0.425, 0.463)
3,580 (2.333, 4.866)
Iaulubee + Edisopptioe
0.246 (0.235, 0.255)
0.325 (0.341. 0,307)
0.487 (0.456, 0.545)
4.602 (3.329, 6.327)
for ÉRISE
0.057 (0.052, 0.064)
0.123 (0.116, 0.135)
0.084 (0.062, 0.134)
3.453 (2.457, 4.782)
fCL ERMSE
0.132 (0.111, 0.151)
0.241 (0.208, 0.285)
0.064 (0.045, 0.091)
1.894 (0.942, 3.108)
Table 5: Emulator performance (including Soboley norm loss) on Lorenz-96 data with noise scale 0.3. The median (25th, 75th percentile) of the evaluation metrics computed on 200 Lorenz-96 test instances (each with 1500 time steps) for the neural operator trained with (1) only RMSE loss RMSE; (2) Sobolev norm loss with dissipative regularization ESobolev Ediisspptive (3) optimal transport (OT) and RMSE loss for + ÉRMSE; and (4) contrastive learning (CL.) and RMSE
loss fCL + {RMSE.
16

---

Training
Histogram Error
Energy Spee. Error
Leading LE Error +
(RMSE
0.390 (0.325, 0.556)
0.290 (0.225, 0,402)
0.101 (0.069, 0.122)
ESobolev + diiesppativee
0,427 (0,289, 0.616)
0.237 (0.204, 0.315)
0.023 (0.012, 0.047)
foT {RMSE
0.172 14,1 0.197)
0.211 10.188, 0.250)
0.094 (0.041, 0.127)
ECL RRSSE
0.193 (0.148, 0.247)
0.176 (0.130, 0.245)
0.108 (0.068, 0.132)
Table 6: Emulator performance (including Sobolev norm loss) on Suramoto-Sivashinsky data with noise scale r = 0.3. The median (25th, 75th percentile) of the evaluation metrics are computed test instances (each with 1000 time steps) for the neural operator trained with (1) only RMSE loss (RMSE; (2) Sobolev norm loss with dissipative regularization (Sooolev + Edissaptive: (3) optimal transport (OT) and RMSE loss for + ÉRMSE; and (4) contrastive learning (CL) and RMSE loss fCL RRMSE.
on 200
B.3 Optimal transport: reduced set of summary statistics
For our optimal transport approach, we test reduced set of summary statistics, which shows how the quality of the summary statistic affects the performance of the method (Table 7). With informative summary statistic, we find even reduced set can still be helpful but, for a non-informative statistic, the optimal transport method fails as expected.
Training statistics
Histogram Error
Energy Spec. Error +
Leading LE Error |
FD Error
(full)
0.057 (0.052, 0.064)
0.123 (0. 11,, 0.135)
0.084 (0.062, 0.134)
3.453 (2.457, 4.782)
S, (partial)
0.090 (0.084, 0.098)
0.198 (0.189, 0,208)
0.263 (0.217, 0.323)
3,992 (2.543, 5,440)
Sy (minimum)
0.221 (0.210, 0.234)
0.221 (0.210, 0.230)
0.276 (0.258, 0.291)
3.204 (2.037, 4.679)
Table 7: Emulator performance for different choices of summary statistics on Lorenz-96 data with noise scale r 0.3. Each neural operator was trained using the optimal transport and RMSE loss using (1) full statistics S(u) == { (d (u441 24-2)u4-11 u;}: (2) partial statistics (u) == {(u+1 14-2)u41]; or (3) minimum statistics S2(u) == {ū}. where ū is the spatial average.
B.4 Contrastive learning: reduced environment diversity
For our contrastive learning approach, we test multi-environment setting with reduced data diversity and find that the contrastive method still performs well under the reduced conditions (Table 8), which demonstrates robustness.
Training
Histogrum Error 1.
Energy Spec. Error
Leading LE Error
Error 1
ARMSE
0.255 (0. 24,448,0 0.263)
0.307 (0.302, 0.315) 0.124 (0. 16, 0.131)
0.459 (0.743, 2.746)
3.879 (2.456 5,076)
for + FRMSE
0.055 (0.050, 0.061)
0.080 (0.045, 0.1)
4.015 (2.401, 5.225) 1.747 (0.792, 2.939)
Eer. FRMSE
0.130 (0.111, 0.152)
0.193 (0. 183, 0.200)
0.031 (0.014, 0.053)
Table Emulator performance with reduced environment diversity (i.e. narrower parameter range) on Lorenz-96 data with noise level r = 0.3. Averaging over 200 testing instances, we show the performance of training the neural operator with (1) only RMSE loss fRMSE: (2) optimal transport (OT) and RMSE loss łoT + RMSE: and (3) contrastive learning (CL) and RMSE loss fCL + {RMSE. We shrink the parameter range for generating the dataset from [10, 18] to [16, 18).
17

---

B.5
Maximum mean discrepancy (MMD) vs. optimal transport loss
We also implement variant of our optimal transport approach that uses maximum mean discrepaney (MMD) as a distributional distance rather than the Sinkhorn divergence. Using the same set of summary statistics, we find that MMD does not perform as well our optimal transport loss for training emulators (Table 9).
Training
Histogram Error
Energy Spec. Error +
Leading LE Error
ÈRMSE:
0,390 (0.325, 0.556)
0.290 (0.225, 0,402)
0.101 (0.069, 0.1
BMMD ARMSE
0.245 (0.218, 0.334)
0.216 (0. 186 0.272)
0.101 (0.058, 0.125)
loT + ERMSE
0.172 (0. 146, 0.197)
0.211 (0.188, 0.250)
0.094 (0.041, 0.127)
fCT. + ÉRMSE
0.193 (0. 148, 0.247)
0.176 (0.130, 0.245)
0.108 (0.068, 0.132)
Sivashinsky data with noise scale r 0.3. The median (25th, 75th percentile) of the evaluation metrics computed test instances (each with 1000 time steps) for the neural operator trained with (1) only RMSE loss fRMS: (2) maximum mean discrepency (MMD) and RMSE loss EMMD + {RMSE: (3) optimal transport (OT) and RMSE loss fot + ÉRMSE; and (4) contrastive learning (CL) and RMSE loss fCL RRMS..
Table 9: Emulator performance (including MMD loss) on
on 200 Srramoto-Sivashinsee
B.6 Additional Lyapunov spectrum evaluation metrics
In the table 10, evaluated the results of Lorenz 96 Lyapunov spectrum error rates and the total number of positive Lyapunov exponents error rates. For the Lyapunov spectrum error, we report the sum of relative absolute errors across the full spectrum: Σ |X; >i//A, where X; is the i-th Lyapunov exponent and d the dimension of the dynamical state. As suggested by [50], we also compare the number of positive Lyapunov exponents (LEs) as an additional statistic to measure the complexity of the chaotic dynamics. We compute the absolute error in the number of positive LEs
Σ
(A > 0 - (A; > 0)}-
Training
Leading LE Error |
Lyapunov Spectrum Error
Total number of positive LEs Error
fIMSE
0.013 (0.006, 0.021)
0.265 ( 100, 0.309)
0.500 (0.000, 1.000)
0.1 EoT + ÉRMSE.
0.050 (0.040, 0.059)
0,248 (0.168, 0.285)
0.000 (0.000, 1.000)
ECt. + ÉRMSE
0.065 (0.058, 0.073)
0.227 10.164, 0.289)
0.000 (0.000, 1.000)
{uMSE
0.170 (U.156, 0,191)
0.612 (0.522, U.727)
4.000 (4.000, 5.000)
0.2 for + fRMSS
0.016 (0.006, 0.030)
0,513 (0.122, 0.590)
3.000 (2.000, 3.000)
fct. ÉRMSE
0.012 (0.006, 0,018)
0.459 (0.138, 0.568)
3.000 (2.000, 3.000)
(RMSE:
0.440 (U:425, 0.463)
0.760 (0.702, U.939)
7,000 (7.000, 8.000)
0.3 for + ÉRMSE fOL. ERMSE
0.084 (0.062, 0.134)
0.661 (0.572, 0.746)
5,000 (4.000, 6.000)
0.064 (0.045, 0.091)
0.654 (0.558, 0.780)
6.000 (5.000, 6.000)
Table 10: Emulator performance on Lyapumov spectrum metrics for Lorenz-96 data. The median (25th, 75th percentile) of the Lyapunov spectrum metrics are computed on 200 Lorenz-96 test instances (each with 1500 time steps) for the neural operator trained with (1) only RMSE loss (RMSE; (2) optimal transport (OT) and RMSE loss for + ÉRMSE; and (3) contrastive learning (CL) and RMSE loss &CT. + ÉRME. In the presence of high noise, OT and CL give lower relative errors on the leading Lyapunov exponent (LE). When evaluating the full Lyapunov spectrum, OT and CL show significant advantages than the baseline. In addition, the lower absolute errors of the total number of the positive Lyapunov exponents (LEs) suggest that OT and CL are able to match the complexity of the true chaotic dynamics.
18

---

B.7 1-step RMSE evaluation results
Evaluating 1-step RMSE only shows short-term prediction performance and is not informative evaluation metric for long-term behavior. Here, we report the 1-step RMSE (Tables 11 and 12) to show that training using our approaches retains similar 1-step RMSE results while significantly improving on long-term statistical metrics (Tables and 2). See Appendix A.2 for additional discussion.
Training
RMSE
FRMSE
0.107 (0. 105, 0.109)
for ÉRMSE
0.108 105, 0.110)
ECL CRMSE
0.109 108, 0.113)
FRMBE
0.202 (0.197, 0,207)
0.2 foT + ÉRMSE
0.207 (0.202, 0.212)
ECT. + ÉRMSE
0.214 (0.203, 0.218)
FRMSE
0.288 (0.282, 0.296)
0.3 foT FRMSE
0,301 (0.293, 0.307)
fc1. ERMSE
0.312 (0.302,0.316)
Table 11: 1-step RMSE performance on Lorenz-96 data. The median (25th, 75th percentile) of
the Lyapunov spectrum metrics are computed on 200 Lorenz-96 test instances (each with 1500 time steps) for the neural operator trained with (1) only RMSE loss {RMSE; (2) optimal transport (OT) and RMSE loss for + ÉRMSE: and (3) contrastive learning (CL.) and RMSE loss fcL. + ERMSE, We see comparable performance on short-term 1-step RMSE.
Training
RMSE
ÊNMSE
0.373 (0.336, 0.421)
foT + ÉRMSE
0.381 (0. 344, 0.430)
Eer. ARMSE 0.402 (0.364, 0.452)
Table 12: 1-step RMSE performance on Kuramoto-Sivashinsk data with noise scale r 0.3. The median (25th, 75th percentile) the evaluation metrics are computed on 200 Kuramoto- Sivashinsky test instances (each with 1000 time steps) for the neural operator trained with (1) only RMSE loss ÉRMSE: (2) optimal transport (OT) and RMSE loss foT + ÉRMSE and (3) contrastive learning (CL) and RMSE loss &CL + ÉRSE. We again comparable performance short-term 1-step RMSE.
C Implementation Details
C.1 Evaluation metrics
To evaluate the long-term statistical behavior of our trained neural operators, we run the neural operator in recurrent fashion for 1500 (Lorenz-96) 1000 Siraaito-Sivvssinsky) time steps and then compute our long- terg evaluation metrics on this autonomously generated time-series.
Histogram error. For distributional distance with pre-specified statistics S, we first compute the histogram of the invariant statistics H()) = {(S,, e1), (S2, 62) C2),.... ...., (Sp, Cn)), where H represents the histogram, and the values of the bins denoted by S, with their corresponding frequencies c-- We then define the L2 histogram error as:
E H. H
Σ
ca -a--
(21)
Note that for a fair comparison across all our experiments, we use the rule of thumb the square root rule decide the number of bins.
Energy spectrum error. We compute the relative mean absolute error of the energy spectrum the over time:
squared norm of the spatial FFT
1
Σ
|| u.
(22)
T
Fu
1
19

---

Leading LE error. The leading Lyapunov exponent (LE) is dynamical invariant that measures how quickly the chaotic system becomes unpredictable. For the leading LE error, we report the relative absolute error i >|/|a| between the model and the ground truth averaged over the test set. We adapted the Julia DynamicalSystem..] package to calculate the leading LE.
FD error. The fractal dimension (FD) is a characterization of the dimension of the attractor. We report the absolute error ID D| between the estimated FD of the model and the ground truth averaged over the test set. We use the Julia DynamicalSystem. package for calculating the fractal dimension.
RMSE. We use 1-step relative RMSE |ur+ar - (o(u,, 0)1l ut/h/ to measure short-term prediction accuracy.
C.2 Time complexity
The optimal transport approach relies on the Sinkhorn algorithm which scales as Onn² log for comparing two distributions n points each (Theorem 2 in [51]. In our experiments, we use 6000 to n 25600 points with no issues, so this approach scales relatively well. The contrastive learning approach requires pretraining but even faster during emulator training since uses fixed, pre-trained embedding network.
Training
Encoder Operator
Total
CRMSE
20
20
Lorenz-96
for + ÉRMSE fCL + EASEE
51
51
22
27
49
fRMSE loT + {RMSE
55
55
Kuramoto-Sivashinkyy
262
262
lCL. + LAMSE
26
56
82
Table 13: Empirical training time. Training time (minutes) with 4 GPUs for 60-dimensional Lorenz-96 and 256-dimensional Kuramoto-Sivashinkkk
C.3 Training of the encoder
Evaluation rule. standard procedure for assessing the performance of an encoder trained with Noise Contrastive Estimation (InfoNCE) loss in an unsupervised manner [32-34], is employing the top-1 accuracy metric. This measures how effectively similar items positioned closer each other compared to dissimilar ones in the embedded space. While this evaluation measure is commonly employed in downstream tasks that focus on classification, is suitably in line with our goals. We aim to learn an encoder that can differentiate whether sequences of trajectories are sampled from different attractors, each characterized by distinct invariant statisties. Therefore, the use of Top-1 1 accuracy to evaluate if two sequences originate from the same trajectory serves our purpose effectively. And we utilize this metric during our training evaluations assess the performance our encoder.
The formal definition of Top-1 accuracy requires us to initially define the softmax function a, which used estimate the likelihood that two samples from different time windows originate from the same trajectory.
exp (Jo (o (.]++K)))
(23)
Σ Dm=L exp
(Ul())) (n) 1.1k)) fo f (m)
HHH+K))
Using the softmax function transform the encoder's output into probability distribution, then define Top-1 accuracy as:
N
1
n arg max (Fo( (U
Acc1
(24)
N
1=1
20

---

where 1(-) is the indicator function representing whether two most close samples in the feature space comes from the same trajectory or not.
Training hyperparameters. We use the ResNet-34 as the backbone of the encoder, throughout all experiments, we train the encoder using the AdamW optimization algorithm [52], with weight decay of 10-s, and set the training duration to 2000 epochs.
For the temperature value balancing the weights of difficult and easy-to-distinguish samples in contrastive learning, we use the same warm-up strategy as in [10]. Initially, we start with a relatively low T value (0.3 in experiments) for the first 1000 epochs. This ensures that samples that are difficult to distinguish get large gradients. Subsequently, we incrementally increase T up to a specified value (0.7 in our case), promoting the grouping of similar samples within the embedded space. From our empirical observations, we have found that this approach leads to an improvement in Top-1 accuracy in our experiments.
The length of the sequence directly influences the Top-1 accuracy. Considering that the sample length L of the training data { { / is finite, excessively increasing the crop length K can have both advantages and disadvantasn vantages. On one hand, it enables the encoder to encapsulate more information; on the other hand, it could lead to the failure of the encoder's training. This is likely to occur if two samples, i.e., f (U I:I+k) and (.J+++))))) from the same trajectory overlap excessively, inhibiting the encoder from learning a meaningful feature space. With this consideration, we' 've empirically chosen the length K of subsequences for the encoder to handle to be approximately 5% of the total length T.
C.4 Lorenz-96
Data generation. To better align with realistic scenarios, we generate training data with random initial conditions drawn from a normal distribution. For the purpose of multi-environment learning, we generate 2000 trajectories for training. Each of these trajectories has a value of (s()) randomly sampled from uniform distribution ranging from 10.0 to 18.0. We discretize these training trajectories into time steps of di 0.1 over a total time of / 2058, yielding 2050 discretized time steps. Moreover, in line with the setup used in [2, 11]. we discard the initial 50 steps, which represent the states during an initial transient period.
Training hyperparameters. We determine the roll-out step during training (i.e., h and haMSe) via the grid search from the set of values {1, 2, 4, 5). Though it is shown in some cases that a larger roll-out number help improve the results [18], we have not observed that in our training. We hypothesized that this discrepancy may be due to our dataset being more chaotic compared to typical cases. Consequently, to ensure a fair comparison and optimal outcomes across all experiments, we have decided to set and ARMSE to
In the case of the optimal transport loss?, described in Section 3.1, the blur parameter Y governs the equilibrium between faster convergence speed and the precision of the Wasserstein distance approximation. We determined the value of y through a grid search, examining values ranging from (0.01, 0.02..... 0.20). Similarly, we decided the weights of the optimal transport loss a through grid search, exploring values from (0.5. 1, 1.3...... ..00). For the experiments conducted using L6renz-96, we selected a and y 0.02.
In the context of the feature loss, the primary hyperparameter we need to consider is its weights, represented as A. Given that we do not have knowledge of the invariant statistics S during the validation phase, we adjust X according to a specific principle: our aim is to reduce the feature loss ffeature (as defined in Eqn. 17) as long as does not adversely affect the RMSE beyond predetermined level (for instance, 10%) when compared with the baseline model, which is exclusively trained on RMSE. As illustrated in Figure we adjust the values of À in systematic grid search from {0.2, 0.4, 0.6, 0.8, 1.0, 1.2). From our observations during the validation phase, we noted that the feature loss was at its lowest with an acceptable RMSE (which is lower than 110% compared to the baseline) when À = 0.8. Therefore, we have reported our results with À set at 0.8.
Results visualization. We present more results with varying noise levels in Figures 5, 6, 7.
21

---

step RMSE Inolsy valication)
statuee loss Inoicy validation)
histogram (no lee testing)
cne-step AMSE (noiseless testing!
0.22
0.40
Nittogram error of invatinnt stateelcs
-0.99
(AASE Inslty welciatiiioo
0.36
0.08
0.88
016
0.32
0.30
0.36
1.2
Figure 4: The trend of feature loss with its weight when the scale of noise is 0.3. The solid lines in the figure represent the evaluation metrics during the validation phase, comparing the outputs of the neural operator the noisy data. In contrast, the dashed lines represent the actual metries we are interested in, comparing the outputs of the neural operator to the clean data and calculating the error of the invariant statistics. In addition, the horizontal solid dashed line correspond to the bar we set for the RMSE, i.e.. 110% of the RMSE when X = O. We observe that, (1) with the increase of from O, the feature loss decreases until À reaches 1.0. (2) The RMSE generally increases with the increase of X. (3) The unseen statistical error generally decreases with the increase of X. We reported the results when À 0.8 as our final result, since the further increase of À does not bring further benefit in decreasing the feature loss, and the result remains in acceptable range in terms of RMSE.
C,5 Kuramoto-Sivashinsky
Data generation. In order to ascertain that operating within a chaotic regime, we set the domain size = 50 and the spatial discretization grids number 256. We select initial conditions randomly from uniform range between [-T, m.. A fourth-order Runge-Kutta method was utilized to perform all simulations. We generate 2000 trajectories for training, with each ((()) being randomly chosen from a uniform distribution within the interval [1.0, 2.6].
Training hy pereparraeeeess. In the case of the optimal transport loss, similar to the discussion in C.4, we search the blur value y and the weights loss from a grid search. And in our experiment, we 0.05. In the case of the feature loss, we again adopt the same rule for deciding
set a = 3 and the value of X, where we compare the trend of feature loss with the relative change of RMSE of the baseline and choose to report the results with the lowest feature loss and acceptable RMSE increase 10% compared the baseline when 0).
22

---

Sampled ectory
Histogram 2d
Histogram 2d
Sampled trajectory
Histogram 2d
Histogram 2d
Truzh
Trm
OT Kong
for VoPP
CL (ours)
CL
Trab
Teah
Energy
104
Energy
101.
10a
10
0
5
10
15
20
25
30
5
10
15
20
25
30
Wavenumber
Wavenumber
Figure 5: Visualization of the predictions when the noise Level = 0.1. We evaluate our method by comparing them to the baseline that is trained solely using RMSE. For two different instances (a) and (b), we visualize the visual comparison of the predicted dynamics (left), two-dimensional histograms of relevant statisties (middle and right). We notice that, with the minimal noise, the predictions obtained from all methods look statistically consistent to the true dynamics.
23

---

Sampled trajectory
Histogram
Histogram
Sampled trajectory
Histogram
Histogram
Truch
PMA RMSE
Of (oun))
01 (our))
CL (ours)
Teah
(ven)
Energy
104
Energy
104
10a
10"
10
15
25
30
15
20
25
30
Wavenumber
Wavenumber
(a)
(b)
Figure 6: Visualization of the predictions when the noise level r 0.2. We evaluate method by comparing them to the baseline that is trained solely using RMSE. For two different instances (a) and (b), visualize the visual comparison of the predicted dynamics (left), two-dimensional histograms of relevant statistics (middle and right). We observe that as the noise level escalates. the degradation of performance in the results of RMSE more rapid compared employs optimal transport and feature loss during training.
method, which
24

---

Sampled trajectory
Histogram
Histogram 2d
Sampled trajectoy
Histogram 2d
Histogram
Truch
RMSE
11 Cour
01 hoo
CL (curs)
OST TOL
ON3 OZ
RMSE 07 jan)
1902
(oen)
Energy
10
Energy
10a
25
Wavenumber
Wavenumber
(a)
(b)
Figure 7: Visualization of the predictions when the noise level r 0.3. We evaluate our method by comparing them to the baseline that is trained solely using RMSE. For two different instances (a) and (b), we visualize the visual comparison the predicted dynamics (left), two-dimensional histograms of relevant statistics (middle and right). We find that, under higher levels of noise, the RMSE results exhibit fewer negative values, as indicated by the blue stripes in the predicted dynamics. This is further confirmed by the energy spectrum, which clearly shows that the RMSE results are significantly deficient in capturing high- high-frequency signals.
25


Reviews:
[{"review_id": "Gje39NtqBP", "reviewer": "Reviewer_GF1X", "summary": "This work considers the problem of learning invariant statistics of chaotic dynamical systems. It is motivated by the observation that training neural operators by standard MSE loss focuses on short-term predictability and thus may miss an attractor’s properties. Two ways of augmenting the MSE loss for capturing invariant stats are discussed: One based on the Wasserstein distance between prob distributions across explicitly provided stats (physics-informed), and one based on contrastive feature learning. Experiments on chaotic benchmarks, the Lorenz-96 and the Kuramoto-Sivashinsky equations, are shown.", "questions": "My primary concern is really the weaknesses in the experimental evaluations as listed above. This sect. was a bit disappointing to read after the exciting start of the paper and presentation of methods.", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 4, "strengths": "The paper is timely and addresses an important issue, namely how we can capture crucial properties of an observed dynamical system beyond just short-term predictions. In general I like the approach, and think this could be a very fruitful addition to the current state. In general, the paper is also well written, and the ideas well motivated. Although incorporating invariant stats into the loss has recently been discussed, this paper offers a different perspective on this topic.", "weaknesses": "On the other hand, the experimental evaluations are, in my view, quite weak, and also a lot of related and relevant literature is not discussed.\n\nThe experiments primarily show that on any of the evaluation measures that particular method wins, for which exactly that measure has been included or accentuated in the loss function: If the loss is only MSE, the only-MSE trained method wins on short-term predictions; if the loss includes invariant statistics, the method wins on exactly those invariant stats. This is not surprising, but would be expected for almost anything explicitly included in the loss term (i.e., the method for which a specific property was included in the loss should have an edge when evaluated on that particular property). \n\nIt would have been much more convincing if methods which include invariant statistics could also reproduce many *other* properties of the attractor that were *not* explicitly included in the loss, like its Lyapunov spectrum or fractal dimension.\nBesides, two of the columns in Table 1 \\& 2 don’t contain any indication of statistical error and are therefore somewhat meaningless in my mind (are the differences significant or negligible?).\nTo see how the methods perform on at least one empirical chaotic problem would have also been nice.\n\nLiterature-wise, the idea of including invariant statistics in the loss has been considered by a number of recent publications, for instance https://arxiv.org/pdf/2304.12865.pdf or https://pubs.aip.org/aip/cha/article-abstract/33/6/063152/2900453/Learning-dynamics-on-invariant-measures-using-PDE?redirectedFrom=fulltext.\nThe fact that purely MSE-based methods have difficulty capturing invariant stats for chaotic systems has recently also been extensively discussed in https://proceedings.neurips.cc/paper_files/paper/2022/hash/495e55f361708bedbab5d81f92048dcd-Abstract-Conference.html and, relatedly, https://arxiv.org/abs/2306.04406, although these authors arrive at different conclusions of how to tackle the problem.\nMy perception is that this problem has received a lot more attention recently than the authors’ related works section makes one believe.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!\n\n**W1**:\n> The experiments primarily show that on any of the evaluation measures that particular method wins…\n\nWe propose emulators with much more consistent long-term behavior and that accurately capture the dynamics of the chaotic system. We make significant gains in long-term behavior, as measured by the evaluation statistics, while retaining very similar short-term prediction performance. Both methods (OT and CL) achieve these benefits, with the contrastive method requiring no prior knowledge and being faster.\n\nFor the CL method, we do not provide any explicit statistics during training. Instead, the contrastive pre-training automatically learns useful invariant statistics. For the OT approach, we show that training using a chosen set of informative summary statistics improves performance on those statistics as well as other statistics (e.g. the energy spectrum, Lyapunov exponent) not used during training.\n\n**W2**:\n> It would have been much more convincing if methods… reproduce… Lyapunov spectrum or fractal dimension.\n\nWe have included estimates of the leading Lyapunov exponent and the fractal dimension in our updated results (Table R1). Note that our method does not aim to regularize dynamical features, such as the Lyapunov exponents, which are distinct from time-invariant statistics. Despite this, we see significant improvements to the model’s leading Lyapunov exponent in high-noise settings. Since our data consists of high-dimensional chaotic attractors, it is very difficult to obtain reliable estimates of the full Lyapunov spectrum and the fractal dimension [1]. However, we do see evidence that the estimated fractal dimension of our CL approach is better than the baseline.\n\n**W3**:\n> Besides, two of the columns in Table 1 & 2 don’t contain any indication of statistical error…\n\nThank you for pointing this out. We have now included quantile intervals in our updated tables. Our results show significantly better performance on long-term statistics while retaining similar short-term prediction performance.\n\n**W4**:\n> To see how the methods perform on at least one empirical chaotic problem would have also been nice.\n\nOur work, like almost all others in this area, uses simulated data for validation. The datasets we currently use are representative of high-dimensional spatiotemporal chaos. It is quite difficult to find publicly available empirical datasets of fully-observed chaotic dynamical systems, especially ones showing spatiotemporal chaos. This is why most papers on spatiotemporal emulators are evaluated on simulated data. One publicly available dataset is global weather data, which we are excited to try but it would require significant resources.\n\n**W5**:\n> Literature-wise, the idea of including invariant statistics in the loss has been considered by a number of recent publications…\n\nWe first note that 3 of the 4 papers ([1], [2], and [4]) should be considered contemporaneous work under the NeurIPS guidelines. That said, we would be happy to cite and add a discussion of all four papers to better clarify our work in light of very recent developments in the field.\n\nTwo of the suggested papers ([3], [4]) focus on training traditional RNNs to emulate chaotic dynamics. Their methods and the ones presented in our submission have similar training protocols for the rMSE loss. Key differences between these papers and our approaches are (a) they only work in the noise-free setting whereas we assume noisy observations, (b) they only use a short-term rMSE loss, and (c) they evaluate their method in the easier single-environment setting. In our experiments, we see that rMSE is a fine loss function in noise-free settings (consistent with the conclusions of [3],[4]), but this loss fails to preserve chaotic attractors in noisy settings.\n\n[1] focuses on reservoir computing—an RNN-like architecture with a single trained output layer—and trains using known dynamical invariants, e.g. Lyapunov exponents. Training using the Lyapunov exponents (a) requires significant prior physical knowledge or (b) empirical estimates of the Lyapunov exponents, which are difficult to estimate stably from data, particularly in the presence of noise and high dimensions. Our updated evaluation results (Table R1) show that our training methods result in accurate estimates of the Lyapunov exponent even when it’s not used during training. \n\nThe method for modeling chaotic systems suggested in [2] requires solving a PDE for the probability density of the state and differs significantly from our proposed approaches, which are based on chosen or learned statistics. In particular, [2] scales poorly to higher dimensional state spaces, like the spatiotemporal dynamics that we consider, because even explicitly representing a high-dimensional probability distribution—let alone solving a high-dimensional PDE using a mesh—suffers from the curse of dimensionality. This is likely why the experiments in [2] focus on low-dimensional (up to 3) dynamical systems since solving a 3D PDE (i.e., a PDE with three “channels” like Lorenz-63) can already be numerically challenging. In contrast, our experiments are on high-dimensional spatiotemporal systems with state space dimensions ranging from 60 to 256.\n\n[1] Constraining Chaos. Platt et al, April 2023.\n\n[2] Learning dynamics on invariant measures using PDE-constrained optimization. Botvinick-Greenhouse et al, Chaos, June 2023.\n\n[3] On the difficulty of learning chaotic dynamics with RNNs. Mikhaeil et al, NeurIPS 2022.\n\n[4] Generalized Teacher Forcing for Learning Chaotic Dynamics. Hess et al, June 2023.\n\n**Q1**:\n> My primary concern is really the weaknesses in the experimental evaluations as listed above.\n\nPlease see the global rebuttal. We hope that our additional evaluation metrics and the above discussion address your concerns!\n"}, {"review_id": "l9eOyVMyw7", "reviewer": "Reviewer_G25b", "summary": "The paper uses neural operators to track the invariant statistics of chaotic systems. It novelly proposes to use the optimal transport loss and contractive learning to match the distribution, so that the learned model correctly track the attractor among the chaotic behavior. The paper uses FNO as a backbone to study the Lorenz and KS equation. Experiments show the proposed methods better track the behaviors such as histogram and spectrum. ", "questions": "I have a few questions:\n- Which method (Sinkhore algorithm and contrastive learning) does better? What are their comparative advantages and disadvantages?\n- How does the training time of the Sinkhore algorithm and contrastive learning? I assume they will be slower than the standard supervised learning. It will be great to report the runtime and training time. Would the wasserstein loss be easier to optimize compared to L2 loss? It will be interesting to show a training curve.\n- How does the proposed methods compared to generative model [1, 2]. For example, in Generative adversarial neural operators they also matches the wasserstein distance.\n\n[1] Rahman, Md Ashiqur, et al. \"Generative adversarial neural operators.\" arXiv preprint arXiv:2205.03017 (2022).\n\n[2] Lim, Jae Hyun, et al. \"Score-based diffusion models in function space.\" arXiv preprint arXiv:2302.07400 (2023).", "limitations": "The paper addressed the limitation.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper proposes to learn the dynamics of chaotic system with the optimal transport loss and contractive learning. The proposed methods are more stable compared to the standard MSE/L2 loss, and they can capture the attractor of the system. Experiments show these methods better track the statistical behaviors such as histogram and spectrum. ", "weaknesses": "The paper only studied Lorenz and 1d Kuramoto–Sivashinsky equation. It will be very interesting to study how these methods scale to 2d Navier-Stokes (Kolmogorov Flow) problem.\n\nIt will be interesting to discuss time complexity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!\n\n**W1**:\n> It will be very interesting to… scale to 2d Navier-Stokes…\n\nWe agree that scaling to 2D Navier–Stokes and even higher-dimensional chaotic problems would be an interesting extension to this work. We note, however, that our current results on 1D spatiotemporal chaos already provide high-dimensional chaotic attractors that allow us to validate our approaches and show the benefits of using chosen or learned statistics to train modern deep learning-based emulators. There is no fundamental impediment to extending our method to this setting; it simply requires using a very different backbone in our architecture and more training time than the rebuttal period allows.\n\n**W2**:\n> It will be interesting to discuss time complexity.\n\nIn the rebuttal PDF Figure R1, we document the training time of the models. The OT approach relies on the Sinkhorn algorithm which scales as $O(n^2\\log(n))$ for comparing two distributions of $n$ points each (Theorem 2 in [3]). In our experiments, we use $n = 6000$ to $n = 25,600$ points with no issues, so this approach scales relatively well. The CL approach requires pretraining but is even faster during emulator training since it uses a fixed, pre-trained embedding network.\n\nWe will gladly include a more detailed discussion of time complexity in our final paper.\n\n**Q1**:\n> Which method (Sinkhore algorithm and contrastive learning) does better? What are their comparative advantages and disadvantages?\n\nPlease see the general rebuttal for a detailed discussion of this topic. We will make sure to include this discussion in our final paper.\n\n**Q2**:\n> How does the training time of the Sinkhore algorithm and contrastive learning? I assume they will be slower than the standard supervised learning. It will be great to report the runtime and training time.\n\nWe show training times in the rebuttal PDF Figure R1. CL is faster than Sinkhorn (OT) during emulator training because CL uses a pre-trained embedding function. However, CL does require a pre-training step.\n\n> Would the wasserstein loss be easier to optimize compared to L2 loss? It will be interesting to show a training curve.\n\nThe Sinkhorn loss is a computationally-efficient proxy for the Wasserstein distance. While the Sinkhorn loss is still slower to compute than the L2 loss (rMSE), it appears to converge faster. We include sample training curves in the rebuttal PDF showing that the Sinkhorn loss converges in fewer iterations than rMSE.\n\n**Q3**:\n> How does the proposed methods compared to generative model [1, 2]. For example, in Generative adversarial neural operators they also matches the wasserstein distance.\n\nBoth cited works are generative models, as you point out. However, our problem is not a generative modeling problem. We are not trying to directly sample from the attractor distribution but instead find and use high quality attractor statistics that allow us to train a better emulator for the dynamical system. As the emulator evolves over time, it may be effectively sampling from the attractor distribution, but the sampling itself is not the end goal. Replacing the emulator with a general-purpose generative model does not accomplish our goal of training a model that accurately captures the time dynamics of the system.\n\nTo address the use of the Wasserstein distance in [1], the generative adversarial neural operator (GANO) uses the Wasserstein distance in the same way as the original Wasserstein GAN uses the Wasserstein distance—as a motivation for the adversarial framework for generative modeling. As noted in [1], “while the cost functional… is well defined, showing that the learned measure is indeed an approximation of [the true distribution]... remains an open problem. We address this issue empirically and perform a set of experiments that demonstrate that GANO produces diverse outputs from the data probability measure.” In other words, it is not clear whether GANs really do compute a Wasserstein approximation and we should really treat the connection to OT as general motivation. In our work, we use the Sinkhorn algorithm to solve an entropy-regularized OT problem, which is a very well-understood approximation theoretically, and then use the Sinkhorn-approximated Wasserstein distance to train an emulator for a dynamical system rather than perform generative modeling.\n\n[1] Rahman, Md Ashiqur, et al. \"Generative adversarial neural operators.\" (2022).\n\n[2] Lim, Jae Hyun, et al. \"Score-based diffusion models in function space.\" arXiv preprint (2023).\n\n[3] Dvurechensky, Pavel, et al. “Computational Optimal Transport: Complexity by Accelerated Gradient\nDescent Is Better Than by Sinkhorn’s Algorithm.” (2018).\n"}, {"review_id": "v4jsyscCHs", "reviewer": "Reviewer_fKLU", "summary": "The authors proposed a training framework to preserve preserve invariant measures of chaotic attractors. First, they identify training standard neural operators using MSE on chaotic dynamics does not work. Then they suggested to train neural operators to preserve the invariant measures of chaotic dynamics and the time-invariant statistics. Importantly, they proposed two approaches to train neural operators to preserve those statistics. One is the optimal transport-based approach, and the other is the contrastive learning approach. They  empirically showed that both approaches capture the true invariant statistics. ", "questions": "What is the advantage of preserving the invariant features in chaotic dynamics? Does it really contribute to long-term forcasting, which suppose to be the objective? If it does, I expect the author to show an improvement in long-term forcasting quatitatively by adopting one of the two proposed approches. \n\nThe proposed method is evaluated on the synthetic data, the Lorenz-96 and the Kuramoto–Sivashinsky chaotic dynamics. I am wondering how the model performs on some real world data, e..g, the climate system. ", "limitations": "The authors have addressed the limitations", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and organized. The paper follows a well-defined structure, with a logical flow of ideas throughout.", "weaknesses": "The experiment part could benefit from more substantial content. Comparison with previous methods should be addressed in the experiment part. See below for the two questions as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!\n\n**W1**:\n> The experiment part could benefit from more substantial content. Comparison with previous methods should be addressed in the experiment part. See below for the two questions as well.\n\nPrior methods for training emulators generally use the standard MSE loss which we show as the baseline in our results. We also include a comparison with a previous approach based on the Sobolev norm (Table 2). Our approaches outperform the Sobolev norm method in our noisy chaotic settings.\n\n**Q1.1**:\n> What is the advantage of preserving the invariant features in chaotic dynamics?\n\nPreserving invariant features of chaotic dynamics means matching the long-term statistical behavior of a chaotic dynamical system. For example, weather models may only be able to predict the exact weather up to two weeks ahead of time. In fact, it is impossible to predict the precise state of the Earth on the time scales (years) required for climate modeling [1], but we would still like to predict important statistical features of the climate, e.g. the average number of hurricanes per year. It is these kinds of statistical features of the chaotic dynamics that we want to preserve for long-term forecasting.\n\n**Q1.2**:\n> Does it really contribute to long-term forecasting, which suppose to be the objective?\n\nYes! Any long-term prediction for the exact state of the system diverges exponentially from the ground truth due to chaos, so the important quantities to consider for long-term forecasting are precisely the statistical features of the dynamics. By matching these statistical features, we are able to train an emulator that truly captures the chaotic dynamics and correctly models the only predictable aspects of the chaotic system.\n\nAnother important point here is that short-term prediction performance does not necessarily imply long-term performance in terms of statistics and model stability. An accurate weather model run for a longer time will not necessarily produce high-quality climate statistics, just as an emulator trained only on short-term MSE may fail to capture the true long-term behavior of the system. The goal of our work is to help rectify this problem by proposing methods that directly tackle long-term forecasting for chaotic systems.\n\n**Q1.3**:\n> If it does, I expect the author to show an improvement in long-term forecasting quantitatively by adopting one of the two proposed approaches.\n\nOur results *do* show a significant improvement (which is often even visually evident in the sample predictions) over the standard MSE-based training method on the relevant long-term statistical behavior. In particular, our metrics include the distributions of key summary statistics and the energy spectrum, both of which represent long-term statistical characterizations of the chaotic system. Our trained deterministic emulators provide much higher quality instances of the system dynamics, which sample the chaotic attractor, and this is the best you can hope for when performing long-term forecasts of chaotic dynamics. The new experiments (see general rebuttal) provide additional evidence for this using new metrics.\n\nThe final paper will better emphasize these points.\n\n**Q2**:\n> The proposed method is evaluated on the synthetic data, the Lorenz-96 and the Kuramoto–Sivashinsky chaotic dynamics. I am wondering how the model performs on some real world data, e..g, the climate system.\n\nWe are also excited to see how our approaches would scale for climate modeling. However, this is outside the scope of this work. Existing climate emulators, such as FourCastNet [2] and ClimaX [3], are large models that cost millions to train, tune, and validate.\n\n[1] What might we learn from climate forecasts? Smith, Leonard A. (2002).\n\n[2] FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators. Pathak, Jaideep et al. (2022).\n\n[3] ClimaX: A foundation model for weather and climate. Nguyen, Tung et al. (2023).\n"}, {"review_id": "YygZdVi4tb", "reviewer": "Reviewer_DhyB", "summary": "This paper proposes two methods to preserve invariant measures of chaotic systems in the multi-environment setting when training neural operators. Given some expert knowledge of the underlying dynamical systems, they propose a new optimal transport loss, which uses this knowledge to match the statistics. Without expert knowledge, they use a contrastive learning approach to learn invariant statistics, which are used to construct the feature loss to preserve these invariant statistics. These two methods show better performance to reproduce the invariant statistics of noisy chaotic systems.", "questions": "1. Is it possible to provide the problem setting, limitations, and advantages of the two methods more clearly, and clearly explain the difference between the two methods, instead of introducing them together?\n2. Whether the performance can be maintained if more limited prior knowledge is utilized for the optimal transport-based approach?\n3. If all those equation constraints have to be used to ensure performance, should it be compared with methods such as physical informed neural operator that utilize similar information instead of the conventional neural operator?\n4. As for the contrastive learning approach, what kind of diversity does the data need to meet to ensure better performance? Could you provide an intuitive explanation?\n", "limitations": "The authors have adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. Long-time prediction of chaotic systems is different, this paper proposed two methods to better represent chaotic systems using neural operators by preserving the invariant statistics of chaotic attractors.\n2. The methods are robust to noise and chaos.", "weaknesses": "1. These two methods are aimed at different problem settings: the optimal transport-based approach needs prior knowledge, while the contrastive learning approach does not; the optimal transport-based approach seems not to depend on the multi-environment setting, whereas the contrastive learning approach does depend on this setting.  These two methods are relatively separate, and the relationship between them is not close.\n2. The two methods are mixed together, so the structure of the paper does not look clear.\n3. It seems that the optimal transport-based approach requires too many equation constraints. For example, in the two examples given in the paper, almost all of the terms in the underlying dynamical equations are used. \n4. As for the contrastive learning approach, higher requirements are put forward for the diversity of data.\n5. The introduction of the Noise Contrastive Estimation (InfoNCE) loss is not described clearly in the main text.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!\n\n**W1**:\n> These two methods are… relatively separate…\n\nWhile the two proposed methods use different machine learning tools, they are very much related by a shared goal and overall approach: to train emulators to capture long-term chaotic behavior by preserving the invariant measure of the attractor. In fact, both methods go about this goal in conceptually similar ways, by matching the statistics learned by the emulator to statistics from the data. The primary difference is how those statistics are identified, whether it is from prior domain knowledge or learned directly from the data. Our contributions consist of both identifying and formulating the shared problem (Section 2) and then proposing two approaches for solving this problem (Sections 3.1 and 3.2). The fact that contrastive learning (CL) can preserve chaotic attractors nearly as well as methods that use prior knowledge is remarkable.\n\nThe multi-environment setting provides an avenue for identifying the relevant statistics for characterizing chaotic attractors. The CL approach uses multiple environments explicitly during pre-training, while the OT approach uses this information implicitly as a part of the prior knowledge used for choosing informative summary statistics. In other words, the best summary statistics for training an emulator are precisely those that can tell the difference between environments/attractors and therefore provide a signal to push the emulator to match the correct attractor.\n\nThe final paper will better emphasize these points.\n\n**W2**:\n> The two methods are mixed together…\n\nThe two methods are presented in separate sections: physics-informed OT (Section 3.1) and unsupervised CL (Section 3.2). Both represent approaches for training an emulator to preserve the invariant measure of the attractor—the problem formulated in Section 2.\n\n**W3**:\n> It seems that the optimal transport-based approach requires too many equation constraints…\n\nWe agree! This was the motivation for developing the CL approach. The OT approach is precisely designed to use existing domain knowledge to choose informative summary statistics and train a better emulator. That said, we are interested in better understanding how the quality and quantity of the chosen statistics influence the trained emulator. To better study this effect, we have performed additional experiments (Table R2 in the rebuttal PDF) with a reduced set of summary statistics including using a minimally informative statistic to demonstrate when the OT approach begins to fail due to a poor choice of statistic.\n\n**W4**:\n> As for the contrastive learning approach, higher requirements… for the diversity of data.\n\nThe multi-environment setting is a very natural setting for many scientific and engineering applications, where different measured trajectories often have varying parameters due to environmental noise or varying control inputs. It also presents a more challenging generalization problem than the single-environment setting. As with any unsupervised representation learning method, CL requires sufficient data diversity, which in this case comes directly from the multi-environment setting. Please also see our response below to Q4.\n\n**W5**:\n> The introduction of the Noise Contrastive Estimation (InfoNCE) loss is not described clearly in the main text.\n\nWe introduce InfoNCE in Section 3.2 eq. 15, cite its origin as a contrastive learning loss, and explain the intuition for its use and its relevance to our problem. \n\n**Q1**:\n> Is it possible to provide the problem setting, limitations, and advantages…\n\nThe two methods share a problem setting as presented in Section 2. Please see the general rebuttal for additional discussion comparing the approaches.\n\n**Q2**\n> Whether the performance can be maintained if more limited prior knowledge…\n\nIn our updated results (Table R2), we show the effect of using a smaller set of summary statistics as well as using a minimally informative statistic for the OT approach. You can often get away with using a fairly limited set of summary statistics if the chosen statistics are highly informative. If there is a real lack of prior knowledge, then our CL approach offers a very compelling alternative.\n\n**Q3**\n> If all those equation constraints have to be used to ensure performance, should it be compared with methods such as physical informed neural operator…\n\nWe believe the reviewer means statistics instead of constraints since we impose no constraints. The OT approach matches summary statistics of the model outputs to those of the data and never treats those statistics as terms in a governing equation. Another important distinction to make here is that methods such as the physics-informed neural operator assume a *known* PDE and then use a neural operator to fit the solution, which can be done even without training data. We are using the neural operator to learn a time evolution operator directly from data without knowing the PDE. (We agree that assuming we know what statistics to preserve with the OT method can be problematic, so we developed the CL method to address this.)\n\n**Q4**:\n> As for the contrastive learning approach, what kind of diversity does the data need to meet to ensure better performance?...\n\nIntuitively, CL uses the data diversity provided by the multi-environment setting to choose invariant statistics by identifying which statistics are informative for distinguishing between the various attractors seen in the data. Empirically, our experiments show that randomly varying just one or a few system parameters is enough to provide the necessary data diversity for CL. In our updated results (Table R3), we show a new experiment that further reduces the range of the only varying parameter F in the Lorenz-96 dataset from [10, 18] to [16, 18]. Even with such minimal diversity, we still see good performance using our CL approach.\n"}, {"review_id": "1HjnLl9SFl", "reviewer": "Reviewer_tYBd", "summary": "This paper mentions that MSE loss is insufficient for learning the solution operator for chaotic dynamics. Two approaches are proposed to solve the problem. One is to use sinkhorn loss, and the other is to use contrastive learning. The authors evaluate the effectiveness of the proposed methods through experiments on several physical systems.", "questions": "- I understood that the prediction in equation (3) is the same as the autoregressive PDE solver proposed in [R1]. It would be good to add such previous studies.\n\n- Since the autoregressive approach performs forecasting iteratively, it seems to me that the numerical error increases in the long-term forecasting setting. Could you add a discussion on this point? Also, is it possible to forecast at arbitrary time intervals $\\Delta t$ during training and testing?\n\n- I think it would be easier to read if the description of the previous studies (optimal transport and contrastive learning) and the proposed method are separated. In section 3, I feel that the novelty of this paper is not described clearly.\n\n- There are other ways (e.g., Maximum Mean Discrepancy) to measure the distance of a distribution besides optimal transport. It would be good to clarify the motivation for using optimal transport, and the same goes for contrastive learning.\n\n[R1] Johannes Brandstetter, Daniel E. Worrall, Max Welling, Message Passing Neural PDE Solvers, ICLR, 2022.", "limitations": "It would be good to discuss in detail the advantages and disadvantages of using Optimal transport and contrastive learning.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The problem of learning solution operators for chaotic dynamics is a challenging and interesting problem, and the authors have an interesting point of view (preserving invariant measures).\n- To my knowledge, the introduction of physics-informed optimal transport and contrastive learning is a new approach.\n", "weaknesses": "- Although the problem setting is important and interesting, the proposed method seems to be a straightforward combination of existing technologies.\n- There is room for improvement in the writing of the paper.\n- There is no reference to previous studies on autoregressive approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!\n\n**W1**:\n> Although the problem setting is important and interesting, the proposed method seems to be a straightforward combination…\n\nOur proposed approaches are not a simple reshuffling of methods to solve a variant of the originally proposed problem (e.g. general representation learning) but rather a completely new use of contrastive learning (CL) and optimal transport (OT) to address a very different problem: training emulators of chaotic systems with good long-term behavior. Our contributions include both framing the problem and then developing new and adapting existing methods to handle this new problem domain.\n\nOur application of CL and OT is deeply motivated by our problem formulation, without which it would be impossible to justify their use in this context. Furthermore, adapting CL for this application is highly non-trivial since it is well outside of its traditional use as a general representation learning method. In the paper, we argue that CL has the right properties to learn exactly the kind of time-invariant statistics that we aim to emulate. We then show that the learned latent space can be adapted into a loss function to train emulators that successfully capture long-term behavior.\n\nThe final paper will emphasize these points.\n\n**W2**:\n> There is room for improvement in the writing of the paper.\n\nPlease let us know if there are any specific writing issues or sections in need of clarification.\n\n**W3/Q1**:\n> There is no reference to previous studies on autoregressive approach.\n> I understood that the prediction in equation (3) is the same as the autoregressive PDE solver… [R1].\n\nWe would be happy to include additional references on emulators. Our work focuses on developing regularizers for training emulators on chaotic systems and not on autoregressive training, and we do not claim autoregressive training to be a novel contribution here.\n\n**Q2**:\n> Since the autoregressive approach performs forecasting iteratively, it seems to me that the numerical error increases in the long-term forecasting setting.\n\nThe key is that, for chaotic dynamics, it does not matter what kind of emulator you choose (autoregressive or not), you will never be able to exactly predict the state of the system in the long term. That is, for any emulator, the mean squared error for long-term predictions will exponentially increase with time. That is why we focus on training emulators that capture long-term statistical behavior, which is the relevant measure of long-term forecasting performance for chaotic systems.\n\n> Also, is it possible to forecast at arbitrary time intervals…\n\nOur emulator does perform 1-step-ahead forecasting for a specific $\\Delta_t$ which can be applied iteratively, but forecast accuracy at a specific time point is not the focus of our work (and indeed an impossible task for chaotic systems). A good analog of our goal would be a climate model which is unable to make an exact prediction for an arbitrary time point but does provide good statistical insights into long-term trends such as the frequency and intensity of hurricanes.\n\n**Q3**:\n> I think it would be easier to read if… the previous studies… and the proposed method are separated.\n\nIn the paper, previous studies are presented in the related work (Section 1.2) and the two proposed approaches are presented in separate sections: physics-informed OT (Section 3.1) and unsupervised CL (Section 3.2).\n\n> In section 3, I feel that the novelty of this paper is not described clearly.\n\nSection 3 presents our approaches to solving the problem formulated in Section 2. The novelty of the work is discussed in our contributions statement (Section 1.1). To reiterate, we show that the standard method for training emulators is insufficient for modeling chaotic dynamics and instead propose to train the emulator to preserve the invariant measure of the chaotic attractor. With this new problem formulation, we propose two new approaches for training the emulator to match the attractor statistics.\n\n**Q4**:\n> There are other ways (e.g., MMD) to measure the distance of a distribution… It would be good to clarify the motivation…\n\nThere are many ways of measuring differences between distributions, including point-wise divergences such as KL, moment-based distances such as MMD, as well as the Sinkhorn distance from optimal transport—which we use in the paper. Using MMD would result in a similar method to our OT approach and act as a drop-in replacement for the Sinkhorn distance for matching the distributions of summary statistics. In fact, we have included a new experiment where we train an emulator using a Gaussian-kernel MMD loss (Table R4). We find that MMD, even after careful hyperparameter tuning, does not perform as well as the Sinkhorn loss.\n\nWe chose the Sinkhorn loss because it has many well-known advantages over point-wise divergences and kernel MMD distances: e.g., for two distributions that do not have overlapping support, the Wasserstein distance underlying the Sinkhorn loss still provides a high-quality gradient signal that makes it useful for optimization. MMD also requires many hyperparameter choices, including the choice of kernel and kernel-specific parameters, that have a significant effect on the properties of the final distance measure.\n\nContrastive learning, on the other hand, allows us to learn a set of invariant statistics from the data. This provides a method that does not rely on prior domain knowledge about the chaotic attractor and instead uses the multi-environment setting to automatically learn informative statistics. This motivation is further discussed in the introduction (Sec. 1), when we propose the method (Sec. 3.2), and in the discussion (Sec. 5).\n\n[R1] Johannes Brandstetter, Daniel E. Worrall, Max Welling, Message Passing Neural PDE Solvers, ICLR, 2022.\n\n**L1**:\n> It would be good to discuss in detail the advantages and disadvantages…\n\nPlease see the discussion in the general rebuttal."}]

Author rebuttal:
We would like to first thank the reviewers and ACs for their helpful comments and questions. We are happy to see that the reviewers generally appreciate the importance of the problem of training better emulators for chaotic dynamics and find the proposed methods and ideas well-motivated. In our response, we provide new experiments, evaluation metrics, and additional discussion to help clarify the properties of our two proposed approaches.

# Comparison of optimal transport and contrastive learning approaches

Several reviewers asked for a more detailed comparison between the two proposed approaches: physics-informed optimal transport (OT) and contrastive learning (CL). **Ultimately, both methods have strong performance according to a variety of metrics, but the CL approach requires no prior physical knowledge and is faster. Both approaches have significant advantages over the classical approach of minimizing rMSE, which fails to preserve important statistical characteristics of the system.**

We present both approaches here because they are both methods for encouraging the emulator to capture the long-term statistical behavior of the chaotic system. In fact, both approaches work by matching statistics computed from the emulator to statistics from the data. The primary conceptual difference between them is that the optimal transport approach takes advantage of prior domain knowledge about the setting to choose informative summary statistics, while the contrastive learning approach learns informative invariant statistics from the data alone.

## Physics-informed OT

### Advantages:
* If prior domain knowledge is available, OT can make good use of it and often performs better than CL when given highly informative summary statistics.
* OT does not require pre-training to learn a distance measure.
### Disadvantages:
* Without domain knowledge, an arbitrary uninformative summary statistic will often have no performance benefit.
* Because OT requires estimating the Wasserstein distance ($O(n^2\log(n))$ operations) for each training step, it is slower than the CL approach during emulator training. However, the scaling behavior is still reasonable, and the computation can be made even faster via sub-sampling.
## Unsupervised CL

### Advantages:
* CL does not require any prior domain knowledge and instead learns informative statistics directly from the data.
* CL is faster than OT during emulator training since it uses a pre-trained encoder for invariant statistics rather than computing distributional distances.
### Disadvantages:
* As an unsupervised representation learning method, CL requires a pre-training step to learn the informative statistics, which are later used for emulator training. *However, this fixed one-time cost is easily amortized.*
* In principle, CL relies on the data diversity of the multi-environment setting, although our new experiments show that we can still obtain good results even with a very minimal diversity of environments. *Also, the multi-environment setting is common in many scientific and engineering applications due to environmental noise or tunable control parameters, so this data diversity is often “free”.*
We have an existing discussion comparing the two approaches in the submitted supplement, and we will include a more detailed discussion in the final paper.

# New experiments and evaluation metrics

We have performed several new experiments that act as additional points of comparison and help us better understand the behavior of our methods under a variety of conditions:

1. For our OT approach (which uses the Sinkhorn loss), we test a reduced set of summary statistics, which shows how the quality of the summary statistic affects the performance of the method (Table R2). With an informative summary statistic, we find even a reduced set can still be helpful but, for a non-informative statistic, the OT method fails as expected.
2. For our CL approach, we test a multi-environment setting with reduced data diversity and find that the contrastive method still performs well under the reduced conditions (Table R3), which demonstrates robustness.
3. We also implement a variant of our OT approach that uses Maximum Mean Discrepancy (MMD) as a distributional distance rather than the Wasserstein distance. Using the same set of summary statistics, we find that MMD does not perform as well as the Wasserstein distance for training emulators (Table R4).

As suggested by the reviewers, we also add new evaluation metrics (Table R1) in the form of:

1. the leading Lyapunov exponent (LLE)—a dynamical quantity that measures how quickly the chaotic system becomes unpredictable; and
2. the fractal dimension—a characterization of the dimension of the attractor.

For the LLE, we report the relative absolute error $|\hat\lambda - \lambda|/|\lambda|$ between the model and the ground truth averaged over the test set. For the fractal dimension, we report the absolute error $|\hat D - D|$ between the estimated fractal dimension of the model and the ground truth averaged over the test set.

We generally find in high-noise settings that both of our approaches give lower LLE errors than an emulator trained on rMSE loss alone, and the contrastive approach has the lowest LLE error. This is despite the fact that our method is only encouraged to match invariant statistics of the final attractor rather than dynamical quantities. We also see evidence that the fractal dimension of the CL approach is closer to the true attractor. However, we note that the fractal dimension is difficult to reliably estimate for high-dimensional chaotic attractors [1].

The new experiments provide useful insights for future applications of our approaches, and the new evaluation metrics shed additional light on the features of chaos preserved by our trained emulators.

[1] Impracticality of a box-counting algorithm for calculating the dimensionality of strange attractors. Greenside, H. S. et al, Phys. Rev. A (1982).

Meta-review: Although some reviewers have some technical concerns such as the weekness of the empirical evaluations in the first review, the authors have provided appropriate feedback in the discussion phase and addressed these concerns adequately. Finally, all reviewers agree that the paper has sufficient merit in the discussed problem. Therefore, I recommend acceptance of this paper for this submission.
Decision: Accept (poster)