SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ruoxi Jiang*
Chicago, IL 60637
roxie62@uchicago.edu
Peter Y. Lu*
lup@uchicago.edu
Elena Orlova
eorlova@uchicago.edu
Rebecca Willett
Department of Statistics and Computer Science
willett@uchicago.edu
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 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.
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 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.
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, 16-19]. 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 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-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 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, 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
(a) Error metrics vs. noise scale  $r$ .
(b) Predictions from emulators.
(c) Energy spectrum.
Figure 1: The impact of noise on 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 Kuramoto-Sivashinsky (KS) system with increasingly noisy initial conditions  $\mathbf{U}_G(\mathbf{u}_0 + \eta)$  as well as with added measurement noise  $\mathbf{U}_G(\mathbf{u}_0 + \eta) + \eta$ . Here,  $\mathbf{U}_G(\cdot)$  refers to the ground truth solution to the differential equation (1) for the KS system given an initial condition, and  $\eta \sim \mathcal{N}(0, r^2 \sigma^2 I)$ , where  $\sigma^2$  is the temporal variance of the trajectory  $\mathbf{U}_G(\mathbf{u}_0)$  and  $r$  is the noise scale. Relative RMSE and Sobolev norm [16], which focus on short-term forecasts, deteriorate rapidly with noise  $\eta$ , 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 a long-term prediction.
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.
Consider a dynamical system with state space  $\mathbf{u} \in \mathcal{U}$  governed by
$$
\frac {\mathrm {d} \mathbf {u}}{\mathrm {d} t} = G (\mathbf {u}, \phi), \tag {1}
$$
where  $G$  is the governing function, and  $\phi \in \Phi$  is a set of parameters that specify the dynamics. We will denote trajectories governed by these dynamics (1) as  $\mathbf{u}(t)$ , particular points on the trajectory as  $\mathbf{u}_t \in \mathcal{U}$ , and a sequence of  $K + 1$  consecutive time points on the trajectory as  $\mathbf{U}_{I:I + K} \coloneqq \{\mathbf{u}_{t_i}\}_{i = I}^{I + K}$ . Our experiments use equally spaced time points with a time step  $\Delta t$ . We refer to settings where  $\phi$  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  $\hat{g}_{\theta}:\mathcal{U}\times \Phi \to \mathcal{U}$  that makes discrete predictions
$$
\hat {\mathbf {u}} _ {t + \Delta t} := \hat {g} _ {\theta} \left(\hat {\mathbf {u}} _ {t}, \phi\right), \tag {2}
$$
where  $\theta$  are the parameters of the neural operator  $\hat{g}_{\theta}$ . In the multi-environment setting, we will have data from a variety of environments  $n \in \{1,2,\dots,N\}$ , and thus our training data  $\{\mathbf{U}_{0:L}^{(n)},\phi^{(n)}\}_{n=1}^N$  consists of trajectories  $\mathbf{U}_{0:L}^{(n)}$  and the corresponding environment parameters  $\phi^{(n)}$ .
Predicted sequence notation. We will denote a sequence of  $h + 1$  autonomously predicted states (using the neural operator  $\hat{g}_{\theta}$ ) with an initial condition  $\mathbf{u}_{t_I}^{(n)}$  as
$$
\hat {\mathbf {U}} _ {I: I + h} ^ {(n)} := \left\{\mathbf {u} _ {t _ {I}} ^ {(n)}, \hat {g} _ {\theta} \left(\mathbf {u} _ {t _ {I}} ^ {(n)}, \phi^ {(n)}\right), \hat {g} _ {\theta} \circ \hat {g} _ {\theta} \left(\mathbf {u} _ {t _ {I}} ^ {(n)}, \phi^ {(n)}\right), \dots , \overbrace {\hat {g} _ {\theta} \circ \cdots \circ \hat {g} _ {\theta}} ^ {h} \left(\mathbf {u} _ {t _ {I}} ^ {(n)}, \phi^ {(n)}\right) \right\}. \tag {3}
$$
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
$$
\hat {\mathbf {U}} _ {I: h: I + K} ^ {(n)} := \hat {\mathbf {U}} _ {I: I + h} ^ {(n)} \oplus \hat {\mathbf {U}} _ {I + h + 1: I + 2 h + 1} ^ {(n)} \oplus \dots \oplus \hat {\mathbf {U}} _ {I + K - h: I + K} ^ {(n)}, \tag {4}
$$
where  $\oplus$  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  $\mathcal{A}$ , we can construct a natural invariant measure (also known as a physical measure)
$$
\mu_ {\mathcal {A}} = \lim  _ {T \rightarrow \infty} \frac {1}{T} \int^ {T} \delta_ {\mathbf {u} _ {\mathcal {A}} (t)} \mathrm {d} t, \tag {5}
$$
where  $\delta_{\mathbf{u}(t)}$  is the Dirac measure centered on a trajectory  $\mathbf{u}_{\mathcal{A}}(t)$  that is in the basin of attraction of  $\mathcal{A}$  [13]. Note that, because  $\mathcal{A}$  is an attractor, any trajectory  $\mathbf{u}_{\mathcal{A}}(t)$  in the basin of attraction of  $\mathcal{A}$  will give the same invariant measure  $\mu_{\mathcal{A}}$ , i.e. the dynamics are ergodic on  $\mathcal{A}$ . Therefore, any time-invariant statistical property  $S_{\mathcal{A}}$  of the dynamics on  $\mathcal{A}$  can be written as
$$
S _ {\mathcal {A}} = \mathbb {E} _ {\mu_ {\mathcal {A}}} [ s ] = \int s (\mathbf {u}) \mathrm {d} \mu_ {\mathcal {A}} (\mathbf {u}) = \lim  _ {T \rightarrow \infty} \frac {1}{T} \int^ {T} s (\mathbf {u} _ {\mathcal {A}} (t)) \mathrm {d} t \tag {6}
$$
for some function  $s(\mathbf{u})$  and a trajectory  $\mathbf{u}_{\mathcal{A}}(t)$  in the basin of attraction of  $\mathcal{A}$ . Conversely, for any  $s(\mathbf{u})$ , (6) gives a time-invariant property  $S_{\mathcal{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
$$
\tilde {\mathbf {u}} _ {\mathcal {A}} (t) = \mathbf {u} _ {\mathcal {A}} (t) + \eta , \quad \eta \sim p _ {\eta} \tag {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  $p_{\eta}$ )
$$
\tilde {\mu} _ {\mathcal {A}} = \lim  _ {T \rightarrow \infty} \frac {1}{T} \int^ {T} \delta_ {\tilde {\mathbf {u}} _ {\mathcal {A}} (t)} \mathrm {d} t = \mu_ {\mathcal {A}} * p _ {\eta} = \lim  _ {T \rightarrow \infty} \frac {1}{T} \int^ {T} p _ {\eta} (\mathbf {u} _ {\mathcal {A}} (t)) \mathrm {d} t \tag {8}
$$
and are therefore given by
$$
\begin{array}{l} \tilde {S} _ {\mathcal {A}} = \mathbb {E} _ {\tilde {\mu} _ {\mathcal {A}}} [ s ] = \int s (\mathbf {u}) \mathrm {d} \tilde {\mu} _ {\mathcal {A}} (\mathbf {u}) \tag {9} \\ = \lim _ {T \to \infty} \frac {1}{T} \int^ {T} s (\tilde {\mathbf {u}} _ {\mathcal {A}} (t)) \mathrm {d} t = \lim _ {T \to \infty} \frac {1}{T} \int^ {T} s (\mathbf {u} _ {\mathcal {A}} (t)) p _ {\eta} (\mathbf {u} _ {\mathcal {A}} (t)) \mathrm {d} t, \\ \end{array}
$$
which can be a good approximation for  $S_{\mathcal{A}} = \mathbb{E}_{\mu_{\mathcal{A}}}[s]$  even in noisy conditions and does not suffer from the exponential divergence of RMSE (Fig. 1).
Figure 2: Our proposed approaches for training neural operators. (a) Neural operators are emulators trained to take an initial state and output future states in a recurrent fashion. To ensure the neural operator respects the statistical properties of chaotic dynamics when trained on noisy data, we 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 an optimal transport loss. (c) In the absence of prior knowledge, we take advantage of self-supervised contrastive learning to automatically learn relevant time-invariant statistics, which can then be used to train neural operators.
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  $\mathbf{s}(\mathbf{u}) = [s^{(1)}(\mathbf{u}), s^{(2)}(\mathbf{u}), \ldots, s^{(k)}(\mathbf{u})]$  between the data  $\mathbf{s}_i = \mathbf{s}(\mathbf{u}_{t_i}) \sim \mu_{\mathbf{s}}$  and the model predictions  $\hat{\mathbf{s}}_j = \mathbf{s}(\hat{\mathbf{u}}_{t_j}) \sim \mu_{\hat{\mathbf{s}}}$ . Here,  $\mu_{\mathbf{s}}$  and  $\mu_{\hat{\mathbf{s}}}$  are the probability measures of the summary statistics for the true chaotic attractor  $\mathcal{A}$  and the attractor  $\hat{\mathcal{A}}$  learned by the neural operator, respectively. Specifically,  $\mu_{\mathbf{s}}(\mathbf{s}') = \int \delta_{\mathbf{s}(\mathbf{u}) - \mathbf{s}'} \mu_{\mathcal{A}}(\mathbf{u}) \mathrm{d}\mathbf{u}$  and  $\mu_{\hat{\mathbf{s}}}(\hat{\mathbf{s}}') = \int \delta_{\mathbf{s}(\hat{\mathbf{u}}) - \hat{\mathbf{s}}'} \mu_{\hat{\mathcal{A}}}(\hat{\mathbf{u}}) \mathrm{d}\hat{\mathbf{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(\mu_{\mathrm{s}}, \mu_{\mathrm{\hat{s}}})$ —a distance metric between probability measures defined in terms of the solution to an optimal transport problem:
$$
\frac {1}{2} W \left(\mu_ {\mathrm {s}}, \mu_ {\hat {\mathrm {s}}}\right) ^ {2} := \inf  _ {\pi \in \Pi \left(\mu_ {\mathrm {s}}, \mu_ {\hat {\mathrm {s}}}\right)} \int c (\mathbf {s}, \hat {\mathbf {s}}) \mathrm {d} \pi (\mathbf {s}, \hat {\mathbf {s}}). \tag {10}
$$