Original: Embed and Emulate: Learning to estimate parameters of dynamical systems with uncertainty quantification
Revised: Embed and Emulate: Learning to estimate parameters of dynamical systems with uncertainty quantification

Original: Abstract
Revised: Abstract

Original: This paper explores learning emulators for parameter estimation with uncertainty estimation of high-dimensional dynamical systems.
Revised: This paper explores learning emulators for parameter estimation with uncertainty estimation of high-dimensional dynamical systems.

Original: We assume access to a computationally complex simulator that inputs a candidate parameter and outputs a corresponding multichannel time series.
Revised: We assume access to a computationally complex simulator that inputs a candidate parameter and outputs a corresponding multichannel time series.

Original: Our task is to accurately estimate a range of likely values of the underlying parameters.
Revised: Our task is to accurately estimate a range of likely values of the underlying parameters.

Original: Standard iterative approaches necessitate running the simulator many times, which is computationally prohibitive.
Revised: Standard iterative approaches necessitate running the simulator many times, which is computationally prohibitive.

Original: This paper describes a novel framework for learning feature embeddings of observed dynamics jointly with an emulator that can replace high-cost simulators for parameter estimation.
Revised: This paper describes a novel framework for learning feature embeddings of observed dynamics jointly with an emulator that can replace high-cost simulators for parameter estimation.

Original: Leveraging a contrastive learning approach, our method exploits intrinsic data properties within and across parameter and trajectory domains.
Revised: Leveraging a contrastive learning approach, our method exploits intrinsic data properties within and across parameter and trajectory domains.

Original: On a coupled 396-dimensional multiscale Lorenz 96 system, our method significantly outperforms a typical parameter estimation method based on predefined metrics and a classical numerical simulator, and with only 1.19% of the baseline’s computation time.
Revised: On a coupled 396-dimensional multiscale Lorenz 96 system, our method significantly outperforms a typical parameter estimation method based on predefined metrics and a classical numerical simulator, and with only 1.19% of the baseline’s computation time.

Original: Ablation studies highlight the potential of explicitly designing learned emulators for parameter estimation by leveraging contrastive learning.
Revised: Ablation studies highlight the potential of explicitly designing learned emulators for parameter estimation by leveraging contrastive learning.

Original:  Introduction
Revised:  Introduction

Original: Physics-based simulators play a vital role in many domains of science and engineering, from energy infrastructure to atmospheric sciences.
Revised: Physics-based simulators play a vital role in many domains of science and engineering, from energy infrastructure to atmospheric sciences.

Original: They are frequently critical for assessing risk and exploring “what if” scenarios, which require running models many times [Snyder et al., 2019, Beusch et al., 2020, Helgeson et al., 2021, Zhao et al., 2021].
Revised: They are frequently critical for assessing risk and exploring “what if” scenarios, which require running models many times [Snyder et al., 2019, Beusch et al., 2020, Helgeson et al., 2021, Zhao et al., 2021].

Original: However, the increasing complexity of operational computer models makes running such simulators a major challenge.
Revised: However, the increasing complexity of operational computer models makes running such simulators a major challenge.

Original: Emulators (also known as surrogate models) are models trained to mimic numerical simulations at a much lower computational cost, particularly for parameters or inputs that have not been simulated.
Revised: Emulators (also known as surrogate models) are models trained to mimic numerical simulations at a much lower computational cost, particularly for parameters or inputs that have not been simulated.

Original: This paper considers the problem of estimating the parameters of a physics simulation that provide the best fit to data.
Revised: This paper considers the problem of estimating the parameters of a physics simulation that provide the best fit to data.

Original: The estimation problem can generally be thought of as a nonlinear inverse problem in which our goal is to estimate parameters ϕ ∈ R k from a noisy multichannel time series Z ∈ R T × d .
Revised: The estimation problem can generally be thought of as a nonlinear inverse problem in which our goal is to estimate parameters ϕ ∈ R k from a noisy multichannel time series Z ∈ R T × d .

Original: Z = H ( ϕ ; Z 0 ) + η , where H ( ϕ ; Z 0 ) represents running is a physics simulator with parameters ϕ and initial condition Z 0 for T time steps, and η ∼ N (0 , Γ ) is observation noise.
Revised: Z = H ( ϕ ; Z 0 ) + η , where H ( ϕ ; Z 0 ) represents running is a physics simulator with parameters ϕ and initial condition Z 0 for T time steps, and η ∼ N (0 , Γ ) is observation noise.

Original: To ease the notation, we drop the initial condition Z 0 in H ( ϕ ; Z 0 ) .
Revised: To ease the notation, we drop the initial condition Z 0 in H ( ϕ ; Z 0 ) .

Original: We are particularly interested in complex simulators for which we do not have analytic expressions for H , evaluating H ( ϕ ) is a computational bottleneck and we cannot readily compute its gradients.
Revised: We are particularly interested in complex simulators for which we do not have analytic expressions for H , evaluating H ( ϕ ) is a computational bottleneck and we cannot readily compute its gradients.

Original: One important application of this problem arises in climate science, where climate scientists have spent decades developing sophisticated physics-based models corresponding to H , often implemented using large-scale software systems that solve complex systems of differential equations
Revised: One important application of this problem arises in climate science, where climate scientists have spent decades developing sophisticated physics-based models corresponding to H , often implemented using large-scale software systems that solve complex systems of differential equations

Original: [Kay et al., 2015].
Revised: [Kay et al., 2015].

Original: In this case, evaluating H can be computationally demanding.
Revised: In this case, evaluating H can be computationally demanding.

Original: Furthermore, tools like physics-informed neural networks [Raissi et al., 2019] are inapplicable because we cannot compute losses that depend upon knowing the form of H .
Revised: Furthermore, tools like physics-informed neural networks [Raissi et al., 2019] are inapplicable because we cannot compute losses that depend upon knowing the form of H .

Original: Estimating the parameters of such models based on observational data is essential for accurate climate forecasting, and we typically seek not only point estimates of parameters, but also quantifiable uncertainty measures that allow us to forecast the full range of possible future outcomes.
Revised: Estimating the parameters of such models based on observational data is essential for accurate climate forecasting, and we typically seek not only point estimates of parameters, but also quantifiable uncertainty measures that allow us to forecast the full range of possible future outcomes.

Original: Parameter uncertainty quantification is vital here [Cleary et al., 2021, Souza et al., 2020, Hansen, 2022], especially when dealing with noisy observations where a small error in ˆ ϕ might lead to a dramatically different forecasts of the dynamics; uncertainty quantification also plays a vital role in the decision-making process to prevent hazardous loss under tail events [Hansen, 2022].
Revised: Parameter uncertainty quantification is vital here [Cleary et al., 2021, Souza et al., 2020, Hansen, 2022], especially when dealing with noisy observations where a small error in ˆ ϕ might lead to a dramatically different forecasts of the dynamics; uncertainty quantification also plays a vital role in the decision-making process to prevent hazardous loss under tail events [Hansen, 2022].

Original: We explore an alternative framework in which we use our simulator to generate training samples of the form ( ϕ i , Z i ) for i = 1 , . . . , n and train a machine learning model for parameter estimation.
Revised: We explore an alternative framework in which we use our simulator to generate training samples of the form ( ϕ i , Z i ) for i = 1 , . . . , n and train a machine learning model for parameter estimation.

Original: Our approach, which we call E MBED & E MULATE , jointly trains two neural networks.
Revised: Our approach, which we call E MBED & E MULATE , jointly trains two neural networks.

Original: The first, f θ , maps an observation Z to a low-dimensional embedding.
Revised: The first, f θ , maps an observation Z to a low-dimensional embedding.

Original: The second, ˆ g θ , emulates the map g θ :
Revised: The second, ˆ g θ , emulates the map g θ :

Original: = f θ ◦ H , and so maps a parameter vector ϕ to the same low-dimensional embedding space.
Revised: = f θ ◦ H , and so maps a parameter vector ϕ to the same low-dimensional embedding space.

Original: The learned ˆ g θ may be used in place of H within an optimization-based framework such as Ensemble
Revised: The learned ˆ g θ may be used in place of H within an optimization-based framework such as Ensemble

Original: Kalman Inversion (EnKI, §2).
Revised: Kalman Inversion (EnKI, §2).

Original: More specifically, our method uses the learned low-dimensional embedding to specify an objective function well-suited to parameter estimation without expert knowledge and the learned emulator as a surrogate for classical numerical methods used to compute H .
Revised: More specifically, our method uses the learned low-dimensional embedding to specify an objective function well-suited to parameter estimation without expert knowledge and the learned emulator as a surrogate for classical numerical methods used to compute H .

Original: Furthermore, inspired by empirical
Revised: Furthermore, inspired by empirical

Original: Bayes methods, our learned network suggests a natural mechanism for specifying the EnKI prior based on the trained network.
Revised: Bayes methods, our learned network suggests a natural mechanism for specifying the EnKI prior based on the trained network.

Original:  Related work
Revised:  Related work

Original: Brief parameter estimation background: A widely-used approach for parameter estimation is Ensemble Kalman Inversion (EnKI; §A.1)
Revised: Brief parameter estimation background: A widely-used approach for parameter estimation is Ensemble Kalman Inversion (EnKI; §A.1)

Original: [Iglesias et al., 2013].
Revised: [Iglesias et al., 2013].

Original: EnKI is a derivative-free optimization framework allows a user to specify a prior distribution over the parameter vector ϕ and results in samples from the corresponding posterior distribution.
Revised: EnKI is a derivative-free optimization framework allows a user to specify a prior distribution over the parameter vector ϕ and results in samples from the corresponding posterior distribution.

Original: As a result, this method provides users with important information about the uncertainty of estimates of ϕ .
Revised: As a result, this method provides users with important information about the uncertainty of estimates of ϕ .

Original: The clear and thought-provoking paper of Schneider et al.
Revised: The clear and thought-provoking paper of Schneider et al.

Original: [2017] highlighted the use of the EnKI in earth system modeling.
Revised: [2017] highlighted the use of the EnKI in earth system modeling.

Original: Despite its success in the perfect-model setting (i.e., where H ( ϕ ) can be computed exactly and perfectly represents the dynamics in Z for some ϕ ), using such methods in practice presents several challenges.
Revised: Despite its success in the perfect-model setting (i.e., where H ( ϕ ) can be computed exactly and perfectly represents the dynamics in Z for some ϕ ), using such methods in practice presents several challenges.

Original: To use the EnKI for parameter estimation, we must specify three key elements: (a) the objective function that EnKI is attempting to minimize; (b) the tool used to compute the forward mapping H (or a surrogate for H ); and (c) a prior p ϕ .
Revised: To use the EnKI for parameter estimation, we must specify three key elements: (a) the objective function that EnKI is attempting to minimize; (b) the tool used to compute the forward mapping H (or a surrogate for H ); and (c) a prior p ϕ .

Original: For instance, Schneider et al.
Revised: For instance, Schneider et al.

Original: [2017] uses a Gaussian prior p ϕ , computes H using Runge-Kutta methods [Dormand and Prince, 1980], and attempts to minimize the Mahalanobis distance
Revised: [2017] uses a Gaussian prior p ϕ , computes H using Runge-Kutta methods [Dormand and Prince, 1980], and attempts to minimize the Mahalanobis distance

Original: [Equation]
Revised: [Equation]

Original: where m ( Z i ) is a vector of first and second moments of different spatial channels of Z i and Σ
Revised: where m ( Z i ) is a vector of first and second moments of different spatial channels of Z i and Σ

Original: i is a diagonal matrix with the j -th diagonal entry Σ
Revised: i is a diagonal matrix with the j -th diagonal entry Σ

Original: i,j,j :=
Revised: i,j,j :=

Original: Var[ m ( H ( ϕ i ))
Revised: Var[ m ( H ( ϕ i ))

Original: j ] .
Revised: j ] .

Original: Using a loss based on moments within the EnKI framework provides critical robustness to the chaotic nature of Z i .
Revised: Using a loss based on moments within the EnKI framework provides critical robustness to the chaotic nature of Z i .

Original: This framework presents two central challenges.
Revised: This framework presents two central challenges.

Original: First, choosing the objective for EnKI to minimize requires non-trivial domain knowledge, and a poor choice may lead to biased parameter estimates or unpredictable sensitivities to certain features in Z .
Revised: First, choosing the objective for EnKI to minimize requires non-trivial domain knowledge, and a poor choice may lead to biased parameter estimates or unpredictable sensitivities to certain features in Z .

Original: Second, high-resolution simulations used to evaluate H for a given ϕ (e.g. Runge-Kutta or other classical numerical solvers) can be computationally demanding.
Revised: Second, high-resolution simulations used to evaluate H for a given ϕ (e.g. Runge-Kutta or other classical numerical solvers) can be computationally demanding.

Original: In particular, the EnKI method uses a collection of particles to represent samples of the posterior, and the forward model H must be calculated for each particle at each iteration.
Revised: In particular, the EnKI method uses a collection of particles to represent samples of the posterior, and the forward model H must be calculated for each particle at each iteration.

Original: For complex optimization landscapes or in high-dimensional settings, the number of particles must be large.
Revised: For complex optimization landscapes or in high-dimensional settings, the number of particles must be large.

Original: This is further complicated by the fact that often the parameter estimation task is conducted repeatedly, e.g., each time new climate data is acquired.
Revised: This is further complicated by the fact that often the parameter estimation task is conducted repeatedly, e.g., each time new climate data is acquired.

Original: Simulation-based inference (SBI): A large body of work in SBI focuses on parameter estimation for physics-based simulators [Beaumont et al., 2009, Papamakarios et al., 2019, Cranmer et al., 2020,
Revised: Simulation-based inference (SBI): A large body of work in SBI focuses on parameter estimation for physics-based simulators [Beaumont et al., 2009, Papamakarios et al., 2019, Cranmer et al., 2020,

Original: Lueckmann et al., 2019, Chen et al., 2020b, Alsing et al., 2019].
Revised: Lueckmann et al., 2019, Chen et al., 2020b, Alsing et al., 2019].

Original: Advanced SBI methods focus on adaptive generation of new training data to approximate the posterior and could be classified into different categories based on how they adaptively choose informative simulations [Lueckmann et al., 2021].
Revised: Advanced SBI methods focus on adaptive generation of new training data to approximate the posterior and could be classified into different categories based on how they adaptively choose informative simulations [Lueckmann et al., 2021].

Original: One approach, likelihood estimation, proposes to approximate an intractable likelihood function [Drovandi et al., 2018].
Revised: One approach, likelihood estimation, proposes to approximate an intractable likelihood function [Drovandi et al., 2018].

Original: For instance, Sequential Neural Likelihood estimation (SNL,
Revised: For instance, Sequential Neural Likelihood estimation (SNL,

Original: Papamakarios et al.
Revised: Papamakarios et al.

Original: [2019]) trains a conditional neural density estimator (e.g., Masked Autoregressive
Revised: [2019]) trains a conditional neural density estimator (e.g., Masked Autoregressive

Original: Flow (MAF)) which models the conditional distribution of data given parameters [Papamakarios et al., 2017]).
Revised: Flow (MAF)) which models the conditional distribution of data given parameters [Papamakarios et al., 2017]).

