H OW N EURAL N ETWORKS E XTRAPOLATE : F ROM F EEDFORWARD TO G RAPH N EURAL N ETWORKS
A BSTRACT
We study how neural networks trained by gradient descent extrapolate , i.e., what they learn outside the support of the training distribution.
Previous works report mixed empirical results when extrapolating with neural networks: while feedforward neural networks, a.k.a.
multilayer perceptrons (MLPs), do not extrapolate well in certain simple tasks, Graph Neural Networks (GNNs) – structured networks with MLP modules – have shown some success in more complex tasks.
Working towards a theoretical explanation, we identify conditions under which MLPs and
GNNs extrapolate well.
First, we quantify the observation that ReLU MLPs quickly converge to linear functions along any direction from the origin, which implies that
ReLU MLPs do not extrapolate most nonlinear functions.
But, they can provably learn a linear target function when the training distribution is sufﬁciently “diverse”.
Second, in connection to analyzing the successes and limitations of GNNs, these results suggest a hypothesis for which we provide theoretical and empirical evidence: the success of GNNs in extrapolating algorithmic tasks to new data (e.g., larger graphs or edge weights) relies on encoding task-speciﬁc non-linearities in the architecture or features.
Our theoretical analysis builds on a connection of over-parameterized networks to the neural tangent kernel.
Empirically, our theory holds across different training settings.
I NTRODUCTION
Humans extrapolate well in many tasks.
For example, we can apply arithmetics to arbitrarily large numbers.
One may wonder whether a neural network can do the same and generalize to examples arbitrarily far from the training data (Lake et al., 2017).
Curiously, previous works report mixed extrapolation results with neural networks.
Early works demonstrate feedforward neural networks, a.k.a.
multilayer perceptrons (MLPs), fail to extrapolate well when learning simple polynomial functions (Barnard & Wessels, 1992; Haley & Soloway, 1992).
However, recent works show Graph
Neural Networks (GNNs) (Scarselli et al., 2009), a class of structured networks with MLP building blocks, can generalize to graphs much larger than training graphs in challenging algorithmic tasks, such as predicting the time evolution of physical systems (Battaglia et al., 2016), learning graph algorithms (Velickovic et al., 2020), and solving mathematical equations (Lample & Charton, 2020).
To explain this puzzle, we formally study how neural networks trained by gradient descent (GD) extrapolate, i.e., what they learn outside the support of training distribution.
We say a neural network extrapolates well if it learns a task outside the training distribution.
At ﬁrst glance, it may seem that neural networks can behave arbitrarily outside the training distribution since they have high capacity (Zhang et al., 2017) and are universal approximators (Cybenko, 1989; Funahashi, 1989;
Hornik et al., 1989; Krkov´a, 1992).
However, neural networks are constrained by gradient descent training (Hardt et al., 2016; Soudry et al., 2018).
In our analysis, we explicitly consider such implicit bias through the analogy of the training dynamics of over-parameterized neural networks and kernel regression via the neural tangent kernel (NTK)
(Jacot et al., 2018).
Hence, MLPs do not extrapolate well in most nonlinear tasks.
But, with appropriate training data,
MLPs can provably extrapolate linear target functions (Theorem 2).
[Figure]
we hypothesize that GNNs can extrapolate well in dynamic programming (DP) tasks if we encode appropriate non-linearities in the architecture (left) and input representation ( domain knowledge or representation learning).
The encoded non-linearities may not be necessary for interpolation, as they can be approximated by MLP modules, but they help extrapolation.
We support the hypothesis theoretically (Theorem 3) and empirically (Figure 6).
Starting with feedforward networks, the simplest neural networks and building blocks of more complex architectures such as GNNs, we establish that the predictions of over-parameterized MLPs with ReLU activation trained by GD converge to linear functions along any direction from the origin.
We prove a convergence rate for two-layer networks and empirically observe that convergence often occurs close to the training data (Figure 1), which suggests ReLU MLPs cannot extrapolate well for most nonlinear tasks.
We emphasize that our results do not follow from the fact that ReLU networks have ﬁnitely many linear regions (Arora et al., 2018;
Hanin & Rolnick, 2019; Hein et al., 2019).
While having ﬁnitely many linear regions implies ReLU MLPs eventually become linear, it does not say whether MLPs will learn the correct target function close to the training distribution.
In contrast, our results are non-asymptotic and quantify what kind of functions MLPs will learn close to the training distribution.
Second, we identify a condition when MLPs extrapolate well: the task is linear and the geometry of the training distribution is sufﬁciently “diverse”.
To our knowledge, our results are the ﬁrst extrapolation results of this kind for feedforward neural networks.
We then relate our insights into feedforward neural networks to GNNs, to explain why GNNs extrapolate well in some algorithmic tasks.
Prior works report successful extrapolation for tasks that can be solved by dynamic programming (DP) (Bellman, 1966), which has a computation structure aligned with GNNs (Xu et al., 2020).
DP updates can often be decomposed into nonlinear and linear steps.
Hence, we hypothesize that GNNs trained by GD can extrapolate well in a DP task, if we encode appropriate non-linearities in the architecture and input representation (Figure 2).
Importantly, encoding non-linearities may be unnecessary for GNNs to interpolate , because the MLP modules can easily learn many nonlinear functions inside the training distribution (Cybenko, 1989; Hornik et al., 1989; Xu et al., 2020), but it is crucial for GNNs to extrapolate correctly.
We prove this hypothesis for a simpliﬁed case using Graph NTK (Du et al., 2019b).
Empirically, we validate the hypothesis on three DP tasks: max degree, shortest paths, and n -body problem.
We show GNNs with appropriate architecture, input representation, and training distribution can predict well on graphs with unseen sizes, structures, edge weights, and node features.
Our theory explains the empirical success in previous works and suggests their limitations: successful extrapolation relies on encoding task-speciﬁc non-linearities, which requires domain knowledge or extensive model search.
From a broader standpoint, our insights go beyond GNNs and apply broadly to other neural networks.
To summarize, we study how neural networks extrapolate.
First, ReLU MLPs trained by GD converge to linear functions along directions from the origin with a rate of O (1 /t ) .
Second, to explain why
GNNs extrapolate well in some algorithmic tasks, we prove that ReLU MLPs can extrapolate well in linear tasks, leading to a hypothesis: a neural network can extrapolate well when appropriate nonlinearities are encoded into the architecture and features.
We prove this hypothesis for a simpliﬁed case and provide empirical support for more general settings.
1. R ELATED WORK
Early works show example tasks where MLPs do not extrapolate well, e.g. learning simple polynomials (Barnard & Wessels, 1992; Haley & Soloway, 1992).
We instead show a general pattern of how
ReLU MLPs extrapolate and identify conditions for MLPs to extrapolate well.
More recent works study the implicit biases induced on MLPs by gradient descent, for both the NTK and mean ﬁeld regimes (Bietti & Mairal, 2019; Chizat & Bach, 2018; Song et al., 2018).
Related to our results, some works show MLP predictions converge to “simple” piecewise linear functions, e.g., with few linear regions (Hanin & Rolnick, 2019;
Maennel et al., 2018;
Savarese et al., 2019; Williams et al., 2019).
Our work differs in that none of these works explicitly studies extrapolation, and some focus only on one-dimensional inputs.
Recent works also show that in high-dimensional settings of the NTK regime, MLP is asymptotically at most a linear predictor in certain scaling limits (Ba et al., 2020;
Ghorbani et al., 2019).
We study a different setting (extrapolation), and our analysis is non-asymptotic in nature and does not rely on random matrix theory.
Prior works explore GNN extrapolation by testing on larger graphs (Battaglia et al., 2018; Santoro et al., 2018; Saxton et al., 2019; Velickovic et al., 2020).
We are the ﬁrst to theoretically study GNN extrapolation, and we complete the notion of extrapolation to include unseen features and structures.
P RELIMINARIES
We begin by introducing our setting.
Let X be the domain of interest, e.g., vectors or graphs.
The task is to learn an underlying function g : X → R with a training set { ( x i , y i ) } ni =1 ⊂ D , where y i =
g
( x i ) and D is the support of training distribution.
Previous works have extensively studied in-distribution generalization where the training and the test distributions are identical (Valiant, 1984;
Vapnik, 2013); i.e., D = X .
In contrast, extrapolation addresses predictions on a domain X that is larger than the support of the training distribution D .
We will say a model extrapolates well if it has a small extrapolation error .
Deﬁnition 1.
(Extrapolation error).
Let f : X → R be a model trained on { ( x i , y i ) } ni =1 ⊂ D with underlying function g : X → R .
Let P be a distribution over X \ D and let (cid:96) : R × R
→ R be a loss function.
We deﬁne the extrapolation error of f as E x ∼P
[ (cid:96) ( f ( x ) , g ( x ))] .
We focus on neural networks trained by gradient descent (GD) or its variants with squared loss .
We study two network architectures: feedforward and graph neural networks.
Graph Neural Networks.
GNNs are structured networks operating on graphs with MLP modules (Battaglia et al., 2018;
Xu et al., 2019).
Let G =
( V, E ) be a graph.
Each node u ∈ V has a feature vector x u , and each edge ( u, v ) ∈
E has a feature vector w ( u,v ) .
GNNs recursively compute node representations h ( k ) u at iteration k
(Gilmer et al., 2017; Xu et al., 2018).
Initially, h (0) u
= x u .
For k = 1 ..K , GNNs update h ( k ) u by aggregating the neighbor representations.
We can optionally compute a graph representation h G by aggregating the ﬁnal node representations.
That is,
[Equation]
The ﬁnal output is the graph representation h G or ﬁnal node representations h ( K ) u depending on the task.
We refer to the neighbor aggregation step for h ( k ) u as aggregation and the pooling step in h G as readout .
Previous works typically use sum-aggregation and sum-readout (Battaglia et al., 2018).
Our results indicate why replacing them may help extrapolation (Section 4).
[Figure]
H OW F EEDFORWARD N EURAL N ETWORKS E XTRAPOLATE
Feedforward networks are the simplest neural networks and building blocks of more complex architectures such as GNNs, so we ﬁrst study how they extrapolate when trained by GD.
Throughout the paper we assume ReLU activation.
Section 3.3 contains preliminary results for other activations.
3. L INEAR E XTRAPOLATION B EHAVIOR OF R E LU MLP S
By architecture, ReLU networks learn piecewise linear functions but what do these regions precisely look like outside the support of the training data?
Figure 1 illustrates examples of how ReLU
MLPs extrapolate when trained by GD on various nonlinear functions.
These examples suggest that outside the training support, the predictions quickly become linear along directions from the origin.
We systematically verify this pattern by linear regression on MLPs’ predictions: the coefﬁcient of determination ( R 2 ) is always greater than 0.99 (Appendix C.2).
That is, ReLU MLPs “linearize” almost immediately outside the training data range.
We formalize this observation using the implicit biases of neural networks trained by GD via the neural tangent kernel (NTK) : optimization trajectories of over-parameterized networks trained by
GD are equivalent to those of kernel regression with a speciﬁc neural tangent kernel, under a set of assumptions called the “NTK regime” (Jacot et al., 2018).
We provide an informal deﬁnition here; for further details, we refer the readers to Jacot et al.
and Appendix A.
Deﬁnition 2.
(Informal) A neural network trained in the NTK regime is inﬁnitely wide, randomly initialized with certain scaling, and trained by GD with inﬁnitesimal steps.
Prior works analyze optimization and in-distribution generalization of over-parameterized neural networks via NTK (Allen-Zhu et al., 2019a;b; Arora et al., 2019a;b; Cao & Gu, 2019; Du et al., 2019c;a; Li & Liang, 2018; Nitanda & Suzuki, 2021).
We instead analyze extrapolation.
Theorem 1 formalizes our observation from Figure 1: outside the training data range, along any direction t v from the origin, the prediction of a two-layer ReLU MLP quickly converges to a linear function with rate O ( 1 t ) .
The linear coefﬁcients β v and the constant terms in the convergence rate depend on the training data and direction v .
The proof is in Appendix B.1.
Theorem 1.
(Linear extrapolation).
Suppose we train a two-layer ReLU MLP f : R d → R with squared loss in the NTK regime.
For any direction v ∈ R d , let x 0 = t v .
As t → ∞ , f ( x 0 + h v ) − f ( x 0 ) → β v · h for any h > 0 , where β v is a constant linear coefﬁcient.
Moreover, given (cid:15) > 0 , for t =
O ( 1 (cid:15) ) , we have | f ( x 0 + h v ) − f ( x 0 )
h − β v | < (cid:15) .
ReLU networks have ﬁnitely many linear regions (Arora et al., 2018; Hanin & Rolnick, 2019), hence their predictions eventually become linear.
In contrast, Theorem 1 is a more ﬁne-grained analysis of how MLPs extrapolate and provides a convergence rate.
While Theorem 1 assumes two-layer networks in the NTK regime, experiments conﬁrm that the linear extrapolation behavior happens across networks with different depths, widths, learning rates, and batch sizes (Appendix C.1 and C.2).
Our proof technique potentially also extends to deeper networks.
Theorem 1 implies which target functions a ReLU MLP may be able to match outside the training data: only functions that are almost-linear along the directions away from the origin.
Indeed, Figure 4a shows ReLU MLPs do not extrapolate target functions such as x (cid:62) A x (quadratic), (cid:80) di
=1 cos(2 π · x ( i ) )
(cos), and (cid:80) d i =1 √ x ( i ) (sqrt), where x ( i ) is the i -th dimension of x .
With suitable hyperparameters,
MLPs extrapolate the L1 norm correctly, which satisﬁes the directional linearity condition.
Figure 4a provides one more positive result: MLPs extrapolate linear target functions well, across many different hyperparameters.
While learning linear functions may seem very limited at ﬁrst, in
Section 4 this insight will help explain extrapolation properties of GNNs in non-linear practical tasks.
Before that, we ﬁrst theoretically analyze when MLPs extrapolate well.
3. W HEN R E LU MLP S P ROVABLY E XTRAPOLATE W ELL
Figure 4a shows that MLPs can extrapolate well when the target function is linear.
However, this is not always true.
In this section, we show that successful extrapolation depends on the geometry of training data.
Intuitively, the training distribution must be “diverse” enough for correct extrapolation.
We provide two conditions that relate the geometry of the training data to extrapolation.
Lemmastates that over-parameterized MLPs can learn a linear target function with only 2 d examples.
Lemma 1.
Let g ( x ) =
β (cid:62) x be the target function for β ∈ R d .
Suppose { x i } ni
=1 contains an orthogonal basis { ˆ x i } di =1 and {− ˆ x i } di =1 .
If we train a two-layer ReLU MLP f on { ( x i , y i ) } ni = with squared loss in the NTK regime, then f ( x ) =
β (cid:62) x for all x ∈ R d .
Lemma 1 is mainly of theoretical interest, as the 2 d examples need to be carefully chosen.
Theorembuilds on Lemma 1 and identiﬁes a more practical condition for successful extrapolation: if the support of the training distribution covers all directions (e.g., a hypercube that covers the origin), the
MLP converges to a linear target function with sufﬁcient training data.
Theorem 2.
(Conditions for extrapolation).
Let g ( x ) =
β (cid:62) x be the target function for β ∈ R d .
Suppose { x i } ni =1 is sampled from a distribution whose support D contains a connected subset S , where for any non-zero w ∈ R d , there exists k >
so that k w ∈ S .
If we train a two-layer ReLU
MLP f : R d → R on { ( x i , y i ) } ni =1 with squared loss in the NTK regime, f ( x ) p −→ β (cid:62) x as n → ∞ .
Experiments: geometry of training data affects extrapolation.
The condition in Theoremformalizes the intuition that the training distribution must be “diverse” for successful extrapolation, e.g., D includes all directions.
Empirically, the extrapolation error is indeed small when the condition of Theorem 2 is satisﬁed (“all” in Figure 4b).
In contrast, the extrapolation error is much larger when the training examples are restricted to only some directions (Figure 4b and Figure 3).
Relating to previous works, Theorem 2 suggests why spurious correlations may hurt extrapolation, complementing the causality arguments (Arjovsky et al., 2019;
Peters et al., 2016; Rojas-Carulla et al., 2018).
When the training data has spurious correlations, some combinations of features are missing; e.g., camels might only appear in deserts in an image collection.
Therefore, the condition for
Theorem 2 no longer holds, and the model may extrapolate incorrectly.
Theorem 2 is also analogous to an identiﬁability condition for linear models, but stricter.
We can uniquely identify a linear function if the training data has full (feature) rank.
MLPs are more expressive, so identifying the linear target function requires additional constraints.
[Figure]
To summarize, we analyze how ReLU MLPs extrapolate and provide two insights: (1) MLPs cannot extrapolate most nonlinear tasks due to their linear extrapolation (Theorem 1); and (2) MLPs extrapolate well when the target function is linear, if the training distribution is “diverse” (Theorem 2).
In the next section, these results help us understand how more complex networks extrapolate.
3. MLP S WITH O THER A CTIVATION F UNCTIONS
Before moving on to GNNs, we complete the picture of MLPs with experiments on other activation functions: tanh σ ( x ) = tanh( x ) , cosine σ ( x ) = cos( x ) (Lapedes & Farber, 1987; McCaughan, 1997; Sopena & Alquezar, 1994), and quadratic σ ( x ) = x 2 (Du & Lee, 2018; Livni et al., 2014).
Details are in Appendix C.4.
MLPs extrapolate well when the activation and target function are similar; e.g., tanh activation extrapolates well when learning tanh, but not other functions (Figure 5).
Moreover, each activation function has different limitations.
To extrapolate the tanh function with tanh activation, the training data range has to be sufﬁciently wide.
When learning a quadratic function with quadratic activation, only two-layer networks extrapolate well as more layers lead to higher-order polynomials.
Cosine activations are hard to optimize for high-dimensional data, so we only consider one/two dimensional cosine target functions.
H OW G RAPH N EURAL N ETWORKS E XTRAPOLATE
Above, we saw that extrapolation in nonlinear tasks is hard for MLPs.
Despite this limitation,
GNNs have been shown to extrapolate well in some nonlinear algorithmic tasks, such as intuitive physics (Battaglia et al., 2016; Janner et al., 2019), graph algorithms (Battaglia et al., 2018; Velickovic et al., 2020), and symbolic mathematics (Lample & Charton, 2020).
To address this discrepancy, we build on our MLP results and study how GNNs trained by GD extrapolate.
4. H YPOTHESIS : L INEAR A LGORITHMIC A LIGNMENT H ELPS E XTRAPOLATION
We start with an example: training GNNs to solve the shortest path problem.
For this task, prior works observe that a modiﬁed GNN architecture with min-aggregation can generalize to graphs larger than those in the training set (Battaglia et al., 2018;
Velickovic et al., 2020):
[Equation]
We ﬁrst provide an intuitive explanation (Figure 2a).
Shortest path can be solved by the Bellman-Ford (BF) algorithm (Bellman, 1958) with the following update: [Equation]
[Equation]
) where w ( v, u ) is the weight of edge ( v, u ) , and d [ k ][ u ] is the shortest distance to node u within k steps.
The two equations can be easily aligned: GNNs simulate the BF algorithm if its MLP modules learn a linear function d [ k − 1][ v ]+ w ( v, u ) .
Since MLPs can extrapolate linear tasks, this “alignment” may explain why min-aggregation GNNs can extrapolate well in this task.
For comparison, we can reason why we would not expect GNNs with the more commonly used sum-aggregation (Eqn. 1) to extrapolate well in this task.
With sum-aggregation, the MLP modules need to learn a nonlinear function to simulate the BF algorithm, but Theorem 1 suggests that they will not extrapolate most nonlinear functions outside the training support.
We can generalize the above intuition to other algorithmic tasks.
Many tasks where GNNs extrapolate well can be solved by dynamic programming (DP) (Bellman, 1966), an algorithmic paradigm with a recursive structure similar to GNNs’ (Eqn. 1) (Xu et al., 2020).
Deﬁnition 3.
Dynamic programming (DP) is a recursive procedure with updates
[Equation]
where Answer [ k ][ s ] is the solution to a sub-problem indexed by iteration k and state s , and DP-Update is a task-speciﬁc update function that solves the sub-problem based on the previous iteration.
From a broader standpoint, we hypothesize that: if we encode appropriate non-linearities into the model architecture and input representations so that the MLP modules only need to learn nearly linear steps, then the resulting neural network can extrapolate well.
Hypothesis 1.
(Linear algorithmic alignment).
Let f : X → R be the underlying function and N a neural network with m MLP modules.
Suppose there exist m linear functions { g i } mi =1 so that by replacing N ’s MLP modules with g
i ’s, N simulates f .
Given (cid:15) > 0 , there exists { ( x i , f ( x i )) } ni =1 ⊂
D (cid:40)
X so that N trained on { ( x i , f ( x i )) } ni =1 by GD with squared loss learns ˆ f with (cid:107) ˆ f − f (cid:107) < (cid:15) .
Our hypothesis builds on the algorithmic alignment framework of (Xu et al., 2020), which states that a neural network interpolates well if the modules are “aligned” to easy-to-learn (possibly nonlinear) functions.
Successful extrapolation is harder: the modules need to align with linear functions.
Applications of linear algorithmic alignment.
In general, linear algorithmic alignment is not restricted to GNNs and applies broadly to neural networks.
To satisfy the condition, we can encode appropriate nonlinear operations in the architecture or input representation (Figure 2).
Learning DP algorithms with GNNs is one example of encoding non-linearity in the architecture (Battaglia et al., 2018;
Corso et al., 2020).
Another example is to encode log-and-exp transforms in the architecture to help extrapolate multiplication in arithmetic tasks (Trask et al., 2018; Madsen & Johansen, 2020).
Neural symbolic programs take a step further and encode a library of symbolic operations to help extrapolation (Johnson et al., 2017; Mao et al., 2019;
Yi et al., 2018).
For some tasks, it may be easier to change the input representation (Figure 2b).
Sometimes, we can decompose the target function f as f = g ◦ h into a feature embedding h and a “simpler” target function g that our model can extrapolate well.
We can obtain h via specialized features or feature transforms using domain knowledge (Lample & Charton, 2020;
Webb et al., 2020), or via representation learning (e.g., BERT) with unlabeled out-of-distribution data in X \ D (Chen et al., 2020; Devlin et al., 2019;
Hu et al., 2020; Mikolov et al., 2013b; Peters et al., 2018).
This brings a new perspective of how representations help extrapolation in various application areas.
For example, in natural language processing, pretrained representations (Mikolov et al., 2013a; Wu & Dredze, 2019) and feature transformation using domain knowledge (Yuan et al., 2020; Zhang et al., 2019) help models generalize across languages, a special type of extrapolation.
In quantitative ﬁnance, identifying the right “factors” or features is crucial for deep learning models as the ﬁnancial markets may frequently be in extrapolation regimes (Banz, 1981; Fama & French, 1993; Ross, 1976).
Linear algorithmic alignment explains successful extrapolation in the literature and suggests that extrapolation is harder in general: encoding appropriate non-linearity often requires domain expertise or model search.
Next, we provide theoretical and empirical support for our hypothesis.
4. T HEORETICAL AND E MPIRICAL S UPPORT
We validate our hypothesis on three DP tasks: max degree, shortest path, and n -body problem, and prove the hypothesis for max degree.
We highlight the role of graph structures in extrapolation.
Theoretical analysis.
We start with a simple yet fundamental task: learning the max degree of a graph, a special case of DP with one iteration.
As a corollary of Theorem 1, the commonly used sum-based GNN (Eqn. 1) cannot extrapolate well (proof in Appendix B.4).
[Figure] [Figure]
Corollary 1.
GNNs with sum-aggregation and sum-readout do not extrapolate well in Max Degree.
To achieve linear algorithmic alignment, we can encode the only non-linearity, the max function, in the readout.
Theorem 3 conﬁrms that a GNN with max-readout can extrapolate well in this task.
Theorem 3.
(Extrapolation with GNNs).
Assume all nodes have the same feature.
Let g and g (cid:48) be the max/min degree function, respectively.
Let { ( G i , g ( G i ) } ni =1 be the training set.
If { ( g ( G i ) , g (cid:48) ( G i ) , g ( G i ) · N max i , g (cid:48) ( G i ) · N min i ) } ni =1 spans R 4 , where N max i and N min i are the number of nodes that have max/min degree on G i , then one-layer max-readout GNNs trained on { ( G i , g ( G i ))
} ni =1 with squared loss in the NTK regime learn g .
Theorem 3 does not follow immediately from Theorem 2, because MLP modules in GNNs only receive indirect supervision.
We analyze the Graph NTK (Du et al., 2019b) to prove Theorem 3 in
Appendix B.5.
While Theorem 3 assumes identical node features, we empirically observe similar results for both identical and non-identical features (Figure 16 in Appendix).
Interpretation of conditions.
The condition in Theorem 3 is analogous to that in Theorem 2.
Both theorems require diverse training data, measured by graph structure in Theorem 3 or directions in
Theorem 2.
In Theorem 3, the condition is violated if all training graphs have the same max or min node degrees, e.g., when training data are from one of the following families: path, C -regular graphs (regular graphs with degree C ), cycle, and ladder.
Experiments: architectures that help extrapolation.
We validate our theoretical analysis with two DP tasks: max degree and shortest path (details in Appendix C.5 and C.6).
While previous works only test on graphs with different sizes (Battaglia et al., 2018; Velickovic et al., 2020), we also test on graphs with unseen structure, edge weights and node features.
The results support our theory.
For max degree, GNNs with max-readout are better than GNNs with sum-readout (Figure 6a), conﬁrming
Corollary 1 and Theorem 3.
For shortest path, GNNs with min-readout and min-aggregation are better than GNNs with sum-readout (Figure 6a).
Experiments conﬁrm the importance of training graphs structure (Figure 7).
Interestingly, the two tasks favor different graph structure.
For max degree, as Theorem 3 predicts, GNNs extrapolate well when trained on trees, complete graphs, expanders, and general graphs, and extrapolation errors are higher when trained on 4-regular, cycles, or ladder graphs.
For shortest path, extrapolation errors follow a U-shaped curve as we change the sparsity of training graphs (Figure 7b and Figure 18 in
Appendix).
Intuitively, models trained on sparse or dense graphs likely learn degenerative solutions.
Experiments: representations that help extrapolation.
Finally, we show a good input representation helps extrapolation.
We study the n -body problem (Battaglia et al., 2016; Watters et al., 2017)
(Appendix C.7), that is, predicting the time evolution of n objects in a gravitational system.
Following previous work, the input is a complete graph where the nodes are the objects (Battaglia et al., 2016).
The node feature for u is the concatenation of the object’s mass m u , position x ( t ) u , and velocity v ( t ) u at time t .
The edge features are set to zero.
We train GNNs to predict the velocity of each object u at time t + 1 .
The true velocity f ( G ; u ) for object u is approximately
[Equation]
where C is a constant.
To learn f , the MLP modules need to learn a nonlinear function.
Therefore,
GNNs do not extrapolate well to unseen masses or distances (“original features” in Figure 6b).
We instead use an improved representation h ( G ) to encode non-linearity.
At time t , we transform the edge features of ( u, v ) from zero to w ( t ) ( u,v ) = m v · (cid:0) x
( t ) v − x ( t ) u (cid:1) / (cid:107) x
( t ) u − x ( t ) v (cid:107) 3 2 .
The new edge features do not add information, but the MLP modules now only need to learn linear functions, which helps extrapolation (“improved features” in Figure 6b).
C ONNECTIONS TO O THER O UT - OF -D ISTRIBUTION S ETTINGS
We discuss several related settings.
Intuitively, from the viewpoint of our results above, methods in related settings may improve extrapolation by 1) learning useful non-linearities beyond the training data range and 2) mapping relevant test data to the training data range.
Domain adaptation studies generalization to a speciﬁc target domain (Ben-David et al., 2010; Blitzer et al., 2008; Mansour et al., 2009).
Typical strategies adjust the training process: for instance, use unlabeled samples from the target domain to align the target and source distributions (Ganin et al., 2016; Zhao et al., 2018).
Using target domain data during training may induce useful non-linearities and may mitigate extrapolation by matching the target and source distributions, though the correctness of the learned mapping depends on the label distribution (Zhao et al., 2019).
Self-supervised learning on a large amount of unlabeled data can learn useful non-linearities beyond the labeled training data range (Chen et al., 2020;
Devlin et al., 2019;
He et al., 2020; Peters et al., 2018).
Hence, our results suggest an explanation why pre-trained representations such as BERT improve out-of-distribution robustness (Hendrycks et al., 2020).
In addition, self-supervised learning could map semantically similar data to similar representations, so some out-of-domain examples might fall inside the training distribution after the mapping.
Invariant models aim to learn features that respect speciﬁc invariances across multiple training distributions (Arjovsky et al., 2019; Rojas-Carulla et al., 2018;
Zhou et al., 2021).
If the model indeed learns these invariances, which can happen in the linear case and when there are confounders or anti-causal variables (Ahuja et al., 2021; Rosenfeld et al., 2021), this may essentially increase the training data range, since variations in the invariant features may be ignored by the model.
Distributional robustness considers small adversarial perturbations of the data distribution, and ensures that the model performs well under these (Goh & Sim, 2010; Sagawa et al., 2020; Sinha et al., 2018; Staib & Jegelka, 2019).
We instead look at more global perturbations.
Still, one would expect that modiﬁcations that help extrapolation in general also improve robustness to local perturbations.
C ONCLUSION
This paper is an initial step towards formally understanding how neural networks trained by gradient descent extrapolate.
We identify conditions under which MLPs and GNNs extrapolate as desired.
We also suggest an explanation how GNNs have been able to extrapolate well in complex algorithmic tasks: encoding appropriate non-linearity in architecture and features can help extrapolation.
Our results and hypothesis agree with empirical results, in this paper and in the literature.
A CKNOWLEDGMENTS
We thank Ruosong Wang, Tianle Cai, Han Zhao, Yuichi Yoshida, Takuya Konishi, Toru Lin, Weihua
Hu, Matt J. Staib, Yichao Zhou, Denny Wu, Tianyi Yang and Dingli (Leo) Yu for insightful discussions.
This research was supported by NSF CAREER award 1553284, NSF III 1900933, and a
Chevron-MIT Energy Fellowship.
This research was also supported by JST ERATO JPMJER and JSPS Kakenhi JP18H05291.
MZ was supported by ODNI, IARPA, via the BETTER Program contract 2019-19051600005.
The views, opinions, and/or ﬁndings contained in this article are those of the author and should not be interpreted as representing the ofﬁcial views or policies, either expressed or implied, of the Defense Advanced Research Projects Agency, the Department of Defense, ODNI,
IARPA, or the U
S. Government.
The U
S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.
A T HEORETICAL B ACKGROUND
In this section, we introduce theoretical background on neural tangent kernel (NTK), which draws an equivalence between the training dynamics of inﬁnitely-wide (or ultra-wide) neural networks and that of kernel regression with respect to the neural tangent kernel.
Consider a general neural network f ( θ , x ) :
X → R where θ ∈ R m is the parameters in the network and x ∈ X is the input.
Suppose we train the neural network by minimizing the squared loss over training data, (cid:96) ( θ ) = 12 (cid:80) ni
=
( f ( θ , x i ) − y i ) 2 , by gradient descent with inﬁnitesimally small learning rate, i.e., d θ ( t )
dt = −∇ (cid:96) ( θ ( t )) .
Let u ( t ) =
( f ( θ ( t ) , x i )) ni
=1 be the network outputs.
u ( t ) follows the dynamics
[Equation]
where H ( t ) is an n × n matrix whose ( i, j ) -th entry is
[Equation]
A line of works show that for sufﬁciently wide networks, H ( t ) stays almost constant during training, i.e., H ( t )
=
H (0) in the limit (Arora et al., 2019a;b; Allen-Zhu et al., 2019a; Du et al., 2019c;a; Li & Liang, 2018; Jacot et al., 2018).
Suppose network parameters are randomly initialized with certain scaling, as network width goes to inﬁnity, H (0) converges to a ﬁxed matrix, the neural tangent kernel (NTK)
(Jacot et al., 2018):
[Equation]
where W is Gaussian.
Therefore, the learning dynamics of sufﬁciently wide neural networks in this regime is equivalent to that of kernel gradient descent with respect to the NTK.
This implies the function learned by a neural network at convergence on any speciﬁc training set, denoted by f NTK ( x ) , can be precisely characterized, and is equivalent to the following kernel regression solution
[Equation]
where NTK train is the n × n kernel for training data, NTK ( x , x i ) is the kernel value between test data x and training data x i , and Y is the training labels.
We can in fact exactly calculate the neural tangent kernel matrix for certain architectures and activation functions.
The exact formula of NTK with ReLU activation has been derived for feedforward neural networks (Jacot et al., 2018), convolutional neural networks (Arora et al., 2019b), and Graph Neural
Networks (Du et al., 2019b).
Our theory builds upon this equivalence of network learning and kernel regression to more precisely characterize the function learned by a sufﬁciently-wide neural network given any speciﬁc training set.
In particular, the difference between the learned function and true function over the domain of X determines the extrapolation error.
However, in general it is non-trivial to compute or analyze the functional form of what a neural network learns using Eqn. 9, because the kernel regression solution using neural tangent kernel only gives point-wise evaluation.
Thus, we instead analyze the function learned by a network in the NTK’s induced feature space , because representations in the feature space would give a functional form.
Lemma 2 makes this connection more precise: the solution to the kernel regression using neural tangent kernel, which also equals over-parameterized network learning, is equivalent to a min-norm solution among functions in the NTK’s induced feature space that ﬁts all training data.
Here the min-norm refers to the RKHS norm.
Lemma 2.
Let φ ( x ) be a feature map induced by a neural tangent kernel, for any x ∈ R d .
The solution to kernel regression Eqn. 9 is equivalent to f NTK ( x ) =
φ ( x ) (cid:62) β NTK , where β NTK is min β (cid:107) β (cid:107)s.t.
φ ( x i ) (cid:62) β
= y i , for i
= 1 , ..., n.
We prove Lemma 2 in Appendix B.6.
To analyze the learned functions as the min-norm solution in feature space, we also need the explicit formula of an induced feature map of the corresponding neural tangent kernel.
The following lemma gives a NTK feature space for two-layer MLPs with
ReLU activation.
It follows easily from the kernel formula described in Jacot et al.
Arora et al. (2019b); Bietti & Mairal (2019).
Lemma 3.
An inﬁnite-dimensional feature map φ ( x ) induced by the neural tangent kernel of a two-layer multi-layer perceptron with ReLU activation function is
[Equation]
where w ( k ) ∼ N ( 0 , I ) , with k going to inﬁnity.
c is a constant, and I is the indicator function.
We prove Lemma 3 in Appendix B.7.
The feature maps for other architectures, e.g., Graph Neural
Networks (GNNs) can be derived similarly.
We analyze the Graph Neural Tangent Kernel (GNTK) for a simple GNN architecture in Theorem 3.
We then use Lemma 2 and 3 to characterize the properties of functions learned by an overparameterized neural network.
We precisely characterize the neural networks’ learned functions in the NTK regime via solving the constrained optimization problem corresponding to the min-norm function in NTK feature space with the constraint of ﬁtting the training data.
However, there still remains many technical challenges.
For example, provable extrapolation (exact or asymptotic) is often not achieved with most training data distribution.
Understanding the desirable condition requires signiﬁcant insights into the geometry properties of training data distribution, and how they interact with the solution learned by neural networks.
Our insights and reﬁned analysis shows in R d space, we need to consider the directions of training data.
In graphs, we need to consider, in addition, the graph structure of training data.
We refer readers to detailed proofs for the intuition of data conditions.
Moreover, since NTK corresponds to inﬁnitely wide neural networks, the feature space is of inﬁnite dimension.
The analysis of inﬁnite dimensional spaces poses non-trivial technical challenges too.
Since different theorems have their respective challenges and insights/techniques, we refer the interested readers to the respective proofs for details.
In Lemma 1 (proof in Appendix B.2), Theorem(proof in Appendix B.3), and Theorem 1 (proof in Appendix B.1) we analyze over-parameterized
MLPs.
The proof of Corollary 1 is in Appendix B.4.
In Theorem 3 we analyze Graph Neural
Networks (proof in Appendix B.5).
B P ROOFS B. P ROOF OF T HEOREM
To show neural network outputs f ( x ) converge to a linear function along all directions v , we will analyze the function learned by a neural network on the training set { ( x i , y i ) } ni =1 , by studying the functional representation in the network’s neural tangent kernel RKHS space.
Recall from Section A that in the NTK regime, i.e., networks are inﬁnitely wide, randomly initialized, and trained by gradient descent with inﬁnitesimally small learning rate, the learning dynamics of the neural network is equivalent to that of a kernel regression with respect to its neural tangent kernel.
For any x ∈ R d , the network output is given by
[Equation]
where NTK train is the n × n kernel for training data, (cid:10) φ ( x ) , φ ( x i ) (cid:11) is the kernel value between test data x and training data x i , and Y is training labels.
By Lemma 2, the kernel regression solution is also equivalent to the min-norm solution in the NTK RKHS space that ﬁts all training data
[Equation]
where the representation coefﬁcient β NTK is
[Equation]
The feature map φ ( x ) for a two-layer MLP with ReLU activation is given by Lemma
[Equation]
where w ( k ) ∼ N ( 0 , I ) , with k going to inﬁnity.
c (cid:48) is a constant, and I is the indicator function.
Without loss of generality, we assume the bias term to be 1 .
For simplicity of notations, we denote each data x plus bias term by, i.e., ˆ x = [ x | 1] (Bietti & Mairal, 2019), and assume constant term is 1 .
Given any direction v on the unit sphere, the network outputs for out-of-distribution data x 0 = t v and x = x 0 + h v = (1 + λ ) x 0 , where we introduce the notation of x and λ for convenience, are given by Eqn.
11 and Eqn.
[Equation]
where we have ˆ x 0 =
[ x 0 | 1] and ˆ x = [(1 + λ ) x 0 | 1] .
It follows that f
[Equation] [Equation]
By re-arranging the terms
[Equation]
we get the following equivalent form of the entries:
[Equation]
Similarly, we have
[Equation] [Equation]
Again
[Equation]
let us denote the part of β NTK corresponding to each w by β w .
Moreover, let us denote the part corresponding to Eqn.
18 by β 1 w and the part corresponding to Eqn.
22 by β 2 w .
Then we have
[Equation]
Note that all β w are ﬁnite constants that depend on the training data.
Next, we show that as t → ∞ , each of the terms above converges in O (1 /(cid:15) ) to some constant coefﬁcient β v that depend on the training data and the direction v .
Let us ﬁrst consider Eqn.
We have
[Equation]
Because β 1 w are ﬁnite constants, it follows that
[Equation] [Equation]
where the right hand side is a constant that depends on training data and direction v .
Next, we show the convergence rate for Eqn. 31.
Given error (cid:15) > 0 , because β 1 (cid:62) w [ v | 0] are ﬁnite constants, we need to bound the following by C · (cid:15) for some constant C ,
[Equation]
Observe that the two terms in Eqn.
33 represent the volume of half-(balls) that are orthogonal to vectors [ x 0 | 1] and [ x 0 | 0] .
Hence, Eqn.
33 is the volume of the non-overlapping part of the two (half)balls, which is created by rotating an angle θ along the last coordinate.
By symmetry, Eqn.
33 is linear in θ .
Moreover, the angle θ = arctan
( C/t )
for some constant C .
Hence, it follows that
[Equation]
In the last inequality
[Equation]
we used the fact that arctan x < x for x > 0 .
Hence, O (1 /t ) < (cid:15) implies t
= O (1 /(cid:15) ) as desired.
Next, we consider Eqn.
[Equation]
Let us ﬁrst analyze the convergence of the following:
[Equation]
The convergence to 0 follows from Eqn. 34.
Now we consider the convergence rate.
The angle θ is at most 1 − 11+ λ times of that in Eqn. 34.
Hence, the rate is as follows
[Equation]
Now we get back to Eqn. 24, which simpliﬁes as the following.
[Equation]
We compare the rate of growth of left hand side and the rate of decrease of right hand side (indicators).
[Equation]
[Equation] [Equation]
Hence, the indicators decrease faster, and it follows that Eqn.
24 converges to 0 with rate O ( 1 (cid:15) ) .
Moreover, we can bound w with standard concentration techniques.
Then the proofs for Eqn. 26 and
Eqn.
27 follow similarly.
This completes the proof.
B. P ROOF OF L EMMA
Overview of proof.
To prove exact extrapolation given the conditions on training data, we analyze the function learned by the neural network in a functional form.
The network’s learned function can be precisely characterized by a solution in the network’s neural tangent kernel feature space which has a minimum RKHS norm among functions that can ﬁt all training data, i.e., it corresponds to the optimum of a constrained optimization problem.
We show that the global optimum of this constrained optimization problem, given the conditions on training data, is precisely the same function as the underlying true function.
Setup and preparation.
Let X = { x 1 , ..., x n } and Y = { y 1 , ..., y n } denote the training set input features and their labels.
Let β g ∈ R
d
denote the true parameters/weights for the underlying linear function g , i.e.,
[Equation]
Recall from Section A that in the NTK regime, where networks are inﬁnitely wide, randomly initialized, and trained by gradient descent with inﬁnitesimally small learning rate, the learning dynamics of a neural network is equivalent to that of a kernel regression with respect to its neural tangent kernel.
Moreover, Lemma 2 tells us that this kernel regression solution can be expressed in the functional form in the neural tangent kernel’s feature space.
That is, the function learned by the neural network (in the ntk regime) can be precisely characterized as
[Equation]
where the representation coefﬁcient β NTK is
[Equation]
An inﬁnite-dimensional feature map φ ( x ) for a two-layer ReLU network is described in Lemma
[Equation]
, where w ( k ) ∼ N ( 0 , I ) , with k going to inﬁnity.
c (cid:48) is a constant, and I is the indicator function.
That is, there are inﬁnitely many directions w with Gaussian density, and each direction comes with two features.
Without loss of generality, we can assume the scaling constant to be 1 .
Constrained optimization in NTK feature space.
The representation or weight of the neural network’s learned function in the neural tangent kernel feature space, β NTK , consists of weight vectors for each x · I (cid:16) w ( k ) (cid:62) x ≥ 0 (cid:17) ∈ R d and w ( k ) (cid:62) x · I (cid:16) w ( k ) (cid:62) x ≥ 0 (cid:17) ∈ R .
For simplicity of notation, we will use w to refer to a particular w , without considering the index ( k ) , which does not matter for our purposes.
For any w ∈ R d , we denote by ˆ β w =
( ˆ β (1) w , ..., ˆ
β ( d ) w ) ∈ R
d the weight vectors corresponding to x · I (cid:0) w (cid:62) x ≥ 0 (cid:1) , and denote by ˆ β (cid:48) w ∈ R d the weight for w (cid:62) x · I (cid:0) w (cid:62) x ≥ 0 (cid:1) .
Observe that for any w ∼ N ( 0 , I )
∈ R d , any other vectors in the same direction will activate the same set of x i ∈ R d .
That is, if w (cid:62) x i ≥ 0 for any w ∈ R d , then ( k · w ) (cid:62) x i ≥ 0 for any k > 0 .
Hence, we can reload our notation to combine the effect of weights for w ’s in the same direction.
This enables simpler notations and allows us to change the distribution of w in NTK features from
Gaussian distribution to uniform distribution on the unit sphere.
More precisely, we reload our notation by using β w and β (cid:48) w to denote the combined effect of all weights ( ˆ β (1) k w , ..., ˆ β ( d ) k w ) ∈ R d and ˆ β (cid:48)
k w ∈ R for all k w with k > 0 in the same direction of w .
That is, for each w ∼ Uni(unit sphere) ∈ R d , we deﬁne β ( j ) w as the total effect of weights in the same direction
[Equation]
where u ∼ N ( 0 , I )
Note that to ensure the β w is a well-deﬁned number, here we can work with the polar representation and integrate with respect to an angle.
Then β w is well-deﬁned.
But for simplicity of exposition, we use the plain notation of integral.
Similarly, we deﬁne β (cid:48) w as reloading the notation of
[Equation]
Here, in Eqn.
48 we have an extra term of (cid:107) u (cid:107) (cid:107) w (cid:107) compared to Eqn. 47 because the NTK features that
Eqn.
48 corresponds to, w (cid:62) x · I (cid:0) w (cid:62) x ≥ 0 (cid:1) , has an extra w (cid:62) term.
So we need to take into account the scaling.
This abstraction enables us to make claims on the high-level parameters β w and β (cid:48) w only, which we will show to be sufﬁcient to determine the learned function.
Then we can formulate the constrained optimization problem whose solution gives a functional form of the neural network’s learned function.
We rewrite the min-norm solution in Eqn.
45 as
[Equation]
where the density of w is now uniform on the unit sphere of R d .
Observe that since w is from a uniform distribution
[Equation]
the probability density function P ( w ) is a constant.
This means every x i is activated by half of the w on the unit sphere
[Equation]
which implies we can now write the right hand side of
Eqn.
50 in the form of left hand side,
i
[Equation]
e., integral form.
This allows us to further simplify Eqn.
50 as
[Equation]
where Eqn. 51 follows from the following steps of simpliﬁcation
[Equation]
Claim 1.
Without loss of generality, assume the scaling factor c in NTK feature map φ ( x ) is 1 .
Then the global optimum to the constraint optimization problem Eqn.
49 subject to Eqn. 51, i.e.,
[Equation]
[Equation] [Equation]
satisﬁes β w + β (cid:48)
w · w = 2 β g for all w .
This claim implies the exact extrapolation we want to prove, i
[Equation]
e., f NTK ( x ) =
g ( x ) .
This is because, if our claim holds, then for any x ∈ R d
[Equation]
Thus, it remains to prove Claim 1.
To compute the optimum to the constrained optimization problem
Eqn. 52, we consider the Lagrange multipliers.
It is clear that the objective Eqn. 52 is convex.
Moreover, the constraint Eqn. 53 is afﬁne.
Hence, by KKT, solution that satisﬁes the Lagrange condition will be the global optimum.
We compute the Lagrange multiplier as
[Equation]
Setting the partial derivative of L ( β , λ ) with respect to each variable to zero gives
[Equation]
It is clear that the solution in Claim 1 immediately satisﬁes Eqn. 58.
Hence, it remains to show there exist a set of λ i for i ∈ [ n ] that satisﬁes Eqn. 56 and Eqn. 57.
We can simplify Eqn. 56 as
[Equation]
where c is a constant.
Similarly
[Equation]
we can simplify Eqn. 57 as
[Equation]
Observe that combining Eqn. 59 and Eqn.
60 implies that the constraint Eqn.
60 can be further simpliﬁed as
[Equation]
It remains to show that given the condition on training data, there exists a set of λ
i so that Eqn.and Eqn.
61 are satisﬁed.
Global optimum via the geometry of training data.
Recall that we assume our training data { ( x i , y i ) } ni =1 satisﬁes for any w ∈ R d , there exist d linearly independent { x w i } di =1 ⊂ X , where
X = { x i } ni =1 , so that w (cid:62) x w i ≥ 0 and
− x w i ∈ X for i
= 1 ..
d , e.g., an orthogonal basis of R d and their opposite vectors.
We will show that under this data regime, we have
(a) for any particular w , there indeed exist a set of λ i that can satisfy the constraints Eqn. 59 and Eqn. 61 for this particular w .
(b) For any w 1 and w 2 that activate the exact same set of { x i } , the same set of λ i can satisfy the constraints Eqn. 59 and Eqn.
61 of both w 1 and w 2 .
(c) Whenever we rotate a w 1 to a w 2 so that the set of x i being activated changed, we can still
ﬁnd
λ i that satisfy constraint of both w 1 and w 2 .
Combining (a), (b) and (c) implies there exists a set of λ that satisfy the constraints for all w .
Hence, it remains to show these three claims.
We ﬁrst prove Claim (a).
For each w , we must ﬁnd a set of λ
i so that the following hold.
[Equation]
Here, β g and w are ﬁxed, and w is a vector on the unit sphere.
It is easy to see that β w is then determined by β g and w , and there indeed exists a solution (solving a consistent linear system).
Hence we are left with a linear system with d linear equations
[Equation]
to solve with free variables being λ i so that w activates x i , i.e., w (cid:62) x i ≥ 0 .
Because the training data { ( x i , y i ) } ni =1 satisﬁes for any w , there exist at least d linearly independent x
i that activate w .
This guarantees for any w we must have at least d free variables.
It follows that there must exist solutions
λ i to the linear system.
This proves Claim (a).
Next, we show that (b) for any w 1 and w 2 that activate the exact same set of { x i } , the same set of λ i can satisfy the constraints Eqn. 59 and Eqn.
61 of both w 1 and w 2 .
Because w 1 and w 2 are activated by the same set of x i , this implies
[Equation]
Since λ i already satisfy constraint Eqn. 59 for w 1 , they also satisfy that for w 2 .
Thus, it remains to show that β w 1 + β
[Equation]
w 1 · w 1 = β w
+ β
[Equation]
w 2 · w 1 assuming β
w 1 = β w 2 , β (cid:48)
w 1 = w (cid:62) 1 β w 1 , and β (cid:48)
w 2 = w (cid:62) 2 β w 2 .
This indeed holds because
[Equation] [Equation]
Here, we used the fact that w 1 and w 2 are vectors on the unit sphere.
This proves Claim (b).
Finally, we show (c) that Whenever we rotate a w 1 to a w 2 so that the set of x i being activated changed, we can still
ﬁnd
λ i that satisfy constraint of both w 1 and w 2 .
Suppose we rotate w 1 to w 2 so that w 2 lost activation with x 1 , x 2 , ..., x p which in the set of linearly independent x i ’s being activated by w 1 and their opposite vectors − x i are also in the training set (without loss of generality).
Then w 2 must now also get activated by − x 1 , − x 2 , ..., − x p .
This is because if w (cid:62) 2 x i < 0 , we must have w (cid:62) 2 ( − x i ) > 0 .
Recall that in the proof of Claim (a), we only needed the λ i from linearly independent x i that we used to solve the linear systems, and their opposite as the free variables to solve the linear system of d equations.
Hence, we can set λ to 0 for the other x i while still satisfying the linear system.
Then, suppose there exists λ i that satisfy
[Equation]
where the x i are the linearly independent vectors that activate w 1 with opposite vectors in the training set, which we have proved in (a).
Then we can satisfy the constraint for β w 2 below
[Equation]
by setting ˆ λ i = − λ i for i = 1 ...p .
Indeed, this gives
[Equation]
Thus, we can also ﬁnd λ i that satisfy the constraint for β w 2 .
Here, we do not consider the case where w 2 is parallel with an x i because such w 2 has measure zero.
Note that we can apply this argument iteratively because the ﬂipping the sign always works and will not create any inconsistency.
Moreover, we can show that the constraint for β (cid:48) w 2 is satisﬁed by a similar argument as in proof of
Claim (b).
This follows from the fact that our construction makes β w 1 = β w 2 .
Then we can follow the same argument as in (b) to show that β w 1 + β (cid:48) w 1 · w 1 = β w 2 + β (cid:48) w 2 · w 1 .
This completes the proof of Claim (c).
In summary, combining Claim (a), (b) and (c) gives that Claim 1 holds.
That is, given our training data, the global optimum to the constrained optimization problem of ﬁnding the min-norm solution among functions that ﬁt the training data satisﬁes β w + β (cid:48) w · w = 2 β g .
We also showed that this claim implies exact extrapolation, i.e., the network’s learned function f ( x ) is equal to the true underlying function g ( x ) for all x ∈ R d .
This completes the proof.
B. P ROOF OF T HEOREM
Proof of the asymptotic convergence to extrapolation builds upon our proof of exact extrapolation, i.e., Lemma 1.
The proof idea is that if the training data distribution has support at all directions, when the number of samples n → ∞ , asymptotically the training set will converge to some imaginary training set that satisﬁes the condition for exact extrapolation.
Since if training data are close the neural tangent kernels are also close, the predictions or learned function will converge to a function that achieves perfect extrapolation, that is, the true underlying function.
Asymptotic convergence of data sets.
We ﬁrst show the training data converge to a data set that satisﬁes the exact extrapolation condition in Lemma 1.
Suppose training data { x i } ni =1 are sampled from a distribution whose support contains a connected set S that intersects all directions, i.e., for any non-zero w ∈ R d , there exists k >
so that k w ∈ S .
Let us denote by S the set of datasets that satisfy the condition in Lemma 1.
In fact, we will use a relaxed condition in the proof of Lemma 1 (Lemma 1 in the main text uses a stricter condition for simplicity of exposition).
Given a general dataset X and a dataset S ∈ S of the same size n , let σ ( X , S ) denote a matching of their data points, i.e., σ outputs a sequence of pairs
[Equation]
Let (cid:96) :
R d × R d → R be the l 2 distance that takes in a pair of points.
We then deﬁne the distance between the datasets d ( X , S ) as the minimum sum of l 2 distances of their data points over all possible matching.
[Equation]
We can then deﬁne a “closest distance to perfect dataset” function D ∗ : X → R which maps a dataset
X to the minimum distance of X to any dataset in S
[Equation]
It is easy to see that for any dataset X
= { x i } ni
=1 , D ∗ ( X ) can be bounded by the minimum of the closest distance to perfect dataset D ∗ of sub-datasets of X of size 2 d .
[Equation]
This is because for any S ∈ S , and any S ⊆ S (cid:48) , we must have S (cid:48) ∈ S because a dataset satisﬁes exact extrapolation condition as long as it contains some key points.
Thus, adding more data will not hurt, i.e., for any X 1 ⊆ X 2 , we always have
[Equation]
Now let us denote by X n a random dataset of size n where each x i ∈
X n is sampled from the training distribution.
Recall that our training data { x i } ni =1 are sampled from a distribution whose support contains a connected set S ∗ that intersects all directions, i.e., for any non-zero w ∈ R d , there exists k >
so that k w ∈ S ∗ .
It follows that for a random dataset X 2 d of size 2 d , the probability that D ∗ ( X 2 d ) > (cid:15) happens is less than 1 for any (cid:15) > 0 .
First there must exist S
= { s i } 2 di =1 ∈ S of size 2 d , e.g., orthogonal basis and their opposite vectors.
Observe that if we scale any s i by k > 0 , the resulting dataset is still in S by the deﬁnition of S .
We denote the set of datasets where we are allowed to scale elements of S 0 by S 0 .
It follows that
[Equation]
where we denote the bound of P ( D ∗ ( X 2 d ) > (cid:15) ) by δ < 1 , and the last step follows from
[Equation]
which further follows from the fact that for any s i ∈ S 0 , by the assumption on training distribution, we can always ﬁnd k > 0 so that k s i ∈ S ∗ , a connected set in the support of training distribution.
By the connectivity of support S ∗ , k s i cannot be an isolated point in S ∗ , so for any (cid:15) > 0 , we must have
[Equation]
x
[Table]
k s i
[Equation]
[Table]
, x
[Table]
∗ f
[Equation]
d
[Equation]
>
Hence, we can now apply Eqn. 62 to bound D ∗ ( X n ) .
Given any (cid:15) > 0 , we have
[Equation]
Here δ < 1 .
This implies D ∗ ( X n ) p −→ 0 , i.e.,
[Equation]
Eqn. 63 says as the number of training samples n → ∞ , our training set will converge in probability to a dataset that satisﬁes the requirement for exact extrapolation.
Asymptotic convergence of predictions.
Let NTK ( x , x (cid:48) ) :
R d × R d → R denote the neural tangent kernel for a two-layer ReLU MLP.
It is easy to see that if x → x ∗ , then NTK ( x , · ) →
NTK ( x ∗ , · ) (Arora et al. (2019b)).
Let NTK train denote the n × n kernel matrix for training data.
We have shown that our training set converges to a perfect data set that satisﬁes conditions of exact extrapolation.
Moreover, note that our training set will only have a ﬁnite number of (not increase with n ) x i that are not precisely the same as those in a perfect dataset.
This is because a perfect data only contains a ﬁnite number of key points and the other points can be replaced by any other points while still being a perfect data set.
Thus, we have NTK train → N ∗ , where N ∗ is the n × n NTK matrix for some perfect data set.
Because neural tangent kernel is positive deﬁnite, we have NTK − 1 train → N ∗ − 1 .
Recall that for any x ∈ R d , the prediction of NTK is
[Equation]
where NTK train is the n × n kernel for training data, NTK ( x , x i ) is the kernel value between test data x and training data x i , and Y is training labels.
Similarly, we have ( NTK ( x , x 1 ) , ..., NTK ( x , x n ))
→ ( NTK ( x , x ∗ 1 ) , ..., NTK ( x , x ∗ n )) , where x ∗ i is a perfect data set that our training set converges to.
Combining this with NTK − 1 train → N ∗ − 1 gives
[Equation]
where f NTK is the function learned using our training set, and f ∗ NTK is that learned using a perfect data set, which is equal to the true underlying function g .
This completes the proof.
B. P ROOF OF C OROLLARY
In order for GNN with linear aggregations
[Equation]
to extrapolate in the maximum degree task, it must be able to simulate the underlying function
[Equation]
Because the max function cannot be decomposed as the composition of piece-wise linear functions, the MLP ( K +1) module in GNN must learn a function that is not piece-wise linear over domains outside the training data range.
Since Theorem 1 proves for two-layer overparameterized MLPs, here we also assume MLP ( K +1) is a two-layer overparameterized MLP, although the result can be extended to more layers.
It then follows from Theorem 1 that for any input and label (and thus gradient), MLP ( K +1) will converge to linear functions along directions from the origin.
Hence, there are always domains where the GNN cannot learn a correct target function.
B. P ROOF OF T HEOREM
Our proof applies the similar proof techniques for Lemma 1 and 2 to Graph Neural Networks (GNNs).
This is essentially an analysis of Graph Neural Tangent Kernel (GNTK), i.e., neural tangent kernel of
GNNs.
We ﬁrst deﬁne the simple GNN architecture we will be analyzing, and then present the GNTK for this architecture.
Suppose G =
( V, E ) is an input graph without edge feature, and x u ∈ R d is the node feature of any node u ∈ V .
Let us consider the simple one-layer GNN whose input is G and output is
[Equation]
G
[Equation]
Note that our analysis can be extended to other variants of GNNs, e.g., with non-empty edge features,
ReLU activation, different neighbor aggregation and graph-level pooling architectures.
We analyze this GNN for simplicity of exposition.
Next, let us calculate the feature map of the neural tangent kernel for this GNN.
Recall from Section A that consider a graph neural network f
( θ , G ) :
G → R where θ ∈ R m is the parameters in the network and G ∈ G is the input graph.
Then the neural tangent kernel is
[Equation]
where θ are the inﬁnite-dimensional parameters.
Hence, the gradients with respect to all parameters give a natural feature map.
Let us denote, for any node u , the degree of u by
[Equation]
It then follows from simple computation of derivative that the following is a feature map of the GNTK for Eqn.
[Equation]
where w ( k ) ∼ N ( 0 , I ) , with k going to inﬁnity.
c is a constant, and I is the indicator function.
Next, given training data { ( G i , y i } ni =1 , let us analyze the function learned by GNN through the min-norm solution in the GNTK feature space.
The same proof technique is also used in Lemmaand 2.
Recall the assumption that all graphs have uniform node feature, i.e., the learning task only considers graph structure, but not node feature.
We assume x v = 1 without loss of generality.
Observe that in this case, there are two directions, positive or negative, for one-dimensional Gaussian distribution.
Hence, we can simplify our analysis by combining the effect of linear coefﬁcients for w in the same direction as in Lemma 1 and 2.
Similarly, for any w , let us deﬁne ˆ β w ∈ R as the linear coefﬁcient corresponding to (cid:80)
[Equation]
· h u in RKHS space, and denote by ˆ β (cid:48) w ∈ R the weight for max
[Equation]
Similarly, we can combine the effect of all ˆ β in the same direction as in Lemma and 2.
We deﬁne the combined effect with β w and β (cid:48) w .
This allows us to reason about w with two directions, + and − .
Recall that the underlying reasoning function, maximum degree, is
[Equation]
We formulate the constrained optimization problem, i
[Equation]
e., min-norm solution in GNTK feature space that ﬁts all training data, as
[Equation]
where G i is the i-th training graph and w ∼ N (0 , 1) .
By combining the effect of ˆ β , and taking the derivative of the Lagrange for the constrained optimization problem and setting to zero, we get the global optimum solution satisfy the following constraints.
[Equation]
where c is some constant, λ i are the Lagrange parameters.
Note that here we used the fact that there are two directions +1 and − 1 .
This enables the simpliﬁcation of Lagrange derivative.
For a similar step-by-step derivation of Lagrange, refer to the proof of Lemma 1.
Let us consider the solution β (cid:48)
+ = 1 and β + = β − = β (cid:48)− = 0 .
It is clear that this solution can ﬁt the training data, and thus satisﬁes Eqn. 71.
Moreover, this solution is equivalent to the underlying reasoning function, maximum degree, g ( G ) = max u ∈ G h u .
Hence, it remains to show that, given our training data, there exist λ
i so that the remaining four constraints are satisﬁes for this solution.
Let us rewrite these constraints as a linear systems where the variables are λ
i
[Equation]
By standard theory of linear systems, there exist λ i to solve Eqn.
73 if there are at least four training data
G i whose following vectors linear independent
[Equation]
Here, N max i denotes the number of nodes that achieve the maximum degree in the graph G i , and
N min i denotes the number of nodes that achieve the min degree in the graph G i .
By the assumption of our training data that there are at least four G i ∼ G with linearly independent Eqn. 74.
Hence, our simple GNN learns the underlying function as desired.
This completes the proof.
B. P ROOF OF L EMMA
Let W denote the span of the feature maps of training data x i , i.e.
W = span ( φ ( x 1 ) , φ ( x 2 ) , ..., φ ( x n )) .
Then we can decompose the coordinates of f NTK in the RKHS space, β NTK , into a vector β 0 for the component of f NTK in the span of training data features W , and a vector β 1 for the component in the orthogonal complement W (cid:62) , i.e., β NTK = β 0 + β 1 .
First, note that since f NTK must be able to ﬁt the training data (NTK is a universal kernel as we will discuss next), i.e.,
[Equation]
Thus, we have φ ( x i ) (cid:62) β
= y i .
Then, β 0 is uniquely determined by the kernel regression solution with respect to the neural tangent kernel
[Equation]
where NTK train is the n × n kernel for training data, (cid:10) φ ( x ) , φ ( x i ) (cid:11) is the kernel between test data x and training data x i , and Y is training labels.
The kernel regression solution f NTK is uniquely determined because the neural tangent kernel NTK train is positive deﬁnite assuming no two training data are parallel, which can be enforced with a bias term (Du et al., 2019c).
In any case, the solution is a min-norm by pseudo-inverse.
Moreover, a unique kernel regression solution f NTK that spans the training data features corresponds to a unique representation in the RKHS space β 0 .
Since β 0 and β 1 are orthogonal, we also have the following
[Equation]
This implies the norm of β NTK is at least as large as the norm of any β such that φ ( x i ) (cid:62) β
NTK =
y i .
Moreover, observe that the solution to kernel regression Eqn. 9 is in the feature span of training data, given the kernel matrix for training data is full rank.
[Equation]
Since β 1 is for the component of f NTK in the orthogonal complement of training data feature span, we must have β 1 = 0 .
It follows that β NTK is equivalent to
[Equation]
as desired.
B. P ROOF OF L EMMA
We ﬁrst compute the neural tangent kernel NTK ( x , x (cid:48) ) for a two-layer multi-layer perceptron (MLP) with ReLU activation function, and then show that it can be induced by the feature space φ ( x ) speciﬁed in the lemma so that NTK ( x , x (cid:48) )
= (cid:10) φ ( x ) , φ ( x (cid:48) )
(cid:11) .
Recall that Jacot et al.
(2018) have derived the general framework for computing the neural tangent kernel of a neural network with general architecture and activation function.
This framework is also described in Arora et al.
(2019b)
; Du et al. (2019b), which, in addition, compute the exact kernel formula for convolutional networks and Graph Neural Networks, respectively.
Following the framework in Jacot et al.
and substituting the general activation function σ with ReLU gives the kernel formula for a two-layer MLP with ReLU activation.
This has also been described in several previous works (Du et al., 2019c; Chizat et al., 2019; Bietti & Mairal, 2019).
Below we describe the general framework in Jacot et al.
(2018) and Arora et al.
(2019b).
Let σ denote the activation function.
The neural tangent kernel for an h -layer multi-layer perceptron can be recursively deﬁned via a dynamic programming process.
Here, Σ ( i ) : R d × R d → R for i = 0 ...
h is the covariance for the i -th layer.
[Equation]
The derivative covariance is deﬁned similarly:
[Equation]
Then the neural tangent kernel for an h -layer network is deﬁned as
[Equation]
where we let ˙Σ ( h ) ( x , x (cid:48) )
= 1 for the convenience of notations.
We compute the explict NTK formula for a two-layer MLP with ReLU activation function by following this framework and substituting the general activation function with ReLU, i.e. σ ( a ) = max(0 , a ) =
a · I ( a ≥ 0) and ˙ σ ( a ) =
I ( a ≥ 0) .
[Equation]
So we can get the NTK via Σ (1) ( x , x (cid:48) ) and ˙Σ (1) ( x , x (cid:48) ) , Σ (0) ( x , x (cid:48) ) .
Precisely, Σ
[Equation]
0) ( x , x (cid:48)
[Equation]
, ∧
[Equation]
Σ (1)
[Equation]
x
[Equation]
x (cid:48) )
[Equation]
c ·
E
[Equation] [Equation]
To sample from N ( 0 , ∧ (1) ) , we let L be a decomposition of ∧ (1) , such that ∧ (1) = LL (cid:62) .
Here, we can see that L =
( x , x (cid:48) )
(cid:62) .
Thus, sampling from N ( 0 , ∧ (1) ) is equivalent to ﬁrst sampling w ∼ N ( 0 , I ) , and output
L w = w (cid:62) ( x , x (cid:48) ) .
Then we have the equivalent sampling ( u, v ) =
( w (cid:62) x , w (cid:62) x (cid:48) ) .
It follows that
[Equation]
It follows from the same reasoning that
[Equation]
The neural tangent kernel for a two-layer MLP with ReLU activation is then
[Equation]
Next
[Equation]
we use the kernel formula to compute a feature map for a two-layer MLP with ReLU activation function.
Recall that by deﬁnition a valid feature map must satisfy the following condition
[Equation]
It is easy to see that the way we represent our NTK formula makes it easy to ﬁnd such a decomposition.
The following inﬁnite-dimensional feature map would satisfy the requirement because the inner product of φ ( x ) and φ ( x (cid:48) ) for any x , x (cid:48) would be equivalent to the expected value in NTK, after we integrate with respect to the density function of w .
[Equation]
where w ( k ) ∼ N ( 0 , I ) , with k going to inﬁnity.
c (cid:48) is a constant, and I is the indicator function.
Note that here the density of features of φ ( x ) is determined by the density of w , i.e. Gaussian.
C E XPERIMENTAL D ETAILS
In this section, we describe the model, data and training details for reproducing our experiments.
Our experiments support all of our theoretical claims and insights.
Overview.
We classify our experiments into the following major categories, each of which includes several ablation studies:
1) Learning tasks where the target functions are simple nonlinear functions in various dimensions and training/test distributions: quadratic, cosine, square root, and l1 norm functions, with MLPs with a wide range of hyper-parameters.
This validates our implications on MLPs generally cannot extrapolate in tasks with nonlinear target functions, unless the nonlinear function is directionally linear out -of- distribution.
In the latter case, the extrapolation error is more sensitive to the hyper-parameters.
2) Computation of the R-Squared of MLP’s learned functions along (thousands of) randomly sampled directions in out-of-distribution domain.
This validates Theorem 1 and shows the convergence rate is very high in practice, and often happens immediately out of training range.
3) Learning tasks where the target functions are linear functions with MLPs.
These validate Theorem 2 and Lemma 1, i.e., MLPs can extrapolate if the underlying function is linear under conditions on training distribution
This section includes four ablation studies : a) Training distribution satisfy the conditions in Theorem 2 and cover all directions, and hence, MLPs extrapolate.
b) Training data distribution is restricted in some directions , e.g., restricted to be positive/negative/constant in some feature dimensions.
This shows when training distribution is restrictive in directions, MLPs may fail to extrapolate.
c) Exact extrapolation with inﬁnitely-wide neural networks , i.e., exact computation with neural tangent kernel (NTK) on the data regime in Lemma 1.
This is mainly for theoretical understanding.
4) MLPs with cosine, quadratic, and tanh activation functions.
5) Learning maximum degree of graphs with Graph Neural Networks.
Extrapolation on graph
structure, number
of
nodes, and node features.
To show the role of architecture for extrapolation, we study the following GNN architecture regimes.
a) GNN with graph-level max-pooling and neighbor-level sum-pooling.
By Theorem 3, this GNN architecture extrapolates in max degree with appropriate training data.
b) GNN with graph-level and neighbor-level sum-pooling.
By Corollary 1, this default GNN architecture cannot extrapolate in max degree.
To show the importance of training distribution, i.e., graph structure in training set, we study the following training data regimes.
a) Node features are identical , e.g., 1 .
In such regimes, our learning tasks only consider graph structure.
We consider training sets sampled from various graph structure, and ﬁnd only those satisfy conditions in Theorem 3 enables GNNs with graph-level max-pooling to extrapolate.
b) Node features are spurious and continuous.
This also requires extrapolation on OOD node features.
GNNs with graph-level max-pooling with appropriate training sets also extrapolate to OOD spurious node features.
6) Learning the length of the shortest path between given source and target nodes, with Graph Neural Networks.
Extrapolation on graph structure, number of nodes, and edge weights.
We
study the following regimes
a) Continuous features.
Edge and node features are real values.
This regime requires extrapolating to graphs with edge weights out of training range.
Test graphs are all sampled from the “general graphs” family with a diverse range of structure.
Regarding the type of training graph structure, we consider two schemes.
Both schemes show a U-shape curve of extrapolation error with respect to the sparsity of training graphs.
a) Speciﬁc graph structure: path, cycle, tree, expander, ladder, complete graphs, general graphs, 4-regular graphs.
b) Random graphs with a range of probability p of an edge between any two nodes.
Smaller p samples sparse graphs and large p samples dense graphs. 7)
Physical reasoning of the n -Body problem in the orbit setting with Graph Neural Networks.
We show that GNNs on the original features from previous works fail to extrapolate to unseen masses and distances.
On the other hand, we show extrapolation can be achieved via an improved representation of the input edge features.
We consider the following extrapolation regimes.
a) Extrapolation on the masses of the objects.
b) Extrapolation on the distances between objects.
We consider the following two input representation schemes to compare the effects of how representation helps extrapolation
a) Original features.
Following previous works on solving n -body problem with GNNs, the edge features are simply set to 0 .
b) Improved features.
We show although our edge features do not bring in new information, it helps extrapolation.
C
L EARNING S IMPLE N ON -L INEAR F UNCTIONS
Dataset details.
We consider four tasks where the underlying functions are simple non-linear functions g : R d → R .
Given an input x ∈ R d , the label is computed by y
= g ( x ) for all x .
We consider the following four families of simple functions g .
a) Quadratic functions g ( x ) =
x (cid:62) A x .
In each dataset, we randomly sample A .
In the simplest case where A =
I , g ( x ) =
(cid:80) di
=1 x 2 i .
a) Cosine functions g ( x ) = (cid:80) di =1 cos (2 π · x i ) .
c) Square root functions g ( x ) =
(cid:80) di =1 √ x i .
Here, the domain X of x is restricted to the space in R d with non-negative value in each dimension.
d) L1 norm functions g ( x ) =
| x | 1 = (cid:80) di =1 | x i | .
We sample each dataset of a task by considering the following parameters
a) The shape and support of training, validation, and test data distributions.
i)
Training, validation, and test data are uniformly sampled from a hyper-cube.
Training and validation data are sampled from [ − a, a ] d with a ∈ { 0 . 5 , 1 . 0 } , i.e., each dimension of x ∈ R
d is uniformly sampled from [ − a, a ] .
Test data are sampled from [ − a, a ] d with a ∈ { 2 . 0 , 5 . 0 , 10 . 0 } .
ii) Training and validation data are uniformly sampled from a sphere, where every point has L 2 distance r from the origin.
We sample r from r ∈ { 0 . 5 , 1 . 0 } .
Then, we sample a random Gaussian vector q in R d .
We obtain the training or validation data x = q / (cid:107) q (cid:107) 2 · r .
This corresponds to uniform sampling from the sphere.
Test data are sampled (non-uniformly) from a hyper-ball.
We ﬁrst sample r uniformly from [0 . 0 , 2 . 0] , [0 . 0 , 5 . 0] , and [0 . 0 , 10 . 0] .
Then, we sample a random Gaussian vector q in R d .
We obtain the test data x = q / (cid:107) q (cid:107) 2 · r .
This corresponds to (non-uniform) sampling from a hyper-ball in R d .
b) We sample 20 , 000 training data, 1 , 000 validation data, and 20 , 000 test data.
c)
We sample input dimension d from { 1 , 2 , 8 } .
d) For quadratic functions, we sample the entries of A uniformly from [ − 1 , 1] .
Model and hyperparameter settings.
We consider the multi-layer perceptron (MLP) architecture.
[Equation]
[Table]
We search the following hyper-parameters for MLPs
a) Number of layers d from { 2 , 4 } .
b) Width of each W ( k ) from { 64 , 128 , 512 } .
c) Initialization schemes.
i)
The default initialization in PyTorch.
ii)
The initialization scheme in neural tangent kernel theory, i
e., we sample entries of W k from N (0 , 1) and scale the output after each W ( k ) by (cid:113) 2 d k , where d k is the output dimension of W ( k ) .
d) Activation function σ is set to ReLU.
We train the MLP with the mean squared error (MSE) loss, and Adam and SGD optimizer.
We consider the following hyper-parameters for training
a) Initial learning rate from { 5 e − 2 , 1 e − 2 , 5 e − 3 , 1 e − 3 } .
Learning rate decays 0 .
5 for every 50 epochs b) Batch size from { 32 , 64 , 128 } .
c)
Weight decay is set to 1 e − 5 .
d) Number of epochs is set to 250 .
Test error and model selection.
For each dataset and architecture, training hyper-parameter setting, we perform model selection via validation set, i.e., we report the test error by selecting the epoch where the model achieves the best validation error.
Note that our validation sets always have the same distribution as the training sets.
We train our models with the MSE loss.
Because we sample test data from different ranges, the mean absolute percentage error (MAPE) loss, which scales the error by the actual value, better measures the extrapolation performance
[Equation]
where A i is the actual value and F i is the predicted value.
Hence, in our experiments, we also report the MAPE.
C. R- SQUARED FOR O UT - OF - DISTRIBUTION D IRECTIONS
We perform linear regression to ﬁt the predictions of MLPs along randomly sampled directions in out-of-distribution regions, and compute the R-squared (or R 2 ) for these directions.
This experiment is to validate Theorem 1 and show that the convergence rate (to a linear function) is very high in practice.
Deﬁnition.
R-squared, also known as coefﬁcient of determination, assesses how strong the linear relationship is between input and output variables.
The closer R-squared is to 1 , the stronger the linear relationship is, with 1 being perfectly linear.
Datasets and models.
We perform the R-squared computation on over 2 , 000 combinations of datasets, test/train distributions, and hyper-parameters, e.g., learning rate, batch size, MLP layer, width, initialization.
These are described in Appendix C.1.
Computation.
For each combination of dataset and model hyper-parameters as described in Section C.1, we save the trained MLP model f : R d → R .
For each dataset and model combination, we then randomly sample 5 , 000 directions via Gaussian vectors N ( 0 , I ) .
For each of these directions w , we compute the intersection point x w of direction w and the training data distribution support (speciﬁed by a hyper-sphere or hyper-cube; see Section C.1 for details).
We then collect 100 predictions of the trained MLP f along direction w (assume w is normalized) with
[Equation]
where r is the range of training data distribution support (see Section C.1).
We perform linear regression on these predictions in Eqn. 75, and obtain the R-squared.
Results.
We obtain the R-squared for each combination of dataset, model and training setting, and randomly sampled direction.
For the tasks of learning the simple non-linear functions, we conﬁrm that more than 96% of the R-squared results are above 0 .
This empirically conﬁrms Theoremand shows that the convergence rate is in fact fast in practice.
Along most directions, MLP’s learned function becomes linear immediately out of the training data support.
C. L EARNING L INEAR F UNCTIONS
Dataset details.
We consider the tasks where the underlying functions are linear g : R d → R .
Given an input x ∈ R d , the label is computed by y
= g ( x )
=
A x for all x .
For each dataset, we sample the following parameters
a) We sample 10 , 000 training data, 1 , 000 validation data, and 2 , 000 test data.
b) We sample input dimension d from { 1 , 2 , 32 } .
c)
We sample entries of A uniformly from [ − a, a ] , where we sample a ∈ { 5 . 0 , 10 . 0 } .
d) The shape and support of training, validation, and test data distributions.
i)
Training, validation, and test data are uniformly sampled from a hyper-cube.
Training and validation data are sampled from [ − a, a ] d with a ∈ { 5 . 0 , 10 . 0 } , i
e., each dimension of x ∈ R
d is uniformly sampled from [ − a, a ] .
Test data are sampled from [ − a, a ] d with a ∈ { 20 . 0 , 50 . 0 } .
ii) Training and validation data are uniformly sampled from a sphere, where every point has L 2 distance r from the origin.
We sample r from r ∈ { 5 . 0 , 10 . 0 } .
Then, we sample a random Gaussian vector q in R d .
We obtain the training or validation data x = q / (cid:107) q (cid:107) 2 · r .
This corresponds to uniform sampling from the sphere.
Test data are sampled (non-uniformly) from a hyper-ball.
We ﬁrst sample r uniformly from [0 . 0 , 20 . 0] and [0 . 0 , 50 . 0] , .
Then, we sample a random Gaussian vector q in R d .
We obtain the test data x = q / (cid:107) q (cid:107) 2 · r .
This corresponds to (non-uniform) sampling from a hyper-ball in R d .
e) We perform ablation study on how the training distribution support misses directions.
The test distributions remain the same as in d).
i)
We restrict the ﬁrst dimension of any training data x i to a ﬁxed number 0 .
1 , and randomly sample the remaining dimensions according to d).
ii)
We restrict the ﬁrst k dimensions of any training data x i to be positive.
For input dimension 32 , we only consider the hyper-cube training distribution, where we sample the ﬁrst k dimensions from [0 , a ] and sample the remaining dimensions from [ − a, a ] .
For input dimensions 1 and 2 , we consider both hyper-cube and hyper-sphere training distribution by performing rejection sampling.
For input dimension 2 , we consider k from { 1 , 2 } .
For input dimension 32 , we consider k from { 1 , 16 , 32 } .
iii)
We restrict the ﬁrst k dimensions of any training data x i to be negative.
For input dimension 32 , we only consider the hyper-cube training distribution, where we sample the ﬁrst k dimensions from [ − a, 0] and sample the remaining dimensions from [ − a, a ] .
For input dimensions 1 and 2 , we consider both hyper-cube and hyper-sphere training distribution by performing rejection sampling.
For input dimension 2 , we consider k from { 1 , 2 } .
For input dimension 32 , we consider k from { 1 , 16 , 32 } .
Model and hyperparameter settings.
For the regression task, we search the same set of hyperparameters as those in simple non-linear functions (Section C.1).We report the test error with the same validation procedure as in Section C.1.
Exact computation with neural tangent kernel
Our experiments with MLPs validate Theoremasymptotic extrapolation for neural networks trained in regular regimes.
Here, we also validate
Lemma 1, exact extrapolation with ﬁnite data regime, by training an inﬁnitely-wide neural network.
That is, we directly perform the kernel regression with the neural tangent kernel (NTK).
This experiment is mainly of theoretical interest.
We sample the same test set as in our experiments with MLPs.
For training set, we sample 2 d training examples according to the conditions in Lemma 1.
Speciﬁcally, we ﬁrst sample an orthogonal basis and their opposite vectors X
= { e i , − e i } di =1 .
We then randomly sample 100 orthogonal transform matrices Q via the QR decomposition.
Our training samples are Q X , i.e., multiply each point in X by Q .
This gives 100 training sets with 2 d data points satisfying the condition in Lemma 1.
We perform kernel regression on these training sets using a two-layer neural tangent kernel (NTK).
Our code for exact computation of NTK is adapted from Arora et al.
Novak et al.
We verify that the test losses are all precisely 0 , up to machine precision.
This empirically conﬁrms
Lemma 1.
Note that due to the differences in hyper-parameter settings in different implementations of NTK, to reproduce our experiments and achieve zero test error, the implementation by Arora et al.
(2020) is assumed.
C. MLP S WITH COSINE , QUADRATIC , AND TANH A CTIVATION
This section describes the experimental settings for extrapolation experiments for MLPs with cosine, quadratic, and tanh activation functions.
We train MLPs to learn the following functions:
a) Quadratic function g ( x ) = x (cid:62) A x , where A is a randomly sampled matrix.
b) Cosine function g ( x ) =
(cid:80) di
=1 cos(2 π · x i ) .
c)
Hyperbolic tangent function g ( x ) =
(cid:80) di
=1 tanh( x i ) .
d) Linear function g ( x ) =
W x + b .
Dataset details.
We use 20,000 training, 1,000 validation, and 20,000 test data.
For quadratic, we sample input dimension d from { 1 , 8 } , training and validation data from [ − 1 , 1] d , and test data from [ − 5 , 5] d .
For cosine, we sample input dimension d from { 1 , 2 } , training and validation data from [ − 100 , 100] d , and test data from [ − 200 , 200] d .
For tanh, we sample input dimension d from { 1 , 8 } , training and validation data from [ − 100 , 100] d , and test data from [ − 200 , 200] d .
For linear, we use a subset of datasets from Appendix C.3: 1 and 8 input dimensions with hyper-cube training distributions.
Model and hyperparameter settings.
We use the same hyperparameters from Appendix C.1, except we ﬁx the batch size to 128, as the batch size has minimal impact on models.
MLPs with cos activation is hard to optimize, so we only report models with training MAPE less than 1.
C. M AX D EGREE
Dataset details.
We consider the task of ﬁnding the maximum degree on a graph.
Given any input graph G =
( V, E ) , the label is computed by the underlying function y
= g ( G ) =
max u ∈ G (cid:80) v ∈N ( u ) 1 .
For each dataset, we sample the graphs and node features with the following parameters
a) Graph structure for training and validation sets.
For each dataset, we consider one of the following graph structure: path graphs, cycles, ladder graphs, 4-regular random graphs, complete graphs, random trees, expanders (here we use random graphs with p = 0 .
8 as they are expanders with high probability), and general graphs (random graphs with p = 0 .
1 to 0 .
9 with equal probability for a broad range of graph structure).
We use the networkx library for sampling graphs.
b) Graph structure for test set.
We consider the general graphs (random graphs with p = 0 .
1 to 0 .
9 with equal probability).
c)
The number of vertices of graphs | V | for training and validation sets are sampled uniformly from [20 ... 30] .
The number of vertices of graphs | V | for test set is sampled uniformly from [50 .. 100] .
d)
We consider two schemes for node features.
i)
Identical features.
All nodes in training, validation and set sets have uniform feature 1 . ii)
Spurious (continuous) features.
Node features in training and validation sets are sampled uniformly from [ − 5 . 0 , 5 . 0] 3 , i.e., a three-dimensional vector where each dimension is sampled from [ − 5 . 0 , 5 . 0] .
There are two schemes for test sets, in the ﬁrst case we do not extrapolate node features, so we sample node features uniformly from [ − 5 . 0 , 5 . 0] 3 .
In the second case we extrapolate node features, we sample node features uniformly from [ − 10 . 0 , 10 . 0] 3 .
e)
We sample 5 , 000 graphs for training, 1 , 000 graphs for validation, and 2 , 500 graphs for testing.
Model and hyperparameter settings.
We consider the following Graph Neural Network (GNN) architecture.
Given an input graph G , GNN learns the output h G by ﬁrst iteratively aggregating and transforming the neighbors of all node vectors h ( k ) u
(vector for node u in layer k ), and perform a max or sum-pooling over all node features h u to obtain h G .
Formally, we have h ( k ) u
= (cid:88)
[Equation]
Here, N ( u ) denotes the neighbors of u , K is the number of GNN iterations, and graph-pooling is a hyper-parameter with choices as max or sum.
h (0) u is the input node feature of node u .
We search the following hyper-parameters for GNNs
a) Number of GNN iterations K is 1 .
b) Graph pooling is from max or sum.
c) Width of all MLPs are set to 256 .
d) The number of layers for MLP ( k ) with k = 1 ..K are set to 2 .
The number of layers for MLP ( K +1) is set to 1 .
We train the GNNs with the mean squared error (MSE) loss, and Adam and SGD optimizer.
We search the following hyper-parameters for training
a) Initial learning rate is set to 0 .
b) Batch size is set to 64 .
c)
Weight decay is set to 1 e − 5 .
d) Number of epochs is set to 300 for graphs with continuous node features, and 100 for graphs with uniform node features.
Test error and model selection.
For each dataset and architecture, training hyper-parameter setting, we perform model selection via validation set, i.e., we report the test error by selecting the epoch where the model achieves the best validation error.
Note that our validation sets always have the same distribution as the training sets.
Again, we report the MAPE for test error as in MLPs.
C. S HORTEST P ATH
Dataset details.
We consider the task of ﬁnding the length of the shortest path on a graph, from a given source to target nodes.
Given any graph G =
( V, E ) , the node features, besides regular node features, encode whether a node is source s , and whether a node is target t .
The edge features are a scalar representing the edge weight.
For unweighted graphs, all edge weights are 1 .
Then the label y
= g ( G ) is the length of the shortest path from s to t on G .
For each dataset, we sample the graphs and node, edge features with the following parameters
a) Graph structure for training and validation sets.
For each dataset, we consider one of the following graph structure: path graphs, cycles, ladder graphs, 4-regular random graphs, complete graphs, random trees, expanders (here we use random graphs with p = 0 .
6 which are expanders with high probability), and general graphs (random graphs with p = 0 .
1 to 0 .
9 with equal probability for a broad range of graph structure).
We use the networkx library for sampling graphs.
b) Graph structure for test set.
We consider the general graphs (random graphs with p = 0 .
1 to 0 .
9 with equal probability).
c)
The number of vertices of graphs | V | for training and validation sets are sampled uniformly from [20 ... 40] .
The number of vertices of graphs | V | for test set is sampled uniformly from [50 .. 70] .
d)
We consider the following scheme for node and edge features.
All edges have continuous weights.
Edge weights for training and validation graphs are sampled from [1 . 0 , 5 . 0] .
There are two schemes for test sets, in the ﬁrst case we do not extrapolate edge weights, so we sample edge weights uniformly from [1 . 0 , 5 . 0] .
In the second case we extrapolate edge weights, we sample edge weights uniformly from [1 . 0 , 10 . 0] .
All node features are [ h, I ( v = s ) , I ( v = t )] with h sampled from [ − 5 . 0 , 5 . 0] .
e) After sampling a graph and edge weights, we sample source s and t by randomly sampling s , t and selecting the ﬁrst pair s , s whose shortest path involves at most 3 hops.
This enables us to solve the task using GNNs with 3 iterations.
f)
We sample 10 , 000 graphs for training, 1 , 000 graphs for validation, and 2 , 500 graphs for testing.
We also consider the ablation study of training on random graphs with different p .
We consider p = 0 .
05 .. 1 . 0 and report the test error curve.
The other parameters are the same as described above.
Model and hyperparameter settings.
We consider the following Graph Neural Network (GNN) architecture.
Given an input graph G , GNN learns the output h G by ﬁrst iteratively aggregating and transforming the neighbors of all node vectors h ( k ) u
(vector for node u in layer k ), and perform a max or sum-pooling over all node features h u to obtain h G .
Formally, we have
[Equation]
Here, N ( u ) denotes the neighbors of u , K is the number of GNN iterations, and for neighbor aggregation we run both min and sum.
h (0) u is the input node feature of node u .
w ( u,v ) is the input edge feature of edge ( u, v ) .
We search the following hyper-parameters for GNNs
a) Number of GNN iterations K is set to 3 .
b) Graph pooling is set to min.
c)
Neighobr aggregation is selected from min and sum.
d) Width of all MLPs are set to 256 .
e)
The number of layers for MLP ( k ) with k = 1 ..K are set to 2 .
The number of layers for MLP ( K +1) is set to 1 .
We train the GNNs with the mean squared error (MSE) loss, and Adam and SGD optimizer.
We consider the following hyper-parameters for training
a) Initial learning rate is set to 0 .
b) Batch size is set to 64 .
c)
Weight decay is set to 1 e − 5 .
d) Number of epochs is set to 250 .
We perform the same model selection and validation as in Section C.5.
C. [Equation] ROBLEM
Task description.
The n-body problem asks a neural network to predict how n stars in a physical system evolves according to physics laws.
That is, we train neural networks to predict properties of future states of each star in terms of next frames, e.g., 0 .
001 seconds.
Mathematically, in an n-body system S = { X i } ni =1 , such as solar systems, all n stars { X i } ni =1 exert distance and mass-dependent gravitational forces on each other, so there were n ( n − 1) relations or forces in the system.
Suppose X i at time t is at position x ti and has velocity v ti .
The overall forces a star X i receives from other stars is determined by physics laws as the following
[Equation]
where G is the gravitational constant, and m i is the mass of star X i .
Then acceralation a ti is determined by the net force F ti and the mass of star m i
[Equation]
Suppose the velocity of star X i at time t is v ti .
Then assuming the time steps dt , i.e., difference between time frames, are sufﬁciently small, the velocity at the next time frame t +1 can be approximated by
[Equation]
Given m i , x ti , and v ti , our task asks the neural network to predict v t
+1 i for all stars X i .
In our task, we consider two extrapolation schemes
a) The distances between stars (cid:107) x t i − x t j (cid:107) 2 are out-of-distribution for test set, i.e., different sampling ranges from the training set.
b) The masses of stars m i are out-of-distribution for test set, i.e., different sampling ranges from the training set.
Here, we use a physics engine that we code in Python to simulate and sample the inputs and labels.
We describe the dataset details next.
Dataset details.
We ﬁrst describe the simulation and sampling of our training set.
We samplevideos of n-body system evolution, each with 500 rollout, i.e., time steps.
We consider the orbit situation: there exists a huge center star and several other stars.
We sample the initial states, i.e., position, velocity, masses, acceleration etc according to the following parameters
a) The mass of the center star is 100 kg .
b) The masses of other stars are sampled from [0 . 02 , 9 . 0] kg .
c)
The number of stars is 3 .
d) The initial position of the center star is (0 . 0 , 0 . 0) .
d)
The initial positions x ti of other objects are randomly sampled from all angles, with a distance in [10 . 0 , 100 . 0] m .
e)
The velocity of the center star is 0 .
f)
The velocities of other stars are perpendicular to the gravitational force between the center star and itself.
The scale is precisely determined by physics laws to ensure the initial state is an orbit system.
For each video, after we get the initial states, we continue to rollout the next frames according the physics engine described above.
We perform rejection sampling of the frames to ensure that all pairwise distances of stars in a frame are at least 30 m .
We guarantee that there are 10 , 000 data points in the training set.
The validation set has the same sampling and simultation parameters as the training set.
We have 2 , 500 data points in the validation set.
For test set, we consider two datasets, where we respectively have OOD distances and masses.
We have 5 , 000 data points for each dataset.
a) We sample the distance OOD test set to ensure all pairwise distances of stars in a frame are from [1 .. 20] m , but have in-distribution masses.
b) We sample the mass OOD test set as follows i)
The mass of the center star is 200 kg , i.e., twice of that in the training set.
ii)
The masses of other stars are sampled from [0 . 04 , 18 . 0] kg , compared to [0 . 02 , 9 . 0]
kg in the training set.
iii)
The distances are in-distribution, i.e., same sampling process as training set.
Model and hyperparameter settings.
We consider the following one-iteration Graph Neural
Network (GNN) architecture, a.k.a.
Interaction Networks.
Given a collection of stars S =
{ X i } ni =1 , our GNN runs on a complete graph with nodes being the stars
X i .
GNN learns the star (node) representations by aggregating and transforming the interactions (forces) of all other node vectors
[Equation]
Here, h v is the input feature of node v , including mass, position and velocity
[Equation]
w ( u,v ) is the input edge feature of edge ( u, v ) .
The loss is computed and backpropagated via the
MSE loss of
[Equation]
where o i denotes the output of GNN for node i , and ans i denotes the true label for node i in the next frame.
We search the following hyper-parameters for GNNs
a) Number of GNN iterations is set to 1 .
b) Width of all MLPs are set to 128 .
c)
The number of layers for MLP (1) is set to 4 .
The number of layers for MLP (2) is set to 2 .
d)
We consider two representations of edge/relations w ( i,j ) .
i)
The ﬁrst one is simply 0 .
ii)
The better representation, which makes the underlying target function more linear, is
[Equation]
We train the GNN with the mean squared error (MSE) loss, and Adam optimizer.
We search the following hyper-parameters for training
a) Initial learning rate is set to 0 .
learning rate decays 0 .
5 for every 50 epochs b) Batch size is set to 32 .
c)
Weight decay is set to 1 e − 5 .
d) Number of epochs is set to 2 , 000 .
D V ISUALIZATION AND A DDITIONAL E XPERIMENTAL R ESULTS D. V ISUALIZATION R ESULTS
In this section, we show additional visualization results of the MLP’s learned function out of training distribution (in black color )
v.s. the underlying true function (in grey color ).
We color the predictions in training distribution in blue color .
In general, MLP’s learned functions agree with the underlying true functions in training range (blue).
This is explained by in-distribution generalization arguments.
When out of distribution, the MLP’s learned functions become linear along directions from the origin.
We explain this OOD directional linearity behavior in Theorem 1.
Finally, we show additional experimental results for graph-based reasoning tasks.
D. E XTRA E XPERIMENTAL R ESULTS
In this section, we show additional experimental results.
all fix pos pos pos
[Figure]
APE
[Figure] [Figure]
[Equation]
APE