SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jihui Jin
jihui@gatech.edu
Etienne Ollivier
Richard Touret
Matthew McKinley
Karim G. Sabra
Justin K. Romberg
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.
Inverse problems arise in many scientific applications where the goal is to reconstruct some unknown signal, image or volume of interest from indirect observations. The forward process, or 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 us consider the task of recovering some signal  $x$  from observations  $y$  that are related by some potentially
non-linear forward model  $F$  via
$$
\boldsymbol {y} = F (\boldsymbol {x}) + \boldsymbol {\eta}, \tag {1}
$$
where  $\eta$  encapsulates noise or other perturbations. Our forward model  $F$  represents a 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 an 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 is typically solved for an in iterative fashion. This process can be potentially computationally expensive due to requiring many applications of the physics based forward model. Computing the adjoint to determine a descent direction increases the computational burden even more.
In this paper, we propose a novel architecture trained to emulate the physics-based forward model. The surrogate model provides a cheaper alternative for both the forward pass and descent direction calculation by leveraging existing auto-grad libraries[30], making iterative solvers feasible for recovering a solution. This work departs from previous works who also aim to emulate the forward model [4, 11, 33] by embedding low-fidelity forward models into the surrogate itself, thereby explicitly incorporating the physics instead of treating it as a black box. More concretely, we propose to use a set of linearizations of the forward model at a set of reference point and train a model to fuse and correct the outputs in an adaptive manner. We then use this trained model in a classical optimization framework to recover the signal given a set of observations.
Concretely, our paper makes the following contributions:
- We propose a novel architecture that learns to emulate the forward model. The model directly embeds physics 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 a simulation of ocean dynamics in the Gulf of Mexico.
2.1 Iterative Inversion
The task of directly inverting some potentially non-linear forward model  $F$  is often non-trivial or mathematically impossible. Instead, a more stable alternative is to iteratively solve for  $x$  given some observations  $y$ . Classically, this is done by formulating the reconstruction as solving a non-linear least squares problem that aims to minimize
$$
\hat {x} = \arg \min  _ {x} \frac {1}{2} \| \boldsymbol {y} - F (\boldsymbol {x}) \| ^ {2}. \tag {2}
$$
The Gauss-Newton method solves Equation (2) by iteratively computing the Jacobian  $J_F$  of  $F$  at the current best estimate  $\boldsymbol{x}^{(k)}$  and solving a set of linear equations to generate an update of the form
$$
\boldsymbol {x} ^ {(k + 1)} = \boldsymbol {x} ^ {(k)} + \left(\boldsymbol {J} _ {F} ^ {\top} \boldsymbol {J} _ {F}\right) ^ {- 1} \boldsymbol {J} _ {F} ^ {\top} (\boldsymbol {y} - F \left(\boldsymbol {x} ^ {(k)}\right)). \tag {3}
$$
The Levenberg-Marquardt algorithm [17, 23] presents a more robust alternative to Gauss-Newton by introducing a (non-negative) damping factor  $\lambda$  leading to an update of the form
$$
\boldsymbol {x} ^ {(k + 1)} = \boldsymbol {x} ^ {(k)} + \left(\boldsymbol {J} _ {F} ^ {\top} \boldsymbol {J} _ {F} + \lambda \mathbf {I}\right) ^ {- 1} \boldsymbol {J} _ {F} ^ {\top} (\boldsymbol {y} - F \left(\boldsymbol {x} ^ {(k)}\right)). \tag {4}
$$
Another approach to address the ill-posed nature of  $F$  is to explicitly introduce a regularization function  $R(\pmb{x})$  as an additional penalty term to Equation (2). The regularizer  $R$  addresses the ill-posed
(b)
Figure 1: (a) Architecture of PETAL for training the forward model. The signal  $\pmb{x}$  is passed through the Embedded Physics module to generate low-fidelity physics-based predictions  $\hat{\pmb{y}}^1, \dots, \hat{\pmb{y}}^N$  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  $\pmb{W}$ . An encoding layer  $\pmb{E}_x$  and decoding layer  $\pmb{D}_x$  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  $\pmb{z} \in \pmb{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  $\ell_2$  norm to encourage smaller norm solutions (similar to Equation 4),  $\ell_1$  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
$$
\boldsymbol {x} ^ {(k + 1)} = \boldsymbol {x} ^ {(k)} + \gamma \boldsymbol {J} _ {F} ^ {\top} (\boldsymbol {y} - F (\boldsymbol {x} ^ {(k)})) - \lambda \nabla R (\boldsymbol {x} ^ {(k)}), \tag {5}
$$
where  $\gamma$  is a tuned step size,  $\lambda$  is a non-negative factor controlling the strength of regularization, and  $J_{F}$  is the Jacobian of  $F$  evaluated at  $\pmb{x}^{(k)}$ .
However all these approaches can be undesirable in practice due to the need to re-compute the Jacobian (or a Jacobian-vector product) at each iteration. Computing a Jacobian-vector product to determine a descent direction often involves solving an auxiliary set of PDEs, which can be computationally expensive if many iterations are required to achieve an acceptable level of accuracy.
Ensemble Kalman filter (EnKF) offers a Bayesian iterative approach to estimating the underlying signal [8, 12]. 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 a 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-linear 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, 13, 25, 29, 37, 38]. Unlike typical supervised learning tasks that attempt to learn a mapping purely from examples, deep learning for inverse problems can leverage our understanding of the physics in the forward model. One common approach to embed the physics is the so-called "Loop Unrolled" architecture heavily inspired by
existing iterative algorithms [1, 2, 9, 10, 35, 40, 41]. The learned model alternates between taking a gradient step computed directly from existing forward models and applying a learned regularizer step. However, such approaches have to apply the forward model multiple times for a 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.
An alternative approach to incorporating physics is termed "Physics Informed Neural Networks (PINN)" [16, 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 a single instance of a 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 a distance comparison (usually through a pre-defined kernel function) with every reference point provided. However, it is designed to only take in 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 a 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  $G_{\theta}$  is trained to emulate the forward model [4, 11]. This is done in a supervised learning fashion, often with a 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  $G_{\theta}$  is substituted for  $F$  in Equation (2) or other similar optimization frameworks. The auto-grad libraries are then used to efficiently compute a gradient with respect to the input  $\pmb{x}$ , making it possible to iteratively solve for the best estimate  $\hat{\pmb{x}}$  that fits some observations  $\pmb{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 [5, 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  $\ell_1$ ,  $\ell_2$ , 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 [15].
Our proposed architecture extends the neural adjoint method in two notable ways. Existing methods are purely data-driven and thus treated like a black box [4, 33], while as our architecture explicitly incorporates knowledge of the physics-based forward model into the learned emulator in the form of linearized approximations. Tangent-linear models are also explored in [11], 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.
(c)
Figure 2: An overview of data from a 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.
The neural adjoint (NA) method is typically decomposed into two stages: training an emulator of the forward model and then using it for inference. We motivate and describe the proposed architecture for the forward model in Section 3.1. Next, we formulate the optimization problem that uses the trained model for inference in Section 3.2. Finally, we propose an augmentation to the optimization formulation to incorporate a learned subspace in Section 3.3.
3.1 Embedding Physics 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 we are given a forward model  $F: \mathbf{x} \rightarrow \mathbf{y}$  that maps our input  $\mathbf{x} \in \mathbb{R}^m$ , where  $m$  denotes the size of the discretization, to some observations  $\mathbf{y} \in \mathbb{R}^n$ , where  $n$  denotes the total number of observations. In our computational ocean tomography examples in Section 4,  $F(\mathbf{x})$  is computed using a ray tracer; in other applications, a PDE might have to be solved. We then train a neural net  $G_{\theta}$ , whose architecture is illustrated in Figure 1 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 a reference point  $\boldsymbol{x}_{\mathrm{ref}}^i$ , we perform a first order Taylor series expansion
$$
\begin{array}{l} \hat {\boldsymbol {y}} \approx F \left(\boldsymbol {x} _ {\text {r e f}} ^ {i}\right) + J _ {F} \left(\boldsymbol {x} _ {\text {r e f}} ^ {i}\right) ^ {\top} \left(\boldsymbol {x} - \boldsymbol {x} _ {\text {r e f}} ^ {i}\right) \tag {6} \\ = \boldsymbol {y} _ {\mathrm {r e f}} ^ {i} + \boldsymbol {A} _ {\mathrm {r e f}} ^ {i} (\boldsymbol {x} - \boldsymbol {x} _ {\mathrm {r e f}} ^ {i}). \\ \end{array}
$$
This linearization approximates the forward model with varying levels of accuracy depending on the proximity of the input  $\pmb{x}$  to the reference point  $\pmb{x}_{\mathrm{ref}}^i$ . Rather than learning the mapping from  $\pmb{x}$  to  $\pmb{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 an ensemble of reference models rather than a single linearization. Only a small subset of linearizations need to be performed once in order to construct the proposed emulator  $G_{\theta}$ , so the additional computational cost is minimal relative to attempting to invert the actual physics based forward model  $F$  for arbitrary measurements  $\pmb{y}$ .
The operation of the architecture acts as follows: Given some input  $\pmb{x}$ , we first pass  $\pmb{x}$  through the embedded physics ensemble  $A_{\mathrm{ref}}^{1}, \dots, A_{\mathrm{ref}}^{N}$  to produce  $N$  predicted observations  $\hat{\pmb{y}}^{1}, \dots, \hat{\pmb{y}}^{N}$  through the application of Equation (6). The predicted observations are then combined through a learned weighted average module to produce the final predicted observation.
A natural way to compute the weights  $w^{i}$  is by calculating the dot product similarity between the input point  $x$  and the reference points  $\pmb{x}_{\mathrm{ref}}^{1}, \dots, \pmb{x}_{\mathrm{ref}}^{N}$  used to generate the linearizations; higher similarity implies that the linearization is a better approximation and thus the prediction is more "trustworthy". We follow an approach similar to attention-based models [36] by learning an embedding space to
perform the dot product and applying a softmax to normalize the weights. Thus for each  $\hat{\pmb{y}}^i$ , we compute the corresponding weight  $w^i$  as
$$
w ^ {i} = \frac {\exp \left\langle \boldsymbol {x} , \boldsymbol {x} _ {\text {r e f}} ^ {i} \right\rangle_ {\boldsymbol {P} _ {x}}}{\sum_ {j} \exp \left\langle \boldsymbol {x} , \boldsymbol {x} _ {\text {r e f}} ^ {j} \right\rangle_ {\boldsymbol {P} _ {x}}}, \tag {7}
$$
where  $P_{x}$  is the learned projection for the dot product space. The predicted  $\hat{\pmb{y}}^i$  get summed together using the learned weights before being passed through a final linear layer  $W$ . Thus the output of the proposed model is
$$
G _ {\theta} (\boldsymbol {x}) = \boldsymbol {W} \sum_ {i} w ^ {i} \left(\boldsymbol {y} _ {\text {r e f}} ^ {i} + \boldsymbol {A} _ {\text {r e f}} ^ {i} \left(\boldsymbol {x} - \boldsymbol {x} _ {\text {r e f}} ^ {i}\right)\right), \tag {8}
$$
where we distinguish components related to the embedded physics module with superscripts. The encoder-decoder layers  $E_{x}$  and  $D_{x}$  are treated as the identity mapping for this section. They are described in further detail in Section 3.3. Note that the learned weights  $w^{i}$  depend on the input  $x$ . To prevent saturation of the learned softmax distribution, we apply spectral normalization to all  $x$  projection layers  $(E_{x}$  and  $P_{x}$ ). The full architecture is outlined in Figure 1a. The model is trained using the mean squared error loss on a dataset of paired example  $(x_{k}, y_{k})$
$$
\min  _ {\theta} \sum_ {k} \| G _ {\theta} (\boldsymbol {x} _ {k}) - \boldsymbol {y} _ {k} \| ^ {2}. \tag {9}
$$
3.2 Performing Inference
In order to perform inference on a set of observations  $\mathbf{y}$ , we solve the following optimization problem that incorporates our trained network
$$
\hat {\boldsymbol {x}} = \underset {\boldsymbol {x}} {\arg \min } \frac {1}{2} \| G _ {\theta} (\boldsymbol {x}) - \boldsymbol {y} \| ^ {2} + R (\boldsymbol {x}). \tag {10}
$$
We solve this iteratively by fixing the weights of the network and computing the gradient with respect to its input  $\pmb{x}$ . Note that this same optimization problem can be set up with the original forward model  $F$ , but computing a gradient is 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 is critical for solving Equation 10. However, these black box models are only trained to match outputs, and thus performing gradient descent can lead to many undesirable 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-based descent directions which may help alleviate these issues (See Appendix for derivations).
Equation (10) can be solved with a variety of optimization algorithms to converge on some locally optimal  $x$ . 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 a regularizer  $R(\cdot)$  as an 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  $\ell_2$  as well as a Sobolev norm ( $\ell_2$  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 we initialize our guess with some estimate. When optimizing from scratch, a reasonable candidate would be the average from the training set. Alternatively, we can leverage an estimated  $\hat{x}$  from other inverse methods by 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 to the iterative nature of the algorithm, but the NA method offers certain trade-offs outlined above that might be more beneficial in practice than the single estimate provided by learned direct inverse methods.
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
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 at 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	—	—
LFM	0.608	0.625	0.602
MLP	0.391 ± 0.005	0.395 ± 0.005	0.393 ± 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) [15]. The PCA basis acts as a de-noiser, removing any artifacts due to the optimization and helps reduce the dimensionality, simplifying the optimization process.