Original: While SNL does not readily scale to high-dimensional data, a recent variant called SNL+ [Chen et al., 2020b] addresses this limitation by learning sufficient statistics (embeddings) of the data based on the infomax principle, and it iteratively updates the network used to compute sufficient statistics and the neural density estimators during sequential sampling.
Revised: While SNL does not readily scale to high-dimensional data, a recent variant called SNL+ [Chen et al., 2020b] addresses this limitation by learning sufficient statistics (embeddings) of the data based on the infomax principle, and it iteratively updates the network used to compute sufficient statistics and the neural density estimators during sequential sampling.

Original: An alternative approach, (sequential) neural posterior estimation ((S)NPE), aims to approximate directly the target posterior (SNPE-A in Papamakarios and Murray [2016], SNPE-B in Lueckmann et al.
Revised: An alternative approach, (sequential) neural posterior estimation ((S)NPE), aims to approximate directly the target posterior (SNPE-A in Papamakarios and Murray [2016], SNPE-B in Lueckmann et al.

Original: [2017], SNPE-C in
Revised: [2017], SNPE-C in

Original: Greenberg et al.
Revised: Greenberg et al.

Original: [2019]), where SNPE-C forces fewer restrictions on the form of prior and posterior by leveraging neural conditional density estimation.
Revised: [2019]), where SNPE-C forces fewer restrictions on the form of prior and posterior by leveraging neural conditional density estimation.

Original: However, in the context of this manuscript’s setting, i.e., estimating multiple ϕ i for multiple different observations Z i at test time, the key technique of SBI, sequential sampling, can require substantial computational investments.
Revised: However, in the context of this manuscript’s setting, i.e., estimating multiple ϕ i for multiple different observations Z i at test time, the key technique of SBI, sequential sampling, can require substantial computational investments.

Original: This idea is discussed further with our experimental results.
Revised: This idea is discussed further with our experimental results.

Original: Learned emulators: Physics models and simulations are pervasive in weather and climate science, astrophysics, high-energy and accelerator physics, and the study of dynamical systems.
Revised: Learned emulators: Physics models and simulations are pervasive in weather and climate science, astrophysics, high-energy and accelerator physics, and the study of dynamical systems.

Original: These models and simulations are used, for example, to infer the underlying physical processes and equations that govern our observations, design new sensors or facilities, estimate errors, understand experimental behavior, and estimate unknown parameters within models.
Revised: These models and simulations are used, for example, to infer the underlying physical processes and equations that govern our observations, design new sensors or facilities, estimate errors, understand experimental behavior, and estimate unknown parameters within models.

Original: Many such physics simulators require expensive computational resources, and are difficult to fully leverage.
Revised: Many such physics simulators require expensive computational resources, and are difficult to fully leverage.

Original: “Learned emulators,” “surrogate models,” or “approximants” are computationally-efficient approximate models that use numerically simulated data to train a machine learning system to mimic these numerical simulations at a much smaller computational cost.
Revised: “Learned emulators,” “surrogate models,” or “approximants” are computationally-efficient approximate models that use numerically simulated data to train a machine learning system to mimic these numerical simulations at a much smaller computational cost.

Original: There have been many recent successes in the development of surrogate model foundations [Brockherde et al., 2017, Raissi et al., 2019, Chattopadhyay et al., 2019, 2020, Raissi et al., 2017, Vlachas et al., 2018, 2020, Brajard et al., 2020, Gagne et al., 2020] and applications [Goel et al., 2008, Mengistu and Ghaly, 2008, Brigham and Aquino, 2007, Kim et al., 2015, Papadopoulos et al., 2018, White et al., 2019, Gentine et al., 2018, Rasp et al., 2018, Cohen et al., 2020, Yuval and O’Gorman, 2020, Xue et al., 2021, Wang et al., 2019].
Revised: There have been many recent successes in the development of surrogate model foundations [Brockherde et al., 2017, Raissi et al., 2019, Chattopadhyay et al., 2019, 2020, Raissi et al., 2017, Vlachas et al., 2018, 2020, Brajard et al., 2020, Gagne et al., 2020] and applications [Goel et al., 2008, Mengistu and Ghaly, 2008, Brigham and Aquino, 2007, Kim et al., 2015, Papadopoulos et al., 2018, White et al., 2019, Gentine et al., 2018, Rasp et al., 2018, Cohen et al., 2020, Yuval and O’Gorman, 2020, Xue et al., 2021, Wang et al., 2019].

Original: Typically, one would directly try to emulate H , and the efficacy of the learned emulator would be evaluated on how well ˆ H ( ϕ ) matches H ( ϕ ) (often using squared error) across a range of ϕ .
Revised: Typically, one would directly try to emulate H , and the efficacy of the learned emulator would be evaluated on how well ˆ H ( ϕ ) matches H ( ϕ ) (often using squared error) across a range of ϕ .

Original: This formulation, however, may be suboptimal when the emulator will be used for parameter estimation.
Revised: This formulation, however, may be suboptimal when the emulator will be used for parameter estimation.

Original: First, as we will detail later, emulating H when Z is very high-dimensional is quite challenging.
Revised: First, as we will detail later, emulating H when Z is very high-dimensional is quite challenging.

Original: (In our experiments, dim( Z ) = 396 , 000 .)
Revised: (In our experiments, dim( Z ) = 396 , 000 .)

Original: Furthermore, in many climate settings, H represents a chaotic process in which very small changes to the initial conditions can result in very large differences later in the process.
Revised: Furthermore, in many climate settings, H represents a chaotic process in which very small changes to the initial conditions can result in very large differences later in the process.

Original: In this setting, training an emulator with a loss akin to (1 /n )
Revised: In this setting, training an emulator with a loss akin to (1 /n )

Original: (cid:80) ni
Revised: (cid:80) ni

Original: =1 ∥
Revised: =1 ∥

Original: Z i − ˆ H ( ϕ i ) ∥may be overly sensitive to noise and initial conditions and not preserve statistical features of Z i that may be essential to parameter estimation.
Revised: Z i − ˆ H ( ϕ i ) ∥may be overly sensitive to noise and initial conditions and not preserve statistical features of Z i that may be essential to parameter estimation.

Original: Inspired by the objective in (1), Cleary et al.
Revised: Inspired by the objective in (1), Cleary et al.

Original: [2021] trained a model ˆ g θ to emulate the moments of parameters using the loss function ℓ moment ( θ ) :
Revised: [2021] trained a model ˆ g θ to emulate the moments of parameters using the loss function ℓ moment ( θ ) :

Original: = (1 /n ) (cid:80) ni
Revised: = (1 /n ) (cid:80) ni

Original: =

Revised: =

Original: ∥ m ( Z i ) − ˆ g θ ( ϕ i ))
Revised: ∥ m ( Z i ) − ˆ g θ ( ϕ i ))

Original: ∥ 2 Σ i to ease the burden of running H .
Revised: ∥ 2 Σ i to ease the burden of running H .

Original: However, ℓ moment ( θ ) is not always the best choice.
Revised: However, ℓ moment ( θ ) is not always the best choice.

Original: First, estimating Σ i for each training sample is an enormous computational burden that far exceeds the cost of generating the original training data (unless the range of candidate ϕ is very small).
Revised: First, estimating Σ i for each training sample is an enormous computational burden that far exceeds the cost of generating the original training data (unless the range of candidate ϕ is very small).

Original: Second, choosing the moment function m (essentially fixing a particular low-dimensional embedding) is a critical design element that requires domain expertise and often must be tuned in practice.
Revised: Second, choosing the moment function m (essentially fixing a particular low-dimensional embedding) is a critical design element that requires domain expertise and often must be tuned in practice.

Original: Finally, even if one is willing to generate accurate estimates of Σ i , one could face additional challenges.
Revised: Finally, even if one is willing to generate accurate estimates of Σ i , one could face additional challenges.

Original: Cleary et al.
Revised: Cleary et al.

Original: [2021] is designed for inferring parameters given one fixed test observation, and they train their emulator for a small domain of parameter space dependent on the test observation at hand.
Revised: [2021] is designed for inferring parameters given one fixed test observation, and they train their emulator for a small domain of parameter space dependent on the test observation at hand.

Original: Extending this framework to a larger domain of parameters ϕ results in large variability among the collection of corresponding Σ i s generated for training, making the loss landscape challenging to navigate.
Revised: Extending this framework to a larger domain of parameters ϕ results in large variability among the collection of corresponding Σ i s generated for training, making the loss landscape challenging to navigate.

Original: Contrastive learning: With the above challenges in mind, we present a novel and generic framework for parameter estimation by jointly optimizing (a) an embedding of the multichannel dynamics used to evaluate the accuracy of a candidate parameter ϕ and (b) an emulator of the dynamics projected into the embedding space.
Revised: Contrastive learning: With the above challenges in mind, we present a novel and generic framework for parameter estimation by jointly optimizing (a) an embedding of the multichannel dynamics used to evaluate the accuracy of a candidate parameter ϕ and (b) an emulator of the dynamics projected into the embedding space.

Original: We leverage a contrastive framework to learn discriminative representations for high-dimensional spatio-temporal data.
Revised: We leverage a contrastive framework to learn discriminative representations for high-dimensional spatio-temporal data.

Original: Contrastive representation learning has exploded in popularity recently for self-supervised visual feature learning that has achieved comparable performance with its supervised counterparts [Caron et al., 2020, Chen et al., 2020a, He et al., 2020, Grill et al., 2020, Zhang and Maire, 2020, Caron et al., 2021].
Revised: Contrastive representation learning has exploded in popularity recently for self-supervised visual feature learning that has achieved comparable performance with its supervised counterparts [Caron et al., 2020, Chen et al., 2020a, He et al., 2020, Grill et al., 2020, Zhang and Maire, 2020, Caron et al., 2021].

Original: The majority of these frameworks operate under the push-pull principle for instance-wise discrimination: images generated from different forms of data augmentation (e.g., cropping, color jittering) are pulled together while other images are pushed away.
Revised: The majority of these frameworks operate under the push-pull principle for instance-wise discrimination: images generated from different forms of data augmentation (e.g., cropping, color jittering) are pulled together while other images are pushed away.

Original: Apart from its common setup in unsupervised representation learning, Khosla et al.
Revised: Apart from its common setup in unsupervised representation learning, Khosla et al.

Original: [2020] also demonstrated selecting positive and negative pairs in a supervised way leveraging label information can boost the performance.
Revised: [2020] also demonstrated selecting positive and negative pairs in a supervised way leveraging label information can boost the performance.

Original: Many efforts focus on understanding the mechanism of contrastive learning [Oord et al., 2018, Wei et al., 2020, Arora et al., 2019, Wang and Isola, 2020, Zimmermann et al., 2021].
Revised: Many efforts focus on understanding the mechanism of contrastive learning [Oord et al., 2018, Wei et al., 2020, Arora et al., 2019, Wang and Isola, 2020, Zimmermann et al., 2021].

Original: Among them,
Revised: Among them,

Original: Zimmermann et al.
Revised: Zimmermann et al.

Original: [2021] show that contrastive loss can be interpreted as the cross-entropy between the latent conditional distribution and ground truth distribution.
Revised: [2021] show that contrastive loss can be interpreted as the cross-entropy between the latent conditional distribution and ground truth distribution.

Original:  Contributions
Revised:  Contributions

Original: In our work, we follow the contrastive framework and learn an embedding network to capture structural information for high-dimensional spatio-temporal data and an emulator that can be used to compute parameter estimates with uncertainty estimates.
Revised: In our work, we follow the contrastive framework and learn an embedding network to capture structural information for high-dimensional spatio-temporal data and an emulator that can be used to compute parameter estimates with uncertainty estimates.

Original: This is effective for several reasons.
Revised: This is effective for several reasons.

Original: First, when the model H is bijective, representations learned using contrastive losses are highly correlated with their underlying parameters [Zimmermann et al., 2021].
Revised: First, when the model H is bijective, representations learned using contrastive losses are highly correlated with their underlying parameters [Zimmermann et al., 2021].

Original: Second, most existing emulating methods focus on approximating dynamical models for a fixed parameter ϕ
Revised: Second, most existing emulating methods focus on approximating dynamical models for a fixed parameter ϕ

Original: and lack generalization capacity [Krishnapriyan et al., 2021].
Revised: and lack generalization capacity [Krishnapriyan et al., 2021].

Original: In part this is due to the difficulties of emulating dynamics with high spatio-temporal resolution.
Revised: In part this is due to the difficulties of emulating dynamics with high spatio-temporal resolution.

Original: We circumvent this bottleneck by training an emulator to output the learned latent representations of the dynamics, which are much lower-dimensional and easier to learn with fewer training samples.
Revised: We circumvent this bottleneck by training an emulator to output the learned latent representations of the dynamics, which are much lower-dimensional and easier to learn with fewer training samples.

Original: Finally, our method is inspired by the Contrastive Language–Image Pretraining (CLIP)
Revised: Finally, our method is inspired by the Contrastive Language–Image Pretraining (CLIP)

Original: [Radford et al., 2021] framework, which is designed for contrastive learning of aligned representations of images and language.
Revised: [Radford et al., 2021] framework, which is designed for contrastive learning of aligned representations of images and language.

Original: Our approach, by analogy, learns aligned representations of dynamics and parameters, allowing us to jointly learn the embedding function and emulator for high-resolution data.
Revised: Our approach, by analogy, learns aligned representations of dynamics and parameters, allowing us to jointly learn the embedding function and emulator for high-resolution data.

Original: More specifically, this paper makes the following key contributions:
Revised: More specifically, this paper makes the following key contributions:

Original: • Incorporation of a contrastive loss function for learning an embedding of the simulation outputs ( Z i ) instead of relying on a known “good” moment function m , leveraging ideas from CLIP [Radford et al., 2021]; • Development of an emulator designed to facilitate parameter estimation with uncertainty quantification.
Revised: • Incorporation of a contrastive loss function for learning an embedding of the simulation outputs ( Z i ) instead of relying on a known “good” moment function m , leveraging ideas from CLIP [Radford et al., 2021]; • Development of an emulator designed to facilitate parameter estimation with uncertainty quantification.

Original: We demonstrate that the proposed emulator can be used effectively within an EnKI framework to generate parameter estimates with accurate posterior distribution estimates.
Revised: We demonstrate that the proposed emulator can be used effectively within an EnKI framework to generate parameter estimates with accurate posterior distribution estimates.

Original: • An empirical evaluation of different methods (EnKI
Revised: • An empirical evaluation of different methods (EnKI

Original: [Schneider et al., 2017], supervised learning, NPE-C Greenberg et al.
Revised: [Schneider et al., 2017], supervised learning, NPE-C Greenberg et al.

Original: [2019], SNL+ [Chen et al., 2020b], and our proposed emulator-based approach) in terms of robustness to noise or errors in training and testing data for a range of sample sizes as well as computational costs.
Revised: [2019], SNL+ [Chen et al., 2020b], and our proposed emulator-based approach) in terms of robustness to noise or errors in training and testing data for a range of sample sizes as well as computational costs.

Original: Compared to using numerical solvers, we show that our method achieves higher accuracy overall in 1.08% of the computation time, even accounting for the computational costs of generating training data.
Revised: Compared to using numerical solvers, we show that our method achieves higher accuracy overall in 1.08% of the computation time, even accounting for the computational costs of generating training data.

