On the Abilities of Mathematical Extrapolation with Implicit Models
Abstract
Deep neural networks excel on a variety of different tasks, often surpassing human abilities.
However, when presented with out-of-distribution data, these models tend to break down even on the simplest tasks.
In this paper, we compare the robustness of implicitly-defined and classical deep learning models on a series of mathematical extrapolation tasks, where the models are tested with out-ofdistribution samples during inference time.
Throughout our experiments, implicit models greatly outperform classical deep learning networks that overfit the training distribution.
We showcase implicit models’ unique advantages for mathematical extrapolation thanks to their flexible and selective framework.
Implicit models, with potentially unlimited depth, not only adapt well to out-of-distribution inputs but also understand the underlying structure of inputs much better.
Introduction
Learning to extrapolate – the ability to infer unknown values that extend the application of a method or conclusion beyond the current scope of the known data – is a core ability of human intelligence and an important development towards general machine intelligence.
Although contemporary neural networks have demonstrated remarkable success in a myriad of domains, they struggle greatly when faced with data points outside of their training distribution [3, 15].
In this work, we investigate the capability of implicitly-defined neural networks [1, 5, 6] to extrapolate on mathematical tasks.
Implicitly-defined neural networks, such as implicit deep learning [6] or deep equilibrium models (DEQ) [1] are a general class of deep learning models that has been proposed as a potential alternative to classical neural networks.
These models do not operate on the premise of explicitly defined layers, but instead have an internal states that are defined via an “equilibrium” (fixed-point) equation and the outputs are determined only implicitly by the underlying equilibrium equation.
Formally, for a given data point u , an implicit model solves the equilibrium equation z = ϕ
( Az + Bu ) , where z is the equilibrium state for an input u , ϕ is an non-linear activation such as ReLU, and matrices
A, B are model parameters.
The prediction is obtained by feeding the equilibrium state z through an affine transformation, ˆ y ( u ) =
Cz + Du , where matrices C, D are also model parameters.
Recent results have shown successes of the implicit models [2, 9, 19].
There has also been emerging work where the equilibrium state is interpreted as a closed-loop feedback system from a neural science perspective [13].
Inspired by these successes, we seek to explore the generalization capabilities of an implicit model when dealing with mathematical and sequential tasks by comparing their capabilities to those of transformers and other architectures specialized for arithmetic computation [18].
Related Work
Prior work on using machine learning to solve logical reasoning tasks have largely focused on developing specialized neural network models to accomplish algorithm learning [7, 8, 11, 12, 17].
Some of these papers take advantage of external memory sources while others allow the network to iterate longer when it comes to more complex inputs.
Deep Equilibrium Models, a specific instance of the variable depth implicit models, have shown superior out of distribution performance. They use a higher number of root-finding iterations before converging for more complex inputs [12].
Neural
Arithmetic Logic Units (NALU) attempt to answer the extrapolation shortcomings of Deep Neural Network’s on arithmetic tasks.
NALU is a novel architecture that is able to explicitly represent mathematical relationships using the neurons of the network [18]. However, even with the improved
NALU (iNALU), instability is seen in training and random reinitializations are required [16]. Transformers are well-suited for addition and subtraction tasks achieving very high accuracy on interpolation experiments [14]. However, when faced with Out of Distribution (OOD) data, testing on longer numbers than the model was trained on, only the transformer with larger than three billion parameters performed well [14]. The same OOD concerns are also observed when experimenting on arithmetic tasks with BART, a de-noising auto-encoder using a transformer-based architecture [20].
On the contrary, when transformers were tested on matrix inversion and eigenvalue decomposition, even with OOD data they provided solutions that were ”roughly correct” and demonstrated some mathematical understanding [4].
This potentially suggests that transformers are better suited for more complex mathematical tasks where there is some possibility of demonstrating mathematical understanding without explicitly solving the problem correctly.
We consider three types of extrapolation tasks: 1) modeling the identity function, 2) performing arithmetic operations, and 3) modeling rolling functions over sequential data.
We evaluate model robustness on out of distribution shifts from the training mean.
We compare implicit models with various classical deep learning models (MLP, LSTM and Transformers) and with Google’s Neural Arithmetic Logic Units (NALU) [18] built for out of distribution arithmetic tasks. ImplicitRNN.
In this paper, we introduce implicit recurrent neural networks (RNNs) designed to receive elements in a sequence one timestep at a time and maintain at least one hidden state through a linear layer to enforce sequence-based memory.
Based on vanilla RNNs, we use a single recurrent layer with input size I and output size O , where the output doubles as our recurrent-hidden state.
We implement this recurrent layer as a single implicit layer with input size I + O , output size O , and a variable hidden size H . Note that the implicit hidden state refers to X and it is not to be confused with the recurrent-hidden state (which is the same as the recurrent layer’s output after the first input of the sequence). We initialize the recurrent-hidden state h to be the zero vector. A single forward pass for input element i of sequence s takes the following steps: (1) concatenate element i and recurrent hidden-state h : x :=
[ i, h ] ; (2) pass x through the implicit layer, producing output o ; (3) set h :
= o . At any step, o can also be used as the layer’s output prediction (i. e. o is the prediction corresponding to input i). Identity function. It has been shown that neural networks struggle to learn the basic task of identity mapping, f ( x ) = x , where models should return the exact input as given [10, 18].
We train on 10,000 data points sampled from a uniform distribution with an input dimension of 10, x train ∈ R 10 ∼ U ( − 5 , 5) , and test on 1,000 data points drawn from multiple shifted uniform distributions, U ( − κ, κ ) , where κ ranges from 10 to 80, for instance x test ∈ R 10 ∼ U ( − 10 , 10) .
We train for 500 epochs for the MLP and 1,000 epochs for both the implicit models and Transformer encoder, all with a learning rate of 0.01.
We compare an implicit model with A ∈ R 4 × 4 , B ∈ R 4 × 10 , C ∈ R 10 × 4 , D ∈ R 10 ×(196 parameters), a 2-layered MLP network of size 10 × 10 × 10 (220 parameters including biases) and a transformer encoder with a single attention head, the smallest encoder we could generate for this task (43,498 parameters).
Arithmetic operations.
We focus on two arithmetic operations: addition and subtraction.
The models take in 10,000 training arrays of length 50.
Replicating the task proposed by Trask et al.
, we randomly select four numbers i < j , k < l from 1 to 50.
For each sample, we construct two new numbers from a given array, ⃗x := ⟨ x 1 , x 2 , · · · , x 50 ⟩ .
We take a = (cid:80) ja = i x a , b = (cid:80) lb = k x b and predict y =
a + b for addition, and y = a − b for subtraction.
The training and testing data follow a uniform distribution where x train ∈ R 50 ∼ U ( − 1 , 1) and {+we expand or shrink our testing distribution by a factor of t ranging from 10 to 10 
 symmetrically such that+} x test ∈ R 50 ∼ U ( − t/ 2 , t/ 2) .
