Title: Training neural operators to preserve invariant measures of chaotic attractors
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 Introduction

Training fast and accurate surrogate models to emulate complex dynamical systems is key to many scientific applications of machine learning, including climate modeling [1], fluid dynamics [2, 3], plasma physics [4], and molecular dynamics [5, 6]. Fast emulators are 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).

One key feature of many of the dynamical systems in these applications is chaos, which is characterized 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 to 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]-properties 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 a long time horizon, which is severely limited by the chaotic dynamics. 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-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 a natural invariant measure on a chaotic attractor that characterizes the time-invariant statistical properties of the dynamics on the attractor [13]. By training a neural operator to preserve this invariant measure-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-based 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 statistics 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-environment 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-environment setting because it requires emulators to generalize over a broader range domain. Most practical use cases for emulators-where the computational costs of generating training data and training an emulator are outweighed by the computational gains at deployment-are in the multi-environment 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 RMSEnamely, 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 a poor signal for training a neural operator. We instead suggest training neural operators to preserve the invariant measures of chaotic attractors and the corresponding time-invariant statistics, which are robust to the combination of noise and chaos. Secondly, we propose a direct optimal transport-based approach to train neural operators to preserve the distribution of a chosen set of summary statistics. Specifically, we use a Sinkhorn 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 a 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.

1.2 Related work

Neural operators. The goal of a 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, 1619]. It 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 time-invariant 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 a 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 an auxiliary input to the neural operator. Like our approach, Dyad learns time-invariant 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 time-invariant features to determine a loss function that can be applied to any neural operator architecture without changes to either the backbone or the inputs.

Long-term predictions. Lu et al. [21] try to 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 a Sobolev norm and dissipative regularization to 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 statistics. 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 a multi-environment setup for chaotic systems, our goal is to learn a neural operator capable of modeling long-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 to model different attractors.

Prior [22] and concurrent work [23] on 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 lowdimensional 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-dimensional attractor. They compute the Wasserstein distance by turning the nonlinear dynamics in the original state space into a linear PDE in the space of distributions and then solving a PDE-constrained 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 FokkerPlanck 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, it is again challenging to scale this method to high-dimensional systems. In contrast, our proposed optimal transport approach focuses on fitting deep learning-based neural operators to high-dimensional chaotic

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Noise scale

r

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Histogram error of invariant statistics

U

G

(

u

0

+

η

) +

η

RMSE,

U

G

(

u

0

+

η

) +

η

Sobolev error,

U

G

(

u

0

+

η

) +

η

U

G

(

u

0

+

η

)

RMSE,

U

G

(

u

0

+

η

)

Sobolev error,

U

G

(

u

0

+

η

)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

RMSE

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

Sobolev error

0

20

40

x

Truth

0

20

40

x

RMSE

t

0

20

40

x

CL (ours)

0

20

40

60

80

100

120

Wavenumber

10

2

10

3

10

4

Energy

Truth

RMSE

OT (ours)

CL (ours)

dynamical systems by first choosing a set of physically relevant summary statistics and then using the computationally efficient Sinkhorn algorithm [31] to compute our optimal transport loss-an entropy-regularized approximation for the (squared) Wasserstein distance-on the distribution of the summary statistics.

Feature loss and contrastive learning. Contrastive learning [32-34] charges an 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 to motivate the encoder to 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 et al. [36] computes MSE on deep features of the classification models as a structural distance metric to train generative models. Tian et al. [37] propose the encoder of contrastive learning as an unsupervised alternative for calculating feature loss.

2 Problem Formulation

Consider a dynamical system with state space u ∈ U governed by

where G is the governing function, and ϕ ∈ Φ is a set of parameters that specify the dynamics. We will denote trajectories governed by these dynamics (1) as u ( t ) , particular points on the trajectory as u t ∈ U , and a sequence of K +1 consecutive time points on the trajectory as U I : I + K := { u t i } I + K i = I . Our experiments use equally spaced time points with a time step ∆ t . We refer to settings where ϕ varies as multi-environment settings .

Our goal is to learn an emulator in a multi-environment setting that approximates the dynamics (1) with a data-driven neural operator ˆ g θ : U × Φ →U that makes discrete predictions

