# The Extrapolation Power of Implicit Models

Juliette Decugis1 Alicia Y. Tsai1,2
Max Emerling1 Ashwin Ganesh1 Laurent El Ghaoui1,2
1UC Berkeley 2VinUniversity
{jdecugis, aliciatsai, memerling, aganesh<EMAIL_ADDRESS>
{alicia.t<EMAIL_ADDRESS>

###### Abstract

In this paper, we investigate the extrapolation capabilities of implicit deep
learning models in handling unobserved data, where traditional deep neural
networks may falter. Implicit models, distinguished by their adaptability in
layer depth and incorporation of feedback within their computational graph,
are put to the test across various extrapolation scenarios: out-of-
distribution, geographical, and temporal shifts. Our experiments consistently
demonstrate significant performance advantage with implicit models. Unlike
their non-implicit counterparts, which often rely on meticulous architectural
design for each task, implicit models demonstrate the ability to learn complex
model structures without the need for task-specific design, highlighting their
robustness in handling unseen data.

##  1 Introduction

Learning to extrapolate, which involves estimating unknown values beyond
observed data, is a fundamental aspect of human intelligence and a significant
step toward achieving general machine intelligence. Despite the remarkable
success of modern neural networks across diverse domains, they often struggle
when faced with unobserved data outside their training distribution. In this
work, we investigate a general class of implicit deep learning models Bai et
al. (2019); Chen et al. (2018a, b); El Ghaoui et al. (2021) that encompasses
classical layered-wise neural networks as special cases. Implicit deep
learning models allow information to propagate both forwardly and backwardly
through closed-form feedback loops, offering a flexible and powerful approach
to learning data representations.

Various formulations have been proposed, including deep equilibrium models
(DEQs) Bai et al. (2019), Neural ODEs Chen et al. (2018b), and implicit models
El Ghaoui et al. (2021). These models don’t have explicitly defined layers and
instead employ state vectors defined via an “equilibrium” (fixed-point)
equation. In this framework, outputs are determined implicitly by the same
equation. Recent results have highlighted the utility of implicit models Bai
et al. (2020); Gu et al. (2020); Tsai et al. (2022). There has also been
emerging work where the equilibrium state is interpreted as a closed-loop
feedback system from a neuroscience perspective Ma et al. (2022). Moreover,
emerging work interprets the equilibrium state as a closed-loop feedback
system from a neuroscience perspective Ma et al. (2022). Unlike traditional
feed-forward models that traverse nodes in a single direction in the
computational graph, inputs of implicit models can revisit or reverse
directions between nodes through closed-loop feedback.

In this paper, we investigate whether implicit deep learning models exhibit
higher extrapolation capabilities, a fundamental skill in human intelligence,
compared to similarly sized non-implicit neural network models. Our
contributions are summarized in the following:

  * •

We demonstrate the extrapolation power of implicit deep learning models across
three distinct domains: well-defined mathematical arithmetic, real-world
earthquake location data, and volatile time series forecasting.

  * •

We conduct analyses and ablation studies on depth adaptability and closed-loop
feedback, revealing that features learned by implicit models are more
generalizable compared to their non-implicit counterparts.

###  1.1 Related Work.

##### Mathematical tasks.

Previous research has primarily focused on developing specialized neural
network models capable of learning algorithms Graves et al. (2014, 2016 10);
Kaiser and Sutskever (2015); Liang et al. (2021); Schwarzschild et al. (2021).
For instance, Neural Arithmetic Logic Units (NALUs) were specifically designed
to represent mathematical relationships within their architecture, with the
goal of improving arithmetic extrapolation Trask et al. (2018). However, these
models were later found to be highly unstable during training Schlör et al.
(2020). Studies by Nogueira et al. demonstrated that transformers performed
effectively for addition and subtraction tasks, achieving high accuracy in
interpolation experiments. However, challenges arose with other Transformer-
based architectures such as BART Wang et al. (2021) and large language models
(LLMs) Wei et al. (2022), which struggled to accurately reproduce functions
when dealing with inputs from a wide distribution range. On the contrary,
Charton showed that Transformers could provide “roughly correct” solutions for
matrix inversion and eigenvalue decomposition tasks, even on out-of-
distribution (OOD) inputs, indicating a notable level of mathematical
understanding.

##### Out-of-distribution generalization.

Only a handful of studies have begun to explore out-of-distribution (OOD)
generalization for implicit deep learning models Liang et al. (2021);
Rackauckas et al. (2021); Kelly et al. (2020). These works showcase the
capabilities of implicit deep learning models on tasks like Blurry MNIST,
sequential tasks Liang et al. (2021), matrix inversion, and graph regression
Anil et al. (2022). Researchers such as Liang et al. and Anil et al. highlight
a characteristic of DEQs known as path independence, where they converge to a
similar fixed point regardless of initialization. This property suggests that
if DEQs iterate for longer before converging on OOD inputs, they could gather
more information about the data and potentially outperform other models. Ramzi
et al. (2023) theorized that this property exists only when testing DEQs on
more complex data than in their training distribution. They demonstrate that
increasing the number of inner iterations can lead to overfitting for
interpolation tasks. We contribute further evidence supporting both hypotheses
by examining well-defined functions and real-world data with out-of-
distribution characteristics.

##### Function extrapolation.

Xu et al. studied how ReLU multi-layer perception (MLPs) and graph neural
networks extrapolate on quadratic, cosine, and linear functions. They
identified specific architectural choices that enhance extrapolation, such as
encoding task-specific non-linearities in model features. Similarly, in the
vision domain, Webb et al. introduced context normalization for more
generalized features. Additionally, Wu et al. demonstrated the advantages of
neural networks with Hadamard products (NNs-Hp) and polynomial networks (PNNs)
for arithmetic extrapolation. In this paper, we hypothesize that implicit
models can “adapt” to distribution changes, thereby eliminating the need for
specific feature transformations for extrapolation.

###  1.2 Background

Given a data input $u$, DEQs are defined as $f_{w}(u,x)$ where $f$ is a
function that depends on model parameters $w$, and $x$ represents the hidden
features. The objective is to find a fixed point $x^{*}$ by iteratively
applying $f$ to $x$ until convergence, such that $x^{*}=f_{w}(u,x^{*})$. In
this paper, we adopt the more generalized formulation presented in El Ghaoui
et al.. This formulation, proven to converge linearly, includes DEQs as a
special case. Formally, for a given data point $u\in{\mathbb{R}}^{p}$, an
implicit model solves the equilibrium equation $x=\phi(Ax+Bu)$, where $x$
represents the equilibrium state for input $u$, $\phi$ denotes a non-linear
activation such as ReLU, and matrices $A$ and $B$ are model parameters.
Prediction is obtained by passing the equilibrium state $x$ through an affine
transformation, yielding $\hat{y}(u)\in{\mathbb{R}}^{q}$, where matrices $C$
and $D$ are also model parameters:

$\hat{y}(u)=Cx+Du,\mbox{ where }x=\phi(Ax+Bu).$ (1)

The “state” vector $x\in{\mathbb{R}}^{n}$ contains the “hidden features” of
the model. These features are not expressible in closed form but are
implicitly defined via the equilibrium equation $x=\phi(Ax+Bu)$.

In this paper, we define convergence as
$\|x_{t}-x_{t-1}\|_{\infty}<\varepsilon$, where $x_{t}$ and $x_{t-1}$
represent the current and previous implicit hidden states, respectively. We
set $\varepsilon=3\times 10^{-6}$ . As demonstrated by El Ghaoui et al.,
implicit models will always converge under certain conditions, such as
projecting the $A$ matrix onto an infinite norm ball, which we ensure in our
experiments. We train implicit models using stochastic gradient descent and
back-propagation. During the forward pass, we aim to approximate the fixed
point $x^{*}$ given a set of weights $A$, $B$, $C$, and $D$, along with a
threshold $\varepsilon$. We denote $x_{N}$ as the approximation of the fixed
point. During the backward pass, implicit differentiation is employed to
compute the Jacobian of $x_{N}$ with respect to the model weights. We define
$\Phi:=\frac{\partial\phi(x_{t})}{\partial x_{t}}$ as a block diagonal matrix
of gradients with $t\in[0,N]$. As demonstrated by El Ghaoui et al. in Appendix
G, the gradient equations for each weight matrix are unique and rely on the
inversion of $(I-\Phi A)^{-1}$, which can be accelerated through approximate
inverse Jacobian methods Fung et al. (2022); Geng et al. (2021).

##  2 Problem Setup

We explore two types of implicit models in this paper: standard implicit
models as defined by Eq. (1), and a variant we term implicitRNN.

##### ImplicitRNN.

The implicitRNN functions similarly to a vanilla RNN, processing sequence
inputs one by one. For each time step $i$, where $s_{i}\in{\mathbb{R}}^{p}$
represents the $i$-th element in a sequence $(s_{1},s_{2},\cdots,s_{t})$. The
model input $u$ for implicitRNN is a concatenation of $s_{i}$ and the previous
hidden state $h_{i-1}$, akin to a vanilla RNN. The implicit prediction rule
for implicitRNN is as follows:

$\displaystyle h_{0}=\mathbf{0};\quad x_{0}=\mathbf{0};\quad
u_{i}=\begin{pmatrix}s_{i}\\\ h_{i-1}\end{pmatrix}$ $\displaystyle
x=\phi(Ax+Bu_{i})\quad\mbox{(equilibrium equation)}$
$\displaystyle\hat{y_{i}}(u_{i})=Cx+Du_{i}\quad\mbox{(prediction equation)}$
$\displaystyle h_{i}=\hat{y_{i}}(u_{i}).$

In this setup, the recurrent layer is replaced with an implicit structure
comprising the equilibrium and prediction equations. We introduce the
implicitRNN to contrast implicit with explicit sequential models, both
retaining representations of data timestep by timestep.

###  2.1 Extrapolate on Mathematical Tasks

We explore three types of functions, ranging from simple to complex: 1)
identity function, 2) arithmetic operations, and 3) rolling functions over
sequential data. To evaluate the extrapolation performance, we use two sets of
data: $u_{\text{train}}\sim P(\mu;\sigma)$ and $u_{\text{test}}\sim
P(\mu+\kappa;\sigma+\kappa)$, where $P$ represents a known distribution, $\mu$
and $\sigma$ are mean and the standard-deviation, and $\kappa$ denote the
hyper-parameters for distribution shift. Our code, models, dataset, and
experiment setup are available on GitHub111The link is redacted for anonymity.
It will be released in our camera-ready version.