For all our tasks, we compare an implicit model with A ∈ R 20 × 20 , B ∈ R 20 × 50 , C ∈ R 1 × 20 , D ∈ R 1 × 50 , a MLP and NALU both of size 50 × 10 × 10 × 1 , and two transformers: a sequential and a depth-wise encoder.
The sequential encoder processes each array as a single sequence and has a single layer with ten attention heads.
On the other hand, the depth-wise encoder processes each element in a given array as a single sequence and has a single layer encoder with one attention head.
Sequence modeling.
We perform three sequence modeling tasks: rolling average, rolling argmax and spiky time series predictions.
The rolling average task consists of predicting for each time step the average of the sequence up to current time step j , (cid:80) ji =1 x i /j .
We train on sequences drawn from a normal distribution, x ∼ N (3 , 1) and test on sequences with a shifted mean such that x ∼ N (3 + t, 1) .
We compare an LSTM (1 × 100 × 100 × 1, 42,210 parameters), an implicit model ( A ∈ R 200 × 200 , B ∈ R 200 × 10 , C ∈ R 10 × 200 , D ∈ R 10 × 10 , 44,100 parameters) and a transformer encoder (with a single layer and 2 attention heads, 42,210 parameters).
The rolling argmax task predicts at each time step the index of the max value seen by the model so far.
We train on sequences x ∼ U (0 , 1) and test on sequences with extrapolation factor t where x ∼ U (0 , t ) .
We compare sequential models who process the sequence one timestep at a time: implicitRNN
( A ∈ R 18 × 18 , B ∈ R 18 × 23 , C ∈ R 22 × 18 , D ∈ R 22 × 23 followed by a single linear layer of size 22 × 10, 1,870 parameters) and LSTM (1 × 19 × 19 × 10, 1872 parameters). We implement a masked transformer decoder (with a single layer and two attention heads, 1920 parameters) which also only has access to previous inputs at a given timestep. We also compared our results to an unmasked transformer decoder (with a single layer and two attention heads, 1,920 parameters) and a regular implicit deep learning model ( A ∈ R 33 × 33 , B ∈ R 33 × 10 , C ∈ R 10 × 33 , D ∈ R 10 × 10 , 1, parameters). In contrast with sequential models, our two transformers and implicit model process the entire sequence at once rather than timestep per timestep.
Finally, for the spiky time series forecasting task, we first generate a time series sequence, then we randomly insert spikes designed from a combination of sine functions (see more details in the appendix). This is not an extrapolation task but rather aims to understand whether our models can predict sudden changes in the data.
We compare an implicitRNN ( A ∈ R 20 × 20 , B ∈ R 20 × 21 , C ∈ R 20 × 20 , D ∈ R 20 × 21 followed by a linear layer 20 × 1,parameters), a masked transformer decoder (single layer and 10 attention heads, 43,529 parameters) and a LSTM (1 × 20 × 20 × 1, 1861 parameters).
Experiments
For the modeling identity task, Figure 1 shows the MSE (mean squarred error) evaluated on the testing set for the MLP, implicit model and transformer.
Our implicit model maintains testing MSE < 5 for training distribution shifts from 0 to 25. Even for very large distribution shifts of up to 40 where x test ∈ R 10 ∼ U ( − 45 , 45) , the implicit model’s testing MSE only grows by a factor of 10.
In comparison, the MLP and transformer encoder testing errors surpass 10 with distribution shifts of only 10 and 5 respectively. We observe the MLP and transformer fail to model the actual identity function and instead replicate patterns observed in the training distribution which leads to increasing error as our testing set shifts away. Specifically, the transformer encoder’s MSE explodes the fastest which may be partly explained by the model’s very large size and therefore higher potential to overfit to a small training set. In contrast, our implicit model converged after only 4 iterations when predicting on the testing data.
We conclude that implicit models, not restricted to a specific number of layers, can limit overfitting through faster convergence when given simple functions. They can therefore effectively model simple mathematical functions such as the identity.
For addition and subtraction tasks, Figure 2 and Figure 3 compare training and out of distribution testing log MSE loss for our five models. In Figure 3, we observe our implicit model outperforms all other models maintaining the lowest testing loss across distribution shifts for both addition and subtraction tasks.
It successfully seems to learn the operations as demonstrated by low training loss and its ability to replicate the operation on out of distribution inputs with testing loss < 1 for distribution shifts < 100 . In Figure 3, transformers, MLP and implicit models’ log MSE loss seem to similarly linearly increase as our log distribution shift factor increases. Further results in our appendix,
Table 2 and Figure 6 however demonstrate the extrapolation advantages of implicit models on even small distribution shifts.
Furthermore as suggested by their high training MSE in Figure 2, transformers seem to underfit the training data which results in their higher extrapolation testing loss.
Implicit models therefore appear a better out of the box model for tasks with fewer training samples. Surprisingly, the NALU model, designed for extrapolation on arithmetic tasks, performs the worst as shown in Figure 3 where its testing loss surpasses 10 10 for an extrapolation shift of only 10. Across our experiments, we weren’t able to replicate robust out of distribution predictions with the NALU model from Trask et al.
We suspect the model’s performance is inflated by hand-crafted evaluation metrics and therefore does not perform well when evaluated using more traditional metrics.
As suggested by Kaiqu Liang et al.
[12], the more selective nature of implicit models may help them generalize better on logical tasks.
For a specific input X, an implicit model’s training only terminates if we find a fixed point representation of X through our equilibrium equation.
During training on both arithmetic operations, we observed our model failed to converge for at least 1/3 of our epochs.
Implicit models would have the ability to filter out internal representations that do not help capture the given arithmetic operation.
On the other hand, MLPs forward pass always terminates in a given number of steps; when the input has gone through each layer. Therefore, the MLP may have a higher chance of overfitting to the training data.
: Testing Log( ) of our five models on testing distribution shifts. We observe the implicit model strongly outperforms all other models on OOD data and maintains a linear increase.
We summarize our results on out of distribution inputs for the three sequence modeling tasks in Figure 4 and 5.
For the rolling argmax task, we observe in Figure 4 that our implicit models maintain the highest and a very stable testing accuracy across distribution shifts.
Both our transformers demonstrate a similar capacity to extrapolate on out of distribution inputs with however a lower accuracy (by at least 5%). In contrast, the LSTM fails to extrapolate as its accuracy drops by almost 50% when evaluated on out of distribution inputs. Figure 5 shows test MSE across distribution shifts for the rolling average task and an example test sequence predictions for the spiky data task.
For the rolling average task, the LSTM and transformer both replicate the training distribution, predicting averages around 3 even in the test set, whereas the implicit model extrapolates to higher values.
For the spiky data predictions, although LSTM and implicit models have similar testing losses, the implicit model seems to have a better understanding of overall sequence structure.
It successfully predicts the specific location and magnitudes of spikes.
Note that this is not an extrapolation task as the training and testing regimes had a similar proportion of spikes.
However, given very few examples of anomalous structure in the data the implicit model performs very well when similar structure appears in the test set.
Across these sequential tasks, we hypothesize that implicit and implicitRNN models here benefit from greater model flexibility.
Specifically for sequential data, the implicit layers within an implicitRNN can run for more iterations when presented with more complex inputs.
Conclusion
Our results showcase implicit models’ superior performance on out of distribution sample points when compared to traditional deep learning models.
This ability is especially apparent on mathematical tasks such as function learning and arithmetic operations.
With potentially infinite depth, implicit models have the ability to grow for more complex inputs and adapt to data shifts.
In all of our experiments, implicit models showcased improved performance over MLPs, LSTMs, transformers and Google’s NALU architecture.
Implicit models therefore appear a better out of the box solution for extrapolation tasks. These results motivate the further study of implicit models as a robust
A Appendix A.1 Activation Function Experiments
For the identity function and arithmetic operation taks, we experimented with 15 different activation functions on our MLP: hardtanh, sigmoid, reLU6, tanh, tanhshrink, hardshrink, leakyrelu, softshrink, softsign, reLU, preLU, multipreLU, softplus, eLU and seLU.
We tried to understand whether specific activations helped the MLP extrapolate as well as our implicit model.
Table 1 summarizes the results of 5 of these activation functions on our identity function task as compared to the implicit deep learning model.
[Table]
Table 2 compares the test MSE for our MLP with ReLU activation, the best MLP across allactivations and our implicit model. We have x train ∈ R 100 ∼ U (1 , 2) and x test ∈ R 100 ∼ U (2 , 5) .
For both operations, the implicit model greatly outperforms the MLP regardless of the activation function.
[Table]
A.2 Arithmetic Operations More Results
For more specific results on the OOD generalization capacities of implicit models, we compare in
Figure 6 the training and validation loss on the addition task of both implicit and MLP models where x train ∈
R 100 ∼ U (1 , 2) and x val ∈ R 100 ∼ U (2 , 5) .
This is therefore one of our small distribution shifts where t = 3.
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 × (sin(2 x )+sin(23 x )+sin(78 x )+sin(100 x )) .
The frequencies were arbitrarily chosen to be between 0 to 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 ) .