PETAL: Physics Emulation Through Averaged Linearizations for Solving Inverse Problems
Authors: Jihui Jin
Etienne Ollivier
Richard Touret
Matthew McKinley
Karim Sabra
Justin Romberg
Venue: NeurIPS 2023 poster (2023)
Keywords: Inverse Problems
Neural Adjoint
Hybrid Machine Learning
Physics
TL;DR: PETAL proposes an explicit approach to embed physics-based linearizations into surrogate modeling for solving inverse problems.

Abstract: Inverse problems describe the task of recovering an underlying signal of interest given observables. Typically, the observables are related via some non-linear forward model applied to the underlying unknown signal. Inverting the non-linear forward model can be computationally expensive, as it often involves computing and inverting a linearization at a series of estimates. Rather than inverting the physics-based model, we instead train a surrogate forward model (emulator) and leverage modern auto-grad libraries to solve for the input within a classical optimization framework. Current methods to train emulators are done in a black box supervised machine learning fashion and fail to take advantage of any existing knowledge of the forward model. In this article, we propose a simple learned weighted average model that embeds linearizations of the forward model around various reference points into the model itself, explicitly incorporating known physics. Grounding the learned model with physics based linearizations improves the forward modeling accuracy and provides richer physics based gradient information during the inversion process leading to more accurate signal recovery. We demonstrate the efficacy on an ocean acoustic tomography (OAT) example that aims to recover ocean sound speed profile (SSP) variations from acoustic observations (e.g. eigenray arrival times) within simulation of ocean dynamics in the Gulf of Mexico.

PETAL: Physics Emulation Through Averaged Linearizations for Solving Inverse Problems
Jihui Jin
Etienne Ollivier Mechanical Engineering Georgia Institute of Technology Atlanta, GA 30332
Electrical and Computer Engineering Georgia Institute of Technology
Atlanta, GA 30332 j ahhu@gatech edu
Richard Touret
Matthew McKinley
Mechanical Engineering Georgia Institute of Technology Atlanta, GA 30332
Mechanical Engineering Georgia Institute Technology Atlanta, GA 30332
Karim G. Sabra
Justin K. Romberg
Mechanical Engineering Georgia Institute of Technology Atlanta, GA 30332
Electrical and Computer Engineering Georgia Institute of Technology Atlanta, GA 30332
Abstract
Inverse problems describe the task recovering an underlying signal of interest given observables. Typically, the observables are related via some non-linear for- ward model applied to the underlying unknown signal. Inverting the non-linear forward model can be computationally expensive, as it often involves computing and inverting a lincarization at a series of estimates. Rather than inverting the physics-based model, we instead train a surrogate forward model (emulator) and leverage modern auto-grad libraries to solve for the input within a classical opti- mization framework. Current methods to train emulators supervised machine learning fashion and fail to take advantage of any existing knowledge of the forward model. In this article, weighted average model that embeds linearizations of the forward model around various reference points into the model itself, explicitly incorporating known physics. Grounding the learned model with physics based linearizations improves the forward modeling accuracy and provides richer physics based gradient infor- mation during the inversion process leading to more accurate signal recovery. We demonstrate the efficacy on an ocean acoustic tomography (OAT) example that aims to recover ocean sound speed profile (SSP) variations from acoustic observa- tions (e.g. eigenray arrival times) within simulation of ocean dynamics in the Gulf of Mexico.
done in black box
propose simple learned
Introduction
Inverse problems arise in many scientific applications where the goal is to reconstruct some unknown signal, image volume of interest from indirect observations. The forward process, the mapping from the data to observations, is typically well-known usually through modeling the physical process. However, inverting the model is often ill-posed or even non-invertible. More formally, let consider the task of recovering some signal æ from observations y that are related by some potentially
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