##### Identify function.

Recent work has shown that neural networks struggle to learn the basic task of
identity mapping, $f(u)=u$, where models should return the exact input as
given He et al. (2016); Trask et al. (2018).

##### Arithmetic operations

We also focus on two arithmetic operations: addition and subtraction with
additional data transformation. Following the task proposed by Trask et al.,
we randomly select four numbers $i,j,k,l$ from 1 to 50, ensuring that $i<j$
and $k<l$. For each sample, we construct two new numbers $a$ and $b$ from the
input $\vec{u}:=\langle u_{1},u_{2},\cdots,u_{50}\rangle$ as follows:
$a=\sum_{a=i}^{j}u_{a}$, $b=\sum_{b=k}^{l}u_{b}$. Finally, we predict $y=a+b$
for addition and $y=a-b$ for subtraction.

##### Rolling functions.

Lastly, we explore two rolling functions over a sequence: average and argmax.
In the rolling average task, we predict the average of the sequence up to the
current time step $j$ for each timestep, calculated as
$\sum_{i=1}^{j}u_{i}/j$. For the rolling argmax task, we predict the index of
the maximum value seen by the model so far for each time step. We convert this
into a classification problem by outputting a one-hot vector of length $L=10$,
representing the index of the predicted maximum input seen so far, where $L$
is the length of the input sequence. Our evaluation focuses on predictions
made for the final element of the sequence.