Original: • Evaluation of the quality of the uncertainty estimates using the Continuous Ranked Probability Score (CRPS) for varying numbers of training samples.
Revised: • Evaluation of the quality of the uncertainty estimates using the Continuous Ranked Probability Score (CRPS) for varying numbers of training samples.

Original: Our empirical results highlight both the improved computational and empirical accuracy of our proposed emulator-centric approach relative to competing methods and, more broadly, the benefits of designing emulators specifically for parameter estimation tasks.
Revised: Our empirical results highlight both the improved computational and empirical accuracy of our proposed emulator-centric approach relative to competing methods and, more broadly, the benefits of designing emulators specifically for parameter estimation tasks.

Original:  Proposed method
Revised:  Proposed method

Original: Our goal is to learn a mapping from observations of a multichannel dynamical system, Z ∈ R T × d , to the underlying system parameters ϕ ∈ R k , where Z = H ( ϕ ) + η for some noise vector η ∈ R T × d .
Revised: Our goal is to learn a mapping from observations of a multichannel dynamical system, Z ∈ R T × d , to the underlying system parameters ϕ ∈ R k , where Z = H ( ϕ ) + η for some noise vector η ∈ R T × d .

Original: We assume we do not have access to an analytical expression for H , but can compute H ( ϕ ) for any ϕ using a computationally complex simulator.
Revised: We assume we do not have access to an analytical expression for H , but can compute H ( ϕ ) for any ϕ using a computationally complex simulator.

Original: We further assume a prior p ϕ over parameter space.
Revised: We further assume a prior p ϕ over parameter space.

Original: Using this prior and the simulator, we generate n training samples as follows: for i = 1 , . . .
Revised: Using this prior and the simulator, we generate n training samples as follows: for i = 1 , . . .

Original: , n , draw ϕ i ∼ p ϕ , and use the simulator to compute Z i
Revised: , n , draw ϕ i ∼ p ϕ , and use the simulator to compute Z i

Original: = H ( ϕ i )
Revised: = H ( ϕ i )

Original: + η i .
Revised: + η i .

Original: We do not explicitly generate noise η i ; rather, the numerical algorithm used to simulate H will generally produce some errors, with faster implementations of H being more error-prone.
Revised: We do not explicitly generate noise η i ; rather, the numerical algorithm used to simulate H will generally produce some errors, with faster implementations of H being more error-prone.

Original: The distribution p ϕ coupled with the noise distribution over η induces the joint distribution p ϕ , Z .
Revised: The distribution p ϕ coupled with the noise distribution over η induces the joint distribution p ϕ , Z .

Original: 4.1 Baseline approach – supervised learning
Revised: 4.1 Baseline approach – supervised learning

Original: Given training samples ( ϕ i , Z i ) for i = 1 , . . .
Revised: Given training samples ( ϕ i , Z i ) for i = 1 , . . .

Original: , n , we train a neural network represented by h θ using a ResNet [He et al., 2016] as the backbone (see details in §B.2).
Revised: , n , we train a neural network represented by h θ using a ResNet [He et al., 2016] as the backbone (see details in §B.2).

Original: θ denotes the learned network parameters and our prediction is ˆ ϕ
Revised: θ denotes the learned network parameters and our prediction is ˆ ϕ

Original: i
Revised: i

Original: =
Revised: =

Original: h θ ( Z i ) .
Revised: h θ ( Z i ) .

Original: We train using the Mean Absolute Percentage Error (MAPE) loss
Revised: We train using the Mean Absolute Percentage Error (MAPE) loss

Original: [Equation]
Revised: [Equation]

Original: As we will see in the experimental results, this approach yields point estimates of ϕ on holdout test data with reasonable accuracy and offers a significant computational advantage over the EnKI method.
Revised: As we will see in the experimental results, this approach yields point estimates of ϕ on holdout test data with reasonable accuracy and offers a significant computational advantage over the EnKI method.

Original: However, it does not offer uncertainty estimates, and its test accuracy is lower than that of our proposed approach (§4.2).
Revised: However, it does not offer uncertainty estimates, and its test accuracy is lower than that of our proposed approach (§4.2).

Original: 4.2 Our approach: joint embedding and emulation via contrastive feature representation
Revised: 4.2 Our approach: joint embedding and emulation via contrastive feature representation

Original: Contrastive feature representation:
Revised: Contrastive feature representation:

Original: Measuring the distance between a pair of dynamics, a necessary task for constructing a loss function used for training, is particularly challenging when the dynamics are chaotic.
Revised: Measuring the distance between a pair of dynamics, a necessary task for constructing a loss function used for training, is particularly challenging when the dynamics are chaotic.

Original: As mentioned above (§2), one alternative in the literature is to use a custom loss function based on summary statistics of Z ; this requires expert knowledge to determine which statistics are relevant, and has demonstrated efficacy only for narrow ranges of parameters ϕ .
Revised: As mentioned above (§2), one alternative in the literature is to use a custom loss function based on summary statistics of Z ; this requires expert knowledge to determine which statistics are relevant, and has demonstrated efficacy only for narrow ranges of parameters ϕ .

Original: We present an alternative framework in which we learn an embedding of Z , denoted f θ ( Z ) , that preserves statistical and structural characteristics of Z most relevant to the parameter estimation task and hence can be used to train an emulator.
Revised: We present an alternative framework in which we learn an embedding of Z , denoted f θ ( Z ) , that preserves statistical and structural characteristics of Z most relevant to the parameter estimation task and hence can be used to train an emulator.

Original: Here θ denotes all the parameters of the neural network defining the embedding function.
Revised: Here θ denotes all the parameters of the neural network defining the embedding function.

Original: Inspired by recent advances in contrastive representation learning, from a collection of parameters and trajectories pairs { ϕ i , Z i } ni =1 , we propose a trajectory encoder f θ which learns a distinguishable representation by minimizing the following variant of the Info
Revised: Inspired by recent advances in contrastive representation learning, from a collection of parameters and trajectories pairs { ϕ i , Z i } ni =1 , we propose a trajectory encoder f θ which learns a distinguishable representation by minimizing the following variant of the Info

Original: Noise
Revised: Noise

Original: [Equation]
Revised: [Equation]

Original: Here f θ ( Z ) ∈ S p − 1 is normalized and lives in a unit hypersphere.
Revised: Here f θ ( Z ) ∈ S p − 1 is normalized and lives in a unit hypersphere.

Original: ˜
Revised: ˜

Original: Z i is selected through data augmentation (see §B.1) or in a supervised way by measuring the distance in the parameter space: ˜
Revised: Z i is selected through data augmentation (see §B.1) or in a supervised way by measuring the distance in the parameter space: ˜

Original: Z i := arg min
Revised: Z i := arg min

Original: Z j ̸ =
Revised: Z j ̸ =

Original: Z i δ ( Z j , Z i ) , where δ is a metric used in calculating distance in parameter space, and τ is a temperature hyperparameter balancing the impact of similar pairs and negative examples [Zhang et al., 2018, Ravula et al., 2021].
Revised: Z i δ ( Z j , Z i ) , where δ is a metric used in calculating distance in parameter space, and τ is a temperature hyperparameter balancing the impact of similar pairs and negative examples [Zhang et al., 2018, Ravula et al., 2021].

Original: In this way, we drive the model to embed data with similar values of the parameter ϕ
Revised: In this way, we drive the model to embed data with similar values of the parameter ϕ

Original: in similar locations while repulsing data with unrelated parameters, resulting in an embedding that respects the latent structure of the parameter and trajectory pairs.
Revised: in similar locations while repulsing data with unrelated parameters, resulting in an embedding that respects the latent structure of the parameter and trajectory pairs.

Original: Emulators: Emulating the dynamics of a simulator by training a deep neural network ˆ g θ ( · ) is a natural approach to easing the computational burden of parameter estimation.
Revised: Emulators: Emulating the dynamics of a simulator by training a deep neural network ˆ g θ ( · ) is a natural approach to easing the computational burden of parameter estimation.

Original: However, recent works in emulating dynamics focus on the fixed-parameter setting [Krishnapriyan et al., 2021] and cannot be easily generalized to multi-parameter settings (i.e. the emulator only approximates H ( ϕ ) for a single fixed ϕ instead of a range of ϕ s, as we need for parameter estimation).
Revised: However, recent works in emulating dynamics focus on the fixed-parameter setting [Krishnapriyan et al., 2021] and cannot be easily generalized to multi-parameter settings (i.e. the emulator only approximates H ( ϕ ) for a single fixed ϕ instead of a range of ϕ s, as we need for parameter estimation).

Original: Learning emulators for a range of parameter values is challenging in part due to the dimension of the dynamics to be emulated (i.e. the dimension of the Z s).
Revised: Learning emulators for a range of parameter values is challenging in part due to the dimension of the dynamics to be emulated (i.e. the dimension of the Z s).

Original: To address this problem, we propose to learn an emulator ˆ g θ of g θ :
Revised: To address this problem, we propose to learn an emulator ˆ g θ of g θ :

Original: =
Revised: =

Original: f θ ◦ H , which represents the mapping from the parameter space ( ϕ ) to the latent representation space
Revised: f θ ◦ H , which represents the mapping from the parameter space ( ϕ ) to the latent representation space

Original: f θ ( Z ) .
Revised: f θ ( Z ) .

Original: which is much lower-dimensional than Z and hence easier to learn with limited training samples.
Revised: which is much lower-dimensional than Z and hence easier to learn with limited training samples.

Original: We say that the trajectory encoder and the emulator are perfectly aligned when
Revised: We say that the trajectory encoder and the emulator are perfectly aligned when

Original: ˆ g θ ( ϕ ) =
Revised: ˆ g θ ( ϕ ) =

Original: f θ ( Z ) , ∀ ( ϕ , Z ) ∼ p ϕ , Z .
Revised: f θ ( Z ) , ∀ ( ϕ , Z ) ∼ p ϕ , Z .

Original: Under perfect alignment, given a test Z , we may seek to minimize
Revised: Under perfect alignment, given a test Z , we may seek to minimize

Original: [Equation]
Revised: [Equation]

Original: instead of (1) with much faster objective function evaluations.
Revised: instead of (1) with much faster objective function evaluations.

Original: Learning the emulator via CLIP:
Revised: Learning the emulator via CLIP:

Original: Pretraining an embedding f θ and then learning an emulator ˆ g θ with an L 2 loss is hard to optimize; the different architectures, necessitated by the very different dimensions of ϕ and Z , make feature alignment challenging and the optimization prone to poor local f θ ( Z )
Revised: Pretraining an embedding f θ and then learning an emulator ˆ g θ with an L 2 loss is hard to optimize; the different architectures, necessitated by the very different dimensions of ϕ and Z , make feature alignment challenging and the optimization prone to poor local f θ ( Z )

Original: ̂ ϕ
Revised: ̂ ϕ

Original: B ack - bo n e Z
Revised: B ack - bo n e Z

Original: M L P
Revised: M L P

Original: Embedding network f θ ℓ ZZ ( θ ; τ ) ℓ MAPE ( θ ) ℓ Z ϕ ( θ ; τ ′ )
Revised: Embedding network f θ ℓ ZZ ( θ ; τ ) ℓ MAPE ( θ ) ℓ Z ϕ ( θ ; τ ′ )

Original: ℓ ϕϕ ( θ ; τ ) ϕ
Revised: ℓ ϕϕ ( θ ; τ ) ϕ

Original: ̂ g θ ( ϕ )
Revised: ̂ g θ ( ϕ )

Original: R e s i
Revised: R e s i

Original: d u a l
Revised: d u a l

Original: M L P
Revised: M L P

Original: Emulator ̂ g θ
Revised: Emulator ̂ g θ

Original: MLP Regression head h θ (a) Network setup ̂
Revised: MLP Regression head h θ (a) Network setup ̂

Original: g θ ( ˜ ϕ 1 )
Revised: g θ ( ˜ ϕ 1 )

Original: ̂ g θ ( ˜ ϕ 2 )
Revised: ̂ g θ ( ˜ ϕ 2 )

Original: ̂ g θ ( ˜ ϕ 3 )
Revised: ̂ g θ ( ˜ ϕ 3 )

Original: ̂ g θ ( ˜ ϕ 4 )
Revised: ̂ g θ ( ˜ ϕ 4 )

Original: ̂ g θ ( ϕ 1 )
Revised: ̂ g θ ( ϕ 1 )

Original: ̂
Revised: ̂

Original: g θ ( ϕ 2 )
Revised: g θ ( ϕ 2 )

Original: ̂ g θ ( ϕ 3 )
Revised: ̂ g θ ( ϕ 3 )

Original: ̂
Revised: ̂

Original: g θ ( ϕ 4 )
Revised: g θ ( ϕ 4 )

Original: Z
Revised: Z

Original: P a r a m e t e r
Revised: P a r a m e t e r