non-linear forward model F via
y == F=x)
(I)
where 77 encapsulates noise other perturbations. Our forward model F represents computational model of the underlying physics of the measurement process. Classical solutions involve modeling the forward process with extremely high accuracy and then attempting to invert a stabilized or linearized variant, which often requires heavy domain knowledge.
Another approach to handle the ill-posed nature of the inversion task is to cast the problem as optimization task and incorporate regularization. A regularizer is a measure of how well the proposed solution fits some known, and often hand-crafted, prior. This term makes the inversion well-posed by biasing towards certain solutions. The inversion typically solved for an in iterative fashion. This process be potentially computationally expensive due requiring many applications of the physics based forward model. Computing the adjoint to determine a descent direction increases the computational burden more.
In this paper, we propose a novel architecture trained to emulate the physics-based forward model. The surrogate model provides cheaper alternative for both the forward pass and descent direction calculation by leveraging existing auto-grad libraries (30), making iterative solvers feasible for recovering solution. This work departs from previous works who also aim to emulate the forward model 4 I 33 by embedding low-fidelity forward models into the surrogate itself, thereby explicitly incorporating the physics instead of treating it black box. More concretely, we propose to use a set linearizations of the forward model at a set of reference point and train a model to fuse and correct the outputs in adaptive manner. We then this trained model in classical optimization framework to recover the signal given a set of observations.
Concretely, our paper makes the following contributions:
We propose novel architecture that learns to emulate the forward model. The model directly embeds physies via linearizations around a subset of reference points.
We introduce a learned encoder/decoder scheme to the neural adjoint method to mitigate artifacts from directly optimizing in the input space.
We demonstrate its efficacy for recovering solutions in a classical optimization framework on an ocean acoustic tomography (OAT) example that aims to recover ocean sound speed profile (SSP) variations from acoustic observations (e.g. eigenray arrival times) within simulation of ocean dynamics the Gulf Mcxico.
2 Related Works
2.1 Iterative Inversion
The task of directly inverting some potentially non-linear forward model F often non-trivial or mathematically impossible. Instead, a more stable alternative is to iteratively solve for æ given some observations y. Classically, this is done by formulating the reconstruction solving non-linear least squares problem that aims to minimize
1 ly =(
a arg min 2
The Gauss-Newton method solves Equation (2) by iteratively computing the Jacobian J of at the current best estimate (k) and solving a set of linear equations to generate an update of the form
a (k+1)
(k) + (JJJF) J ( (()
(3)
The Meenberrg-Maqquarttt algorithm 17 23
presents more robust alternative Gauss-Newton by introducing a (non-negative) damping factor À leading to an update of the form
x (k+1) = x (k) (JIJ XL) -(
Another approach address the ill-posed nature of F is explicitly introduce regularization function R(x) as an additional penalty term to Equation 2 The regularizer R addresses the ill iol-posed

---

Dx
Xref
Are
Ex
Unlike typical Embeedded Plysics
Xyef
A
Emberdeed Plysier
PP
D1
PN
E <'SPPs
Softmax
g
pv
Softmax
W
W
Σ
w W
Σ
Learned Weighted Average
Learned Weighted Average
Dx
W
W
(a)
(b)
Figure 1: (a) Architecture of PETAL for training the forward model. The signal æ is passed through the Embedded Physics module to generate low-fidelity physics- based predictions ÿ "'....."" G that are then merged in an adaptive manner through attention mechanisms with the Learned Weighted Average module before being corrected with a final linear output layer W. An encoding layer Ex and decoding layer D2 are simultaneously learned. (b) The architecture rearranged for inversion. The decoding layer is moved to the input to allow for inversion in the learned lower dimensional subspace z E Z. Learned modules are denoted in green.
nature by stabilizing the inversion process and biasing the solution towards those with expected or desired structure. Some common examples include the l, norm to encourage smaller norm solutions (similar to Equation 4 & to promote sparsity, and total variation (TV) to induce homogeneous regions separated by sharp boundaries. For differentiable R, the augmented optimization problem can be solved via steepest descent, computing updates of the form
(k+1) (k) + J ()) (V
(5)
where y is tuned step size, À is non-negative factor controlling the strength of regularization, and J. is the Jacobian of F evaluated at (*).
However all these approaches be undesirable in practice due to the need to re-compute the Jacobian (or Jacobian-vector product) at each iteration. Computing a Jacobian-vector product to determine descent direction often involves solving auxiliary set of PDEs, which can be computationally expensive many iterations are required to achieve an acceptable level of accuracy.
Ensemble Kalman filter (EnKF) offers Bayesian iterative approach to estimating the underlying signal [8. [2]. Rather than estimating a single instance of the signal, EnKF recursively updates an estimated distribution of the state, assuming Gaussian probability. A sample mean and covariance is calculated using an ensemble of samples that are updated iteratively via Bayesian updates. However, such method still ultimately relies on applying the physics based forward model to each member of the ensemble, which can be prohibitively expensive for complex non- annlinnar forward operators.
2.2 Learned Inversion
An increasingly popular approach for solving inverse problems takes advantage of machine learning. Deep learning has achieved tremendous success in areas of natural language processing, computer vision and other tasks, in part due to the availability of large labelled datasets, Recent works attempt
to tackle inverse problems using these data-driven methods [7
3 25
29
37
38
supervised learning tasks that attempt to learn a mapping purely from examples, deep learning for inverse problems can leverage our understanding of the physics the forward model. One common approach to embed the physics is the so-called "Loop Unrolled" architecture heavily inspired by

---

The learned model alternates between taking gradient step computed directly from existing forward models and applying learned regularizer step. However, such approaches have to apply the forward model multiple times for single forward pass through the architecture, making the method infeasible for more complex non-linear simulators. Our approach bypasses this obstacle by using more computationally tractable approximations of the forward model.
existing iterative algorithmss [ 10, 35 40
An alternative approach to incorporating physics is termed "Physics Informed Neural Networks (PINN)" [6 32 These methods incorporate the governing partial differential equations directly in the loss, guiding a parameterized model towards physics-obeying solutions. However, an important distinction between PINNs and more generalized machine learning for inverse problems is that each model is trained for single instance of PDE solve given some boundary/initial conditions and must be re-trained for any new conditions or observables (in the case of inverse problems). Every training iteration involves a PDE misfit calculation, which can be expensive and ill-posed, making scaling to larger dimensions difficult. Unlike PINNs, the linearizations we use only need to be computed once before training rather than at each iteration of the training process.
Gaussian process regression (GPR) is another approach that leverages a set of references. GPR attempts to model some unknown function f given data points by assuming a multivariate normal distribution prior and interpolating between the observed samples [39]. GPR is similar to the proposed model in that it does distance comparison (usually through a pre-defined kernel function) with every reference point provided. However, is designed only take input-output pairs and scales poorly with the number of examples, making it difficult to take advantage of large training sets. Our proposed model only needs to retain a subset of the training set for comparison purposes, but can still leverage large datasets for training more accurate predictor.
2.3 Neural Adjoint
The neural adjoint method 33 was proposed to tackle inverse problems with more computationally intensive forward models. A parameterized model Go is trained to emulate the forward model [ This is done in a supervised learning fashion, often with simple mean-squared error loss. Not only does this provide a cheaper/quicker alternative to the physics- based forward model, if trained with existing auto-grad libraries such as PyTorch (30), this also allows for efficient computation of the gradient with respect to the input, bypassing the need to explicitly solve for the adjoint when calculating the gradient.
Once trained, the parameters are fixed and the model Ge is substituted for F in Equation 2 or other similar optimization frameworks. The auto-grad libraries are then used efficiently compute a gradient with respect to the input æ, making it possible to iteratively solve for the best estimate that fits some observations y. Existing works primarily focused on lower dimensional examples (on the order of 5-15) where the test set was drawn from the same distribution as the training set 22 28 31 33), thus a simple "boundary-loss" regularizer was often sufficient to produce accurate results. Direct learned inversion methods are much faster than this iterative based method, but yield only a single estimate and are susceptible to overfitting to the training set. The neural adjoint method allows for exploration of the solution space with different initializations and the incorporation of various regularizers such as L,, f2, or even gradient based (21), to guide the optimization towards specific solutions. In addition, one can also restrict the optimization to some pre-determined basis that better represents the data while reducing the dimensionality (5)..
Our proposed architecture extends the neural adjoint method in two notable ways. Existing methods are purely nata-driven and thus treated like a black box 4 13), while as our architecture explicitly incorporates knowledge of the physics- linearized approximations. Tangent- -inear models are also explored in [II,, but on the surrogate in order to verify the accuracy of the learned model, while as the proposed architecture uses linearizations of the physics based forward model in the construction of the surrogate. In addition, we also propose to jointly learn a subspace to improve the accuracy of the optimization stage.
forward model into the learned emulator in the form of

---

31
101
800
500
M400
rwage (1m)
2000
2500
3000
W
68%
NPW
6FW
10:0
Loophadee
cound speed (mil)
(a)
(h)
(c)
Figure 2: An overview of data from month long simulation of the (a) Gulf of Mexico. (b) Example 2D slices of the sound speed profile (top) and its deviation from the mean (bottom). The 10 slices used for experiments are separated by white lines. (c) The average SSP as a function of depth. Note that most of the fluctuations occur near the surface.
3 Method
The neural adjoint (NA) method is typically decomposed into two stages: training emulator of the forward model and then using it for inference. motivate and describe the proposed architecture for the forward model in Section 3.1 Next, we formulate the optimization problem that uses the Finally, we propose an augmentation to the optimization
trained model for inference Section 3.2
formulation to incorporate a learned subspace in Section 3.3
3.1 Embedding Physies in the Learned Forward Model
The neural adjoint (NA) method aims to learn an accurate emulator of the forward model to replace existing physics simulators. More formally, assume that are given forward model F y that maps our input E R where m denotes the size of the discretization, to some observations our computational ocean tomography
y R", where denotes the total number of observations. examples in Section 4 F F(x) is computed using ray tracer; in other applications, a PDE might have be solved. We then train neural Go. whose architecture is illustrated in Figure and described in detail below, to approximate the mapping F.
We first motivate the design of the architecture by discussing the more classical approach of linearizing the forward model F. Given reference point whet, we perform first order Taylor series expansion
ÿ F( + JF(x(ee) (x æfet)
(6)
yier + Aler wia)).
This linearization approximates the forward model with varying levels of accuracy depending on the proximity of the input to the reference point a Rather than learning the mapping from æ to y in a pure data-driven fashion, we propose to leverage a set of these linearizations that already perform the mapping. We mitigate the modeling inaccuracies by using ensemble of reference models rather than single linearization. Only small subset of linearizations need to be performed once in order construct the proposed emulator Go. the additional computational cost minimal relative to attempting to invert the actual physics based forward model for arbitrary measurements
The operation of the architecture acts follows: Given some input æ, first pass through the embedded physics ensemble A ..., A, produce N predicted observations ),.,, y.M through the application of Equation (6). The predicted observations are then combined through a learned weighted average module to produce the final predicted observation.
natural way to compute the weights wi by calculating the dot product similarity between the input point æ and the reference points a Tref:... ..., .रe used to generate the linearizations; higher similarity implies that the linearization better approximation and thus the prediction more "trustworthy". We follow approach similar to attention- based models (36) by learning embedding space to
5

---

perform the dot product and applying a softmax to normalize the weights. Thus for each y', we compute the corresponding weight w² as
exp (x, x Feel)
"² Σj exp (x, Fret) Px
where Pz is the learned projection for the dot product space, The predicted O' get summed together using the learned weights before being passed through final linear layer W. Thus the output of the proposed model is
Go(x)= =W Σ w/ (y + Aje - Fict)),
(8)
where we distinguish components related to the embedded physics module with superscripts. The encoder-decoder layers Ex and Dr are treated as the identity mapping for this section. They are described in further detail Section 3.3. Note that the learned weights w' depend on the input æ. To prevent saturation of the learned softmax distribution, we apply spectral normalization to all æ projection layers (Ex and Pz). The full architecture is outlined Figure la The model is trained using the mean squared error loss on dataset of paired example (Ix, 5k)
min Σ [Ga(ma)) -2
3.2 Performing Inference
In order perform inference on set of observations y. we solve the following optimization problem that incorporates our trained network
1 c = arg min 2 [[)a(z)) - yl R(x).
(10)
We solve this iteratively by fixing the weights of the network and computing the gradient with respect to its input æ. Note that this same optimization problem can be set up with the original forward model F, but computing a gradient often non-trivial and computationally expensive. By training a forward model approximation, we can leverage existing auto-grad libraries (30 to efficiently compute the gradient. Having an accurate descent direction critical for solving Equation TO However, these black box models only trained match outputs, and thus performing gradient descent can lead to many undesireable local minima. Due to the construction of our emulator, a convex combination of the gradients from using the individual linear forward model approximations (slightly modulated by the learned weights) arises in the calculation, providing some physics-basee descent directions which may help alleviate these issues (See Appendix for derivations).
Equation (TPp can be solved with a variety of optimization algorithms to converge on some locally optimal ar. Since we are substituting the forward model with an approximation, we can account for any inaccuracies by introducing a tolerance level. Once the observation loss drops below a pre-determined level, the optimization terminates early.
Note that we incorporated regularizer R(-) as additional cost term. The regularizer encourages certain properties (e.g. smaller values) and helps guide the optimization towards particular solutions. In our experiments, we used & well a Sobolev norm (&g norm performed on the discrete x and y gradient of our input x:).
Finally, it should be noted that the iterative nature of this method requires that initialize guess with some estimate. When optimizing from scratch, a reasonable candidate would be the average from the training first initializing with that estimate and then refining it with the NA procedure. We note that there is an increase in computation time compared to directly learning the inverse due the iterative nature of the algorithm, but the NA method offers certain trade-off's outlined above that might be more beneficial in practice than the single estimate provided by learned direct inverse methods.
A Alternatively, we can leverage an estimated å from other inverse methods by
3.3 Learning a Subspace for Reconstruction
Directly inverting on the input space of neural networks often generates artifacts in the final result or gets trapped in local minima due to the highly non-convex model. One way to address this issue is by

---

1480
1000
dspth (m)
1900
2000
2000 00
500
1000
1000
2000
2000
3000
3000
4000
4500
5000
range [m]
Figure 3: An example OAT forward model set up with 20 sources and 20 receivers. The direct and surface bounce paths are denoted in blue and red respectively. The average SSP influencing the arrival times the receivers can be seen on the right.
Table 1: RMSE (m/s) of inversion with various initializations.
Model
Avg Init
LFM Init Tik Init
Tik
0.760 0.608
LFM
0.625
0.602 0.393 + 0.005
MLP
0.391 = 0.005
0.395 + 0.005
PETAL (ours)
0.364 = 0.018
0.357 0.014
0.357 + 0.014
optimizing in some lower dimensional subspace, such as one determined by principle component analysis (PCA) [13]. The PCA basis acts de-noiser, removing any artifacts due to the optimization and helps reduce the dimensionality, simplifying the optimization process.
Rather than pre-computing basis, instead propose to jointly learn linear projection and reconstruction layer along with the forward model. The model as described Subsection [.I] can be augmented with linear encoder Ex and decoder D layer. The encoder layer projects our input onto a learned subspace, producing the latent code z. This code is then passed through the decoder layer reconstruct à An additional input reconstruction loss in the form of mean squared error loss between æ and è is included during training time, forcing the model to not only learn to approximate the forward model but also learn linear subspace of the input. We leave learning more complex non-linear encoder-decoder schemes for future work.
During inference, then optimize in this subspace. More concretely, rearrange the proposed architecture so that the æ Decoder layer becomes the first input layer as shown in Figure [b] The optimization variable is passed through this layer to produce estimated & that is then passed through the rest of the model as described in Section 3.1 Our optimization framework thus becomes
1
2 arg min
|G D==)) v1l R(Dx2),
(11)
2
where we recover our final estimate with D
Experimental Set Up
Although the method described in this article should be generalizable to any forward model, we demonstrate on an ocean acoustic tomography problem. Sound speed variations in the ocean are essential for accurate predictions of sound propagation and the various acoustic applications that rely on these predictions (6 [4] 24]. Typically, the ocean sound speed is estimated using empirical formulas based on the temperature, salinity and density. However, this would require dense sampling both spatially and temporally throughout the whole volume of interest. Alternatively, the fundamental relationship between acoustic observations (e.g. arrival time measurements) can be leveraged to indirectly estimate the volumetric bypassing the need to densely sample.
variability of the ocean sound speed profiles (SSPs),
Ocean acoustic tomography (OAT) aims to reconstruct the SSP variations (with respect to a known reference environment) within an ocean slice given the changes in acoustic measurements from the
7

---

Tik
LFM
MLP
PETAL
GT
Depth (km)
1
0
0
5
5
5
Range (km)
(a)
Tik
LFM
MLP
PETAL
GT
Depth (km)
5 0
0
0
5
-2
Range (km)
(b)
Figure 4: Example visualizations of predicted sound speed profiles for each method. Tik manages to recover general structure, but often gets the magnitude wrong. LFM introduces many artifacts during the optimization process due the ill-posed matrix. MLP only captures coarse features but fails to capture subtler details like PETAL (ours).
propagating acoustic waves between multiple pairs of sources and receivers 26 The "forward model" that computes arrival time measurements between source-receiver pairs given an SSP is fundamentally non-linear and would require solving the wave equation. However, SSP fluctuations typically small compared to baseline values (typically 1%). Modeling assumptions can be made to simplify (e.g. ray-tracing methods). Classical OAT methods also further simplify by linearizing the relationship between SSP perturbations and the measured variations in arrival times of stable "paths" propagating between source and receiver pairs, providing a more numerically tractable solution
(inverting linearized model) 3 26 27 34
We perform our experiments on high fidelity month long simulation of the Gulf Mexico as seen in Figure 2 19 20). We restrict experiments to 10 2D range-dependent slices within the data cube (Figure 2b The first 1000 time samples are used for training. the next 200 for validation and the remaining 239 for testing for each slice. This particular train/test/split in time was selected to mimic existing conditions (i.e. collect data for set amount of time to train models and then deploy on future samples). Note that this creates slightly more difficult problem due to the temporally changing dynamics in the test/validation set not present in the training set. The testset further divided based the variability within each slice. This is measured by the average deviation of the last 239 samples from the average of the first 1000.
The ray tracer BELLHOP simulates forward model of sound waves being emitted from each of the sources. The receivers aim to detect peaks of the observed pressure wave, corresponding to the "arrival time" of the direct path and the surface bounce path of the sound wave. The simulation is an approximation of the wave equation. Rather than solving the full partial differential equations, the wave is modeled equation solve along the path of the ray between the source and receiver.
eigen-ray traversing through space, reducing it ordinary differential
We experiment with two source-receiver configurations. "dense" set with 20 sources and receivers approximately every 50 m in depth and km apart in range as shown in Figure 3 We also construct the more challenging "sparse" set up by sub-sampling in depth by factor of 2. To counter the

---

differing bathymetry between the 10 slices, we restrict ourselves to the upper 1000 m portion of the ocean (where a majority of the SSP variation lies as shown in Figure (2] and consider only direct or surface bounce arrival time paths. The 2D SSP is discretized into range by depth grid of 11 x 231.
Thus, the linearized forward models (LFMM) are of dimension 800 x 2541 (or 200 x 2541 in the sparse case), yielding an ill-posed matrix for inversion. We construct our proposed model with 10 reference SSPs: the last available SSP in the trainset (time 1000) for each of the 10 slices. Once trained, the weights are fixed and the ssp æ is solved for iteratively given some observations y. We perform the optimization in the learned subspace and with Eg as well as Sobolev regularization. All optimized models are given a budget of 1000 gradient descent iterations with learning rate of 50.
Results
We compare our proposed model against three baselines. First we compare ourselves against the pseudo-inverse performed in the PCA space and using Tikhonov Regularization as proposed in [15. hereby referred to as "Tik". We select the last available SSP in the train set (time 1000) for each respective slice as the reference point for linearization when evaluating on the test set. We perform PCA on the first 1000 SSPs for each respective slice to construct the basis. Next, we compare against using the linearized forward model (LFM) in iterative method using the same regularization as the proposed model, but with the linearization around a single reference as the forward model. Similar to the Tik method, we linearize around the last available SSP in the training set for each respective slice. Finally, we compare ourselves with a simple multi-layer perceptron (MLP) trained to emulate the forward model. The MLP does not incorporate any physical information and is simply trained to map SSPs to arrival times in a black box manner. All iterative methods are initialized with: the average SSP, the Tik solution and the LFM solution.
The full results are summarized in Table [ When provided no initialization (avg), the proposed method performs the best at 0.364 + 0.018 m/s RMSE. MLP achieves the second best average performance at 0.391 + 0.005 m/s RMSE, where the loss performance likely due to its inability to emulate the forward model as accurately the proposed model. Despite using the same forward model set up. LFM (optimized with different regularization) is able to outperform Tik. We hypothesize that this is due to the basis computed by applying PCA to the training set failing to generalize to the dynamics of the test set. However, when allowed to optimize further (LFM init), degradation occurs due to the inaccuracy of the forward model approximation, increasing from 0.608 to 0.625 m/s error.
MLP also fails to refine the solution for both the LFM as well as the Tik initialization, suggesting that the learned non-linear forward model gets caught in local minimas desipte having better initialization. Our proposed model achieved better results when initialized with the more accurate estimates, dropping the RMSE to 0.357 for both LFM and Tik initializations.
Visualizations of the recovered SSPs can be seen Figure 4 All models are able recover the background levels and other coarse features. Tik is able to recover some finer details, but often fails to achieve the correct magnitudes. LFM introduces many artifacts due to the ill- ill-posed forward matrix, an issue mitigated by using a subspace in Tik and PETAL. MLP is able to capture some of the finer details, but also suffers from artifacts introduced during the non-convex optimization.
Table 2 breaks down the accuracy into further sub categories. Note that the level of variability does not always correlate with the performance of each method. For examples, PETAL performs the worst on the dense case with low variability with average initialization at 0.365 m/s RMSE. However, when provided with the Tik initialization (that performed better on the low variability scenario), is able to achieve the best RMSE of 0.339 m/s. The sparse source/receiver scenario shows a degradation in performance across all methods as expected given the less observations available to each model. However, trends are still upheld with learned surrogates MLP and PETAL outperforming classical linearization based solutions and PETAL performing the best overall.
6 Ablations
In this section, we explore which design decisions contribute to the overall success of PETAL. More specifically, we try to determine whether the learned weighted averaging, learned transformation of the predicted arrival times the learned encoding/decoding of the SSP are all necessary components.

---

Table RMSE (m/s) of inversion across different scenarios.
Low Variability
Med Variability
High Variability
Model
Avg
Tik
Avg
Tik
Avg
Tik
Tik
0.724
0.967
0.926
Sparse S/R Pairs (10 x 10)
LFM
0.692
0.670
0.674
0.677
0.727
0.728
MLP
0.465
0.456
0.477
0.494
0,463
0.466
PETAL (Ours)
0.450
0.444
0.423
0.431
0.416
0.420
Tik
0.647
0.773
0.881
Dense S/R Pairs (20 20)
LFM
0.620
0.597
0,584
0,580
0.617
0.630
MLP
0.384
0.378
0,406
0,409
0.424
0.428
PETAL (Ours)
0.365
0.339
0.360
0.346
0.361
0.374
Table 3: Ablation study of PETAL
Model
Weighted Avg
AT Transform
SSP Subspace
RMSE (m/s)
PETAL (ours)
0.343
A-LFM
X
X
X
0.585
AXX
WA-LFM
X
0.577
WA-LFM + Dec
0.508
WAN
X
0.372
The baseline will be the proposed model which incorporates all three design choices. We compare this against (1) the average of all reference linearizations (A-LFM), (2) weighted average linearizations (WA-LFM), (3) weighted average combined optimized in the learned subspace (WA-LFM + Dec), and (4) a learned weighted average network (WAN). full summary as well as the average RMSE on test set can be found in T
This study shows that each component essential in the success of the proposed model. A-LFM performs the worst overall, though still noticeably better than the single LFM presented in Table suggesting that incorporating more references crucial for representing the varied dynamics present in the test set. Simply adjusting the weights of the average (WA-LFM) as learned by a fully trained PETAL model already leads to an improvement in performance, dropping the RMSE from 0.585 to 0.577. Incorporating the learned SSP subspace improves even further, dropping the RMSE to 0.508. Learning an AT transform allows the surrogate model better approximate the true model, leading to more dramatic improvement in RMSE at 0.372 for WAN. And finally, incorporating all three components leads to the best performance overall at an RMSE of 0.343
Conclusions
In this study, we propose novel architecture to embed physics in learned surrogate models by incorporating linearizations around a set of reference points. We also include an encoder-decoder structure learn subspace optimize in when solving inverse problems, mitigating issues arising from the non-convex optimization. We demonstrate the efficacy of our approach on an Ocean Acoustic Tomography example, outperforming classical methods well learned forward models that do not incorporate any known physics. We validate the necessity of each component in an ablation study, confirming that each contributes to the success of the proposed model.
Acknowledgments and Disclosure of Funding
This work was sponsored by the Ocean Acoustics program at the Office of Naval Research (Code 322) and the Task Force Ocean. We would also like to thank Dr. Guangpeng Liu and Dr. Annalisa Bracco (EAS, GaTech) for sharing Gulf of Mexico simulations of the SSPs.
10

---

References
[1] Jonas Adler and Ozan Öktem. Learned primal-dual reconstruction. IEEE transactions on medical imaging, 1726):1322-1332, 2018.
[2] Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
[3] D Behringer, T Birdsall, M Brown, B Cornuelle, R Heinmiller, R Knox, K Metzger, W Munk, J Spiesberger, R Spindel, D Webb, P Worcester, and C Wunsch. A demonstration of ocean acoustic tomography. Nature, 29966879))121-125,, 1982.
[4] Matthew Chantry, Sam Hatfield, Peter Dueben, Inna Polichtchouk, and Tim Palmer. Machine learning emulation of gravity wave drag in numerical weather forecasting, Journal of Advances in Modeling Earth Systems, 237):=2221I 2021.
[5] Yang Deng, Simiao Ren, Kebin Fan, Jordan M Malof, and Willie J Padilla. Neural-adjoint method for the inverse design of all-dielectric metasurfaces. Optics Express, 2965)77526-754, 2021.
[6] David R Dowling and Karim G Sabra. Acoustic remote sensing. Annual Review of Fluid Mechanics, 47:221-243, 2015.
[7] Nicholas C Durofchalk, Jihui Jin, Heriberto Vazquez, Kay L Gemba, Justin Romberg, and Karim G Sabra. Data driven source localization using a library of nearby shipping sources of opportunity. JASA Express Letters, ((12):24400) 2021.
[8] Geir Evensen. Sequential data assimilation with a nonlinear quasi-geostrophing model using monte carlo methods forecast error statisties. Journal of Geophysical Research: Oceans, 99(C5):1014331011 1994.
[9] Zalan Fabian and Mahdi Soltanolkotabi. Humus-net: Hybrid unrolled multi-scale network architecture for accelerated mri reconstruction. arXiv preprint arXiv:2203.082113 2022.
[10] Peimeng Guan, Jihui Jin, Justin Romberg, and Mark A Davenport. Loop unrolled shallow cquilibrium regularizer (luser)-a memory-efficient inverse problem solver. arXiv preprint vrXiv:2210.04987 2022.
[11] Sam Hatfield, Matthew Chantry, Peter Dueben, Philippe Lopez, Alan Geer, and Tim Palmer. Building tangent-lineat and adjoint models for data assimilation with neural networks. Journal of Advances in Modeling Earth Systems, :30)) 2021.
[12] Peter L Houtekamer and Herschel Mitchell. Data assimilation using an ensemble kalman filter technique. Monthly Weather Review, 126(3):796-8)) 1998.
[13] Zhaoqiong Huang. Ji Xu, Zaixiao Gong, Haibin Wang, and Yonghong Yan. Source localization using deep neural networks in shallow water environment. The journal of the acoustical society of America, 143(5):2922229) 2932, 2018.
[14] Finn B Jensen, William A Kuperman, Michael B Porter, Henrik Schmidt, and Alexandra Tolstoy. Computational Ocean Acoustics, volume 794. Springer, 2011.
[15] Jibui Jin, Priyabrata Saha, Nicholas Durofchalk, Saibal Mukhopadhyay, Justin Romberg, and Karim G Sabra. Machine learning approaches for ray-based ocean acoustic tomography. The Journal of the Acoustical Society of America, 152(6))3768-3788 2022.
[16] George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6):422-440, 2021.
[17] Kenneth Levenberg. A method for the solution of certain non-linear problems in least squares. Quarterly of applied mathematics, 2(2):164-168, 1944.
[18] Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 4503-3) 1989.