We compare implicit models with neural networks specially designed to excel on
each task: MLPs for simple functions, LSTMs for sequential data Hochreiter and
Schmidhuber (1996), NALU for OOD arithmetic task Trask et al. (2018), and
various Transformers-based models Vaswani et al. (2017). The details of each
task and the model architectures are provided in the Appendix. We optimize all
models using grid search and 5-fold cross-validation on in-distribution
inputs. To ensure a fair comparison, we set the model dimension of the
implicit models to have a similar number of parameters as their non-implicit
counterparts.

###  2.2 Extrapolate on Noisy Real-world Data

In addition to mathematical extrapolation where data is generated from a known
underlying function, we consider extrapolation on real-world problems where
data is typically noisy and often lacks known data generation functions. We
focus on two real-world applications known for predicting unseen events:
oscillating time series forecasting and earthquake location prediction

##### Oscillating time series forecasting.

We first consider a simpler synthetic version known as spiky time series
forecasting, where spikes are inserted into a time series derived from a
combination of sine functions at random intervals (see Appendix A.3). In the
second case, we forecast AMC stock data volatility, particularly focusing on
the drastic increase in average volatility observed at the beginning of 2021
as illustrated in Figure 1. Volatility, in this context, refers to the
expected amount by which a security’s price might change suddenly, serving as
a measure of financial risk for an asset. Our prediction task involves
forecasting AMC’s volatility over the next 10 minutes based on its volume-
weighted average trade price (VWAP) for each of the past 60 minutes.
Volatility is calculated as the variance of the VWAP prices222Normally,
volatility is calculated as the standard deviation of returns. However, we
wish to amplify the changes in distribution between our train and test set, so
we use raw prices and the variance instead. during the forecast period. Our
training data goes from 2/1/2015 to 12/31/2020 and our validation set from
1/1/2021 to 12/31/2021. Our training data spans from 2/1/2015 to 12/31/2020,
with the validation set covering 1/1/2021 to 12/31/2021. We deliberately
refrain from making our data stationary through differencing or examining
returns, as we aim to evaluate our models’ ability to adapt to changes in
price distribution. We compare the performance of implicit models against
simple linear regression and non-implicit neural network baselines, with
details of the baseline models included in Table 4 in the Appendix.

Figure 1: Time series of a 21-day rolling average of AMC stock volatility
plotted on a log scale, highlights a drastic volatility increase at the
beginning of our validation cutoff.

##### Earthquake location prediction.

The earthquake location prediction is a well-established problem in seismology
Smith et al. (2021); Saad et al. (2021). It involves predicting the location
($X$, $Y$, and $Z$ coordinates) and the seismic wave travel time ($T$) of an
earthquake based on seismographs recorded from nearby seismometers.
Successfully solving this problem can have significant humanitarian
implications, allowing for early warnings before the potentially destructive
second wave(s) of an earthquake. However, the problem remains challenging due
to the sparsity of the event observations Chuang et al. (2023).

Figure 2: Left: Geometric visualization of one set of training features
($x_{i}$, $y_{i}$, $z_{i}$, $p_{i}$, $\theta_{i}$) and its corresponding
labels ($X$, $Y$, $Z$, $T$). The triangles correspond to stations and the star
corresponds to a source. Right: The map shows the training set region colored
in blue, roughly corresponding to the Pacific Ring of Fire. The two red areas
are the testing set regions for $k=3$.

We follow the methodology outlined by Chuang et al., generating seismograph
samples recorded by five stations with one anchor station designated as the
reference for all seismic wave arrival times. As illustrated in the left panel
of Figure 2, our input features include a station’s coordinates ($x$, $y$,
$z$), event-station back-azimuths ($\theta$), and relative wave travel times
($p$) with respect to the anchor station. Training data is synthetically
generated within a range of $90^{\circ}$E and $-90^{\circ}$E, roughly
corresponding to the Pacific Ring of Fire, as shown in Figure 2 (further
details are provided in Appendix A.4). Testing is conducted on regions shifted
from $10^{\circ}$E to $90^{\circ}$E beyond this Ring of Fire. As a comparison
baseline, we use EikoNet, a deep learning model introduced by Smith et al.
specifically designed for earthquake location prediction. While earthquakes
primarily occur in certain active tectonic boundaries, an extrapolated
earthquake location prediction system can aid in detecting earthquakes in new
areas, whether natural or human-made (e.g., caused by mining, oil, and gas
activities). Such a system is also valuable for explosion monitoring,
providing universal mapping capabilities.

##  3 Experimental Results

###  3.1 Mathematical Extrapolation

##### Identify function.

The test mean squared error (MSE) for the identify function task is shown in
Figure 3. The implicit model maintains the lowest test MSE (below 5) for test
data with a distribution shift from 0 to 25. Even for very large distribution
shifts of up to 200, where $u_{\text{test}}\in{\mathbb{R}}^{10}\sim
U(-205,205)$, the implicit model outperforms the Transformers encoder model by
a factor of $10^{5}$ and the MLP by a factor of $10^{3}$. This disparity
highlights the implicit model’s robustness in handling distribution shifts
compared to non-implicit models, which tend to overfit to the training
distribution and exhibit increased error with larger distribution shifts. The
identity function task exposes the spurious features learned by both the MLP
and Transformers models. Our implicit model reaches equilibrium after only 4
training iterations. In contrast, similar-sized MLPs and transformer encoders
struggle with overfitting, highlighting the implicit model’s ability to limit
overfitting through faster convergence when dealing with simpler functions
like identity mapping.

##### Arithmetic operations.

The results for the addition and subtraction arithmetic tasks are depicted in
Figure 4. The implicit model not only outperforms various Transformer encoders
but also NALU, a model specifically designed for mathematical operations. As
shown by Wei et al., Transformers models require much larger model sizes
($10^{23}$) to perform an arithmetic task effectively. In contrast, an
implicit model with only $7,650$ parameters successfully learns these
operations with a testing loss of $<1$ for shifts $<100$. Throughout the
experiment, we consistently observed that implicit models outperform non-
implicit models on extrapolation tasks, especially with limited training
samples. Surprisingly, the specialized NALU model achieves the worst testing
loss, exceeding $10^{10}$ for an extrapolation shift of only 10, as shown in
Figure 4. Across our experiments, we couldn’t produce robust out-of-
distribution predictions with the NALU model.