Original: E n c od e r ̂ g θ ϕ ( ϕ 1 , ˜ ϕ 1 )( ϕ 2 , ˜ ϕ 1 ) ( ϕ 1 , ˜ ϕ 2 )( ϕ 2 , ˜ ϕ 2 ) ( ϕ 1 , ˜ ϕ 3 )( ϕ 2 , ˜ ϕ 3 ) ( ϕ 1 , ˜ ϕ 4 )( ϕ 2 , ˜ ϕ 4 ) ( ϕ 3 , ˜ ϕ 1 )( ϕ 4 , ˜ ϕ 1 ) ( ϕ 3 , ˜ ϕ 2 )( ϕ 4 , ˜ ϕ 2 ) ( ϕ 3 , ˜ ϕ 3 )( ϕ 4 , ˜ ϕ 3 ) ( ϕ 3 , ˜ ϕ 4 )( ϕ 4 , ˜ ϕ 4 ) ( Z 1 , ϕ 1 )( Z 1 , ϕ 2 ) ( Z 2 , ϕ 1 )( Z 2 , ϕ 2 ) ( Z 3 , ϕ 1 )( Z 3 , ϕ 2 ) ( Z 4 , ϕ 1 )( Z 4 , ϕ 2 ) ( Z 1 , ϕ 3 )( Z 1 , ϕ 4 ) ( Z 2 , ϕ 3 )( Z 2 , ϕ 4 ) ( Z 3 , ϕ 3 )( Z 3 , ϕ 4 ) ( Z 4 , ϕ 3 )( Z 4 , ϕ 4 ) ℓ ϕϕ ( θ ; τ ) ( Z 1 ,˜
Revised: E n c od e r ̂ g θ ϕ ( ϕ 1 , ˜ ϕ 1 )( ϕ 2 , ˜ ϕ 1 ) ( ϕ 1 , ˜ ϕ 2 )( ϕ 2 , ˜ ϕ 2 ) ( ϕ 1 , ˜ ϕ 3 )( ϕ 2 , ˜ ϕ 3 ) ( ϕ 1 , ˜ ϕ 4 )( ϕ 2 , ˜ ϕ 4 ) ( ϕ 3 , ˜ ϕ 1 )( ϕ 4 , ˜ ϕ 1 ) ( ϕ 3 , ˜ ϕ 2 )( ϕ 4 , ˜ ϕ 2 ) ( ϕ 3 , ˜ ϕ 3 )( ϕ 4 , ˜ ϕ 3 ) ( ϕ 3 , ˜ ϕ 4 )( ϕ 4 , ˜ ϕ 4 ) ( Z 1 , ϕ 1 )( Z 1 , ϕ 2 ) ( Z 2 , ϕ 1 )( Z 2 , ϕ 2 ) ( Z 3 , ϕ 1 )( Z 3 , ϕ 2 ) ( Z 4 , ϕ 1 )( Z 4 , ϕ 2 ) ( Z 1 , ϕ 3 )( Z 1 , ϕ 4 ) ( Z 2 , ϕ 3 )( Z 2 , ϕ 4 ) ( Z 3 , ϕ 3 )( Z 3 , ϕ 4 ) ( Z 4 , ϕ 3 )( Z 4 , ϕ 4 ) ℓ ϕϕ ( θ ; τ ) ( Z 1 ,˜

Original: Z 1 )( Z 1 ,˜
Revised: Z 1 )( Z 1 ,˜

Original: Z 2 ) ( Z 2 ,˜ Z 1 )( Z 2 ,˜ Z 2 ) ( Z 3 ,˜ Z 1 )( Z 3 ,˜ Z 2 ) ( Z 4 ,˜
Revised: Z 2 ) ( Z 2 ,˜ Z 1 )( Z 2 ,˜ Z 2 ) ( Z 3 ,˜ Z 1 )( Z 3 ,˜ Z 2 ) ( Z 4 ,˜

Original: Z 1 )( Z 4 ,˜
Revised: Z 1 )( Z 4 ,˜

Original: Z 2 ) ( Z 1 ,˜ Z 3 )( Z 1 ,˜ Z 4 ) ( Z 2 ,˜ Z 3 )( Z 2 ,˜ Z 4 ) ( Z 3 ,˜ Z 3 )( Z 3 ,˜ Z 4 ) ( Z 4 ,˜ Z 3 )( Z 4 ,˜ Z 4 )
Revised: Z 2 ) ( Z 1 ,˜ Z 3 )( Z 1 ,˜ Z 4 ) ( Z 2 ,˜ Z 3 )( Z 2 ,˜ Z 4 ) ( Z 3 ,˜ Z 3 )( Z 3 ,˜ Z 4 ) ( Z 4 ,˜ Z 3 )( Z 4 ,˜ Z 4 )

Original: f θ (˜ Z 1 ) f
Revised: f θ (˜ Z 1 ) f

Original: θ (˜ Z 2 )
Revised: θ (˜ Z 2 )

Original: f
Revised: f

Original: θ (˜ Z 3 )
Revised: θ (˜ Z 3 )

Original: f
Revised: f

Original: θ (˜ Z 4 )
Revised: θ (˜ Z 4 )

Original: f θ ( Z 1 ) f
Revised: f θ ( Z 1 ) f

Original: θ ( Z 2 ) f
Revised: θ ( Z 2 ) f

Original: θ ( Z 3 ) f
Revised: θ ( Z 3 ) f

Original: θ ( Z 4 )
Revised: θ ( Z 4 )

Original: ̂ ϕ̂ ϕ̂ ϕ̂ ϕℓ ZZ
Revised: ̂ ϕ̂ ϕ̂ ϕ̂ ϕℓ ZZ

Original: ( θ ; τ ) ℓ Z ϕ
Revised: ( θ ; τ ) ℓ Z ϕ

Original: ( θ ; τ ′ )
Revised: ( θ ; τ ′ )

Original: Trajectory Encoder f θ (
Revised: Trajectory Encoder f θ (

Original: [Figure]
Revised: [Figure]

Original: losses and minimizers.
Revised: losses and minimizers.

Original: Alternatively, inspired by the recent success of CLIP [Radford et al., 2021, Ravula et al., 2021] in cross-domain representation learning, we use the following variant of InfoNCE loss to align the representation space between two networks:
Revised: Alternatively, inspired by the recent success of CLIP [Radford et al., 2021, Ravula et al., 2021] in cross-domain representation learning, we use the following variant of InfoNCE loss to align the representation space between two networks:

Original: [Equation]
Revised: [Equation]

Original: Intra-parameter domain contrastive loss:
Revised: Intra-parameter domain contrastive loss:

Original: In order to fully exploit data potential within parameter space and better guide the learning of the representation and preserve similarities within the parameter space, we add an intra-parameter contrastive loss and find empirically that it helps accelerate the training process:
Revised: In order to fully exploit data potential within parameter space and better guide the learning of the representation and preserve similarities within the parameter space, we add an intra-parameter contrastive loss and find empirically that it helps accelerate the training process:

Original: [Equation]
Revised: [Equation]

Original: where ˜ ϕ is selected through data augmentation by perturbing ϕ with small amount of noise (see §B.1) or by finding nearest neighbors in the training set.
Revised: where ˜ ϕ is selected through data augmentation by perturbing ϕ with small amount of noise (see §B.1) or by finding nearest neighbors in the training set.

Original: Regression head: Zimmermann et al.
Revised: Regression head: Zimmermann et al.

Original: [2021] shows that under certain assumptions (see §A.2), the feature encoder f θ implicitly learns to invert the data generating process H up to an affine transformation, implying we can estimate ϕ via h ( f θ ( Z )) for some linear function h .
Revised: [2021] shows that under certain assumptions (see §A.2), the feature encoder f θ implicitly learns to invert the data generating process H up to an affine transformation, implying we can estimate ϕ via h ( f θ ( Z )) for some linear function h .

Original: Inspired by their analysis, we add a regression head h θ (a single fully-connected linear layer) in our framework, and we ensure the affine relationship between h θ and f θ by letting them share the same backbone up to the last layers before the output.
Revised: Inspired by their analysis, we add a regression head h θ (a single fully-connected linear layer) in our framework, and we ensure the affine relationship between h θ and f θ by letting them share the same backbone up to the last layers before the output.

Original: Unified training procedure: To this end, we formulate the final loss of our E MBED & E MULATE method as
Revised: Unified training procedure: To this end, we formulate the final loss of our E MBED & E MULATE method as

Original: [Equation]
Revised: [Equation]

Original: ( θ ) , where λ ZZ , λ ϕϕ , λ Z ϕ and λ MAPE control relative loss weights.
Revised: ( θ ) , where λ ZZ , λ ϕϕ , λ Z ϕ and λ MAPE control relative loss weights.

Original: We show our basic setup in Fig. 1(a) and loss framework in Fig. 1(b).
Revised: We show our basic setup in Fig. 1(a) and loss framework in Fig. 1(b).

Original:  Experiments
Revised:  Experiments

Original: We conduct a numerical case study on the multiscale Lorenz-96 (L96) model [Lorenz, 1996], which is a common test model for climate models with both “fast” (high-frequency) and “slow” (lowfrequency) components and other geophysical applications [Majda and Harlim, 2012, Law and Stuart, 2012, Law et al., 2016, Brajard et al., 2020].
Revised: We conduct a numerical case study on the multiscale Lorenz-96 (L96) model [Lorenz, 1996], which is a common test model for climate models with both “fast” (high-frequency) and “slow” (lowfrequency) components and other geophysical applications [Majda and Harlim, 2012, Law and Stuart, 2012, Law et al., 2016, Brajard et al., 2020].

Original: It is a prototypical turbulent dynamical system “designed to mimic baroclinic turbulence in the midlatitude atmosphere” [Majda et al., 2010].
Revised: It is a prototypical turbulent dynamical system “designed to mimic baroclinic turbulence in the midlatitude atmosphere” [Majda et al., 2010].

Original: Its dynamics exhibit strong energy-conserving non-linearities, and for some settings of ϕ , it can exhibit strong chaotic turbulence.
Revised: Its dynamics exhibit strong energy-conserving non-linearities, and for some settings of ϕ , it can exhibit strong chaotic turbulence.

Original: Code is available at: https://github.com/roxie62/Embed-and-Emulate
Revised: Code is available at: https://github.com/roxie62/Embed-and-Emulate

Original: The governing equations of the L96 system are defined as follows:
Revised: The governing equations of the L96 system are defined as follows:

Original: [Equation]
Revised: [Equation]

Original: where X t ∈ R K denotes the slow variable at the t -th time stamp and Y t ∈ R KJ denotes the fast variable.
Revised: where X t ∈ R K denotes the slow variable at the t -th time stamp and Y t ∈ R KJ denotes the fast variable.

Original: We use Z t :=
Revised: We use Z t :=

Original: [ X t , Y t ] ∈ R K ( J +1) to denote the system state at the t -th time stamp.
Revised: [ X t , Y t ] ∈ R K ( J +1) to denote the system state at the t -th time stamp.

Original: We choose K = 36 and J = 10 throughout experiments in this section, as in Schneider et al.
Revised: We choose K = 36 and J = 10 throughout experiments in this section, as in Schneider et al.

Original: We set our baseline following Schneider et al.
Revised: We set our baseline following Schneider et al.

Original: Within the EnKI framework, they use
Revised: Within the EnKI framework, they use

Original: Runge-Kutta Dormand and Prince [1980] to compute the forward model H and optimize (1).
Revised: Runge-Kutta Dormand and Prince [1980] to compute the forward model H and optimize (1).

Original: After analyzing the physics information of the L96 system, they define the moment function as: m ( Z ) :=
Revised: After analyzing the physics information of the L96 system, they define the moment function as: m ( Z ) :=

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ¯ Y
Revised: ¯ Y

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: Y ⟩
Revised: Y ⟩

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Y
Revised: Y

Original: [Equation]
Revised: [Equation]

Original: ⟩ T ] , where ⟨·⟩ T denotes an empirical average over T time steps, ¯ Y denotes an average across the J Y channels, ⟨ ¯ Y ⟩ T =
Revised: ⟩ T ] , where ⟨·⟩ T denotes an empirical average over T time steps, ¯ Y denotes an average across the J Y channels, ⟨ ¯ Y ⟩ T =

Original: (1 /JT )
Revised: (1 /JT )

Original: (cid:80) Tt =1 (cid:80)
Revised: (cid:80) Tt =1 (cid:80)

Original: Jj =1 Y j,kt ∈ R K , and other quantities are computed similarly (see details in §B.1).
Revised: Jj =1 Y j,kt ∈ R K , and other quantities are computed similarly (see details in §B.1).

Original: m ( Z t ) ∈ R 5 K .
Revised: m ( Z t ) ∈ R 5 K .

Original: For the SBI methods, to suit the scenario we care about, i.e., estimating multiple ϕ
Revised: For the SBI methods, to suit the scenario we care about, i.e., estimating multiple ϕ

Original: i for multiple different observations
Revised: i for multiple different observations

Original: Z i at test time, we compare with the non-adaptive counterpart of the algorithms referenced in these papers.
Revised: Z i at test time, we compare with the non-adaptive counterpart of the algorithms referenced in these papers.

Original: For instance, we set the round number of SNL+ [Chen et al., 2020b] to be one.
Revised: For instance, we set the round number of SNL+ [Chen et al., 2020b] to be one.

Original: And use the fixed training dataset for SNL+ and NPE-C
Revised: And use the fixed training dataset for SNL+ and NPE-C

Original: [Greenberg et al., 2019].
Revised: [Greenberg et al., 2019].

Original: We conduct experiments and demonstrate the efficacy when measuring (1) the accuracy of averaged point estimates, (2) the accuracy of uncertainty quantification, (3) the empirical computational complexity.
Revised: We conduct experiments and demonstrate the efficacy when measuring (1) the accuracy of averaged point estimates, (2) the accuracy of uncertainty quantification, (3) the empirical computational complexity.

Original: For all of our experiments (including SNL+ and NPE-C), we use ResNet-34 [He et al., 2016] as the backbone of the trajectory encoder.
Revised: For all of our experiments (including SNL+ and NPE-C), we use ResNet-34 [He et al., 2016] as the backbone of the trajectory encoder.

Original: We apply average pooling at the last layer to generate a 512-d hidden vector.
Revised: We apply average pooling at the last layer to generate a 512-d hidden vector.

Original: For contrastive learning, we project the 512-d hidden vector into a 128-d feature vector as the output of the trajectory encoder.
Revised: For contrastive learning, we project the 512-d hidden vector into a 128-d feature vector as the output of the trajectory encoder.

Original: For the regression head, we map the 512-d hidden vector to a 4-d parameter vector.
Revised: For the regression head, we map the 512-d hidden vector to a 4-d parameter vector.

Original: It is important to note that activation functions are not used in the regression head to ensure an affine mapping between two projections.
Revised: It is important to note that activation functions are not used in the regression head to ensure an affine mapping between two projections.

Original: We parameterize the emulator network with a ResNet backbone, and to enhance the representation power of the network, we use five residual connection blocks (see §B.2).
Revised: We parameterize the emulator network with a ResNet backbone, and to enhance the representation power of the network, we use five residual connection blocks (see §B.2).

Original: We find this increases the alignment between the two representation spaces.
Revised: We find this increases the alignment between the two representation spaces.

Original: Within each skip block, there is a residual connection between the input layer and the output.
Revised: Within each skip block, there is a residual connection between the input layer and the output.

Original: The implementation details for training are shown in §B.1.
Revised: The implementation details for training are shown in §B.1.

Original: 5.1 Higher quality estimates with lower sample cost
Revised: 5.1 Higher quality estimates with lower sample cost

Original: In this section, we evaluate our method with different training sizes in the perfect-model setting where no noise is injected in observations.
Revised: In this section, we evaluate our method with different training sizes in the perfect-model setting where no noise is injected in observations.

Original: These experiments demonstrate the usefulness of our method in realistic scenarios where forward models are prohibitively expensive and only limited quantities of training data are available.
Revised: These experiments demonstrate the usefulness of our method in realistic scenarios where forward models are prohibitively expensive and only limited quantities of training data are available.

Original: We train our E MBED & E MULATE method and the supervised regression baselines when the training size equals 500 and 1 , 000 .
Revised: We train our E MBED & E MULATE method and the supervised regression baselines when the training size equals 500 and 1 , 000 .

Original: The training samples are sampled uniformly from ϕ min =
Revised: The training samples are sampled uniformly from ϕ min =

Original: [ F min , h min , c min , b min ] =
Revised: [ F min , h min , c min , b min ] =

Original: [ − 5 , 0 , 0 . 1 , 0] to ϕ max =
Revised: [ − 5 , 0 , 0 . 1 , 0] to ϕ max =

Original: [ F max , h max , c max , b max ] =
Revised: [ F max , h max , c max , b max ] =

Original: [20 , 5 , 25 , 25] , and each simulation is of length 100, with dt = 0 .
Revised: [20 , 5 , 25 , 25] , and each simulation is of length 100, with dt = 0 .

Original: We then sample 200 testing samples uniformly from a narrower range of [ − 3 , 0 . 5 , 2 , 2] to [18 , 4 . 5 , 23 , 23] .
Revised: We then sample 200 testing samples uniformly from a narrower range of [ − 3 , 0 . 5 , 2 , 2] to [18 , 4 . 5 , 23 , 23] .

Original: As assumed in Schneider et al.
Revised: As assumed in Schneider et al.