where θ are the parameters of the neural operator ˆ g θ . In the multi-environment setting, we will have data from a variety of environments n ∈ { 1 , 2 , . . . , N } , and thus our training data { U ( n ) 0: L , ϕ ( n ) } N n =1 consists of trajectories U ( n ) 0: L and the corresponding environment parameters ϕ ( n ) .

Predicted sequence notation. We will denote a sequence of h +1 autonomously predicted states (using the neural operator ˆ g θ ) with an initial condition u ( n ) t I as

We will often use a concatenated sequence (with total length K +1 ) of these autonomous prediction sequences (each with length h +1 ), which we will denote as

where ⊕ is concatenation.

Chaotic dynamical systems and invariant measures of chaotic attractors. We focus on chaotic dynamical systems that have one or more chaotic attractors, which exhibit a sensitive dependence on initial conditions characterized by a positive Lyapunov exponent [13]. For each chaotic attractor A , we can construct a natural invariant measure (also known as a physical measure)

where δ u ( t ) is the Dirac measure centered on a trajectory u A ( t ) that is in the basin of attraction of A [13]. Note that, because A is an attractor, any trajectory u A ( t ) in the basin of attraction of A will give the same invariant measure µ A , i.e. the dynamics are ergodic on A . Therefore, any time-invariant statistical property S A of the dynamics on A can be written as

for some function s ( u ) and a trajectory u A ( t ) in the basin of attraction of A . Conversely, for any s ( u ) , (6) gives a time-invariant property S A of the dynamics.

In this work, we assume each sampled trajectory in the data is from a chaotic attractor and therefore, has time-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

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 p η )

and are therefore given by

which can be a good approximation for S A = E µ A [ s ] even in noisy conditions and does not suffer from the exponential divergence of RMSE (Fig. 1).

̂

̂

̂

̂

(1)

(2)

(3)

ℓ

RMSE

=

1

K

+1

K

∑

t

=0

∥

U

(

n

)

I

+

t

-

U

(

n

)

I

+

t

∥

2

/∥

U

(

n

)

I

+

t

∥

2

ℓ

OT

=

ℓ

OT

(

S

(

n

)

I

:

I

+

K

,

S

(

n

)

I

:

I

+

K

)

ℓ

CL

=

∑

l

cos

(

f

l

ψ

(

U

(

n

)

I

:

I

+

K

)

,

f

l

ψ

(

U

(

n

)

I

:

I

+

K

)

)

f

ψ

shared

weights

f

ψ

f

ψ

Attract

Random Crops

U

(

n

)

0:

L

{

U

(

m

)

0:

L

}

N

m

=1,

m

≠

n

Repulse

Hyperspherical

Embedding

(c) Unsupervised

Contrastive Learning

(b) Optimal Transport

U

(

n

)

I

:

I

+

K

U

(

n

)

I

:

I

+

K

g

θ

(

u

(

n

)

t

I

,

ϕ

(

n

)

)

U

(

n

)

I

:

I

+

K

(a) Neural Operator

Training

U

(

n

)

I

+1:

I

+

K

Training Loss

S

(

n

)

I

:

I

+

K

S

(

n

)

I

:

I

+

K

distribution

matching

du

u

?u

?u

du

1.x

dxdt

Fixed

dr2

Summary

statistic

3 Proposed Approaches

3.1 Physics-informed optimal transport

We propose an optimal transport-based loss function to match the distributions of a set of summary statistics s ( u ) = [ s (1) ( u ) , s (2) ( u ) , . . . , s ( k ) ( u ) ] between the data s i = s ( u t i ) ∼ µ s and the model predictions ˆ s j = s (ˆ u t j ) ∼ µ ˆ s . Here, µ s and µ ˆ s are the probability measures of the summary statistics for the true chaotic attractor A and the attractor ˆ A learned by the neural operator, respectively. Specifically, µ s ( s ′ ) = ∫ δ s ( u ) -s ′ µ A ( u ) d u and µ ˆ s (ˆ s ′ ) = ∫ δ s (ˆ u ) -ˆ s ′ µ ˆ A (ˆ u ) dˆ u . By choosing the set of summary statistics using expert domain knowledge, this approach allows us to directly guide the neural operator toward preserving important physical properties of the system.