---

[19] Guangpeng Liu, Annalisa Bracco, and Uta Passow. The influence of mesoscale and subme- soscale circulation on sinking particles in the northern gulf of mexico. Elementa: science of the Anthropocene, 6, 2018.
[20] Guangpeng Liu, Annalisa Bracco, and Alexandra Sitar. Submesoscale mixing across the mixed layer in the gulf of mexico. Frontiers in Marine Science, 8:615066, 2021.
[21] Xiaofeng Liu, Yalan Song, and Chaopeng Shen. Bathymetry inversion using deep-learning- based surrogate for shallow water equations solvers. arXiv preprint arXiv:2203.022 022222,, 2022.
[22] Wei Ma, Feng Cheng, and Yongmin Liu. Deep-learing-enbbled on-demand design of chiral metamaterials. ACS nano, 12(6):6326-6334, 2018
[23] Donald W Marquardt. An algorithm for least-squares estimation of nonlinear parameters. Journal of the society for Industrial and Applied Mathematics, 11(2):431-44 -411, 1963.
[24] Herman Medwin and Clarence S Clay. Fundamentals of Acoustical Oceanography. Elsevier, 1998.
[25] Ali Mousavi and Richard G Baraniuk. Learning to invert: Signal recovery via deep convolutional networks. In 2017 IEEE international conference acoustics, speech and signal processing (ICASSP), pages 2272 22762276. IEEE, 2017.
[26] Walter Munk and Carl Wunsch. Ocean acoustic tomography: A scheme for large scale monitor- ing. Deep Sea Research Part A. Oceanographic Research Papers, 26(2): 1979.
[27] Walter H Munk and Peter F Worcester. Ocean acoustic tomography. Oceanography, 1(1):8-10, 1988.
(28) Christian C Nadell, Bohao Huang, Jordan M Malof, and Willie J Padilla. Deep learning for accelerated all-dielectric metasurface design. Optics express, 27(20):25523-2753 2019.
[29] Haiqiang Niu, Emma Ozanich, and Peter Gerstoft. Ship localization in santa barbara chan- nel using machine learning classifiers. The journal of the acoustical society of America, EL45))EL455-BL 2017.
[30] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic diffcrentiation in pytorch. 2017.
[31] John Peurifoy, Yichen Shen, Li Jing, Yi Yang, Fidel Cano-Renteria, Brendan G DeLacy, John D Joannopoulos, Max Tegmark, and Marin Soljačić. Nanophotonic particle simulation and inverse design using artificial neural networks. Science advances, 4(6):eaar4206, 2018
[32] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physies-informed neural networks:
A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
[33] Simiao Ren, Willie Padilla, and Jordan Malof. Benchmarking deep inverse models over time, and the neural-adjoint method. Advances in Neural Information Processing Systems, 33:38-48, 2020.
[34] EK Skarsoulis, GA Athanassoulis, and Uwe Send. Ocean acoustic tomography based on peak arrivals. The Journal of the Acoustical Society of America, 100(2):797-817 -13, 1996.
[35] Jian Sun, Huibin Li, Zongben Xu, et al. Deep admm-net for compressive sensing mri. Advances in neural information processing systems, 29, 2016.
[36] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[37] Bo Wang and Huafeng Liu. Fbp-net for direct reconstruction of dynamic pet images. Physics in Medicine & Biology, 65(23)2 2020.
12