Figure 3: Test MSE for the identity function task. MSE for MLP and
Transformers model increases as the distribution shift hyper-parameter
$\kappa$ increases. Figure 4: Test Log(MSE) for the arithmetic operations. The
implicit model strongly outperforms all other models on OOD data.

##### Rolling functions.

In Figure 6, we present the performance of the two sequence modeling tasks. In
the rolling average task (left plot of Figure 6), the LSTM and Transformers
exhibit similar behavior while the implicit model maintains close to constant
loss regardless of the test distribution. For the rolling argmax, we use a
Transformer encoder-decoder architecture, where the target sequence is the
right-shifted argmax labels. This architecture has an inherent advantage, as
simply outputting the final element in this right-shifted sequence will
provide the argmax of the input sequence up to but not including the current
element. Given that argmax labels of a sequence are uniformly distributed,
this setup implies an expected accuracy of 90% or of $1-1/L$, where $L=10$ is
the sequence length. Surprisingly, we observe that the standard implicit model
and implicitRNN both outperform the LSTM and Transformer models. Additionally,
the implicitRNN outperforms the standard implicit model, indicating the
potential benefits of a rolling latent representation in tasks requiring
awareness of relative positions within a sequence.

Comparing different Transformers variants highlights the potential for small
Transformers to overfit to their positional encodings (PE) on simpler tasks.
The Transformer without positional encodings, lacking the ability to
differentiate between different positions in the sequence, outperforms the
other two architectures, consistently hovering around the 90% accuracy
benchmark. This performance suggests that it effectively learns to always
output the final value of its target sequence. Overall, these results
demonstrate a very competitive performance of the standard implicit model and
implicitRNN over Transformers and LSTMs in extrapolation tasks, particularly
in scenarios where an understanding of relative positions within a sequence is
crucial.

It is worth emphasizing that the implicitRNN structure resembles that of a
vanilla RNN, with the standard recurrent cell replaced by an implicit layer.
The substantial performance improvement over the LSTM suggests an exciting
potential of the implicit layer to learn effective memory-gating mechanisms.
However, further experimentation, specifically with longer sequence lengths,
would be required to verify this conjecture. For further insights and
experimental findings, please refer to Appendix A.2.

Figure 5: For rolling tasks, implicit models maintain close to constant loss
($\downarrow$) and accuracy ($\uparrow$) across shifts.

Figure 6: The implicitRNN most accurately models spike magnitudes.

###  3.2 Oscillating Time Series Forecasting

In our prior experiments, we tested implicit models on clearly defined
mathematical operations. Now, we shift focus to real-world noisy data, where
the underlying function is unknown. This transition allows us to explore
whether the extrapolation benefits of implicit models extend to such complex,
real-world scenarios.

In the spiky time series forecasting task, our objective is for models to
effectively capture sudden and short-lived distribution changes in our
generated data. Such extreme events, resembling those seen in stock prices,
sales volumes, or storage capacity predictions, often challenge forecasting
models that struggle to extrapolate to these unprecedented patterns. The
ability to accurately adapt to synthetic spikes in this task carries
significant real-life implications. Table 1 demonstrates that implicitRNN
achieves a threefold reduction in test MSE compared to non-implicit models in
the spiky time series task. Furthermore, Figure 6 showcases the implicit
model’s capacity to accurately capture the locations and magnitudes of these
spikes. In contrast, the transformer decoder and LSTM models often output the
average of the time series. This tendency of standard models to revert to the
mean also hints at their potential weakness in simulating time-series data
effectively. Encouraged by these positive results, we applied implicit models
to predict the sharp rise in AMC stock volatility observed in 2021. We report
the mean absolute percentage error (MAPE)333The Mean Absolute Percentage Error
(MAPE) metric assesses a model’s capability to predict changes in magnitudes
relative to their sizes. Consequently, it is not directly comparable across
distributions. Despite the higher Root Mean Square Error (RMSE) on the
validation data (for the implicit model: train RMSE = 0.001784, validation
RMSE = 0.266366), the validation MAPE is lower since the average volatility in
the validation set is orders of magnitude larger. in Table 1, where the
implicit model outperforms other baselines by a factor of 1.67.

| Spiky Data | AMC Stock
---|---|---
| (MSE) | (MAPE)
Model |  Train |  Test |  Train |  Test
Transformer |  0.061 |  0.012 |  12.2 |  6.51
ImplicitRNN |  0.015 |  0.004 |  2.61 |  1.71
LSTM |  0.011 |  0.011 |  10.5 |  5.46
MLP |  - |  - |  5.51 |  2.87
Linear regression |  - |  - |  7.19 |  3.94

Table 1:  Train and test metrics ($\downarrow$) in forecasting time series
with sudden changes for synthetic (spiky time series) and real-world data (AMC
stock volatility). Our architectures vary between tasks (see Table 4) but the
implicit model outperforms all fine-tuned models.

####  3.2.1 Earthquake Location Prediction

In the earthquake location prediction task, our implicit model demonstrates an
improvement of $0.25\mathrm{e}{-3}$ in the in-distribution test loss compared
to EikoNet, a model specifically designed for seismic data. Notably, in the
left panel of Figure 7, as $k$ increases, we observe a progressive improvement
in the performance of the implicit model in terms of the extrapolated test
MSE. By the time $k$ reaches 2, the implicit model has surpassed EikoNet, and
at $k$ = 9, the implicit model’s test loss is $1.59\mathrm{e}{-2}$ better than
that of EikoNet. This improvement translates to an average enhancement of
$11^{\circ}$ in longitude and $2^{\circ}$ in latitude. Further refinement
could be achieved by limiting the source latitude along with longitude during
training, potentially leading to a more substantial improvement in latitude
extrapolation. However, as depicted in the right panel of Figure 7, the
implicit model encounters greater difficulty in predicting time and depth,
exhibiting deteriorating performance with increasing $k$. Specifically, at $k$
= 9, the implicit model’s average performance worsens by 9.2 seconds and 409
km. Future research should explore whether training an implicit model
exclusively on time and depth labels can enhance its performance in these
aspects. This two-pass approach resembles traditional location prediction
software such as HYPOINVERSE (USGS), where constraints on depth and time are
notably challenging.