Optimal transport-Wasserstein distance. We match the distributions of summary statistics using the Wasserstein distance W ( µ s , µ ˆ s ) -a distance metric between probability measures defined in terms of the solution to an optimal transport problem:

We use a quadratic cost function c ( s , ˆ s ) = 1 2 ∥ s -ˆ s ∥ 2 (i.e., W is the 2-Wasserstein distance), and the map π ∈ Π( µ s , µ ˆ s ) must be a valid transport map between µ s and µ ˆ s (i.e., a joint distribution for s , ˆ s with marginals that match µ s and µ ˆ s ) [38].

In the discrete setting with distributions represented by samples S = { s i } L i =1 and ˆ S = { ˆ s j } L j =1 , the Wasserstein distance is given by

where the cost matrix C ij = 1 2 ∥ s i -ˆ s j ∥ 2 . The set of valid discrete transport maps Π consists of all matrices T such that ∀ i, j , T ij ≥ 0 , ∑ j T ij = 1 , and ∑ i T ij = 1 .

Entropy-regularized optimal transport-Sinkhorn 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 in a convex relaxation of the original optimal transport problem:

where h ( T ) = -∑ i,j T ij log T ij is the entropy of the transport map. This entropy-regularized optimal transport problem can be solved efficiently using the Sinkhorn algorithm [31].

As γ → 0 , the entropy-regularized Wasserstein distance W γ → W 0 = W reduces to the exact Wasserstein distance (11). For γ > 0 , we can further correct for an entropic bias to obtain the Sinkhorn divergence [39, 40]

which gives us our optimal transport loss. Combined with relative root mean squared error (RMSE)

for short-term prediction consistency [2, 16, 18], our final loss function is

where S ( n ) I : I + K := { s ( u ) | u ∈ U ( n ) I,I + K } , ˆ S ( n ) I : h : I + K := { s (ˆ u ) | ˆ u ∈ ˆ U ( n ) I : h : I + K } , and α > 0 is a hyperparameter. Note that ˆ S ( n ) I : h : I + K and ˆ U ( n ) I : h RMSE : I + K implicitly depend on weights θ .

3.2 Contrastive feature learning

When there is an absence of prior knowledge pertaining to the underlying dynamical system, or when the statistical attributes are not easily differentiable, we propose an alternative contrastive learning-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 to train an encoder f ψ ( U ) (with parameters ψ ) to compute relevant invariant statistics of the dynamics from sequences U with fixed length K + 1 . We do not explicitly train on the environment parameters ϕ but rather use a general-purpose contrastive learning approach that encourages the encoder f ψ to learn a variety of time-invariant features that are able to distinguish between sequences from different trajectories (and therefore different ϕ ).

A 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 as 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 ) I : I + K , U ( n ) J : J + K from the same trajectory U ( n ) 0: L both sample the same chaotic attractor, i.e. their statistics should be similar, so we treat any such pair of sequences as positive pairs. Two sequences U ( n ) I : I + K , U ( m ) H : H + K from different trajectories are treated as negative pairs. This allows us to formulate the InfoNCE loss as:

̸

The term in the numerator enforces alignment of the positive pairs which ensures we obtain timeinvariant statistics, while the term in 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 a series of features of the encoder network f ψ [42]:

where f l ψ gives the output the l -th layer of the neural network. The combined loss that we use for training the neural operator is given by

where λ > 0 is a hyperparameter.

4 Experiments

We evaluate our approach on the 1D chaotic Kuramoto-Sivanshinsky (KS) system and a finitedimensional 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://github.com/roxie62/neural_operators_for_chaos .

Experimental setup. Our data consists of noisy observations u ( t ) with noise η ∼ N (0 , r 2 σ 2 I ) , where σ 2 is the temporal variance of the trajectory and r is a scaling factor. Baselines. We primarily consider the baseline as training with RMSE [2]. We have additional baselines in Appendix B, including Gaussian denoising and a 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 metrics and other measures that characterize the chaotic attractor. See Appendix C.1 for details.

4.1 Lorenz-96

As is a 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

Its dynamics exhibit strong energy-conserving non-linearity, and for a large F ≥ 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 i dt , ( u i +1 -u i -2 ) u i -1 , u i } . We generate 2000 training data points with each ϕ ( n ) randomly sampled from a 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 are presented in Table 1, and predictions and invariant statistics are shown in Fig. 3 (refer to C.4 for more visualizations).