Original: [2017], each testing sample is of length 1 , 000 , with dt = 0 .
Revised: [2017], each testing sample is of length 1 , 000 , with dt = 0 .

Original: The EnKI prior used for baselines (i.e. minimizing (1)) is a Gaussian distribution with means at the middle of the range of the testing instances, diagonal covariance matrix, and variances broad enough that all test samples are within 2 σ of the mean, denoted p ϕ , fixed (see §B.1).
Revised: The EnKI prior used for baselines (i.e. minimizing (1)) is a Gaussian distribution with means at the middle of the range of the testing instances, diagonal covariance matrix, and variances broad enough that all test samples are within 2 σ of the mean, denoted p ϕ , fixed (see §B.1).

Original: Within the context of the E MBED & E MULATE approach, we adopt an empirical Bayes approach: for each test instance, we compute ˆ ϕ using the regression head and use these values as the prior means when using EnKI to minimize (4); we denote this prior p ϕ , empB .
Revised: Within the context of the E MBED & E MULATE approach, we adopt an empirical Bayes approach: for each test instance, we compute ˆ ϕ using the regression head and use these values as the prior means when using EnKI to minimize (4); we denote this prior p ϕ , empB .

Original: The impact of the empirical Bayes approach is explored in §5.3.
Revised: The impact of the empirical Bayes approach is explored in §5.3.

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: Reported in minutes, total time for E MBED & E MULATE , the supervised regression, or NPE-C are 1 .

Revised: 19% of EnKI with Runge-Kutta.

Revised: All neural approaches are trained with 4 GPUs and tested with

Revised: GPU.

Revised: Both training data generation and EnKI w/o Learning are run with 32 CPU Cores.

Original: We first evaluate the accuracy of point estimates using the mean and median absolute percentage error (MAPE and MdAPE), see §B.1.
Revised: We first evaluate the accuracy of point estimates using the mean and median absolute percentage error (MAPE and MdAPE), see §B.1.

Original: In Table 1 and Table 2, we see that E MBED & E MULATE evaluated with the regression head ( h θ ) and EnKI guarantees similar performance in terms of averaged accuracy.
Revised: In Table 1 and Table 2, we see that E MBED & E MULATE evaluated with the regression head ( h θ ) and EnKI guarantees similar performance in terms of averaged accuracy.

Original: However, the regression head ( h θ ) is slightly better, which may be explained by our empirical observation that when estimates from h θ with E MBED & E MULATE are far from the true parameter values, then using this estimate as the prior mean for EnKI can worsen the estimate.
Revised: However, the regression head ( h θ ) is slightly better, which may be explained by our empirical observation that when estimates from h θ with E MBED & E MULATE are far from the true parameter values, then using this estimate as the prior mean for EnKI can worsen the estimate.

Original: Despite this challenge, the EnKI has the advantage of providing uncertainty estimates, which are discussed below and reflected in Table 3.
Revised: Despite this challenge, the EnKI has the advantage of providing uncertainty estimates, which are discussed below and reflected in Table 3.

Original: Compared with other methods, it is clear that E MBED & E MULATE yields a significant improvement in accuracy, especially for the parameter affecting high-frequency dynamics (e.g., c ).
Revised: Compared with other methods, it is clear that E MBED & E MULATE yields a significant improvement in accuracy, especially for the parameter affecting high-frequency dynamics (e.g., c ).

Original: In particular,
Revised: In particular,

Original: NPE-C
Revised: NPE-C

Original: [Greenberg et al., 2019] performs worse with a smaller training set size and requires a longer training time to converge; SNL+ [Chen et al., 2020b] using MCMC or rejection sampling suffers from increasing evaluation time as the number of test samples increases.
Revised: [Greenberg et al., 2019] performs worse with a smaller training set size ; SNL+ [Chen et al., 2020b] using MCMC or rejection sampling suffers from increasing evaluation time as the number of test samples increases.

Original: Moreover, when compared to the classical method of running EnKI using a predefined moment function and expensive numerical solvers, E MBED & E MULATE yields better performance in terms of accuracy and computation time.
Revised: Moreover, when compared to the classical method of running EnKI using a predefined moment function and expensive numerical solvers, E MBED & E MULATE yields better performance in terms of accuracy and computation time.

Original: Last, our E MBED & E MULATE framework also performs well relative to supervised regression for smaller training sets.
Revised: Last, our E MBED & E MULATE framework also performs well relative to supervised regression for smaller training sets.

Original: We then evaluate our results based on the continuous ranked probability score (CRPS)
Revised: We then evaluate our results based on the continuous ranked probability score (CRPS)

Original: [Hersbach, 2000, Zamo and Naveau, 2018, Pappenberger et al., 2015].
Revised: [Hersbach, 2000, Zamo and Naveau, 2018, Pappenberger et al., 2015].

Original: CRPS is an important metric used in quantifying uncertainty (e.g., in weather forecasts).
Revised: CRPS is an important metric used in quantifying uncertainty (e.g., in weather forecasts).

Original: It measures the accuracy of estimated posterior distributions and is defined as
Revised: It measures the accuracy of estimated posterior distributions and is defined as

Original: [Equation]
Revised: [Equation]

Original: where ˆ ϕ ∗ j represents the j -th component of the estimate parameter vector, ϕ ∗ j represents the jth component of the true parameter vector, C is the cumulative density function of the ensemble estimates with C ( y )
Revised: where ˆ ϕ ∗ j represents the j -th component of the estimate parameter vector, ϕ ∗ j represents the jth component of the true parameter vector, C is the cumulative density function of the ensemble estimates with C ( y )

Original: =
Revised: =

Original: P ( ϕ ∗ j ≤ y ) , and U is the Heaviside step function.
Revised: P ( ϕ ∗ j ≤ y ) , and U is the Heaviside step function.

Original: For a deterministic estimate from supervised regression, CRPS is equal to the mean absolute error (MAE).
Revised: For a deterministic estimate from supervised regression, CRPS is equal to the mean absolute error (MAE).

Original: Table 3 shows that E MBED & E MULATE achieves high-accuracy estimates of uncertainty.
Revised: Table 3 shows that E MBED & E MULATE achieves high-accuracy estimates of uncertainty.

Original: n
Revised: n

Original: F ↓ h ↓ c ↓ b ↓
Revised: F ↓ h ↓ c ↓ b ↓

Original: EnKI w/o Learning 0.
Revised: EnKI w/o Learning 0.

Original: h θ w/ E MBED & E MULATE 0.
Revised: h θ w/ E MBED & E MULATE 0.

Original: EnKI
Revised: EnKI

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Supervised Regression 0.
Revised: Supervised Regression 0.

Original: NPE-C 1.
Revised: NPE-C 0.

Original: SNL+ 1.
Revised: SNL+ 1.

Original: h θ w/ E MBED & E MULATE 0.
Revised: h θ w/ E MBED & E MULATE 0.

Original: EnKI w/
Revised: EnKI w/

Original: [Table]
Revised: [Table]

Original: MBED &
Revised: MBED &

Original: [Table]
Revised: [Table]

Original: MULATE 0.
Revised: MULATE 0.

Original: Supervised Regression 0.
Revised: Supervised Regression 0.

Original: NPE-C 1.
Revised: NPE-C 0.

Original: [Table]
Revised: [Table]

Original: 5.2 Visualizing uncertainty with noisy observations
Revised: 5.2 Visualizing uncertainty with noisy observations

Original: In this subsection we go beyond the perfect-model setting and evaluate our method in the realistic scenarios where observations are noisy and obtained through: Z = H ( ϕ ) + η  where η ∼ N (0 , r Γ ) , where Γ is the temporal covariance of the trajectory Z and r is a scaling value.
Revised: In this subsection we go beyond the perfect-model setting and evaluate our method in the realistic scenarios where observations are noisy and obtained through: Z = H ( ϕ ) + η , where η ∼ N (0 , r Γ ) , where Γ is the temporal covariance of the trajectory Z and r is a scaling value.

Original: Specifically, we set ϕ
Revised: Specifically, we set ϕ

Original: =
Revised: =

Original: [10 , 1 , 10 , 10] following Schneider et al.
Revised: [10 , 1 , 10 , 10] following Schneider et al.

Original: [2017] to ensure we are in the chaotic regime.
Revised: [2017] to ensure we are in the chaotic regime.

Original: We use model trained with n = 4 , 000 data samples and compare the results obtained in both the noiseless and noisy cases.
Revised: We use model trained with n = 4 , 000 data samples and compare the results obtained in both the noiseless and noisy cases.

Original: For this visualization, both methods use the fixed prior from §5.1 based on the range of test values, so any differences observed are not caused by differences in the prior, but rather differences in the choice of objective and the method of computing the forward model
Revised: For this visualization, both methods use the fixed prior from §5.1 based on the range of test values, so any differences observed are not caused by differences in the prior but rather differences in the choice of objective and the method of computing the forward model

Original: As shown in Fig.
Revised: As shown in Fig.

Original: E MBED & E MULATE is more robust to noise than the baseline method with numerical solvers and predefined moments.
Revised: E MBED & E MULATE is more robust to noise than the baseline method with numerical solvers and predefined moments.

Original: Reconstructed posterior distributions learned with
Revised: Reconstructed posterior distributions learned with

Original: E MBED & E MULATE in Fig.
Revised: E MBED & E MULATE in Fig.

Original: 2 are more consistent with increasing noise levels, especially for the parameters F and b .
Revised: 2 are more consistent with increasing noise levels, especially for the parameters F and b .

Original: 5.3 Ablation study: role of the regression head
Revised: 5.3 Ablation study: role of the regression head

Original: In this section we empirically verify the utility of the regression head h θ in our E MBED & E MULATE framework.
Revised: In this section, we empirically verify the utility of the regression head h θ in our E MBED & E MULATE framework.

Original: We use n = 4000 training samples in two settings: first, we use the full E MBED &
Revised: We use n = 4000 training samples in two settings: first, we use the full E MBED &

Original: E MULATE model of §5.1; second, we discard the regression head (i.e. fix h θ = 0 ) while keeping intra- and inter-domain contrastive losses.
Revised: E MULATE model of §5.1; second, we discard the regression head (i.e. fix h θ = 0 ) while keeping intra- and inter-domain contrastive losses.

Original: For both trained models, we run the EnKI.
Revised: For both trained models, we run the EnKI.

Original: We run the
Revised: We run the

Original: EnKI using components from the E MBED & E MULATE framework for two choices of prior: first, we use p ϕ , fixed , and second, we use p ϕ , empB .
Revised: EnKI using components from the E MBED & E MULATE framework for two choices of prior: first, we use p ϕ , fixed , and second, we use p ϕ , empB .

Original: Note the empirical Bayes prior is only possible with the regression head.
Revised: Note the empirical Bayes prior is only possible with the regression head.

Original: These priors are detailed in §5.1 and §B.1.
Revised: These priors are detailed in §5.1 and §B.1.

Original: 9 10 11 F 0.
Revised: 9 10 11 F 0.

Original: M o m e n t s a n d R K 0.8 1.0 1.2 h      7.5 10.0 12.5 c 0.
Revised: M o m e n t s a n d R K 0.8 1.0 1.2 h      7.5 10.0 12.5 c 0.

Original: 12 b 0.
Revised: 12 b 0.

Original: 9 10 11 F 0.
Revised: 9 10 11 F 0.

Original: E m b e d a n d E m u l a t e 0.8 1.0 1.2 h      7.5 10.0 12.5 c 0.
Revised: E m b e d a n d E m u l a t e 0.8 1.0 1.2 h      7.5 10.0 12.5 c 0.

Original: 12 b 0.
Revised: 12 b 0.

Original: r=0r=0.1r=0.
Revised: r=0r=0.1r=0.

Original: Truth
Revised: Truth

Original: Table 4 shows that having the regression component of the loss complement the contrastive losses yields a substantial improvement in parameter estimation accuracy using E MBED & E MULATE .
Revised: Table 4 shows that having the regression component of the loss complement the contrastive losses yields a substantial improvement in parameter estimation accuracy using E MBED & E MULATE .

Original: In other words, explicitly training the embedding function and emulator for the parameter estimation task yields an emulator that is far more accurate during parameter estimation than training a generic embedding and generic emulator.
Revised: In other words, explicitly training the embedding function and emulator for the parameter estimation task yields an emulator that is far more accurate during parameter estimation than training a generic embedding and generic emulator.

Original: Furthermore, this table illustrates the efficacy of our empirical
Revised: Furthermore, this table illustrates the efficacy of our empirical

Original: Bayes procedure – i.e., using our learned regressor to alter the prior used by EnKI.
Revised: Bayes procedure – i.e., using our learned regressor to alter the prior used by EnKI.

Original: The medians are slightly improved while the means are strongly improved, suggesting that the empirical Bayes procedure is particularly helpful in the tails.
Revised: The medians are slightly improved while the means are strongly improved, suggesting that the empirical Bayes procedure is particularly helpful in the tails.

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: that having the regression component of the loss complement the contrastive losses yields a substantial improvement in parameter estimation accuracy using E MBED & E MULATE .
Revised: that having the regression component of the loss complement the contrastive losses yields a substantial improvement in parameter estimation accuracy using E MBED & E MULATE .

Original:  Conclusions
Revised:  Conclusions

Original: The proposed E MBED & E MULATE framework trains an emulator of a complex simulation to facilitate parameter estimation.
Revised: The proposed E MBED & E MULATE framework trains an emulator of a complex simulation to facilitate parameter estimation.

Original: Unlike generic emulation methods, which can lead to poor parameter estimates and require expert domain knowledge to construct, our method (a) leverages a contrastive learning framework coupled with a regression head to jointly learn a low-dimensional embedding of simulator outputs that can be used to construct an objective function for parameter estimation and (b) yields an emulator that can be used within an optimization framework such as the EnKI to produce accurate parameter estimates in 1.19% of the computation time of an approach using classical numerical methods, even accounting for the time required to generate training samples.
Revised: Unlike generic emulation methods, which can lead to poor parameter estimates and require expert domain knowledge to construct, our method (a) leverages a contrastive learning framework coupled with a regression head to jointly learn a low-dimensional embedding of simulator outputs that can be used to construct an objective function for parameter estimation and (b) yields an emulator that can be used within an optimization framework such as the EnKI to produce accurate parameter estimates in 1.19% of the computation time of an approach using classical numerical methods, even accounting for the time required to generate training samples.

Original: We explore our approach in the context of earth system modeling as described by Schneider et al.
Revised: We explore our approach in the context of earth system modeling as described by Schneider et al.

Original: [2017] and Cleary et al.
Revised: [2017] and Cleary et al.

Original: [2021] and hypothesize that these tools can facilitate improved climate forecasts that account for uncertainties and cover the full range of possible outcomes.
Revised: [2021] and hypothesize that these tools can facilitate improved climate forecasts that account for uncertainties and cover the full range of possible outcomes.

Original: The social impacts of improved climate forecasting are positive if acted upon.
Revised: The social impacts of improved climate forecasting are positive if acted upon.

Original: While learned emulators are gaining in popularity, we as a community still know little about which systems are more or less challenging to emulate or how to design task-specific emulators more generally.
Revised: While learned emulators are gaining in popularity, we as a community still know little about which systems are more or less challenging to emulate or how to design task-specific emulators more generally.

