SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rather than pre-computing a basis, we instead propose to jointly learn a linear projection and reconstruction layer along with the forward model. The model as described in Subsection 3.1 can be augmented with a linear encoder  $E_{x}$  and decoder  $D_{x}$  layer. The encoder layer projects our input  $x$  onto a learned subspace, producing the latent code  $z$ . This code is then passed through the decoder layer to reconstruct  $\hat{x}$ . An additional input reconstruction loss in the form of a mean squared error loss between  $x$  and  $\hat{x}$  is included during training time, forcing the model to not only learn to approximate the forward model but also to learn a linear subspace of the input. We leave learning more complex non-linear encoder-decoder schemes for future work.
During inference, we then optimize in this subspace. More concretely, we rearrange the proposed architecture so that the  $x$  Decoder layer becomes the first input layer as shown in Figure 1b. The optimization variable  $z$  is passed through this layer to produce an estimated  $\hat{x}$ , that is then passed through the rest of the model as described in Section 3.1. Our optimization framework thus becomes
$$
\hat {z} = \arg \min  _ {z} \frac {1}{2} \left\| G _ {\theta} \left(\boldsymbol {D} _ {\boldsymbol {x}} z\right) - \boldsymbol {y} \right\| ^ {2} + R \left(\boldsymbol {D} _ {\boldsymbol {x}} z\right), \tag {11}
$$
where we recover our final estimate with  $\hat{\pmb{x}} = D_{\pmb{x}}\hat{\pmb{z}}$
4 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, 14, 24]. Typically, the ocean sound speed is estimated using empirical formulas based on the temperature, salinity and density. However, this would require a 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 spatio-temporal variability of the ocean sound speed profiles (SSPs), bypassing the need to densely sample.
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
(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 to 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 are 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 a linearized model) [3, 26, 27, 34].
We perform our experiments on a high fidelity month long simulation of the Gulf of Mexico as seen in Figure 2 [19, 20]. We restrict our 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 a set amount of time to train models and then deploy on future samples). Note that this creates a slightly more difficult problem due to the temporally changing dynamics in the test/validation set not present in the training set. The testset is further divided based on 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 a 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 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.
We experiment with two source-receiver configurations, a "dense" set up with 20 sources and receivers approximately every  $50\mathrm{m}$  in depth and  $5\mathrm{km}$  apart in range as shown in Figure 3. We also construct the more challenging "sparse" set up by sub-sampling in depth by a factor of 2. To counter the
differing bathymetry between the 10 slices, we restrict ourselves to the upper  $1000\mathrm{m}$  portion of the ocean (where a majority of the SSP variation lies as shown in Figure 2c) and consider only direct or surface bounce arrival time paths. The 2D SSP is discretized into a range by depth grid of  $11\times 231$
Thus, the linearized forward models (LFM) are of dimension  $800 \times 2541$  (or  $200 \times 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  $\pmb{x}$  is solved for iteratively given some observations  $\pmb{y}$ . We perform the optimization in the learned subspace and with  $\ell_2$  as well as Sobolev regularization. All optimized models are given a budget of 1000 gradient descent iterations with a learning rate of 50.
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 an 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 1. When provided no initialization (avg), the proposed method performs the best at  $0.364 \pm 0.018 \, \text{m/s}$  RMSE. MLP achieves the second best average performance at  $0.391 \pm 0.005 \, \text{m/s}$  RMSE, where the loss in performance is likely due to its inability to emulate the forward model as accurately as 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 \, \text{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 despite 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 in Figure 4. All models are able to 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-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\mathrm{m / s}$  RMSE. However, when provided with the Tik initialization (that performed better on the low variability scenario), it is able to achieve the best RMSE of  $0.339~\mathrm{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.
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 or the learned encoding/decoding of the SSP are all necessary components.
Table 2: RMSE (m/s) of inversion across different scenarios.
Low Variability	Med Variability	High Variability
Model	Avg	Tik	Avg	Tik	Avg	Tik
Sparse S/R Pairs (10 × 10)	Tik	0.724	—	0.967	—	0.926	—
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
Dense S/R Pairs (20 × 20)	Tik	0.647	—	0.773	—	0.881	—
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
WA-LFM	✓	X	X	0.577
WA-LFM + Dec	✓	X	✓	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 of linearizations (WA-LFM), (3) weighted average combined optimized in the learned subspace (WA-LFM + Dec), and (4) a learned weighted average network (WAN). A full summary as well as the average RMSE on a test set can be found in Table 3.
This study shows that each component is 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 1, suggesting that incorporating more references is 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 to better approximate the true model, leading to a 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
In this study, we propose a 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 to learn a subspace to 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 as well as 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.
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.
[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-2331, 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 a fast iterative shrinkage thresholding network for inverse problems in imaging. IEEE Transactions on Medical Imaging, 40(5):1329-1339, 2021.
[41] Jian Zhang and Bernard Ghanem. Ista-net: Iterative shrinkage-thresholding algorithm inspired deep network for image compressive sensing. 2017.
Table 4: Architecture details of final proposed PETAL model.
Layer	Input Dim	Output Dim	Spectral Norm?
x Encoder	2541	1000	✓
Px	1000	1000	✓
W	800	800	×
x Decoder	1000	2541	✓
A Training the Forward Model
All experiments were performed on a 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, it is able to learn a complex non-linear model due to the attention-inspired mechanism. The exact details of each linear sub-component can be found in Table 4. The embedded physics layer takes in an input of dimension  $BS \times N_x$ , where  $BS$  is the batch size and  $N_x = 2541$  the size of the vectorized SSP, and outputs a tensor of dimension  $BS \times N_{\mathrm{ref}} \times N_y$  where  $N_{\mathrm{ref}}$  refers to the number of linearized reference models and  $N_y = 800$  is the number of observations.
The weighted average layer inspired by attention layers takes in a query of dimension  $BS \times 1 \times H_x$ , where  $H_x$  is the hidden dimension for the latent space of input SSPs, a key of  $BS \times N_{\mathrm{ref}} \times H_x$  and a value of  $BS \times N_{\mathrm{ref}} \times N_y$ . It outputs a predicted arrival time vector  $\hat{y}$  of dimension  $BS \times N_y$ .
The model was trained using ADAMW with a learning rate of 1e-5 for 500 epochs. The learning rate was dropped by a factor of 0.2 at epoch 300.
The model was trained to minimize the MSE of the arrival time prediction as well as a MSE on the SSP reconstruction. The selected model achieved training RMSE of 4.88e-2 and validation RMSE of 1.02e-1. It achieve an unnormalized (removing training normalization) AT RMSE of 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 as 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 a training RMSE of  $3.62\mathrm{e - }2$  and validation RMSE of  $1.22\mathrm{e - }1$ . The unnormalized AT validation RMSE was  $6.08\mathrm{e - }4$  (higher than PETAL). The model was trained using Adam for 250 epochs with a learning rate of 1e-5.
B Optimization Framework
The neural adjoint method is an iterative method to recover an SSP  $\mathbf{x}$  given some observations  $\mathbf{y}$ . All models are optimized using Pytorch's Stochastic Gradient Descent with a learning rate of 50 for 1000 epochs.
We use two forms of regularization: an  $\ell_2$  penalty on  $x$  with a scale of 1e-7 and a Sobolev penalty  $(\ell_2$  on the discrete x and y gradient) with a scale of 1e-4.
The optimization is performed in batches. We set an early cutoff rate of 1e-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 is lower than the final (normalized) MSE AT loss of any of the models, so the assumption is that any further optimization beyond this point will just overfit to the model.
Figure 5: RMSE of different models vs number of epochs optimized given (a) average, (b) LFM, and (c) Tik initializations.
Figure 6: Boxplot of RMSE of different models.
B.1 Results
The results of NA given different initializations for each of the forward models can 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 6.
B.2 Robustness to Unseen Slices
In this section we explore the robustness of surrogate models to 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 7 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 in 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 all slices dropping from 0.33715 (when evaluated only on slice 10) to 0.33736 for our proposed model PETAL.
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	Val	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 a (simplification) of the PETAL model as
$$
\begin{array}{l} \hat {\boldsymbol {y}} = \boldsymbol {W} \left(\sum_ {i} w ^ {i} \hat {\boldsymbol {y}} ^ {i}\right) \\ = \sum_ {i} w ^ {i} \boldsymbol {W} \left(\boldsymbol {A} _ {\text {r e f}} ^ {i} \boldsymbol {x} + \boldsymbol {b} ^ {i}\right). \tag {12} \\ \end{array}
$$
Note that by construction, the weights  $w^{i}$  sum up to 1. If we include this in a simple MSE loss we get
$$
L = \frac {1}{2} \| \hat {\boldsymbol {y}} - \boldsymbol {y} \| ^ {2}. \tag {13}
$$
Computing a gradient w.r.t.  $x$  gives
$$
\frac {\partial L}{\partial \boldsymbol {x}} = \sum_ {i} \sum_ {j} \frac {\partial w ^ {i}}{\partial \boldsymbol {x}} w ^ {j} \boldsymbol {W} \left(\boldsymbol {A} ^ {i} \boldsymbol {x} + \boldsymbol {b} ^ {i}\right) \left(\boldsymbol {W} \boldsymbol {A} ^ {j} \boldsymbol {x} + \boldsymbol {W} \boldsymbol {b} ^ {j} - \boldsymbol {y}\right) + w _ {i} w _ {j} \boldsymbol {A} ^ {i \top} \boldsymbol {W} ^ {\top} \left(\boldsymbol {W} \boldsymbol {A} ^ {j} \boldsymbol {x} + \boldsymbol {W} \boldsymbol {b} ^ {j} - \boldsymbol {y}\right), \tag {14}
$$
where the right term reduces to a convex combination of the gradient of the linearized physics based forward models, modulated by some matrix  $\mathbf{W}$ , when  $i = j$ .
D Extensions to Interpolation: Additional Experiments
In this section, we begin a preliminary study on PETAL for interpolation (i.e. a 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 heuristics to select appropriate reference models. Instead, we perform a projected  $k$ -means approach to select our reference points. The standard
$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 to be embedded in the model. For our experiments, we used  $k = 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 a 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 a huge degradation in performance. We hypothesize that this is due to the nature of the data not being suitable for convolutions. The input SSPs are not translationally invariant depth-wise. The output arrival times are spatially small  $(20 \times 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 to 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 a total of 5 linear layers separated by leaky ReLU. 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 RLUs. The  $x$  decoder also uses 3 convolutional layers with kernel size  $3 \times 7$  and 32 filters. Finally, the output layer consists of a 3 layer MLP with hidden dimension 800.
The performance of each model is shown in Table 6. We experiment with gradient descent (SGD) and limited memory Broyden-Fletcher-Goldfarb-Shanno algorithm (LBFGS)[18] as the optimization algorithm. For both MLP and PETAL, LBFGS reached a local minima 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 a neural net.