4.2 Kuramoto-Sivashinsky

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 u ∂u ∂x , models dispersion effects through ∂ 2 u ∂x 2 , and manages discontinuities by introducing hyper-viscosity via ∂ 4 u ∂x 4 :

0

25

50

75

0

20

40

x

Sampled trajectory

Truth

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

Histogram 2d

u i

Histogram 2d

0

25

50

75

0

20

40

x

RMSE

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

25

50

75

0

20

40

x

OT (ours)

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

25

50

75

t

0

20

40

x

CL (ours)

u i

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

du i

/dt i

u i

(a) Lorenz-96

0

50

100

0

50

100

x

Sampled trajectory

∂u/∂x

Histogram 2d

∂u/∂x

Histogram 2d

0

50

100

0

50

100

x

∂u/∂x

∂u/∂x

0

50

100

0

50

100

x

∂u/∂x

∂u/∂x

0

50

100

t

0

50

100

x

∂u/∂t

∂u/∂x

∂

2

u/∂x

2

∂u/∂x

Experimental setup. For the KS system, we define s ( u ) := { ∂u ∂t , ∂u ∂x , ∂ 2 u ∂x 2 } [47]. We generate 2000 training data points, with each ϕ ( n ) 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 2 and visualize the predictions and invariant statistics in Fig. 3.

5 Discussion and Limitations

Wehave demonstrated two approaches for training neural operators on chaotic dynamics by preserving the invariant measures of the chaotic attractors. The optimal transport approach uses knowledge of the underlying system to construct a physics-informed objective that matches relevant summary statistics, while the contrastive learning approach uses a multi-environment setting to directly learn relevant invariant statistics. In both cases, we see a significant improvement in the ability of the emulator to reproduce the invariant statistics of chaotic attractors from noisy data when compared with traditional RMSE-only training that focuses on short-term forecasting.

Wealso find, in high-noise settings, that both of 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. We also see 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 to reliably estimate for high-dimensional chaotic attractors [48].

Limitations. Because we rely on invariant measures, our current approach is limited to trajectory data from attractors, i.e. we assume that the dynamics have reached an attractor and are not in a transient phase. We also cannot handle explicit time dependence, including time-dependent 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 statistics also depends on the diversity of the environments present in the multi-environment 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 a wide range of applications, we foresee no direct negative societal impacts.

Acknowledgments and Disclosure of Funding

The authors gratefully acknowledge the support of DOE grant DE-SC0022232, AFOSR grant FA955018-1-0166, and NSF grants DMS-2023109 and DMS-1925101. In addition, Peter Y. Lu gratefully acknowledges the support of the Eric and Wendy Schmidt AI in Science Postdoctoral Fellowship, a Schmidt Futures program.

References

Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. Fourcastnet: A global datadriven high-resolution weather model using adaptive fourier neural operators, 2022.

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 on Learning Representations , 2021. URL https://openreview.net/forum?id=c8P9NQVtmnO .

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.

A. Mathews, M. Francisquez, J. W. Hughes, D. R. Hatch, B. Zhu, and B. N. Rogers. Uncovering turbulent plasma dynamics via deep learning from partial observations. Phys. Rev. E , 104: 025205, Aug 2021. doi: 10.1103/PhysRevE.104.025205. URL https://link.aps.org/doi/ 10.1103/PhysRevE.104.025205 .

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-10186, Aug 2021. ISSN 0009-2665. doi: 10.1021/acs.chemrev.0c01111.

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.1038/s41467-023-36329-y.

Yuming Chen, Daniel Sanz-Alonso, and Rebecca Willett. Autodifferentiable ensemble kalman filters. SIAM Journal on Mathematics of Data Science , 4(2):801-833, 2022.

Peter Y. Lu, Samuel Kim, and Marin Soljaˇ ci´ c. Extracting interpretable physical parameters from spatiotemporal systems using unsupervised learning. Phys. Rev. X , 10:031056, Sep 2020. doi: 10.1103/PhysRevX.10.031056. URL https://link.aps.org/doi/10.1103/PhysRevX.10. 031056 .