Figure 7: Extrapolation comparison between EikoNet and implicit model on the
location prediction task as the extrapolation factor increases. Left: The
implicit model has the general edge in terms of MSE loss. Right: Breaking down
the prediction values, we observe that the implicit model only outperforms
EikoNet in longitude and latitude predictions.

##  4 Analysis

In our analysis, we pinpoint two key properties that contribute to the
effective extrapolation capabilities of implicit models, even with limited
datasets. The first property is ”depth adaptability,” meaning these models are
not constrained by a fixed number of layers. The second property, “closed-loop
feedback”, allows inputs to revisit the same node within a single pass through
the model, enhancing its learning and adaptability.

##### Depth adaptability.

During a forward pass through an implicit model, the process concludes either
upon convergence to a fixed point $x^{*}$ or upon reaching a predetermined
iteration limit. We investigated whether the complexity of the input affects
the number of iterations required for the model to achieve equilibrium. Across
our experiments, we noticed that as our model familiarized itself with
parameter matrices, the iteration count stabilized for inputs within the
distribution. On average, implicit models converged in about 15 iterations for
learning the addition operation, 30 iterations for subtraction, and 175
iterations for the rolling argmax task. We can interpret the number of
iterations required for a given task as an indication of the model’s perceived
difficulty with that task. Consistent with the findings of Liang et al., we
observed that as our input deviated further from the training distribution,
the number of iterations needed to converge increased. This suggests that the
input $U$ underwent more transformations by the model parameters (as shown in
the bottom two panels of Figure 8). Therefore, we interpret that implicit
models dynamically “grow” in depth to adapt their feature space to
extrapolated inputs, iterating longer compared to in-distribution inputs. This
adaptability allows them to benefit from low depth for in-distribution inputs
(reducing overfitting) and higher depth for out-of-distribution inputs
(reducing underfitting).

Figure 8: Left: Relationship between model performance during training and the
number of iterations. Middle and right: Growth of implicit models as the input
complexity increases in the arithmetic tasks and rolling argmax. Figure 9: A
matrix and node diagram representation of one iteration through the implicit
model: $Ax+Bu$. In red, we have the weights that induce feedback in $A$, and
in yellow, the non-feedback weights.

##### Closed-loop feedback.

Introduced by Wiener, refers to a system’s ability to self-correct based on
its outputs. This concept, elaborated by Patten and Odum, involves returning a
portion of the output, which partly influences the input at the next step. In
neural networks, feedback is encoded through the backward pass. Implicit
models, as depicted in Figure 9, inherently incorporate feedback within a
single iteration via the feedback connection in the parameter matrix $A$
(lower-triangular part). We observe that inputs in implicit models do not only
move forward from one set of nodes to the next, as in a standard feed-forward
neural network. Instead, they may revisit the same nodes or even move
backward, allowing the model to correct itself from one layer to the next
within a single iteration. This is in contrast to non-implicit models, which
correct themselves after a pass through all layers. Ma et al. highlighted
closed-loop feedback as a strength of implicit models and noted its similarity
to human neural networks. Therefore, we evaluated the utility of closed-loop
feedback in extrapolation through ablation studies on the lower triangular
part of the model’s $A$ matrix. We illustrate the absence of feedback in an
implicit model by considering a $3\times 3$ example with a strictly upper
triangular $A$ matrix. At time $t+1$, the iteration equation simplifies to
$x^{t+1}=\phi(Ax^{t}+Bu)$. Here, we can disregard the activation $\phi$ and
$Bu$ terms, focusing solely on the $x^{t}$ term as it carries the outputs from
the previous iteration. This setup allows us to demonstrate how an implicit
model without feedback operates, contrasting it with the typical behavior of a
truly implicit models.

$\displaystyle x^{t+1}\approx Ax^{t}$
$\displaystyle=\begin{pmatrix}x^{t+1}_{0}\\\ x^{t+1}_{1}\\\ x^{t+1}_{2}\\\
\end{pmatrix}=\begin{pmatrix}\star&W_{0}&W_{1}\\\ \star&\star&W_{2}\\\
\star&\star&\star\\\ \end{pmatrix}\begin{pmatrix}x^{t}_{0}\\\ x^{t}_{1}\\\
x^{t}_{2}\\\ \end{pmatrix}$ (2) $\displaystyle=\begin{pmatrix}\star\cdot
x^{t}_{0}+W_{0}x^{t}_{1}+W_{1}x^{t}_{2}\\\ \star\cdot x^{t}_{0}+\star\cdot
x^{t}_{1}+W_{2}x^{t}_{2}\\\ \star\cdot x^{t}_{0}+\star\cdot
x^{t}_{1}+\star\cdot x^{t}_{2}\\\ \end{pmatrix}.$

Closed-loop feedback corresponds to $x_{i}^{t}$ being used to generate
$x_{i}^{t+1}$. Equation 2 illustrates how the $\star$ weights encode this
feedback mechanism. Specifically, the $i$-th state at time step $t+1$, denoted
as $x_{i}^{t+1}$, depends not only on its past output $x_{i}^{t}$ directly
through weights on the diagonal but also indirectly through weights on the
lower triangular matrix. As an example of this indirect dependence, consider
$x_{2}^{t+1}$, which depends not only on $x_{2}^{t}$ but also on $x_{1}^{t}$
due to the lower triangular $\star$ weights. Furthermore, both $x_{1}^{t}$ and
$x_{2}^{t}$ depend on $x_{0}^{t-1}$, highlighting the chain of dependencies
facilitated by the feedback mechanism.