---

[38] Yun Wang and Hua Peng. Underwater acoustic source localization using generalized regression neural network. The Journal of the Acoustical Society of America, 143(4):2321-233 2018.
[39] Christopher Williams and Carl Rasmussen. Gaussian processes for regression. Advances in neural information processing systems, 8, 1995.
[40] Jinxi Xiang, Yonggui Dong, and Yunjie Yang. Fista-net: Learning fast iterative shrinkage thresholding network for inverse problems in imaging. IEEE Transactions on Medical Imaging, 40(5):1329-1333) 2021.
[41] Jian Zhang and Bernard Ghanem. Ista-net: Iterative shrinkage-thresholling algorithm inspired deep network for image compressive sensing. 2017.
13

---

Table 4: Architecture details of final proposed PETAL model.
Layer
Input Dim
Output Dim
Spectral Norm?
a Encoder
2541
1000
Pz 1000
1000
W
800
800
Decoder
1000
2541
/
A Training the Forward Model
All experiments were performed on GeForce RTX 2080 Super.
A.1 Data Preparation
We normalize our data using the training set "pixel-wise" average and standard deviation for training only.
A.2 PETAL
The proposed PETAL model only uses linear layers throughout. However, is able learn complex non-linear model due to the attention-inspired mechanism. The exact details of each linear sub-component can be found in Table The embedded physics layer takes in an input of dimension BS x NN. where BS is the batch size and N₂ 2541 the size of the vectorized SSP, and outputs tensor of dimension BS x Nief N2 where Nict refers to the number of linearized reference models and N₂ 800 is the number of observations.
The weighted average layer inspired by attention layers takes in query of dimension BS x 1 H where H is the hidden dimension for the latent space of input SSPs, a key of BS x N x H, and a value of BS x Nee Ny. outputs a predicted arrival time vector of dimension BS Ny-
The model was trained using ADAMW with a learning rate of le-5 for 500 epochs. The learning rate was dropped by factor of 0.2 epoch 300.
The model was trained to minimize the MSE of the arrival time prediction well MSE on the SSP reconstruction. The selected model achieved training RMSE of 4. 4.88e-2 and validation RMSE of 1.02e-1. It achievse an unnormalized (removing training normalization) AT RMSE of 4. 4.98e-4 and SSP RMSE of 5.37e-2 on the validation set.
A.3 MLP
We experimented with both encoder-decoder like structures well as models without the bottle- neck layers. The final best performing model had 4 hidden layers of dim 1500 with leaky ReLU non-linearities. It achieved training RMSE of 3.62e-2 and validation RMSE of 1.22 1.22e-1. The unnormalized AT validation RMSE was 6,08e-4 (higher than PETAL). The model was trained using Adam for 250 epochs with learning rate of le-5.
B Optimization Framework
The neural adjoint method an iterative method to recover an SSP æ given some observations y. All models are optimized using Pytorch's 's Stochastic Gradient Descent with a learning rate of 50 for 1000 epochs.
We use two forms of regularization: an lz penalty on æ with a scale of 1e-7 and a Sobolev penalty (lg the discrete and y gradient) with scale of le-4.
The optimization is performed batches. We set an early cutoff rate of le-2 such that for any sample, if the forward model observation loss drops below this value, we cut off the gradient to that sample. This value lower than the final (normalized) mse AT loss of any the models, so the assumption that any further optimization beyond this point will just overfit to the model.
14