Rui Wang, Robin Walters, and Rose Yu. Meta-learning dynamics forecasting using task inference. Advances in Neural Information Processing Systems , 35:21640-21653, 2022.

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 in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=6scShPCpdDu .

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-229, Mar 2021. ISSN 2522-5839. doi: 10.1038/ s42256-021-00302-5.

Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. arXiv preprint arXiv:2205.13671 , 2022.

Alfredo Medio and Marji Lines. Nonlinear Dynamics: A Primer . Cambridge University Press, Cambridge, 2001. doi: 10.1017/CBO9780511754050.

Peter Davidson. Turbulence: An Introduction for Scientists and Engineers . Oxford University Press, 06 2015. ISBN 9780198722588. doi: 10.1093/acprof:oso/9780198722588.001.0001.

J. R. Dorfman. An Introduction to Chaos in Nonequilibrium Statistical Mechanics . Cambridge Lecture Notes in Physics. Cambridge University Press, 1999. doi: 10.1017/ CBO9780511628870.

Zongyi Li, Miguel Liu-Schiaffini, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, 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 https://openreview.net/forum?id=1C36tFZn7sR .

Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural pde solvers. In International Conference on Machine Learning , pages 2241-2256. PMLR, 2022.

Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatiotemporal-scale generalized pde modeling. arXiv preprint arXiv:2209.15616 , 2022.

Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations , 2022. URL https:// openreview.net/forum?id=vSix3HPYKSU .

Matthieu Kirchmeyer, Yuan Yin, Jérémie Donà, Nicolas Baskiotis, Alain Rakotomamonjy, and Patrick Gallinari. Generalizing to new physical systems via context-informed dynamics model. In International Conference on Machine Learning , pages 11283-11301. PMLR, 2022.

Fei Lu, Kevin K Lin, and Alexandre J Chorin. Data-based stochastic model reduction for the kuramoto-sivashinsky equation. Physica D: Nonlinear Phenomena , 340:46-57, 2017.

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://proceedings.neurips.cc/paper_files/paper/ 2022/file/495e55f361708bedbab5d81f92048dcd-Paper-Conference.pdf .

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 International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 13017-13049. PMLR, 2023. URL https://proceedings.mlr.press/v202/ hess23a.html .

Jason A Platt, Stephen G Penny, Timothy A Smith, Tse-Chun Chen, and Henry D I Abarbanel. Constraining chaos: Enforcing dynamical invariants in the training of reservoir computers. Chaos , 33(10), 1 October 2023. URL http://dx.doi.org/10.1063/5.0156999 .

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning , volume 70 of Proceedings of Machine Learning Research , pages 214-223. PMLR, 06-11 Aug 2017. URL https://proceedings.mlr.press/ v70/arjovsky17a.html .

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.

Patrick T. Komiske, Eric M. Metodiev, and Jesse Thaler. Metric space of collider events. Phys. Rev. Lett. , 123:041801, Jul 2019. doi: 10.1103/PhysRevLett.123.041801. URL https: //link.aps.org/doi/10.1103/PhysRevLett.123.041801 .

Peter Y. Lu, Rumen Dangovski, and Marin Soljaˇ ci´ c. Discovering conservation laws using optimal transport and manifold learning. Nat. Commun. , 14(1):4744, 7 August 2023. URL http://dx.doi.org/10.1038/s41467-023-40325-7 .

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-310, 2023. doi: 10.1137/21M1421337.

Jonah Botvinick-Greenhouse, 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.1063/5.0149673 .

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 26, Red Hook, NY, 2013. Curran Associates, Inc. URL https://proceedings.neurips.cc/paper/2013/file/ af21d0c97db2e27e13572cbf59eb343d-Paper.pdf .

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.

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.

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.

Xiao Zhang and Michael Maire. Self-supervised visual representation learning from hierarchical grouping. Advances in Neural Information Processing Systems , 33:16579-16590, 2020.

Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part II 14 , pages 694-711. Springer, 2016.

Yunjie Tian, Lingxi Xie, Xiaopeng Zhang, Jiemin Fang, Haohang Xu, Wei Huang, Jianbin Jiao, Qi Tian, and Qixiang Ye. Semantic-aware generation for self-supervised visual representation learning. arXiv preprint arXiv:2111.13163 , 2021.