We compared a standard implicit model (with feedback loops) against an ablated
implicit model (without feedback loops) where the upper-triangularity of $A$
was enforced during training. The standard implicit models correspond to those
used in our experiments, as described in Section 2. The results of the
mathematical tasks for implicit models with and without feedback are presented
in Table 2. It is evident that the feedback loops play a crucial role in
helping the models achieve significantly lower testing loss, especially on
inputs with distribution shifts. Figure 10 presents the ablation results for
both models across distribution shifts in arithmetic and rolling sequential
tasks. Ablating the feedback harms model performance for subtraction but not
for addition. The regular model, with its feedback loops, has twice as many
weights, which could potentially lead to overfitting on simpler tasks.
Feedback loops appear to increase stability to distribution shifts for rolling
average tasks, whereas they only provide better performance in calculating the
rolling argmax of a sequence. We also note that besides superior overall
performance, the presence of closed-loop feedback seems to make the model more
resistant to distribution shifts: notably, we observe that its loss increases
slower in the subtraction and rolling average experiments. Our analysis
demonstrates that implicit models possess the capability to adapt their
architecture by learning the necessary depth, node connections, and correcting
themselves based on past predictions within a single training iteration,
thanks to their closed-loop feedback mechanism.

Task |  Feedback |  No feedback
---|---|---
Add. |  3.06 |  4.07
Sub. |  0.953 |  1.80
R. Max. |  93.7% |  91.8%
R. Avg. |  $110-4$ |  $110-3$

Table 2:  Testing metrics (MSE $\downarrow$ for the arithmetic operations and
rolling average, accuracy $\uparrow$ for rolling argmax) for a distribution
shift of 100 comparing an implicit model trained with and without closed-loop
feedback. Figure 10: Implicit models benefit from closed-loop feedback
specifically in harder extrapolation tasks.

Implicit models represent a significant advancement over standard feed-forward
neural network layers by relying on the convergence of an equilibrium equation
to extract features from input data. In this paper, we have demonstrated their
remarkable ability to learn useful function representations for extrapolation
tasks, even in scenarios with limited training data. Our experiments
consistently show that implicit models outperform non-implicit baselines when
handling out-of-distribution inputs, temporal and geographical shifts. The
adaptive nature of training implicit models allows them to explore and
identify optimal architectures with minimal hand-engineering, resulting in a
robust inner function representation for extrapolation tasks involving out-of-
distribution data. These results highlight the potential of implicit models
and motivate further research into their use as more interpretable solutions
for wider scenarios with limited data set.

## References

  * Anil et al. [2022] Cem Anil, Ashwini Pokle, Kaiqu Liang, Johannes Treutlein, Yuhuai Wu, Shaojie Bai, Zico Kolter, and Roger Grosse.  Path independent equilibrium models can better exploit test-time computation.  In Advances in Neural Information Processing Systems (NeurIPS 2022), 2022.
  * Bai et al. [2019] Shaojie Bai, J. Zico Kolter, and Vladlen Koltun.  Deep equilibrium models.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * Bai et al. [2020] Shaojie Bai, Vladlen Koltun, and J. Zico Kolter.  Multiscale deep equilibrium models.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 5238–5250. Curran Associates, Inc., 2020.
  * Charton [2022] François Charton.  What is my math transformer doing? – three results on interpretability and generalization, 2022.
  * Chen et al. [2018a] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * Chen et al. [2018b] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  Advances in neural information processing systems, 31, 2018.
  * Chuang et al. [2023] Lindsay Y. Chuang, Jesse Williams, Louisa Barama, Zhigang Peng, and Andrew V. Newman.  Deep learning and masksembles for sparse network earthquake location and uncertainty estimation.  Manuscript in preparation., 2023.
  * El Ghaoui et al. [2021] Laurent El Ghaoui, Fangda Gu, Bertrand Travacca, Armin Askari, and Alicia Tsai.  Implicit deep learning.  SIAM Journal on Mathematics of Data Science, 3(3):930–958, 2021.
  * Fung et al. [2022] Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, and Wotao Yin.  Jfb: Jacobian-free backpropagation for implicit networks.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6648–6656, 2022.
  * Geng et al. [2021] Zhengyang Geng, Xin-Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin.  On training implicit models.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 24247–24260. Curran Associates, Inc., 2021.
  * Graves et al. [2014] Alex Graves, Greg Wayne, and Ivo Danihelka.  Neural turing machines.  CoRR, abs/1410.5401, 2014.
  * Graves et al. [2016 10] Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, Adrià Puigdomènech Badia, Karl Moritz Hermann, Yori Zwols, Georg Ostrovski, Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukcuoglu, and Demis Hassabis.  Hybrid computing using a neural network with dynamic external memory.  Nature., 538(7626), 2016-10.
  * Gu et al. [2020] Fangda Gu, Heng Chang, Wenwu Zhu, Somayeh Sojoudi, and Laurent El Ghaoui.  Implicit graph neural networks.  In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS’20, Red Hook, NY, USA, 2020. Curran Associates Inc.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Identity mappings in deep residual networks.  In European conference on computer vision, pages 630–645. Springer, 2016.
  * Hochreiter and Schmidhuber [1996] Sepp Hochreiter and Jürgen Schmidhuber.  Lstm can solve hard long time lag problems.  In M.C. Mozer, M. Jordan, and T. Petsche, editors, Advances in Neural Information Processing Systems, volume 9. MIT Press, 1996.
  * Kaiser and Sutskever [2015] Łukasz Kaiser and Ilya Sutskever.  Neural gpus learn algorithms, 2015.
  * Kelly et al. [2020] Jacob Kelly, Jesse Bettencourt, Matthew James Johnson, and David Duvenaud.  Learning differential equations that are easy to solve, 2020.
  * Kennett et al. [1995] B. L. N. Kennett, E. R. Engdahl, and R. Buland.  Constraints on seismic velocities in the Earth from traveltimes.  Geophysical Journal International, 122(1):108–124, 07 1995.
  * Liang et al. [2021] Kaiqu Liang, Cem Anil, Yuhuai Wu, and Roger Grosse.  Out-of-distribution generalization with deep equilibrium models.  ICML 2021 Workshop on Uncertainty and Robustness in Deep Learning, 2021.
  * Ma et al. [2022] Yi Ma, Doris Tsao, and Heung-Yeung Shum.  On the principles of parsimony and self-consistency for the emergence of intelligence.  Frontiers of Information Technology & Electronic Engineering, pages 1–26, 2022.
  * Nogueira et al. [2021] Rodrigo Nogueira, Zhiying Jiang, and Jimmy Lin.  Investigating the limitations of the transformers with simple arithmetic tasks.  CoRR, abs/2102.13019, 2021.
  * Patten and Odum [1981] B. Patten and E. Odum.  The cybernetic nature of ecosystems.  118(6), December 1981.
  * Rackauckas et al. [2021] Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Dominic Skinner, Ali Ramadhan, and Alan Edelman.  Universal differential equations for scientific machine learning, 2021.
  * Ramzi et al. [2023] Zaccharie Ramzi, Pierre Ablin, Gabriel Peyré, and Thomas Moreau.  Test like you train in implicit deep learning.  2023\.
  * Saad et al. [2021] Omar M. Saad, Ali G. Hafez, and M. Sami Soliman.  Deep learning approach for earthquake parameters classification in earthquake early warning system.  IEEE Geoscience and Remote Sensing Letters, 18(7):1293–1297, 2021.
  * Schlör et al. [2020] Daniel Schlör, Markus Ring, and Andreas Hotho.  inalu: Improved neural arithmetic logic unit.  Frontiers in Artificial Intelligence, 3, 2020.
  * Schwarzschild et al. [2021] Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein.  Can you learn an algorithm? generalizing from easy to hard problems with recurrent networks.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 6695–6706. Curran Associates, Inc., 2021.
  * Smith et al. [2021] Jonathan D. Smith, Kamyar Azizzadenesheli, and Zachary E. Ross.  EikoNet: Solving the eikonal equation with deep neural networks.  IEEE Transactions on Geoscience and Remote Sensing, 59(12):10685–10696, dec 2021.
  * Trask et al. [2018] Andrew Trask, Felix Hill, Scott E Reed, Jack Rae, Chris Dyer, and Phil Blunsom.  Neural arithmetic logic units.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * Tsai et al. [2022] Alicia Y Tsai, Juliette Decugis, Laurent El Ghaoui, and Alper Atamtürk.  State-driven implicit modeling for sparsity and robustness in neural networks.  arXiv preprint arXiv:2209.09389, 2022.
  * [31] USGS.  Hypoinverse earthquake location.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30, pages 5998–6008. Curran Associates, Inc., 2017.
  * Wang et al. [2021] Cunxiang Wang, Boyuan Zheng, Yuchen Niu, and Yue Zhang.  Exploring generalization ability of pretrained language models on arithmetic and logical reasoning.  In Lu Wang, Yansong Feng, Yu Hong, and Ruifang He, editors, Natural Language Processing and Chinese Computing, pages 758–769, Cham, 2021. Springer International Publishing.
  * Webb et al. [2020] Taylor W. Webb, Zachary Dulberg, Steven M. Frankland, Alexander A. Petrov, Randall C. O’Reilly, and Jonathan D. Cohen.  Learning representations that support extrapolation.  CoRR, abs/2007.05059, 2020.
  * Wei et al. [2022] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus.  Emergent abilities of large language models.  Transactions on Machine Learning Research, 2022.  Survey Certification.
  * Wiener [1948] N. Wiener.  Cybernetics.  MIT Press, Cambridge, Mass, 1948.
  * Wu et al. [2022] Yongtao Wu, Zhenyu Zhu, Fanghui Liu, Grigorios Chrysos, and Volkan Cevher.  Extrapolation and spectral bias of neural nets with hadamard product: a polynomial net study.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  * Xu et al. [2021] Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S. Du, Ken ichi Kawarabayashi, and Stefanie Jegelka.  How neural networks extrapolate: From feedforward to graph neural networks.  2021\.