---

Avg Results
LFM Results
Tik Init Results
(ours)
0.6
PETAL (ours)
(aurs)
1.2
LFM
LFM
LFM
RMSE (m/))
MLP
RMSE (m/s)
MLP
RMSE (m//s)
MLP
1.0
0.5
0.5
0.4
0.4
0.4
250
500
750
1000
250
500 Epochs
750
1000
0
250
750
1000
Epochs
Epachs
(a)
(b)
(c)
Figure 5: RMSE of different models vs number of epochs optimized given (a) average, (b) LFM, and (c) Tik initializations.
Boxplot of RMSEs on Test Set
1.25
RMSE (m/s)
1.00
0.75
0.50
0.25
Tik
LFM
MLP
PETAL
Method
Figure 6: Boxplot RMSE of different models.
B.1 Results
The results of NA given different initializations for each of the forward models be seen in Figure 5) Although further iterations might yield higher performance, the overall RMSE already begins to plateau around 1000 epochs. For some models (particularly LFM), the performance already starts to degrade. The distribution of errors after 1000 iterations can be found in Figure
B.2 Robustness to Unseen Slices
In this section we explore the robustness of surrogate models unseen slices. We perform this experiment by training the surrogate models on only slices 1-9 (with the same train/val/test split) and then evaluating on the entirety of slice 10. The performance can be seen in Figure and Table
5 We refer to the subsets of slice 10 as "Train" "Val", and "Test" for convenience, referring to the temporal split of the data, but no samples from slice 10 were available during train time. We select the linearization around the last available SSP the times corresponding to the train set for LFM.
Both trained surrogate models greatly outperform LFM in the subset of the slice overlapping in time with the trainset, suggesting that there are some shared dynamics across space that can be learned. Notably, most models begin to degrade in the times corresponding to the validation and test set, highlighting the difficulty in capturing dynamic shifts over time. However, the learned models still remained more robust to this shift and the performance only degraded slightly compared to when trained with slices dropping from 0.33715 (when evaluated only on slice 10) to 0.33736 for our proposed model PETAL.
15

---