Cédric Villani. Optimal Transport: Old and New . Springer Berlin Heidelberg, Berlin, Heidelberg, 2009. ISBN 978-3-540-71050-9. doi: 10.1007/978-3-540-71050-9_6.

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://proceedings.mlr.press/v89/feydy19a.html .

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://proceedings.mlr.press/v119/ janati20a.html .

Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning , pages 9929-9939. PMLR, 2020.

Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 586-595, 2018.

Dirk Leendert van Kekem. Dynamics of the lorenz-96 model: Bifurcations, symmetries and waves, 2018. https://research.rug.nl/en/publications/ dynamics-of-the-lorenz-96-model-bifurcations-symmetries-and-waves .

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.

Andrew J Majda and John Harlim. Filtering Complex Turbulent Systems . Cambridge University Press, 2012.

James M Hyman and Basil Nicolaenko. The kuramoto-sivashinsky equation: a bridge between pde's and dynamical systems. Physica D: Nonlinear Phenomena , 18(1-3):113-126, 1986.

Pantelis R Vlachas, Georgios Arampatzis, Caroline Uhler, and Petros Koumoutsakos. Multiscale simulations of complex systems by learning their effective dynamics. Nature Machine Intelligence , 4(4):359-366, 2022.

H. S. Greenside, A. Wolf, J. 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/PhysRevA.25.3453. URL https://link.aps.org/doi/10.1103/PhysRevA.25. 3453 .

Bowen Zhu, Jingwei Hu, Yifei Lou, and Yunan Yang. Implicit regularization effects of the sobolev norms in image processing. arXiv preprint arXiv:2109.06255 , 2021.

Chuanfu Wang, Chunlei Fan, and Qun Ding. Constructing discrete chaotic systems with positive lyapunov exponents. International Journal of Bifurcation and Chaos , 28(07):1850084, 2018.

Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's algorithm. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 1367-1376. PMLR, 10-15 Jul 2018. URL https://proceedings.mlr.press/v80/dvurechensky18a.html .

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id= Bkg6RiCqY7 .

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 physics-informed optimal transport approach requires a choice of summary statistics, so it is much more dependent on 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 L 1 histogram distance of the summary statistics S (Table 1 and 2), which is 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 statistics

Comparing our results in Tables 1 and 2, we find that emulators trained using our approaches perform significantly better on the long-term evaluation metrics (e.g. L 1 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 a 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 statistics ˆ S A of trajectories generated by our trained neural operators are closer to the true statistics S A than the statistics directly computed using the noisy training data ˜ S A , i.e., | ˆ S A -S A | < | ˜ S A -S A | , 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.

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 a 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 distribution to establish a convolution matrix that's applied to the original image. The fundamental idea involves substituting the noisy pixel with a 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 circumstances, the Gaussian function collapses to a single point, leading to the elimination of the 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 high-frequency signals that could be affected by the blurring preprocessing.

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 derivatives, 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 a high-frequency domain, as minimizing the Sobolev norm then fails to accurately capture relevant statistics, as shown in Tables 5 and 6.

B.3 Optimal transport: reduced set of summary statistics

For our optimal transport approach, we test a reduced set of summary statistics, which shows how the quality of the summary statistic affects the performance of the method (Table 7). With an informative summary statistic, we find even a reduced set can still be helpful but, for a non-informative statistic, the optimal transport method fails as expected.

B.4 Contrastive learning: reduced environment diversity

For our contrastive learning 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 8), which demonstrates robustness.

B.5 Maximum mean discrepancy (MMD) vs. optimal transport loss