Original: There are further opportunities for exploring emulators for parameter estimation using optimization frameworks beside the EnKI.
Revised: There are further opportunities for exploring emulators for parameter estimation using optimization frameworks beside the EnKI.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: We thank the anonymous reviewers and area chair for their helpful comments.
Revised: We thank the anonymous reviewers and area chair for their helpful comments.

Original: We thank Owen Melia and Xiao Zhang for helpful discussion and proofreading our paper.
Revised: We thank Owen Melia and Xiao Zhang for helpful discussion and proofreading our paper.

Original: RJ and RW are partially supported by DOE DE-AC02-06CH113575, DE-SC0022232, AFOSR FA9550-18-1-0166, NSF OAC-1934637,
Revised: RJ and RW are partially supported by DOE DE-AC02-06CH113575, DE-SC0022232, AFOSR FA9550-18-1-0166, NSF OAC-1934637,

Original: NSF DMS-2023109, and NSF DGE-2022023.
Revised: NSF DMS-2023109, and NSF DGE-2022023.

Original: A Algorithms and analysis A.1 Ensemble Kalman Inversion
Revised: A Algorithms and analysis A.1 Ensemble Kalman Inversion

Original: Ensemble Kalman Inversion (EnKI) gained its popularity in addressing Bayesian inverse problems since its proposal [Iglesias et al., 2013].
Revised: Ensemble Kalman Inversion (EnKI) gained its popularity in addressing Bayesian inverse problems since its proposal [Iglesias et al., 2013].

Original: It is a derivative-free optimization method for objective functions of the form [Chada et al., 2020] :
Revised: It is a derivative-free optimization method for objective functions of the form [Chada et al., 2020] :

Original: [Equation]
Revised: [Equation]

Original: where y measure is the measured data, g is a given map and R is a symmetric positive definite matrix representing the data measurement precision (we explain how different approaches in our paper using different y measure and R below).
Revised: where y measure is the measured data, g is a given map and R is a symmetric positive definite matrix representing the data measurement precision (we explain how different approaches in our paper using different y measure and R below).

Original: The EnKI estimates a posterior distribution of ϕ , which is approximated using an ensemble of particles.
Revised: The EnKI estimates a posterior distribution of ϕ , which is approximated using an ensemble of particles.

Original: At each iteration, the EnKI consists of two steps: the prediction step and the analysis step.
Revised: At each iteration, the EnKI consists of two steps: the prediction step and the analysis step.

Original: In the prediction step, we apply the forward model g to each particle of the ensemble.
Revised: In the prediction step, we apply the forward model g to each particle of the ensemble.

Original: In the analysis step, each particle is updated by using artificially perturbed observations and forward-mapped particles.
Revised: In the analysis step, each particle is updated by using artificially perturbed observations and forward-mapped particles.

Original: The EnKI is summarized in Algorithm 1.
Revised: The EnKI is summarized in Algorithm 1.

Original: Algorithm 1 EnKI ( y measure , g, p ϕ , α, R , M, N )
Revised: Algorithm 1 EnKI ( y measure , g, p ϕ , α, R , M, N )

Original: Prediction propagate the ensemble of particles to data space { g (cid:0) ˆ ϕ
Revised: Prediction propagate the ensemble of particles to data space { g (cid:0) ˆ ϕ

Original: ( i,m ) (cid:1) }
Revised: ( i,m ) (cid:1) }

Original: Mm =1 by applying the forward model g M times.
Revised: Mm =1 by applying the forward model g M times.

Original: From the ensemble, calculate the empirical means and covariance matrices:
Revised: From the ensemble, calculate the empirical means and covariance matrices:

Original: [Equation]
Revised: [Equation]

Original: Analysis step : calculate the Kalman gain matrix
Revised: Analysis step : calculate the Kalman gain matrix

Original: [Equation]
Revised: [Equation]

Original: artificially perturb the data y ( i,m ) measure =
Revised: artificially perturb the data y ( i,m ) measure =

Original: y measure + η
Revised: y measure + η

Original: ( i,m ) where η ( i,m ) ∼ N (0 , α − 1 R ) , and update the ensemble of particles:
Revised: ( i,m ) where η ( i,m ) ∼ N (0 , α − 1 R ) , and update the ensemble of particles:

Original: [Equation]
Revised: [Equation]

Original: 4: end for
Revised: 4: end for