##  Appendix A Appendix

###  A.1 Model Specificites

We used ReLU activation function for our implicit models, transformers have a
dropout of 0.1 and a layer norm epsilon of $\epsilon=10^{-5}$. See Tables 3
and 4 on the next page for model architecture and experiment specificities.

Task |  Training Distribution |  Testing Distribution
---|---|---
Identity Function |  $u_{\text{train}}\in{\mathbb{R}}^{10,000\times 10}\sim U(-5,5)$ |  $u_{\text{test}}\in{\mathbb{R}}^{3,000\times 10}\sim U(-\kappa,\kappa)$, where $\kappa$ ranges from 10 to 80
Arithmetic Operations |  $u_{\text{train}}\in{\mathbb{R}}^{10,000\times 50}\sim U(-1,1)$ |  $u_{\text{test}}\in{\mathbb{R}}^{3,000\times 50}\sim U(-\kappa/2,\kappa/2)$, $\kappa$ ranges from $10$ to $10^{5}$
Rolling Average |  $u_{\text{train}}\in{\mathbb{R}}^{10,000\times 10}\sim\mathcal{N}(3,1)$ |  $u_{\text{test}}\in{\mathbb{R}}^{3,000\times 10}\sim\mathcal{N}(3+\kappa,1)$, $\kappa$ ranges from 5 to 100
Rolling Argmax |  $u_{\text{train}}\in{\mathbb{R}}^{10,000\times 10}\sim U(0,1)$ |  $u_{\text{test}}\in{\mathbb{R}}^{3,000\times 10}\sim U(0,\kappa)$, $\kappa$ ranges from $10^{1}$ to $10^{5}$
Earthquake Location |  720,576 ($X$, $Y$, $Z$) locations sampled between (90, -90)∘E, 30 features |  20,016 samples in each extrapolation region $(90-10\kappa,100-10\kappa)\cup(-100+10\kappa,-90+10k)$, $\kappa$ ranges from 1 to 9
Table 3:  Details of training and out-of-distribution test set for each extrapolation task.  Task |  Baseline model |  Implicit models |  Transformers
---|---|---|---
Identity Function |  MLP: $10\times 9\times 9\times 10$ |  Regular: $A\in{\mathbb{R}}^{4\times 4},B\in{\mathbb{R}}^{4\times 10},C\in{\mathbb{R}}^{10\times 4},D\in{\mathbb{R}}^{10\times 10}$ |  Encoder-decoder: $10\times 10\times 5$, 5 attention heads
Arithmetic Operations |  • MLP: $50\times 10\times 10\times 1$ |  Regular: $A\in{\mathbb{R}}^{20\times 20},B\in{\mathbb{R}}^{20\times 50},C\in{\mathbb{R}}^{1\times 20},D\in{\mathbb{R}}^{1\times 50}$ |  • Sequential encoder: 1 layer, 10 attention heads, feedforward dim 50 - processes each array as a single sequence
|  • NALU: $50\times 10\times 10\times 1$ |  |  • Depth-wise encoder: 1 layer, 1 attention head, feedforward dim 500, max PE length 50 - processes each element in a given array as a single sequence
Rolling Average |  LSTM: 1$\times$18$\times$18$\times$1 |  Regular: $A\in{\mathbb{R}}^{32\times 32},B\in{\mathbb{R}}^{32\times 10},C\in{\mathbb{R}}^{10\times 32},D\in{\mathbb{R}}^{10\times 10}$ |  Encoder-decoder: $10\times 10\times 5\times 10$, $5$ attention heads
Rolling Argmax |  LSTM: 1$\times$21$\times$21$\times$10 |  • Regular: $A\in{\mathbb{R}}^{36\times 36},B\in{\mathbb{R}}^{36\times 10},C\in{\mathbb{R}}^{10\times 36},D\in{\mathbb{R}}^{10\times 10}$ |  • Masked encoder-decoder: 1 encoder layer, 1 decoder layer, 2 attention heads, feedforward dim 10, max PE length 10
|  |  • RNN: $A\in{\mathbb{R}}^{21\times 21},B\in{\mathbb{R}}^{21\times 23},C\in{\mathbb{R}}^{22\times 21},D\in{\mathbb{R}}^{22\times 23}$ |  • Unmasked encoder-decoder: 1 encoder layer, 1 decoder layer, 2 attention heads, feedforward dim 10, max PE length 10
|  |  |  • Unmasked encoder-decoder without PE: 1 encoder layer, 1 decoder layer, 2 attention heads, feedforward dim 10
Spiky Time Series |  LSTM: 1$\times$20$\times$20$\times$1 |  RNN: $A\in{\mathbb{R}}^{20\times 20},B\in{\mathbb{R}}^{20\times 21},C\in{\mathbb{R}}^{20\times 20},D\in{\mathbb{R}}^{20\times 21}$ with a 20$\times$1 linear layer |  1x10 linear layer (expansion) $\xrightarrow[]{}$ masked decoder (1 layer, 2 attention heads, feedforward dim 2048, max PE length 10) $\xrightarrow[]{}$ 10x1 linear layer (contraction)
Volatility Prediction |  • LSTM: $1\times 38\times 38\times 1$ |  • Regular: $A\in{\mathbb{R}}^{53\times 53},B\in{\mathbb{R}}^{53\times 60},C\in{\mathbb{R}}^{1\times 53},D\in{\mathbb{R}}^{1\times 60}$ |  Sequential encoder (1 layer, 1 attention head, feedforward dim 2048, max PE length 60) $\xrightarrow[]{}$ 60x1 linear layer
|  • SGD Linear Regression |  • RNN: $A\in{\mathbb{R}}^{37\times 37},B\in{\mathbb{R}}^{37\times 41},C\in{\mathbb{R}}^{40\times 37},D\in{\mathbb{R}}^{40\times 41}$ with a 40$\times$1 linear layer |
|  • MLP: $60\times 50\times 27\times 27\times 27\times 10\times 1$ |  |
Earthquake Location Prediction |  EikoNet: $270\times 32\times 128\times 128\times 128\times 32\times 4$ (42,500) |  Regular: $A\in{\mathbb{R}}^{190\times 190},B\in{\mathbb{R}}^{190\times 270},C\in{\mathbb{R}}^{4\times 190},D\in{\mathbb{R}}^{4\times 270}$ (42,680) |
Table 4:  Details of the explicit and implicit network architectures used in
our experiments.