We also implement a variant of our optimal transport approach that uses maximum mean discrepancy (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 as our optimal transport loss for training emulators (Table 9).

B.6 Additional Lyapunov spectrum evaluation metrics

In the table 10, we evaluated the results of Lorenz 96 on 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: ∑ d i | ˆ λ i -λ i | /λ i , where λ i is the i -th Lyapunov exponent and d is 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 ∑ d i | 1 ( ˆ λ i > 0 -1 ( λ i > 0) | .

B.7 1-step RMSE evaluation results

Evaluating on 1-step RMSE only shows short-term prediction performance and is not an 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 1 and 2). See Appendix A.2 for additional discussion.

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 a recurrent fashion for 1500 (Lorenz-96) or 1000 (Kuramoto-Sivashinsky) time steps and then compute our long-term evaluation metrics on this autonomously generated time-series.

Histogram error. For distributional distance with a pre-specified statistics S , we first compute the histogram of the invariant statistics H ( S ) = { ( S 1 , c 1 ) , ( S 2 , c 2 ) , . . . , ( S B , c B ) } , where H represents the histogram, and the values of the bins are denoted by S b with their corresponding frequencies c b . We then define the L 1 histogram error as:

Note that for a fair comparison across all our experiments, we use the rule of thumb-the square root rule to decide the number of bins.

Energy spectrum error. We compute the relative mean absolute error of the energy spectrum-the squared norm of the spatial FFT F [ u t ] -averaged over time:

Leading LE error. The leading Lyapunov exponent (LE) is a dynamical invariant that measures how quickly the chaotic system becomes unpredictable. For the leading LE error, we report the relative absolute error | ˆ λ -λ | / | λ | between the model and the ground truth averaged over the test set. We adapted the Julia DynamicalSystem.jl 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 | ˆ D -D | between the estimated FD of the model and the ground truth averaged over the test set. We use the Julia DynamicalSystem.jl package for calculating the fractal dimension.

RMSE. We use 1-step relative RMSE ∥ u t +∆ t -ˆ g θ ( u t , ϕ ) ∥ 2 / ∥ u t +∆ t ∥ 2 to measure short-term prediction accuracy.

C.2 Time complexity

The optimal transport 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 [51]). In our experiments, we use n = 6000 to n = 25600 points with no issues, so this approach scales relatively well. The contrastive learning approach requires pretraining but is even faster during emulator training since it uses a fixed, pre-trained embedding network.

C.3 Training of the encoder

Evaluation rule. A 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 are positioned closer to each other compared to dissimilar ones in the embedded space. While this evaluation measure is commonly employed in downstream tasks that focus on classification, it 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 statistics. Therefore, the use of Top-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 to assess the performance of our encoder.

The formal definition of Top-1 accuracy requires us to initially define the softmax function σ , which is used to estimate the likelihood that two samples from different time windows originate from the same trajectory,