Performance on unseen slice
Train
Val
Test
PETAL (ours)
1.0
MLP
RMSE (m/s)
0.8
LFM
0.6
0.4
0.2
0
10
15
20
25
30
Time (days)
Figure 7: Performance of models on unseen slice. Both trained forward models perform well on the subset of the slice overlapping in time with the trainset, suggesting that dynamics are shared throughout the region.
Table 5: Average RMSE (m/s) of inversion on unseen slice.
Model
Train
Test
LFM
0,405
0.447
0.583
MLP
0.196
0.288
0.402
PETAL (ours)
0.149
0.217
0.337
C
Gradient of PETAL
Define (simplification) of the PETAL model
w w w'D')
(12)
= Σ w W (4 ( x b¹).
Note that by construction, the weights wt sum up 1. If include this in simple MSE loss we get
L v1
(13)
Computing gradient w.r.t x gives
0L
ΣΣ Ou" dr uP W W(A'z+b' Wh' - -y) + A W (WA Wb y),
dx
(14)
where the right term reduces to convex combination of the gradient of the linearized physics based forward models, modulated by some matrix W, when =
D Extensions to Interpolation: Additional Experiments
In this section, begin preliminary study PETAL for interpolation (i.e. random split of the available data). By removing the dynamics shift between the training and the validation set, we are able to explore more complex models beyond linear components.
D.1 Selecting Reference Points
Due to the random split, we can no longer rely on heuristies to select appropriate reference models. Instead, we perform projected k-means approach to select our reference points. The standard
16

---

k-means algorithm is applied to all SSPs in the training set. Once the algorithm terminates, the closest point in the training set to each mean is selected as the reference point be embedded in the model. For our experiments, we used l = 20.
D.2 Extending the layers
We maintain the overall structure of PETAL by retaining an SSP encoder-decoder, an embedded physics module, a learned weighted average module, and a final layer for the predicted arrival time. However, we extend the learned components by replacing the single linear layer used in the main manuscript with MLPs.
In addition to replacing the final linear layer W with series of linear layers followed by a leaky ReLU activation, we also include a skip connection between the output of the learned weighted average module and the final arrival times prediction. We also observed that including an auxiliary loss in the form of mean squared error between the ground truth arrival times and the direct output of the weighted average layer (prior to passing through the final layers) was beneficial.
We experimented with replacing linear layers with convolutional layers, but noticed huge degradation in performance. We hypothesize that this is due the nature of the data not being suitable for convolutions. The input SSPs are not translationally invariant depth-wise. The output arrival times spatially small (20 x 20), and thus the choice of padding can easily introduce errors at the edge, degrading the final prediction.
However, we did notice introducing convolutional layers to have a beneficial effect when added the end of the SSP decoder. The convolutional filters play the role of filtering/smoothing the predicted output to create more realistic reconstructions.
D.3 Results
We compare all our results against a Gaussian process regression (GPR) baseline. We construct a dataset by randomly sampling 1000 SSPs from the 1439 available from each of the 10 slices, totaling 10,000 total training samples. We use a random subset of 1000 points to train GPR due to memory limitations of the method. Two models (MLP and PETAL) are trained to approximate the forward model and then used in an optimization framework for inference.
The MLP consists of 4 hidden layers with dimension 1500 for total of 5 linear layers separated by leaky ReL.U. Batch normalization was found to degrade performance and thus was omitted for the final model.
The final PETAL architecture uses a 2 layer MLP with hidden dimension 1000 for the x encoder and x decoder separated by leaky ReLUs. The æ decoder also uses 3 convolutional layers with kernel size 3 x and 32 filters. Finally, the output layer consists of 3 layer MLP with hidden dimension 800.
The performance of each model is shown in Table We experiment with gradient descent (SGD) and limited memory Broyden-Fleecher- Fletcher-Goldfarr -oldfarb-Shanno Shanno algorithm (LBFGS) as the optimization algorithm. For both MLP and PETAL, LBFGS reached a local minima sig significantly faster than SGD, However, it did not always achieve the most accurate reconstruction, likely getting trapped in local minimas due to the non-convex nature of optimizing neural net. MLP achieved higher accuracy for emulating the forward model. However, this higher performance did not translate to the inference stage. MLP was able to perform well when given a more accurate starting point (GPR init), but failed to outperform GPR with the average intialization, getting caught in local minimas. PETAL was able to outperform GPR regardless of initialization. This suggests that the structure embedded into PETAL's architecture aids in the optimization process when determining descent directions. Examples of reconstructed SSPs can be found in Figure g Note the artifacts and hallucinations introduced by optimizing the MLP surrogate.
Fletcher Goldfarb
E
Limitations
Our proposed model was evaluated on noise-less simulations, both with respect to measurements and sensor/receiver placement, which is not true in practice for data collected the real world. We also did not explore the selection process of the reference points to linearize around, assuming that the
17

---

Table RMSE (m/s) of random data split
Model
Optimization
Avg Init
GPR Init
GPR
0.176
PETAL
SGD
0.144
0.128
PETAL
LBFGS
0.142
0.129
MLP
SGD
0.224
0.146
MLP
LBFGS
0.231
0.164
Depat (km))
04
0.0
-0.4
0
Range [km]
Range (km)
Range (km)
Range (km)
Range (km)
Range (km)
(a)
GT: 0.000
GPR: 0.065
PETAL: 0.053
PFTAL: 0.038
M.PP: D.071
M.IP: 0.072
0.8
06
Dapth
-0.4
Range [kml
Range (km)
Range (km)
Range (km)
Range (km)
Range (km)
GT: 0.000
GPR: 0.071
S-PETAL: 0.049
L-PETA:: 0.046
S-MLP: 0.065
L-MLP: 0.063
0.6
02
-0.2
-0.4
Range
Range (km)
Ranige
Range
Range
Range (km)
Figure 8: Example visualizations of predicted sound speed profiles for each method. The "S-" prefix denotes models optimized with SGD and "L-" with LBFGS.
18

---

chosen subset sufficiently represented the data. However, section B.2 suggests that the selection of reference points is somewhat robust unseen dynamics.
19


Reviews:
[{"review_id": "eI9ukCBIoj", "reviewer": "Reviewer_exJb", "summary": "This paper further improves a method used in reference [30]. The goal is to reconstruct the inverse problem, through a \"gray-box\" model. For the forward solver, a parametrized model is proposed: many linearizations centered at different (reference) points of the nonlinear forward problem are ensembled. The ensemble weights depend on a learned similarity measure. The inverse problem is solved using the usual way (iterative methods with a regularization) with no parameterization involved. \n\nOverall, I think this is a decent submission and lean towards acceptance.\n", "questions": "- line 68: while the meaning of \"TV\" is obvious in the context, it is better to define it before using the abbreviation for the first time.\n- What does it mean that \"direct inversion methods\" yield a \"single estimate\"?\n- line 133: $n$ the total number of \"observations\"?\n- line 162: What is the spectral \"norm\", based on the context right before it should be \"normalization\", not a norm, please confirm this.\n- What does \"embedding physics\" in Section 3 mean? ", "limitations": "The authors did not discuss any limitations.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "- This paper is clear-written and well-motivated.\n- There is a full-blown systematic ablation study.", "weaknesses": "- The contribution versus that of [30] seems limited.\n- Certain presentations are somewhat focused on a specific inverse problem (ocean tomography), for example, what is a \"ray tracer\" on line 134? What is an \"eigenray\"? What precisely is the forward model in this case?\n- The biggest weakness is the assumption that the Jacobian of the forward operator can be straightforwardly obtained in equation (6). Using EIT as an example, how to measure the change in the nonlinear forward operator with respect to the conductivity coefficients is still quite unsolved. \n- The error for fitting the forward problem is not reported. \n- The reason why, or an empircal study of, not updating the parameters of the forward problem solver during the iterative procedure to solve the inverse problem is not mentioned, at all.\n- Missing a whole line of reference on EnKF for inverse problems.\n- The baseline for comparison is not that strong.\n\n\nThere are some grammatical errors:\n- line 23: \"well known\" -> \"well-known\".\n- line 53: there should be an article before \"simulation\".\n- line 93: \"physics obeying\" -> \"physics-obeying\".\n- line 104: \"physics based\" -> \"physics-based\" and a few other places as well.\n- line 105: a minor nitpicking, \"Pytorch\" should be \"PyTorch\".\n- line 154: \"dot-product\" should be \"dot product\" here, different from that in \"dot-product attention\".\n- line 173: \"undesireable local minimas\" should be \"undesirable local minima\", same problem on line 279.\n- line 190: \"trade offs\" -> \"trade-offs\".\n- line 193: \"generate\" -> \"generates\".\n- line 214: \"is essential\" -> \"are essential\".\n- line 224: \"a SSP\" -> \"an SSP\".\n- line 232: \"range dependent\" -> \"range-dependent\".\n- line 281 and 312: \"out-perform\" should be one word without the hyphen.\n- line 314: \"each contribute to ...\" -> \"each contributes to\".\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and thoughtful review. We are delighted that you found our work well-motivated. We respond to your highlighted issues below:\n\n>The contribution versus that of [30] seems limited.  \n\nWhile our work is an extension of [30], we believe that tackling much larger dimensional problems (around dimension 10 vs 1000) and providing the tools to do so in the form of a learned basis is a notable contribution. This also departs from other previous work such as [13] that used principal component analysis to generate a basis. \nFurthermore, our primary contribution lies in a suggestion for a model architecture that takes advantage of a specific type of low-fidelity physics model (linearizations around references). [30] describes a methodology assuming a generic black-box architecture that is not necessarily tailored to solving inverse problems. The proposed model can even be the foundation for more complex surrogate models (additional layers to further refine the output). \n\n> Certain presentations are somewhat focused on a specific inverse problem (ocean tomography), for example, what is a \"ray tracer\" on line 134? What is an \"eigenray\"? What precisely is the forward model in this case?\n\nTo expand on section 4 (line 212), the forward model simulates emitting a sound wave from each of the 20 sources. The receivers aim to detect peaks of the observed pressure wave, corresponding to the \"arrival time\" of the direct path and the surface bounce path of the sound wave, leading to a toal of 800 observables. The simulation is an approximation of the wave equation for the pressure waves. Rather than solving the full partial differential equations, we model the wave as an eigen-ray traversing through space, reducing it to an ordinary differential equation solve along the path of the ray between the source and receiver. \n\nThe ray tracer refers to BELLHOP, a program used in the ocean acoustics community. This program is available freely at the Ocean Acoustics Library (http://oalib.hlsresearch.com/AcousticsToolbox/). We will add this to the final revision for clarity.\n\n> The biggest weakness is the assumption that the Jacobian of the forward operator can be straightforwardly obtained in equation (6). Using EIT as an example, how to measure the change in the nonlinear forward operator with respect to the conductivity coefficients is still quite unsolved.\n\nThe reviewer brings up a very important critique. Although we use explicit Jacobians in our model formulation, it is not a requirement. The methodology outlined in our paper can also take advantage of implicit formulations, since we only need the ability to apply the Jacobian to arbitrary vectors, as is the case in any problem that uses the adjoint state method (computing Jacobian-vector products). \n\nWe would also like to point out that the full Jacobian can be calculated using the same adjoint calculation by applying different vectors (e.g. elementary). Although this is expensive and would scale with the dimension of the output, it can be treated as an offline cost and would only need to be done for a small subset (the reference points selected). However, we do acknowledge that the need to apply the Jacobian in some form (explicit or implicit) does somewhat limit our method to a certain subclass of problems. \n\n> The error for fitting the forward problem is not reported.\n\nPETAL achieves a (normalized by training set statistics) RMSE of 4.88e-2 and MLP 3.62e-2. The (normalized) validation RMSEs are 1.02e-1 and 1.22e-1, respectively. We report a validation (unnormalized) rmse of 4.98e-4 seconds for arrival times for PETAL and 6.08e-4 seconds for MLP in the appendix (A.2 and A.3). We will revise the final draft to include the training numbers.  \n\n> The reason why, or an empircal study of, not updating the parameters of the forward problem solver during the iterative procedure to solve the inverse problem is not mentioned, at all.\n\nWe followed the techniques practiced in the underwater acoustic community (a single linearization used across a range of samples) for computational reasons. We did a preliminary study on re-linearizing at intermediate points (manuscript submitted but not available for release yet) and found that on the full dense setup, it lead to minimal improvement. However, it did improve the estimation results further on sparser source/receiver pair set ups ($10\\times10$ and $3\\times10$).\n\n> Missing a whole line of reference on EnKF for inverse problems.\n\nWe thank the reviewer for bringing this to our attention. We will make sure to include the relevant references in the final revision. \n\n> Grammatical errors/suggestions: \n\nWe will edit these in the final revision. Thank you for pointing these out.\n\n> What does it mean that \"direct inversion methods\" yield a \"single estimate\"?\n\nNon-bayesian models that focus on learning a mapping from $y$ to $x$ only produce a single estimate of $x$ for each observation $y$. This contrasts from optimization based methods that can produce different answers with different initializations\n\n> Spectral \"norm\"\n\nThe reviewer is correct. We have updated the manuscript to reflect that it is a \"normalization\".\n\n> What does \"embedding physics\" in Section 3 mean?\n\nWe use embedding physics in this context to mean directly incorporating physics-derived (low-fidelity approximation) models into the machine learning architecture. \n\n"}, {"review_id": "eoQvCJycIK", "reviewer": "Reviewer_yV52", "summary": "A surrogate model based on multiple linearizations of the target simulator is proposed. It is then used for solving inverse problems by gradient descent.", "questions": "1. What is common or different from the previous studies [4, 10, 30]? The authors mention that their work departs from them, but a concrete explanation of the relation is missing.\n\n2. In the proposed model, how do you choose $x_ \\text{ref}$ in general?\n\n3. What is the role of $W$ in Eq. (3)? Isn't this redundant with $P_y$?\n\n----\n\nPost-rebuttal: The questions above are basically solved after the rebuttal.\n", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The idea of using linearizations for surrogate model is interesting.\n- The explanation about the better gradients of such a model is reasonable.\n- The method is examined with a concrete example.\n- The rationale of why the proposed model's gradient can be better is very nice.", "weaknesses": "- The difference from previous studies [4,10,30] is unclear, which is the main reason why I lean to rejection currently.\n- The experiment is conducted only for a single forward model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. We respond to the highlighted questions below:\n\n> What is common or different from the previous studies [4, 10, 30]? The authors mention that their work departs from them, but a concrete explanation of the relation is missing.\n\n- [4] ML Emulation of gravity wave  \nThis work primarily focuses on building surrogate models. Although they include the surrogate as part of a pipeline that includes physics-based components, the learned model only serves to replace a specific stage in a data-driven manner. Our proposed methodology is similar in that it is attempting to approximate a physics-based process (“forward model”), but differs in that it incorporates low-fidelity versions (linerizations at different reference points) of the forward model directly into the surrogate, so it is not a purely data-driven method. \n- [10] Tangent-linear and adjoint models  \nThis work proposes a surrogate model in the form of a multi-layer perceptron net (MLP) and discusses tests to verify the accuracy of both its output as well as the gradient for inversion. They construct/test the tangent-linear (linearization of the forward model around a reference point) as well as the adjoint (linearization applied for gradient calculation), but **for the learned surrogate model**. We use an ensemble of tangent-linear models (referred to as jacobians in the manuscript) of the **true physics-based forward model** at a number of reference points rather than the learned surrogate when constructing our architecture. We also compute the adjoint of the surrogate when performing inversion, but do so implicitly through the use of auto-grad libraries and note that this specific approach is not our primary contribution, but that of [30].   \n- [30] Benchmarking Inverse Problems and the Neural Adjoint Method  \nWhile our work is an extension of [30], we believe that tackling much larger dimensional problems (around dimension 10 vs 1000) and providing the tools to do so in the form of a learned basis for improved reconstruction is a notable contribution. This also departs from other previous work such as [13] that used principal component analysis to generate a basis. \nFurthermore, our primary contribution lies in a suggestion for a model architecture that takes advantage of a specific type of low-fidelity physics model (linearizations around references). [30] describes a methodology assuming a generic black-box architecture that is not necessarily tailored to solving inverse problems. Our proposed model can even be the foundation for more complex surrogate models (additional layers to further refine the output). \n\n> In the proposed model, how do you choose $x_{ref}$ in general?  \n\nCurrently we are employing a heuristic approach based on current practices (i.e. selecting the most recently available point for linearization). We are currently exploring a k-means clustering approach to select a reference set for better coverage, but are currently not ready to comment on such an approach.\n\n> What is the role of $W$ in Eq. (3)? Isn't this redundant with $P_y$? \n\nThe reviewer is correct. Notation was done to reflect architecture design choices in the code related to attention layers, but are otherwise indeed redundant. We will simplify the notation to a single matrix in the final revision.\n\n> The experiment is conducted only for a single forward model.\n\nWe agree with the reviewer that empirically showing the proposed model on a variety of forward models would best highlight its potential to solve inverse problems. We are actively investigating other potential models and datasets, but do not have available results to share just yet. However, we have performed additional experiments on the simulation of the Gulf of Mexico as shared in the overall rebuttal. More specifically, we repeat a subset of experiments on a different region of the ocean and we simulate a different forward model configuration on the region presented in the manuscript. Both new experiments match the trends of the original experiment with our proposed model performing the best. \n"}, {"review_id": "VEuxfu25Vx", "reviewer": "Reviewer_1TNX", "summary": "The paper proposes a new neural network architecture to simulate the evolution of a given physical system that can be seen as a generalized weighted combination of linearizations around a set reference pairs of state and measurements $(x_{ref, i}, y_{ref, i})_{i=1}^{N}$. The paper proposes to use this method when solving inverse problems in physics through MAP (maximum a posteriori) estimations. It then shows that using this architecture they can achieve smaller RMSE w.r.t the state $x$ that was used to generate the observation $y$ with different regularizations (priors) for the problem of ocean acoustic tomography when compared to other forward models.", "questions": "* If $P_y$ is a matrix, then I do not see why the choice of separating $W$ and $P_y$ in equation 8, since $WP_Y$ is just another matrix of dimension $2431 \\times 800$, that could be directly parameterized.\n\n* It is known that self-normalisation of weights that correspond to inner products in big dimensions lead to weight collapse (after self normalisation one of the reference ($x_ref, y_ref$) can have a weight of almost one while the others 0). In this case, the inverse problem would be basically based upon $A_ref$ for this given pair. Is this something that was seen during the optimization for solving the inverse problem? If this is the case, the model would be basically selecting one of the reference solutions, and could be simplified.\n\n* How sensitive is the model to the reference solution that were chosen? How would it be impacted if another set was chosen? \n\n*The MLP model seems to be overfitting (lines 275 and 276). If i'm not mistaken, standard techniques to reduce overfitting have not been tried (dropout, batch normalisation, etc...). How would those impact the performance?\n\n\n", "limitations": "NA", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "The paper proposes a reasonable model that is based on reference set. The presentation of the inverse problem applications to the physics domain is good and the application to ocean acoustic tomography is original and not often found amongst the NeurIPS community.", "weaknesses": "One of the main weakness of the paper is the lack of support of the claim that the proposed architecture is a efficient tool for modelling general forward problems in physics. There are two main problems with the validation of the method:\n\n1. Lack of replicates for the validation metrics:\n - Neither table 1 or 2 contain confidence intervals. Those kind of metrics are expected to have aggregated results from several replications of initializations of the network before training as well as several initializations of the optimization problem for solving the inverse problem. Without it, it is quite hard to know if the RMSE differences are signficative in say MLP vs Petal in table 1 or PETAL vs WAN in table 2.\n\n2. Comparison with other baselines:\nI think the comparison with other possible forward models is lacking. Indeed I'd expect a comparison to a more expressive meta model than a MLP, by either training a more complex network that better exploits the 2d nature of the problem. Other methods that take advantage of some reference set, such as Gaussian processes could have been exploited as well.\nI'd like to know as well if the proposed method scales to other domains than Ocean Acoustic Tomography. This would be of great interest in reproducibility, since the dataset used is not available and therefore the community has no way of replicating the experiments.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough review. We respond to their concerns below:\n> Neither table 1 or 2 contain confidence intervals.\n\nWe repeated the model training as the reviewer suggested and have updated Table 1 in the author rebuttals attachment to reflect the confidence intervals. The reviewers suggestion was helpful in further understanding our model. We observe that PETAL seem more sensitive to model initialization than MLP with its higher standard deviation, but still manages to outperform other baselines.\n\n> Comparison with other baselines. Indeed I'd expect a comparison to a more expressive meta model than a MLP, by either training a more complex network that better exploits the 2d nature of the problem. \n\nWe had similar thoughts to the reviewer for data-driven surrogate models with respect to exploiting the 2d nature of the problem. Beyond simple MLPs, we also tested fully convolutional models, both 1d and 2d, as well as hybrid convolutional/linear models. We also experimented with encoder-decoder structures, mixing and matching linear/convolutional/deconvolutional layers. Unfortunately, all models underperformed during the first stage of attempting to learn the forward model with the exception of MLP and PETAL. We conjecture that this is partially due to the changing statistics between training and validation set, making it more difficult to train a generalizable model. We did an experiment with a random split of train/validation/test, and all models achieved very good performance. However, this does not reflect a realistic data split, hence the current split used in the manuscript.\n\n> Other methods that take advantage of some reference set, such as Gaussian processes could have been exploited as well.\n\nWe thank the reviewer for bringing this to our attention. We will make sure to thoroughly investigate this alternative method for future work.\n\n> I'd like to know as well if the proposed method scales to other domains than Ocean Acoustic Tomography. This would be of great interest in reproducibility, since the dataset used is not available and therefore the community has no way of replicating the experiments.\n\nWe agree with the reviewer that empirically showing the proposed model on a variety of forward models would best highlight its potential to solve inverse problems. We are actively investigating other potential models and datasets, but do not have available results to share just yet. However, we have performed additional experiments on the simulation of the Gulf of Mexico as shared in the overall rebuttal. More specifically, we repeat a subset of experiments on a different region of the ocean and we simulate a different forward model configuration on the region presented in the manuscript. Both new experiments match the trends of the original experiment with our proposed model performing the best.\n\n> If $P_y$ is a matrix, then I do not see why the choice of separating...\n\nThe reviewer is correct. Notation was done to reflect architecture design choices in the code related to attention layers, but are otherwise indeed redundant. We will simplify the notation to a single matrix in the final revision.\n\n> It is known that self-normalisation of weights that correspond to inner products in big dimensions lead to weight collapse (after self normalisation one of the reference ($x_{ref}$, $y_{ref}$) can have a weight of almost one while the others 0). In this case, the inverse problem would be basically based upon $A_{ref}$ for this given pair. Is this something that was seen during the optimization for solving the inverse problem? If this is the case, the model would be basically selecting one of the reference solutions, and could be simplified.\n\nThe reviewer makes an insightful observation. Indeed, we found similar results in our initial experiments where the predicted weights collapsed into favoring a single $A_{ref}$ across a number of examples. However, we addressed this by applying spectral normalization to the projection matrix. This led to outputs being at an order of magnitude that did not collapse to a delta after the soft-max function. Instead we observe a weight distribution that changes depending on the input. \n\n> How sensitive is the model to the reference solution that were chosen? How would it be impacted if another set was chosen?\n\nFor the scope of this manuscript, we did not experiment extensively with the selection of reference models, instead relying on a heuristic (the most recently available point in the trainset). However, we did conduct one experiment in the appendix (B.2) where we construct and train the model with references only from 9 slices before testing it on the 10th slice. There was very little degradation in performance, suggesting (at least in this instance) some robustness to the reference set. We are currently exploring a k-means clustering approach to selecting the reference model but are not ready to comment at this point. \n\n> The MLP model seems to be overfitting (lines 275 and 276). If i'm not mistaken, standard techniques to reduce overfitting have not been tried (dropout, batch normalisation, etc...). How would those impact the performance?\n\nDuring experiments, we did test various methods to reduce overfitting including batch normalization and dropout. However, this led to a degradation in performance on the validation set, so we chose to not include them in the final model. We note that some of the issues arise from the nature of our dataset and how we chose to split it (sequentially in time). \n\n"}, {"review_id": "tX4wg67lby", "reviewer": "Reviewer_Bbps", "summary": "This paper tackles the area of inverse problems using a neural network architecture. Specifically this follows the neural adjoint paradigm, whereby the algorithm learns a forward model and then uses it for inference. The authors expand on this existing methodology by incorporating knowledge of the physics-based forward model and using a latent space in the optimization stage.", "questions": "1. Please include more details regarding the hyper-parameters of your PETAL network.\n2. Please include details about the MLP used in the benchmarking.\n3. What would the performance of the MLP be if the embedded physics is also provided to it?\n4. I am not convinced based on Table 1 that PETAL is more robust to initialization compared to the MLP, as it is claimed on page 8 line 281. It could also be chance, given how small the differences are between the various initializations for both PETAL and the MLP. Please provide further evidence for this claim.\n5. How many simulations of the Gulf of Mexico have been done? If only one, I would suggest further benchmarking on either other simulations of this Gulf or even more beneficial, on other inverse problems.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is clearly written and well motivated, providing background and alternative methodologies. The paper also has clear mathematical notation and is benchmarked on a simulated system that appears grounded in reality, i.e. the sounds speeds in the Gulf of Mexico. The authors also attempt to tease out and explain the effects of the various architecture components using ablation studies.", "weaknesses": "The paper doesn’t discuss the hyper-parameters of the networks used and their effects on the results. When discussing the performance of their algorithm against the MLP, it is unclear what this MLP actually is and while the MLP doesn’t appear to perform as well, it is rather a close second. Hence, the extra complexity of the PETAL algorithm and the benefit of the embedded physics is not as transparent as it is stated in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful feedback. In particular, we are glad the reviewer found our work clearly written and well motivated. We attempt to address some of their issues below:\n\n> Hyper-parameters of PETAL and MLP\n\nWe include the sizes of the various layers included in PETAL in Appendix A.2 as well as training hyperparameters. The details of the MLP are included in A.3. We are happy to include any additional detail suggested by the reviewer.\n\n> What would the performance of the MLP be if the embedded physics is also provided to it?\n\nThe suggestion of providing embedded physics to an MLP model is a interesting. If we consider the case of stacking the output across all the reference linearized models and feeding this into an MLP, I would expect the performance to be potentially worse than PETAL. The primary reason is that the weights learned to merge the different outputs would be fixed, while as PETAL produces adaptive weights depending on the proximity of the input to each specific reference point. Fixed weights might be more accurate if each of the models had some general approximation error, but that is not the case. These linearizations are more accurate within a small neighborhood of the reference point used and thus their accuracy should vary depending on the input. \n\nWe can also consider the case of using PETAL’s learned adaptive weighting scheme to combine the outputs of the embedded physics module before feeding it into an MLP, but I would argue that this is simply an extension of the proposed model with a non-linear final output layer instead of the current linear one. \n\n> I am not convinced based on Table 1 that PETAL is more robust to initialization compared to the MLP, as it is claimed on page 8 line 281. It could also be chance, given how small the differences are between the various initializations for both PETAL and the MLP. Please provide further evidence for this claim.\n\nWe have extended the experiments, and included additional tables of results. Reviewing Table 1, we can see that MLP lead to degradation in accuracy, despite being initialized with the LFM/Tik estimate which is more accurate than the training set average (approximately 1 m/s RMSE), suggesting that it gets caught in a local minima. Alternatively, PETAL is able to refine the answer further across different realizations (different trained weight initialization, but fixed initialization for optimization for inference purposes).\n\nFurther examination of Table 3 also shows similar trends. PETAL is able to achieve better results when given a better starting point in all three variability scenarios for the dense $20\\times20$ forward model and for the low variability scenario for the spare $10\\times10$ configuration. However, it does fail to achieve as good of performance for the medium and high variability scenario. This is likely due to the more challenge forward model setting with less observations making the optimization problem more difficult to solve. Alternatively, MLP shows more sensitivity to its input, only succesfully achieving better results with a better initialization in the low variability regime. \n\n> How many simulations of the Gulf of Mexico have been done? If only one, I would suggest further benchmarking on either other simulations of this Gulf or even more beneficial, on other inverse problems.\n\nWhile we have not finished designing experiments for other domains, we have performed additional experiments on different Ocean Acoustic examples. We experimented with a sparse source/receiver pair configuration ($10\\times10$ instead of $20\\times20$) as well as on a different portion of the ocean for the full $20\\times20$ configuration. All models followed similar trends (i.e. PETAL outperformed MLP which outperformed linearizations). Please see the attachment for summarized results. \n\n\n"}]

Author rebuttal:
We thank all the reviewers for their time and thoughtful feedback. We have included an additional set of experiments as well as a more detailed breakdown of the scenario presented in the manuscript.

For Table 1, we retrained all neural net models 10 different initialized weights before performing the optimization for inference with fixed input initializations. We have included the standard deviation for MLP and PETAL to showcase the confidence intervals. Note that the proposed models upper end still falls below that of other models.

For Table 2, we perform initial experiments on a different region of the Gulf of Mexico simulation. We still use 10 slices with the same train/validation/test split. 

For Table 3, we divide the original dataset into three scenarios based on the variability seen in the test set (last 239 samples) for each slice. Slices are sorted based on the average root mean square error (RMSE) from the average of the first 1000 samples (corresponding to the train set) for each respective slice. This approximates how much the slice deviates over time. The low variability slices had around 1.6 m/s deviation, the medium 2.5 m/s, and the high 3.4 m/s.

We then repeat the experiments for a different configuration of the forward model. Instead of using $20\times20$ source/receiver pairs, we reduce this to a $10\times10$ configuration, leading to signficantly less observations. We would expect performance to degrade in this sparser scenario, which is reflected in the experimental results. However, general trends still hold with the trained surrogates out-performing the classical baselines and the proposed PETAL model performing the best. 

One interesting thing to observe is that although there is some degradation in performance as variability increases, this does not always hold true. For example, PETAL actually performs better in the higher variability scenario. We hypothesize that this is due to being trained on all 10 slices rather than each individual slice and thus learning a more generalized forward model. 

Meta-review: This paper describes a new method for solving inverse problems that rely on complex physics-based forward operators. The paper proposes to train a differentiable surrogate forward emulator using an interesting learned model that uses linearizations around reference points. 
The authors did a great job in their rebuttal by including additional results for fewer source/receiver pairs in a much more under-sampled regime. The fewer observations lead to graceful degradation as the additional experiments show. 
One reviewer remains concerned about limited evaluation to more general physics problem but overall I think there is sufficient novelty in this paper.
Decision: Accept (poster)