###  A.2 Additional Experimental Results

We provide more in-depth results on the OOD generalization capacities of
implicit models for a specific small distribution shift in Figure 11. We
compare the training and validation loss on the addition task of both implicit
and MLP models where $u_{\text{train}}\in{\mathbb{R}}^{100}\sim U(1,2)$ and
$u_{\text{val}}\in{\mathbb{R}}^{100}\sim U(2,5)$. Even with this small
distribution shift, we observe a large improvement.

Figure 11: Top and middle: The MLP test loss explodes whereas the implicit
model achieves testing loss close to 0. Bottom: For rolling argmax prediction
with extrapolation factor t = 10, our implicitRNN performs similarly on
interpolated and extrapolated data.

###  A.3 Spiky Data Generation

Both the LSTM and the implicit model were trained on 7000 data points and
tested on 3000 data points. The training regime featured 20 spiky regions of
100 data points each. The testing regime featured a proportionate amount of
spiky regions. The data points in the spiky regions were sampled from
$y=5\times(\sin(2x)+\sin(23x)+\sin(78x)+\sin(100x))$. We arbitrarily choose
frequencies in $[0,100]$ to generate a sufficiently spiky pattern. The
magnitude of the spiky regions is at most 20. For the non-spiky regimes, the
data points were sampled from $y=\sin(x)$ with added noise
$\epsilon\sim\mathcal{N}(0,0.25)$.

###  A.4 Earthquake Data Generation

To generate samples of seismic waves between specific longitudes, based on the
methods presented by Chuang et al., we used a 1D velocity model called Ak135
from the Python library obspy.taup. Obspy is a Python framework used to
process seismological data. Kennett et al. demonstrate the accuracy of this
model compared to real-world data (see specifically Figure 6). A 1D velocity
model assumes the P-wave travel time (the duration the P-wave takes to travel
from point A to point B) only depends on two attributes: the distance between
the source and the receiver station and the depth of the source. We use this
model to create a travel time lookup table based on these two attributes. We
then generate source locations from a mesh that spans the entire globe while
adding perturbation to each latitude and longitude pair. We generate station
locations using the source-station distances we have from the lookup table and
place the stations in random orientations (azimuths) from the source.