Using the softmax function to transform the encoder's output into a probability distribution, we then define Top-1 accuracy as:

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 a weight decay of 10 -5 , 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 τ value (0.3 in our experiments) for the first 1000 epochs. This ensures that samples that are difficult to distinguish get large gradients. Subsequently, we incrementally increase τ 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 { U ( n ) 0: L } N n =1 is finite, excessively increasing the crop length K can have both advantages and disadvantages. 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 ( n ) I : I + K ) and f ψ ( U ( n ) J : J + K ) , 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 our 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 ϕ ( n ) randomly sampled from a uniform distribution ranging from 10 . 0 to 18 . 0 . We discretize these training trajectories into time steps of dt = 0 . 1 over a total time of t = 205 s , 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 h RMSE ) via the grid search from the set of values { 1 , 2 , 3 , 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 h and h RMSE to 1.

In the case of the optimal transport loss 2 , described in Section 3.1, the blur parameter γ governs the equilibrium between faster convergence speed and the precision of the Wasserstein distance approximation. We determined the value of γ through a grid search, examining values ranging from { 0 . 01 , 0 . 02 , . . . , 0 . 20 } . Similarly, we decided the weights of the optimal transport loss α through a grid search, exploring values from { 0 . 5 , 1 , 1 . 5 , . . . , 3 . 0 } . For the experiments conducted using Lorenz-96, we selected α = 3 and γ = 0 . 02 .

In the context of the feature loss, the primary hyperparameter we need to consider is its weights, represented as λ . Given that we do not have knowledge of the invariant statistics S during the validation phase, we adjust λ according to a specific principle: our aim is to reduce the feature loss ℓ feature (as defined in Eqn. 17) as long as it does not adversely affect the RMSE beyond a predetermined level (for instance, 10% ) when compared with the baseline model, which is exclusively trained on RMSE. As illustrated in Figure 4, we adjust the values of λ in a 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.

0.0

0.2

0.4

0.6

0.8

1.0

1.2

lambda

0.14

0.16

0.18

0.20

0.22

histogram error of invariant statistics

one-step RMSE (noisy validation)

feature loss (noisy validation)

one-step RMSE (noiseless testing)

0.28

0.30

0.32

0.34

0.36

0.38

0.40

one-step RMSE (noisy validation)

0.15

0.16

0.17

0.18

0.19

0.20

feature loss (noisy validation)

C.5 Kuramoto-Sivashinsky

Data generation. In order to ascertain that we are operating within a chaotic regime, we set the domain size L = 50 and the spatial discretization grids number d = 256 . We select initial conditions randomly from a uniform range between [ -π, π ] . A fourth-order Runge-Kutta method was utilized to perform all simulations. We generate 2000 trajectories for training, with each ϕ ( n ) being randomly chosen from a uniform distribution within the interval [1 . 0 , 2 . 6] .

Training hyperparameters. In the case of the optimal transport loss, similar to the discussion in C.4, we search the blur value γ and the weights of loss α from a grid search. And in our experiment, we set α = 3 and γ = 0 . 05 . In the case of the feature loss, we again adopt the same rule for deciding the value of λ , 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 ( 110% compared to the baseline when λ = 0 ).

0

20

40

60

80

0

20

40

x

Sampled trajectory

Truth

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

Histogram 2d

u i

Histogram 2d

0

20

40

60

80

0

20

40

x

RMSE

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

0

20

40

x

OT (ours)

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

t

0

20

40

x

CL (ours)

u i

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

du i

/dt i

u i

0

5

10

15

20

25

30

Wavenumber

10

3

10

4

Energy

Truth

RMSE

OT (ours)

CL (ours)

(a)

0

20

40

60

80

0

20

40

x

Sampled trajectory

Truth

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

Histogram 2d

u i

Histogram 2d

0

20

40

60

80

0

20

40

x

RMSE

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

0

20

40

x

OT (ours)

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

t

0

20

40

x

CL (ours)

u i

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

du i

/dt i

u i

0

5

10

15

20

25

30

Wavenumber

10

3

10

4

Energy

Truth

RMSE

OT (ours)

CL (ours)

(b)

0

20

40

60

80

0

20

40

x

Sampled trajectory

Truth

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

Histogram 2d

u i

Histogram 2d

0

20

40

60

80

0

20

40

x

RMSE

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

0

20

40

x

OT (ours)

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

t

0

20

40

x

CL (ours)

u i

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

du i

/dt i

u i

0

5

10

15

20

25

30

Wavenumber

10

3

10

4

Energy

Truth

RMSE

OT (ours)

CL (ours)

(a)

0

20

40

60

80

0

20

40

x

Sampled trajectory

Truth

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

Histogram 2d

u i

Histogram 2d

0

20

40

60

80

0

20

40

x

RMSE

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

0

20

40

x

OT (ours)

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

t

0

20

40

x

CL (ours)

u i

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

du i

/dt i

u i

0

5

10

15

20

25

30

Wavenumber

10

3

10

4

Energy

Truth

RMSE

OT (ours)

CL (ours)

(b)

0

20

40

60

80

0

20

40

x

Sampled trajectory

Truth

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

Histogram 2d

u i

Histogram 2d

0

20

40

60

80

0

20

40

x

RMSE

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

0

20

40

x

OT (ours)

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

t

0

20

40

x

CL (ours)

u i

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

du i

/dt i

u i

0

5

10

15

20

25

30

Wavenumber

10

3

10

4

Energy

Truth

RMSE

OT (ours)

CL (ours)

(a)

0

20

40

60

80

0

20

40

x

Sampled trajectory

Truth

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

Histogram 2d

u i

Histogram 2d

0

20

40

60

80

0

20

40

x

RMSE

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

0

20

40

x

OT (ours)

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

u i

0

20

40

60

80

t

0

20

40

x

CL (ours)

u i

(

u

i

+1

-

u

i

-

2

)

u

i

-

1

du i

/dt i

u i

0

5

10

15

20

25

30

Wavenumber

10

3

10

4

Energy

Truth

RMSE

OT (ours)

CL (ours)

(b)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nIt would be good to discuss in detail the advantages and disadvantages of using Optimal transport and contrastive learning.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

comments: 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.

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.

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. 

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.

LIMITATIONS:
The authors have addressed the limitations
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nI 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).\n\nLIMITATIONS:\nThe paper addressed the limitation.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nOn 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.\n\nQUESTIONS:\nMy 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-24"}