Original: Output: { ˆ ϕ ( N
Revised: Output: { ˆ ϕ ( N

Original: [Equation]
Revised: [Equation]

Original: m
Revised: m

Original: [Equation]
Revised: [Equation]

Original: When using E MBED & E MULATE to minimize the objective function (4), y measure corresponds to the embedding of the raw trajectory f θ ( Z ) , the forward model g corresponds to the emulator ˆ g θ , and
Revised: When using E MBED & E MULATE to minimize the objective function (4), y measure corresponds to the embedding of the raw trajectory f θ ( Z ) , the forward model g corresponds to the emulator ˆ g θ , and

Original: R corresponds to an identity matrix.
Revised: R corresponds to an identity matrix.

Original: When using the baseline method to minimize the objective function (1), y measure corresponds to the moment vector m
Revised: When using the baseline method to minimize the objective function (1), y measure corresponds to the moment vector m

Original: ( Z ) , g corresponds to the m ◦ H where
Revised: ( Z ) , g corresponds to the m ◦ H where

Original: H is computed using Runge-Kutta method, and R is a diagonal matrix with j -th diagonal entry
Revised: H is computed using Runge-Kutta method, and R is a diagonal matrix with j -th diagonal entry

Original: R j,j :=
Revised: R j,j :=

Original: Var[ m ( Z ) j ] .
Revised: Var[ m ( Z ) j ] .

Original: A.2 Contrastive learning and inverse problems
Revised: A.2 Contrastive learning and inverse problems

Original: In this section, we briefly reviewed the assumptions and results of Zimmermann et al.
Revised: In this section, we briefly reviewed the assumptions and results of Zimmermann et al.

Original: [2021] and explain how we relate the analysis to our work in addressing the parameter estimation problem.
Revised: [2021] and explain how we relate the analysis to our work in addressing the parameter estimation problem.

Original: Recent literature studies the reason why contrastive learned representations can be successfully extended to multiple downstream tasks [Arora et al., 2019, Wang and Isola, 2020, Zimmermann et al., 2021].
Revised: Recent literature studies the reason why contrastive learned representations can be successfully extended to multiple downstream tasks [Arora et al., 2019, Wang and Isola, 2020, Zimmermann et al., 2021].

Original: Among them, a recent work from Zimmermann et al.
Revised: Among them, a recent work from Zimmermann et al.

Original: [2021] points out that under certain assumptions, the representations f θ learned in the contrastive framework to minimize ℓ ZZ ( θ ; τ ) (shown in (3)) inverts the underlying data generating process H − 1 ( Z ) of observed data up to an affine transformation.
Revised: [2021] points out that under certain assumptions, the representations f θ learned in the contrastive framework to minimize ℓ ZZ ( θ ; τ ) (shown in (3)) inverts the underlying data generating process H − 1 ( Z ) of observed data up to an affine transformation.

Original: Assumptions: Zimmermann et al.
Revised: Assumptions: Zimmermann et al.

Original: [2021] firstly make assumptions on the data generating process .
Revised: [2021] firstly make assumptions on the data generating process .

Original: They assume the generator H : Φ → Z is an injective function with Φ ∈ S k − 1 in a unit hypersphere and the distribution of ϕ ∈ Φ is uniform.
Revised: They assume the generator H : Φ → Z is an injective function with Φ ∈ S k − 1 in a unit hypersphere and the distribution of ϕ ∈ Φ is uniform.

Original: Second, they assume the conditional distribution p (˜ ϕ | ϕ ) of a “positive parameter” ˜
Revised: Second, they assume the conditional distribution p (˜ ϕ | ϕ ) of a “positive parameter” ˜

Original: ϕ given ϕ is a von Mises-Fisher (vMF) distribution with p (˜ ϕ | ϕ )
Revised: ϕ given ϕ is a von Mises-Fisher (vMF) distribution with p (˜ ϕ | ϕ )

Original: =
Revised: =

Original: c − 1 p exp( κ ⟨ ˜ ϕ , ϕ ⟩ ) where c p is the normalizing constant and κ is the concentration parameter.
Revised: c − 1 p exp( κ ⟨ ˜ ϕ , ϕ ⟩ ) where c p is the normalizing constant and κ is the concentration parameter.

Original: Third, they assume the learned representation f θ : Z → S k − 1 is in the unit hypersphere where the dimension of the representation is equal to the dimension of the parameters.
Revised: Third, they assume the learned representation f θ : Z → S k − 1 is in the unit hypersphere where the dimension of the representation is equal to the dimension of the parameters.

Original: They further assume the learned conditional distribution of ˜ ϕ given ϕ with the contrastive trained neural network is a vMF distribution with q θ (˜ ϕ | ϕ ) =
Revised: They further assume the learned conditional distribution of ˜ ϕ given ϕ with the contrastive trained neural network is a vMF distribution with q θ (˜ ϕ | ϕ ) =

Original: C q ( ϕ ) − 1 exp( ⟨ f θ ( H (˜ ϕ ) , f θ ( H ( ϕ )) /τ ⟩ ) where C q ( ϕ ) is the partition function.
Revised: C q ( ϕ ) − 1 exp( ⟨ f θ ( H (˜ ϕ ) , f θ ( H ( ϕ )) /τ ⟩ ) where C q ( ϕ ) is the partition function.

Original: In the context of the parameter estimation problem of this paper, the relationship proved by Zimmermann et al.
Revised: In the context of the parameter estimation problem of this paper, the relationship proved by Zimmermann et al.

Original: [2021], that f θ ( Z )
Revised: [2021], that f θ ( Z )

Original: =
Revised: =

Original: AH − 1 ( Z ) =
Revised: AH − 1 ( Z ) =

Original: A ϕ where A is a full rank matrix, is clearly a desirable property as, if A were known, we could simply invert
Revised: A ϕ where A is a full rank matrix, is clearly a desirable property as, if A were known, we could simply invert

Original: A to obtain an unbiased estimate of ϕ .
Revised: A to obtain an unbiased estimate of ϕ .

Original: Empirically, we find that directly adding a linear regression head on top of f θ (which effectively serves as a mechanism for learning A − 1 ; see Fig. 1) provides faster convergence and better performance (as shown in §5.3).
Revised: Empirically, we find that directly adding a linear regression head on top of f θ (which effectively serves as a mechanism for learning A − 1 ; see Fig. 1) provides faster convergence and better performance (as shown in §5.3).

Original: B Implementation details B.1 Experiment details
Revised: B Implementation details B.1 Experiment details

Original: Lorenz 96 dataset: Initial conditions of the ODE/PDE systems affect the values of dynamical variables at future times.
Revised: Lorenz 96 dataset: Initial conditions of the ODE/PDE systems affect the values of dynamical variables at future times.

Original: To better fit in real-world scenarios, we simulate the training data used for
Revised: To better fit in real-world scenarios, we simulate the training data used for

Original: E MBED & E MULATE and the supervised regression, as well as the testing data used for evaluations of all methods, with random initial conditions sampled from the normal distributions.
Revised: E MBED & E MULATE and the supervised regression, as well as the testing data used for evaluations of all methods, with random initial conditions sampled from the normal distributions.

Original: During EnKI evaluations of the baseline we are required to specify initial conditions.
Revised: During EnKI evaluations of the baseline we are required to specify initial conditions.

Original: We follow Schneider et al.
Revised: We follow Schneider et al.

Original: [2017] and set the initial conditions of all simulations in the EnKI evaluations as random samples drawn from the observation.
Revised: [2017] and set the initial conditions of all simulations in the EnKI evaluations as random samples drawn from the observation.

Original: Note that, unlike the baseline method, E MBED & E MULATE directly learns the structural embeddings of ϕ and does not require initial conditions.
Revised: Note that, unlike the baseline method, E MBED & E MULATE directly learns the structural embeddings of ϕ and does not require initial conditions.

Original: Beyond the setup of initial conditions, there are two stages to generate our train and test data.
Revised: Beyond the setup of initial conditions, there are two stages to generate our train and test data.

Original: In the first stage, we generate a collection of ϕ i , and for each one, we run a Runge-Kutta solver to compute a corresponding Z i .
Revised: In the first stage, we generate a collection of ϕ i , and for each one, we run a Runge-Kutta solver to compute a corresponding Z i .

Original: In the second stage, we filter out some of the { ϕ i , Z i } generated during the first stage –whenever the
Revised: In the second stage, we filter out some of the { ϕ i , Z i } generated during the first stage –whenever the

Original: Runge-Kutta iterations are unsuccessful, generating NaN values, and whenever the Z i are degenerate, resulting in a standard deviation below 5e-5.
Revised: Runge-Kutta iterations are unsuccessful, generating NaN values, and whenever the Z i are degenerate, resulting in a standard deviation below 5e-5.

Original: Data processing : (a) Trajectory cropping: during the training of E MBED & E MULATE and the supervised regression, we randomly crop a collection of length-250 trajectories from lengthtrajectories, where each crop Z crop ∈ R 250 × 396 .
Revised: Data processing : (a) Trajectory cropping: during the training of E MBED & E MULATE and the supervised regression, we randomly crop a collection of length-250 trajectories from lengthtrajectories, where each crop Z crop ∈ R 250 × 396 .

Original: (b) “positive” data selection: E MBED & E MULATE trained with contrastive learning objectives tries to pull “anchor” data Z (or ϕ ) towards “positive” data
Revised: (b) “positive” data selection: E MBED & E MULATE trained with contrastive learning objectives tries to pull “anchor” data Z (or ϕ ) towards “positive” data

Original: ˜ Z in (3) (or ˜ ϕ
Revised: ˜ Z in (3) (or ˜ ϕ

Original: in (6)) and at the same time push “negative” data away.
Revised: in (6)) and at the same time push “negative” data away.

Original: To identify a “positive” trajectory ˜ Z (or parameter ˜ ϕ ), we use a two-step approach.
Revised: To identify a “positive” trajectory ˜ Z (or parameter ˜ ϕ ), we use a two-step approach.

Original: First, we select it in a supervised manner by measuring the distance between the “anchor” data Z i (or ϕ i ) and a sample Z j (or ϕ j ) in the parameter space.
Revised: First, we select it in a supervised manner by measuring the distance between the “anchor” data Z i (or ϕ i ) and a sample Z j (or ϕ j ) in the parameter space.

Original: We find the index ˜
Revised: We find the index ˜

Original: i of the nearest neighbor of the “anchor” data in the parameter space with ˜ i :=
Revised: i of the nearest neighbor of the “anchor” data in the parameter space with ˜ i :=

Original: arg min i ̸ =
Revised: arg min i ̸ =

Original: j δ ( ϕ i , ϕ j ) , where δ is a metric function defined as: δ ( ϕ i , ϕ j ) :
Revised: j δ ( ϕ i , ϕ j ) , where δ is a metric function defined as: δ ( ϕ i , ϕ j ) :

Original: = 12 (cid:8) APE( ϕ i ; ϕ j )
Revised: = 12 (cid:8) APE( ϕ i ; ϕ j )

Original: + APE( ϕ j ; ϕ i ) (cid:9) (see APE in §B.3).
Revised: + APE( ϕ j ; ϕ i ) (cid:9) (see APE in §B.3).

Original: However, since the training data is sampled from finite grids with limited size, the distance between the “anchor” data and its nearest neighbor δ ( ϕ i , ϕ ˜ i ) could be very large.
Revised: However, since the training data is sampled from finite grids with limited size, the distance between the “anchor” data and its nearest neighbor δ ( ϕ i , ϕ ˜ i ) could be very large.

Original: We apply a threshold filter and only accept a candidate nearest neighbor as “positive” data when the empirical distance δ ( ϕ i , ϕ ˜ i ) is below some threshold (e.g., 0 .
Revised: We apply a threshold filter and only accept a candidate nearest neighbor as “positive” data when the empirical distance δ ( ϕ i , ϕ ˜ i ) is below some threshold (e.g., 0 .

Original: 45 for n = 500 , 0 .
Revised: 45 for n = 500 , 0 .

Original: 4 for n = 1 , 000 and 4 , 000 ).
Revised: 4 for n = 1 , 000 and 4 , 000 ).

Original: Second, for the cases where no sample passes the threshold filter, we set ˜ Z and ˜ ϕ
Revised: Second, for the cases where no sample passes the threshold filter, we set ˜ Z and ˜ ϕ

Original: in two different ways.
Revised: in two different ways.

Original: For ˜ Z , we simply set it in an unsupervised way by randomly making another crop on Z i .
Revised: For ˜ Z , we simply set it in an unsupervised way by randomly making another crop on Z i .

Original: For ˜ ϕ , we set it by randomly adding a small amount of noise on ϕ i with ˜ ϕ ij :
Revised: For ˜ ϕ , we set it by randomly adding a small amount of noise on ϕ i with ˜ ϕ ij :

Original: = ϕ ij + ξ ϕ ij , j = 1 , 2 , 3 , 4 with a certain probability (e.g., 50% ), where ξ is a normal variable with mean 0 and a small standard deviation (e.g., 0.04).
Revised: = ϕ ij + ξ ϕ ij , j = 1 , 2 , 3 , 4 with a certain probability (e.g., 50% ), where ξ is a normal variable with mean 0 and a small standard deviation (e.g., 0.04).

Original: Training Hyperparameters: We train E MBED & E MULATE and the supervised regression using the
Revised: Training Hyperparameters: We train E MBED & E MULATE and the supervised regression using the

Original: AdamW optimizer.
Revised: AdamW optimizer.

Original: For both methods, we linearly warm up the learning rate to 0.01 at the beginning of the training and then gradually decay the learning rate with a cosine decay scheduler.
Revised: For both methods, we linearly warm up the learning rate to 0.01 at the beginning of the training and then gradually decay the learning rate with a cosine decay scheduler.

Original: We use λ = 1e-5 for weight decay.
Revised: We use λ = 1e-5 for weight decay.

Original: We train NPE-C
Revised: We train NPE-C

Original: [Greenberg et al., 2019] and SNL+ [Chen et al., 2020b] with learning rate 0.001 as we find lower learning rate leads to better performance.
Revised: [Greenberg et al., 2019] and SNL+ [Chen et al., 2020b] with learning rate 1e-4 and 5e-4 as we find lower learning rate leads to better performance.

Original: We train all the methods with 1 , 000 epochs except for NPE-C, which requires longer time to converge (e.g., 2 , 000 ).
Revised: We train all the methods with 1 , 000 epochs.

Original: We use batch size 1000 for large sized training data with n ≥ 1 , 000 .
Revised: We use batch size 1000 for large sized training data with n ≥ 1 , 000 .

Original: For the small size training data with n = 500 , we set it as 500 .
Revised: For the small size training data with n = 500 , we set it as 500 .

Original: Temperature values in contrastive learning balance the influence between barely distinguishable samples and easily distinguishable “negative” samples [Ravula et al., 2021].
Revised: Temperature values in contrastive learning balance the influence between barely distinguishable samples and easily distinguishable “negative” samples [Ravula et al., 2021].

Original: They are commonly chosen to be less than one where smaller values indicate a larger influence of barely distinguishable “hard” samples.
Revised: They are commonly chosen to be less than one where smaller values indicate a larger influence of barely distinguishable “hard” samples.

Original: For E MBED & E MULATE , we initialize the temperature values τ in (3) and (6) and τ ′ in (5) at 0.1.
Revised: For E MBED & E MULATE , we initialize the temperature values τ in (3) and (6) and τ ′ in (5) at 0 .

Original: We keep τ and τ ′ fixed at 0 .
Revised: We keep τ and τ ′ fixed at 0 .

Original: 1 for the first 500 epochs to make sure “hard” samples get large gradients and are updated sufficiently.
Revised: 15 for the first 500 epochs to make sure “hard” samples get large gradients and are updated sufficiently.

Original: We then adopt the “heating-up” strategy from Zhang et al.
Revised: We then adopt the “heating-up” strategy from Zhang et al.

Original: [2018] and linearly increase τ in (3) and (6) to 0.5 to increase the impact of “easier” samples.
Revised: [2018] and linearly increase τ in (3) and (6) to 0.5 to increase the impact of “easier” samples.

Original: We keep τ ′ in (5) fixed all the time so that the learned emulator ˆ g θ is capable of distinguishing embeddings of “hard” samples from the embedding network f θ , and vice versa.
Revised: We keep τ ′ in (5) fixed all the time so that the learned emulator ˆ g θ is capable of distinguishing embeddings of “hard” samples from the embedding network f θ , and vice versa.

Original: In training a contrastive learning target, we usually need sufficient amount of “negative” samples to approximate the uniform repulsion.
Revised: In training a contrastive learning target, we usually need sufficient amount of “negative” samples to approximate the uniform repulsion.

Original: We follow He et al.
Revised: We follow He et al.

Original: [2020] and construct a first-in-first-out queue with a rolling updating scheme.
Revised: [2020] and construct a first-in-first-out queue with a rolling updating scheme.

Original: Practically, the memory bank size should be no less than the training size.
Revised: Practically, the memory bank size should be no less than the training size.

Original: In our experiments, we set the size of the memory bank as 5 , 000 when the training size n = 4 , 000 ; the size of the memory bank as 2 , 000 when the training size n = 1 , 000 ; and the size of the memory bank as 1 , 500 when the training size n = 500 .
Revised: In our experiments, we set the size of the memory bank as 5 , 000 when the training size n = 4 , 000 ; the size of the memory bank as 2 , 000 when the training size n = 1 , 000 ; and the size of the memory bank as 1 , 500 when the training size n = 500 .

Original: All the hyperparameters are chosen using a grid search with a reserved validation set consisting ofsamples.
Revised: All the hyperparameters are chosen using a grid search with a reserved validation set consisting ofsamples.

Original: The range of values searched over are as follows:
Revised: The range of values searched over are as follows:

Original: • The initialized learning rates for E MBED & E MULATE and the supervised regression were selected from the set { 0 . 001 , 0 . 01 } .
Revised: • The initialized learning rates for E MBED & E MULATE and the supervised regression were selected from the set { 0 . 001 , 0 . 01 } .

Original: • The initialized temperature values of τ in (3) and (6) and τ ′ in (5) were selected from the set { 0 . 10 , 0 . 15 , 0 . 20 } .
Revised: • The initialized temperature values of τ in (3) and (6) and τ ′ in (5) were selected from the set { 0 . 10 , 0 . 15 , 0 . 20 } .

Original: The heated up maximum value of τ in (3) and (6) were selected from the set { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 } .
Revised: The heated up maximum value of τ in (3) and (6) were selected from the set { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 } .

Original: Priors Initialization of the EnKI:
Revised: Priors Initialization of the EnKI:

Original: As Schneider et al.
Revised: As Schneider et al.

Original: [2017], for the initialization of the EnKI prior, we use a normal prior for ( ϕ 1 , ϕ 2 , ϕ 4 ) =
Revised: [2017], for the initialization of the EnKI prior, we use a normal prior for ( ϕ 1 , ϕ 2 , ϕ 4 ) =

Original: ( F, h, b ) and a log-normal prior for ϕ 3 =
Revised: ( F, h, b ) and a log-normal prior for ϕ 3 =

Original: c in order to enforce positivity.
Revised: c in order to enforce positivity.

Original: For the baseline method minimizing (1) in §5.1 and §5.2, we use a fixed prior p ϕ , fixed .
Revised: For the baseline method minimizing (1) in §5.1 and §5.2, we use a fixed prior p ϕ , fixed .

Original: For the normal variables, we choose the mean at the middle of the range of testing instances: µ ϕ , fixed ( F, h, b ) =
Revised: For the normal variables, we choose the mean at the middle of the range of testing instances: µ ϕ , fixed ( F, h, b ) =

Original: (7 . 5 , 2 . 5 , 12 . 5) and variances broad enough so that all test samples are within 2 σ of the mean: σ 2 ϕ , fixed ( F, h, b ) =
Revised: (7 . 5 , 2 . 5 , 12 . 5) and variances broad enough so that all test samples are within 2 σ of the mean: σ 2 ϕ , fixed ( F, h, b ) =

Original: For the log-normal variable log c , we set its mean µ ϕ , fixed (log c ) =
Revised: For the log-normal variable log c , we set its mean µ ϕ , fixed (log c ) =

Original: log(11 . 5) and variance σ 2 ϕ , fixed (log c )
Revised: log(11 . 5) and variance σ 2 ϕ , fixed (log c )

Original: = 0 .
Revised: = 0 .

Original: 15 (i.e., a mean value of 12 . 5 for c ).
Revised: 15 (i.e., a mean value of 12 . 5 for c ).

Original: For E MBED & E MULATE , we can choose between the instance-wise prior p ϕ , empB and the fixed prior defined above.
Revised: For E MBED & E MULATE , we can choose between the instance-wise prior p ϕ , empB and the fixed prior defined above.

Original: Specifically, we alter the mean of p ϕ , empB with the empirical estimate provided by the regression head and use a relatively small variance.
Revised: Specifically, we alter the mean of p ϕ , empB with the empirical estimate provided by the regression head and use a relatively small variance.

Original: σ 2 ϕ , emp ( F, h, b ) = (18 . 0 , 1 . 125 , 18 . 0) for normal variables and σ 2 ϕ , emp (log c ) = 0 .
Revised: σ 2 ϕ , emp ( F, h, b ) = (18 . 0 , 1 . 125 , 18 . 0) for normal variables and σ 2 ϕ , emp (log c ) = 0 .

Original: 075 for log-normal variable c .
Revised: 075 for log-normal variable c .

Original: Setup of the EnKI:
Revised: Setup of the EnKI:

Original: For experiments in §5.1, we set the ensemble size M = 100 , step size α = 0 .and
Revised: For experiments in §5.1, we set the ensemble size M = 100 , step size α = 0 .and

Original: iteration number N = 50 for the baseline method with p ϕ , fixed and E MBED & E MULATE with p ϕ , empB .
Revised: iteration number N = 50 for the baseline method with p ϕ , fixed and E MBED & E MULATE with p ϕ , empB .

Original: For the study of noise impact in §5.2 when evaluating the baseline and E MBED &
Revised: For the study of noise impact in §5.2 when evaluating the baseline and E MBED &

Original: E MULATE , we use p ϕ , fixed .
Revised: E MULATE , we use p ϕ , fixed .

Original: For the ablation study in §5.3 when evaluating E MBED & E MULATE and no regression head contrastive learning with p ϕ , fixed , we set a large ensemble size M = 10 , 000 to prevent potential collapse in a local minimum and N = 100 to ensure convergence.
Revised: For the ablation study in §5.3 when evaluating E MBED & E MULATE and no regression head contrastive learning with p ϕ , fixed , we set a large ensemble size M = 10 , 000 to prevent potential collapse in a local minimum and N = 100 to ensure convergence.

Original: Arguably, this could slow down the computational times.
Revised: Arguably, this could slow down the computational times.

Original: However, empirically we only increased the time from 2 seconds to 16 seconds per instance, which is significantly lower than the times consumed by the baseline method.
Revised: However, empirically we only increased the time from 2 seconds to 16 seconds per instance, which is significantly lower than the times consumed by the baseline method.

Original: Computational resources: Experiments of E MBED & E MULATE and the supervised regression were performed on a system with 4x Nvidia A40 GPUs, 2 AMD EPYC 7302 CPUs, and 128GB of RAM.
Revised: Computational resources: Experiments of E MBED & E MULATE and the supervised regression were performed on a system with 4x Nvidia A40 GPUs, 2 AMD EPYC 7302 CPUs, and 128GB of RAM.

Original: The baseline methods and data generating process were performed on 2x Xeon Gold 6130 CPUs.
Revised: The baseline methods and data generating process were performed on 2x Xeon Gold 6130 CPUs.

Original: B.2 Network Architecture
Revised: B.2 Network Architecture

Original: [Figure]
Revised: [Figure]

Original: Emulator ˆ g θ : To enhance the representation power of the emulator
Revised: Emulator ˆ g θ : To enhance the representation power of the emulator

Original: ˆ g θ and enlarge the alignment between two learned representation spaces, we design the emulator using residual blocks as in Fig.
Revised: ˆ g θ and enlarge the alignment between two learned representation spaces, we design the emulator using residual blocks as in Fig.

Original: Embedding network f
Revised: Embedding network f

Original: θ : We use Resnet

Revised: θ : We use Resnet

Original: [He et al., 2016] as the backbone of the trajectory branch of E MBED & E MULATE in Fig.
Revised: [He et al., 2016] as the backbone of the trajectory branch of E MBED & E MULATE in Fig.

Original: 1(a)
Revised: 1(a)

Original: and the supervised regression.
Revised: and the supervised regression.

Original: We apply average pooling at the last layer to generate a 512-d hidden vector.
Revised: We apply average pooling at the last layer to generate a 512-d hidden vector.

Original: To respect the temporal periodic pattern of Lorenz 96, we replace the standard zero padding with the circular padding on the temporal axis.
Revised: To respect the temporal periodic pattern of Lorenz 96, we replace the standard zero padding with the circular padding on the temporal axis.

Original: For the regression head h
Revised: For the regression head h

Original: θ , we project the 512-d hidden vector into a 4-d output using only a single fully connected layer.
Revised: θ , we project the 512-d hidden vector into a 4-d output using only a single fully connected layer.

Original: For the embedding network, we project the 512-d hidden vector into a 128-d feature vector using a similar structured branch as projection residual (out) in Fig. 3(d) except that the residual block skips two layers instead of one.
Revised: For the embedding network, we project the 512-d hidden vector into a 128-d feature vector using a similar structured branch as projection residual (out) in Fig. 3(d) except that the residual block skips two layers instead of one.

Original: B.3 Performance metrics
Revised: B.3 Performance metrics

Original: Absolute Percentage Error (APE): : APE is used in selecting “positive” samples of E MBED &
Revised: Absolute Percentage Error (APE): : APE is used in selecting “positive” samples of E MBED &

Original: E MULATE and is calculated on pairs of training samples:
Revised: E MULATE and is calculated on pairs of training samples:

Original: [Equation]
Revised: [Equation]

Original: here k is the dimension of ϕ , i, j represent indexes of two data samples, and ϵ is a small value to address overflow issues.
Revised: here k is the dimension of ϕ , i, j represent indexes of two data samples, and ϵ is a small value to address overflow issues.

Original: Mean Absolute Percentage Error (MAPE): MAPE is a measure of prediction accuracy over the entire testing data and is expressed as a ratio with the formula:
Revised: Mean Absolute Percentage Error (MAPE): MAPE is a measure of prediction accuracy over the entire testing data and is expressed as a ratio with the formula:

Original: [Equation]
Revised: [Equation]

Original: where ˆ ϕ ∗ j represents the j -th component of the estimate parameter vector, ϕ ∗ j represents the j -th component of the true parameter vector (e.g., [ ϕ ∗ 1 , ϕ ∗ 2 , ϕ ∗ 3 , ϕ ∗ 4 ] =
Revised: where ˆ ϕ ∗ j represents the j -th component of the estimate parameter vector, ϕ ∗ j represents the j -th component of the true parameter vector (e.g., [ ϕ ∗ 1 , ϕ ∗ 2 , ϕ ∗ 3 , ϕ ∗ 4 ] =

Original: [ F, h, c, b ] for Lorenz 96), and ϵ is a small value to address overflow issues.
Revised: [ F, h, c, b ] for Lorenz 96), and ϵ is a small value to address overflow issues.

Original: Continuous Ranked Probability Score(CRPS): As described in §5.1, CRPS is a probability metric to evaluate the performance of a probabilistic estimation.
Revised: Continuous Ranked Probability Score(CRPS): As described in §5.1, CRPS is a probability metric to evaluate the performance of a probabilistic estimation.

Original: CRPS is computed as a quadratic measurement of cumulative distribution function (CDF) between the predicted distribution C and the empirical distribution of the observation.
Revised: CRPS is computed as a quadratic measurement of cumulative distribution function (CDF) between the predicted distribution C and the empirical distribution of the observation.

Original: CRPS is defined as [Hersbach, 2000, Zamo and Naveau, 2018, Pappenberger et al., 2015]:
Revised: CRPS is defined as [Hersbach, 2000, Zamo and Naveau, 2018, Pappenberger et al., 2015]:

Original: [Equation]
Revised: [Equation]

Original: where ϕ ∗ j and ˆ ϕ ∗ j are defined above, C is the CDF of the ensemble estimation with C ( y j ) =
Revised: where ϕ ∗ j and ˆ ϕ ∗ j are defined above, C is the CDF of the ensemble estimation with C ( y j ) =

Original: P ( ˆ ϕ ∗ j ≤ y j ) , and U is the CDF of the Heaviside step function with U ( y j − ϕ ∗ j ) = 1 { y j ≥ ϕ ∗ j } .
Revised: P ( ˆ ϕ ∗ j ≤ y j ) , and U is the CDF of the Heaviside step function with U ( y j − ϕ ∗ j ) = 1 { y j ≥ ϕ ∗ j } .

Original: Practically estimating C is difficult since it has no analytic form.
Revised: Practically estimating C is difficult since it has no analytic form.

Original: In our case, using EnKI with ensemble size M , we can estimate (10) using the empirical distribution calculated from the particles { ˆ ϕ ( m ) ∗ j } Mm =1 1 :
Revised: In our case, using EnKI with ensemble size M , we can estimate (10) using the empirical distribution calculated from the particles { ˆ ϕ ( m ) ∗ j } Mm =1 1 :

Original: [Equation]
Revised: [Equation]

Original: B.4 Visualizing the objective functions
Revised: B.4 Visualizing the objective functions

Original: f 2 4 6 8 10 12 14 16b 4 0.5 0.0 0.
Revised: f 2 4 6 8 10 12 14 16b 4 0.5 0.0 0.

Original: h 0.6 0.8 1.0 1.2 1.
Revised: h 0.6 0.8 1.0 1.2 1.

Original: c 4 0.5 0.0 0.
Revised: c 4 0.5 0.0 0.

Original: [Figure]
Revised: [Figure]

Original: All objective values greater equal
Revised: All objective values greater equal

Original: , but intentionally negative
Revised: , but intentionally negative

Original: “star” points between different marginal heatmaps using the function are aligned with post-processing.
Revised: “star” points between different marginal heatmaps using the function are aligned with post-processing.

Original: Unlike the value of the objective function (4) are bounded between [0 , 4] values of (1) have upper bound and contain extremely large numbers
Revised: Unlike the value of the objective function (4) are bounded between [0 , 4] values of (1) have upper bound and contain extremely large numbers

Original: In this section, we visualize the objective function of (1) with predefined moments function and
Revised: In this section, we visualize the objective function of (1) with predefined moments function and

Original: Runge-Kutta and (4) of E MBED & E MULATE with emulator f θ
Revised: Runge-Kutta and (4) of E MBED & E MULATE with emulator f θ

Original: and ˆ g θ .
Revised: and ˆ g θ .

Original: We use the same instance in §5.2 as the observation with the ground truth ϕ
Revised: We use the same instance in §5.2 as the observation with the ground truth ϕ

Original: =
Revised: =

Original: [10 , 1 , 10 , 10] and the noiseless Z is of length 10 , 000 .
Revised: [10 , 1 , 10 , 10] and the noiseless Z is of length 10 , 000 .

Original: Fig.
Revised: Fig.

Original: 4 shows the marginal heatmaps of the objective functions when holding two parameters fixed at the ground truth values and sampling the other two parameters uniformly in a fixed range.
Revised: 4 shows the marginal heatmaps of the objective functions when holding two parameters fixed at the ground truth values and sampling the other two parameters uniformly in a fixed range.

Original: We can see that the objective function learned with E MBED & E MULATE can accurately reflect the dynamics in the parameter space with smooth curves.
Revised: We can see that the objective function learned with E MBED & E MULATE can accurately reflect the dynamics in the parameter space with smooth curves.

Original: In contrast, the predefined objective function is insensitive to the parameter c with a flat curve showing in the marginal heatmap of ( h, c ) pairs.
Revised: In contrast, the predefined objective function is insensitive to the parameter c with a flat curve showing in the marginal heatmap of ( h, c ) pairs.

Original: C Further Experiments
Revised: C Further Experiments

Original: In this section, we study the 1-d Kuramoto-Sivashinsky equation (KSE).
Revised: In this section, we study the 1-d Kuramoto-Sivashinsky equation (KSE).

Original: Kuramoto-Sivashinsky equation is known for its chaotic dynamics and can be applied in describing a variety of physical phenomena, including flows in pipes, plasmas, chemical reaction dynamics, etc.
Revised: Kuramoto-Sivashinsky equation is known for its chaotic dynamics and can be applied in describing a variety of physical phenomena, including flows in pipes, plasmas, chemical reaction dynamics, etc.

Original: [Hyman and
Revised: [Hyman and

Original: Nicolaenko, 1986]
Revised: Nicolaenko, 1986]

Original: The Kuramoto-Sivashinsky equation is a fourth-order partial differential equation with the form:
Revised: The Kuramoto-Sivashinsky equation is a fourth-order partial differential equation with the form:

Original: [Equation]
Revised: [Equation]

Original: where V xt is the state of KSE, λ 2 , λ 4 and λ nonlinear are the coefficients of second-order derivative, fourth-order derivative and nonlinear term.
Revised: where V xt is the state of KSE, λ 2 , λ 4 and λ nonlinear are the coefficients of second-order derivative, fourth-order derivative and nonlinear term.

Original: In all, we define ϕ :=
Revised: In all, we define ϕ :=

Original: [ λ 2 , λ 4 , λ nonlinear ] , and our goal is to learn ϕ from sequences of observations V ∈ R T × d .
Revised: [ λ 2 , λ 4 , λ nonlinear ] , and our goal is to learn ϕ from sequences of observations V ∈ R T × d .

Original: Choosing the moments function: We choose the first order statistics of V to define the moments function: m ( V ) =
Revised: Choosing the moments function: We choose the first order statistics of V to define the moments function: m ( V ) =

Original: [ ⟨ V ⟩ ] where ⟨·⟩ denotes empirical average over time.
Revised: [ ⟨ V ⟩ ] where ⟨·⟩ denotes empirical average over time.

Original: Note here we have tried multiple combinations of moments, including second-order and third-order statistics but found no choice of moments performs better than using only the first-order.
Revised: Note here we have tried multiple combinations of moments, including second-order and third-order statistics but found no choice of moments performs better than using only the first-order.

Original: Preparation of dataset: We set L = 32 and d = 256 to ensure we are in a chaotic regime.
Revised: Preparation of dataset: We set L = 32 and d = 256 to ensure we are in a chaotic regime.

Original: Initial conditions are randomly sampled from a uniform range [ − π, π ] .
Revised: Initial conditions are randomly sampled from a uniform range [ − π, π ] .

Original: All simulations performed here are completed with a Runge-Kutta 4-stage method as Pachev et al.
Revised: All simulations performed here are completed with a Runge-Kutta 4-stage method as Pachev et al.

Original: Training data are sampled uniformly for all three components of ϕ ranging from min = 0 .
Revised: Training data are sampled uniformly for all three components of ϕ ranging from min = 0 .

Original: 1 to max = 10 . 0 with length
Revised: 1 to max = 10 . 0 with length

Original: T = 400 , dt = 0 .
Revised: T = 400 , dt = 0 .

Original: 5 and size n = 500 .
Revised: 5 and size n = 500 .

Original: 100 testing data are sampled in a narrower range from min = 0 .
Revised: 100 testing data are sampled in a narrower range from min = 0 .

Original: 5 and max = 9 .
Revised: 5 and max = 9 .

Original: 5 with length 800 and dt = 0 .
Revised: 5 with length 800 and dt = 0 .

Original: Setup of the EnKI:
Revised: Setup of the EnKI:

Original: We set the fixed Gaussian prior p ϕ , fixed with mean at the middle of the testing range and variance broader enough to cover all testing instances in 2 σ : µ ϕ ∗ j , fixed = 5 , σ 2 ϕ ∗ j , fixed = 6 . 25 , j = 1 , 2 , 3 .
Revised: We set the fixed Gaussian prior p ϕ , fixed with mean at the middle of the testing range and variance broader enough to cover all testing instances in 2 σ : µ ϕ ∗ j , fixed = 5 , σ 2 ϕ ∗ j , fixed = 6 . 25 , j = 1 , 2 , 3 .

Original: We use the same variance for p ϕ , empB and alter its mean with the empirical estimate provided by the regression head.
Revised: We use the same variance for p ϕ , empB and alter its mean with the empirical estimate provided by the regression head.

Original: For both baseline and E MBED & E MULATE , we set ensemble size
Revised: For both baseline and E MBED & E MULATE , we set ensemble size

Original: M = 100 , iteration number N = 40 , step size α = 0 .
Revised: M = 100 , iteration number N = 40 , step size α = 0 .

Original: Setup of the training: We train both E MBED & E MULATE and the supervised regression for 1, epochs.
Revised: Setup of the training: We train both E MBED & E MULATE and the supervised regression for 1, epochs.

Original: We set the memory bank size of E MBED & E MULATE for 1,000.
Revised: We set the memory bank size of E MBED & E MULATE for 1,000.

Original: We set the dimension of learned representation, i.e., f θ ( V ) and ˆ g θ ( ϕ ) to be 120.
Revised: We set the dimension of learned representation, i.e., f θ ( V ) and ˆ g θ ( ϕ ) to be 120.

Original: We use the same learning rate and cosine decay scheduler.
Revised: We use the same learning rate and cosine decay scheduler.

Original: We initialize τ, τ ′ = 0 .
Revised: We initialize τ, τ ′ = 0 .

Original: 1 , and linearly warm up τ ′ from 500 to 900 epochs with the maximum value 0 .
Revised: 1 , and linearly warm up τ ′ from 500 to 900 epochs with the maximum value 0 .

Original: λ 2 ↓ λ 4 ↓ λnonlinear ↓
Revised: λ 2 ↓ λ [Table] λ

Original: [Table]
Revised: [Table]

Original: prior p ϕ , empB ; (d) Supervised regression; (e) NPE-C [Greenberg et al., 2019].
Revised: prior p ϕ , empB ; (d) Supervised regression; (e) NPE-C [Greenberg et al., 2019].

Original: Results: We first evaluate the accuracy of point estimates using the mean absolute percentage error (MAPE) and the median absolute percentage error (MdAPE).
Revised: Results: We first evaluate the accuracy of point estimates using the mean absolute percentage error (MAPE) and the median absolute percentage error (MdAPE).

Original: In Table 5, we see that E MBED &
Revised: In Table 5, we see that E MBED &

Original: E MULATE significantly outperforms EnKI using the predefined moments function and the expensive numerical solver.
Revised: E MULATE significantly outperforms EnKI using the predefined moments function and the expensive numerical solver.

Original: This experiment demonstrates that designing moments function like Schneider et al.
Revised: This experiment demonstrates that designing moments function like Schneider et al.

Original: [2017] did for the Lorenz 96 system needs physical information of the dynamical system and in applications (e.g., climatology) this requires extra domain expertise.
Revised: [2017] did for the Lorenz 96 system needs physical information of the dynamical system and in applications (e.g., climatology) this requires extra domain expertise.

Original: We then evaluate the quantified uncertainty using the continuous ranked probability score (CRPS).
Revised: We then evaluate the quantified uncertainty using the continuous ranked probability score (CRPS).

Original: In Table 6, we see that E MBED &
Revised: In Table 6, we see that E MBED &

Original: E MULATE outperforms the baselines of the EnKI without learning and the supervised regression.
Revised: E MULATE outperforms the baselines of the EnKI without learning and the supervised regression.

Original: This
Revised: This

Original: [Table] [Table]
Revised: [Table] [Table]

Original: experiment demonstrates that E MBED & E MULATE with the EnKI is capable of providing a reliable probabilistic estimation with low errors than the EnKI without learning which requires domain knowledge to perform well, the supervised regression which only provides points estimates [Table] and
Revised: experiment demonstrates that E MBED & E MULATE with the EnKI is capable of providing a reliable probabilistic estimation with low errors than the EnKI without learning which requires domain knowledge to perform well, the supervised regression which only provides points estimates, and

Original: NPE-C which lacks competence when trying to estimate multiple different observations at test time.
Revised: NPE-C which lacks competence when trying to estimate multiple different observations at test time.

Original: Lastly, we show the empirical computational time needed for E MBED & E MULATE , the supervised regression, NPE-C, and the EnKI with Runge-Kutta for different stages.
Revised: Lastly, we show the empirical computational time needed for E MBED & E MULATE , the supervised regression, NPE-C, and the EnKI with Runge-Kutta for different stages.

Original: From Table 7, we see thatboth E MBED & E MULATE and the supervised regression only requires 2 .
Revised: From Table 7, we see that E MBED & E MULATE , the supervised regression, and NPE-C only require 1 .

Original: 15% (or below) of the time needed by the EnKI with Runge-Kutta.
Revised: 23% (or below) of the time needed by the EnKI with Runge-Kutta.