How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks
Authors: Keyulu Xu
Mozhi Zhang
Jingling Li
Simon Shaolei Du
Ken-Ichi Kawarabayashi
Stefanie Jegelka
Venue: ICLR 2021 Oral (2021)
Keywords: extrapolation
deep learning
out-of-distribution
graph neural networks
deep learning theory

Abstract: 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 sufficiently 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-specific 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.

Published as a conference paper at ICLR 2021
How NEURAL NETWORKS EXTRAPOLATE: FROM FEEDFORWARD TO GRAPH NEURAL NETWORKS
Keyulu Xut, Mozhi Zhang", Jingling Lit, Simon S. Dus, Ken-ichi Kawarabayashi",
Stefanie Jegelkal
"Massachusetts Institute of Technology (MIT)
Unniversity of Maryland
University of Washington
National Institute of Informatics (NII)
keyuluêmit. edu, keyuluxćcsail Liit. edu
ABSTRACT
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 feedfor- ward 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 theoretical explanation,
identify conditions under which MLPs and GNNs extrapolate well. First, we quantify the observation that ReI.U MLPss quickly converge linear functions along any direction from the origin, which implies that ReLU MLPs do not extrapolate most nonlinear functions. But, they can provably learn linear target function when the training distribution is sufficiently "diverse". Second, in connection to analyzing the successes and limitations of GNNs, these results suggest a hypothesis for which we provide theoretical and empirical evi- dence: the success of GNNs in extrapolating algorithmic tasks to new data (e.g.. larger graphs or edge weights) relies on encoding task-specific non-linearities the architecture or features. Our theoretical analysis builds a connection of over- over-paramen parameterizee networks to the neural tangent kernel. Empirically, our theory holds across different training settings.
INTRODUCTION
Humans extrapolate well in many tasks. For example, we can apply arithmetics to arbitrarily large numbers. One may wonder whether 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 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 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 task outside the training distribution. At first 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; Kurkova, 1992). However, neural networks are constrained by gradient descent training (Hardt al., 2016; Soudry et al., 2018). In analysis, we explicitly consider such implicit bias through the analogy of the training dynamics of over over-parameterized neural networks and kernel regression via the neural langent kernel (NTK) (Jacot et al., 2018).

---

Published as a conference paper at ICLR 2021
Figure I: How ReLU MLPs extrapolate. We train MLPs to learn nonlinear functions (grey) and plot their predictions both within (blue) and outside (black) the training distribution. MLPs converge quickly to linear functions outside the training data range along directions from the origin (Theorem 1). Hence, MLPs do not extrapolate well in most nonlinear tasks. But, with appropriate training data, MLPs can provably extrapolate linear target functions (Theorem 2).
f: hard for extrapolation
GNN Architectures
DP Algorithm
fG)
(Target Function)
A
(M!""(M"" (ag- bi-- wiy 1))
MLP has to learn non-linear steps
d[Rl[o] min,
G: input graph
g: easier for extrapolation
h: input transform
A
min- MLP (M 1./--1 (0)
dlk 11l 10(2.20
h
MG)
g
AG)
MLP Jeams linear steps
(a) Network architecture
(b) Input representation
Figure 2: How GNNs extrapolate. Since MLPs can extrapolate well when learning linear functions, we hypothesize that GNNs can extrapolate well dynamic programming (DP) tasks if we encode appropriate oonaaiietsss in the architecture (left) and input representation (right; through domain knowledge or representation learning). The encoded non-lincarities may not be necessary for interpolation, 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 wer-parameterizer 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 finitely many linear regions (Arora et al., 2018; Hanin & Rolnick, 2019; Hein et al., 2019). While having finitely many linear regions implies ReLU MLPs eventually become linear, it does not say whether MLPs will learn the correct target function close 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 sufficiently "diverse". To our knowledge, our results are the first extrapolation results of this kind for feedforward neural networks.
We then relate insights into feedforward neural networks 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 computation structure aligned with GNNs (Xu et al., 2020). DP updates can often be decomposed into nonlinear and linear steps. Hence, hypothesize that GNNs trained by GD can extrapolate well in DP task, if we encode appropriate non-linearities in the architecture and inpul representation (Figure 2). Importantly, encoding non-line linearinasssies may be unnecessary for GNNs 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 simplified case using Graph NTK (Du et al., 2019b). Empirically, we validate the hypothesis on three DP tasks: max degree, shortest paths, and ro-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 encoding task-specific non-line non-lincarities, which requires domain knowledge or extensive model search. From a broader standpoint, our insights go beyond GNNs and apply broadly other neural networks.
2

---

Published as a conference paper at ICLR 2021
To summarize, 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, prove that ReLU MLPs can extrapolate well in linear tasks, leading to a hypothesis: a neural network can extrapolate well when appropriate non- linearities are encoded into the architecture and features. We prove this hypothesis for simplified case and provide empirical support for more general settings.
1.1 RELATED WORK
Early works show example tasks where MLPs do not extrapolate well, e.g. learning simple polynomi- als (Barnard & Wessels, 1992; Haley & Soloway, 1992). We instead show 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 field 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 gigh-dimensiona settings of the NTK regime, MLP is asymptotically at most linear predictor in certain scaling limits (Ba 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 al., 2018; Saxton et al., 2019; Velickovic al., 2020). We the first to theoretically study GNN extrapolation, and we complete the notion of extrapolation to include unseen features and structures.
2 PRELIMINARIES
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 with training set W:)) DD where Y: (xx,) 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 domain X' that is larger than the support of the training distribution D. We will say a model extrapolates well if it has small extrapolation error.
X Definition 1. (Extrapolation error). Let f
be a model trained on {(2,, y1) } underlying function g X + R.. Let P be distribution over X' 'D and let E IR x IR + R be a loss function. We define the extrapolation error of f as (E-p((((<)) g(x))).
C D with
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 mod- ules (Battaglia al., 2018; Xu et al., 2019). Let = (V. E) be graph. Each node a E has feature vector Iu, and each edge (u, v) E has feature vector (...))) GNNs recursively compute node representations (k) at iteration & (Gilmer et al., 2017; Xu et al., 2018). Initially, ( &w. For k 1..K, GNNs update 1e by aggregating the neighbor representations. We can optionally compute a graph representation he by aggregating the final node representations. That is,
h (k)
Σ
(LP($)
(k-1) W(en))) hG MLP(K (- 1) (k-1) h -.-0=1 (11
(Σ Ex(R)
(1)
(EN((a
KEG
The final output is the graph representation hG or final node representations (K) depending on the task. We refer to the neighbor aggregation step for (a) aggregation and the pooling step in he as readout. Previous works typically use aun-aggregation and sum-readout (Battaglia et al., 2018). Our results indicate why replacing them may help extrapolation (Section 4).
3

---

Published as a conference paper at ICLR 2021
Figure 3: Conditions for ReLU MLPs to extrapolate well. We train MLPs to learn linear functions (grey) with different training distributions (blue) and plot dittr--distribution predictions (black). Following Theorem 2, MLPs extrapolate well when the training distribution (blue) has support in all directions (first panel), but not otherwise: in the two middle panels, some dimensions of the training data are constrained to be positíve (red arrows): in the last panel, one dimension is fixed constant.
10
10"
MAPE
10
MAPE
10"
10"
10"
10
10 0
quadratic
linear
fixt
neg1
neg16
neg32
(a) Different target functions
(b) Different training distributions for linear target
Figure 4: Extrapolation performance of ReLU MLPs. We plot the distributions of MAPE (mean absolute percentage error) of MIPs trained with various hyperparameters (depth, width, learning rate, batch size). (a) Learning different target functions; (b) Different training distributions for learning linear target functions: "all" covers all directions, "fixl" " has dimension fixed to constant, and "negd" has d dimensions constrained to negative values. ReLU MLPs generally do not extrapolate well unless the target function linear along each direction (Figure 4a), and extrapolate linear target functions if the training distribution covers sufficiently many directions (Figure 4b).
3 How FEEDFORWARD NEURAL NETWORKS EXTRAPOLATE
Feedforward networks are the simplest neural networks and building blocks of more complex architectures such as GNNs, so we first study how they extrapolate when trained by GD. Throughout the paper, assume ReLU activation. Section 3.3 contains preliminary results for other activations.
3.1 LINEAR EXTRAPOLATION BEHAVIOR OF RELU MLPS
By architecture, ReLU networks learn piecewise linear functions, but what do these regions precisely look like outside the support of the training data? Figure 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 coefficient of determination (R²2 is always greater than 0.99 (Appendix C.2). That is, ReLU MLPs "lincarize" 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 specific neural tangent kernel, under a set of assumptions called the "NTK regime" (Jacot et al., 2018). We provide an informal definition here; for further details, we refer the readers to Jacot et al. (2018) and Appendix
Definition 2. (Informal) A neural network trained the NTK regime is infinitely wide, randomly initialized with certain scaling, and trained by GD with infinitesimal steps.
Prior works analyze optimization and in-distribution generalization of wver-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 to from the origin, the prediction of two-layer ReL.U MIP quickly converges to linear
4

---

Published as a conference paper at ICLR 2021
function with rate ). The linear coefficients 3, 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 + R with squared loss in the NTK regime. For any direction v RR, ler Do tv. As co, fxo + hwv) - f(o) + Se h for any h O, where Bu is a constant linear coefficient. Moreover, given E > 0. t = ) we have {(mo (Botth)) hw) f(mo) -Byl <.
ReL.U networks have finitely many linear regions (Arora et al., 2018; Hanin Rolnick, 2019), hence their predictions evensually become linear. In contrast, Theorem is a more fine-grained analysis of how MIPs extrapolate and provides a convergence rate. While Theorem assumes two-layer networks in the NTK regime, experiments confirm that the linear extrapolation behavior happens across networks with different depths, widths, learning rates, and batch sizes (Appendix C.I and C.2). Our proof technique potentially also extends to deeper networks.
Theorem 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 (2 cos(27 (1) (cos), and E i= V (a) (sqrt), where (+)) is the i-th dimension of æ. With suitable hyperparameters, MLPs extrapolate the LI norm correctly, which satisfies the directional linearity condition.
shows ReLU MLPs do not extrapolate target functions such as x T Ax (quadratic), Ei= cos
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 first, in Section 4 this insight will help explain extrapolation properties of GNNs in non-linear practical tasks. Before that, we first theoretically analyze when MLPs extrapolate well.
3.2 WHEN RELU MLPs PROVABLY EXTRAPOLATE WELL
Figure 4a shows that MLPs can extrapolate well when the target function is linear. However, this is not always true. 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. Lemma states that over- peer--arramterizeed MIPs can learn linear target function with only 2d examples.
Lemma 1. Let (xx) BT x be the target function for /3 E Rd. Suppose {2; l2 contains an orthogonal basis { and {- If we train two-layer ReLU MLP f on {(x,, v:) with squared loss in the NTK regime, then f f(x = B for all x: E Rdd
Lemma is mainly of theoretical interest, as the 2d examples need to be carefully chosen. Theorem 2 builds on Lemma and identifies a more practical condition for successful extrapolation: if the support of the training distribution covers all directions (e.g.. hypercube that covers the origin), the MLP converges to a linear target function with sufficient training data.
Theorem 2. (Conditions for extrupolation). Let g(x) = BT x be the target function for B E Rdd Suppose (x:) Pi 1 is sampled from distribution whose support 'D contains a connected subset S, where for any non-zero w RRd,, there exists > so that kw S. If we train a two-layer ReLU
MLP f: Rd + R on (x),, with squared loss in the NTK regime, f(x) P, 3 x as n + co.
geometry of training data affects extrapolation. The condition in Theorem 2 formalizes 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 satisfied ("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).
Experiments:
Relating to previous works, Theorem 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 image collection. Therefore, the condition for Theorem no longer holds, and the model may extrapolate incorrectly. Theorem 2 is also analogous to an identifiability condition for linear models, but stricter. We can uniquely identify linear function the training data has full (feature) rank. MLPs are more expressive, so identifying the linear target function requires additional constraints.
5

---

Published as a conference paper at ICLR 2021
10"
10'
10²
MAPE
MAPE
MAPE
10"
10"
10"
10"
10"
quad
600
tanh
Insar
10
quid
tanh
Inher
quad2
quad4
cos
(a) tanh activation
(b) cosine activation
(c) quadratic activation
Figure 5: Extrapolation performance of MLPs with other activation. MLPs can extrapolate well when the activation is "similar" to the target function. When learning quadratic with quadratic activation, 2-layer networks (quad-2) extrapolate well, but 4-layer networks (quad-4) do not.
To summarize, 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.3 MLPs WITH OTHER ACTIVATION FUNCTIONS
Before moving on to GNNs, we complete the picture of MLPs with experiments on other activation functions: tanh (xx tanh(1) (x). cosine ((r) = cos(x)) (Lapedes & Farber, 1987; McCaughan, 1997; Sopena Alquezar, 1994), and quadratic (x) x² (Du & Lee, 2018; Livni 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, 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 sufficiently wide. When learning a quadratic function with quadratic activation, only two-layer networks extrapolate well 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.
How GRAPH NEURAL NETWORKS EXTRAPOLATE
Above, 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 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.1 HYPOTHESIS: LINEAR ALGORITHMIC ALIGNMENT HELPS EXTRAPOLATION
We start with example: training GNNs to solve the shortest path problem. For this task, prior works observe that modified GNN architecture with min-aggregation can generalize to graphs larger than those in the training set (Battaglia et al., 2018; Velickovic et al., 2020):
(k) min MLP(k) (h (k-1), (k-1)) w(..))
vEN(uw
We first provide intuitive explanation (Figure 2a). Shortest path can be solved by the Bellman-Ford (BF) algorithm (Bellman, 1958) with the following update:
[u]][[]] min dlk 1][r] + w(v u),
vEN(u)
where w(u u) is the weight edge (u, u). and [u] is the shortest distance to node u within steps. The two equations can be easily aligned: GNNs simulate the BF algorithm if its MLP modules learn a linear function dik - 1[[b] [[]] w(v, u). Since MLPs can extrapolate linear tasks, this "alignment" may explain why min-aggregation GNNs extrapolate well in this task.

---

Published as a conference paper at ICLR 2021
For comparison, we reason why would not expect GNNs with the more commonly used sum -ana-gggegttion (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 suggests that they will not extrapolate most nonlinear functions outside the training support.
We can generalize the above intuition 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. (Xu al., 2020).
Definition 3. Dynamic programming (DP) is recursive procedure with updates
Answer(k [A]] DP- DP-UUpd (Answer(K: ][[5') 8' 1...m),
where Answer [s] is the solution sub teob--roblem indexed by iteration & and state s, and DP- DP-Update is a task-specific update function that solves the sub-problem based on the previous iteration.
From a broader standpoint, we hypothesize that: if we encode appropriate non-linearitics into the model architecture and input representations 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 neural network with m MLP modules. Suppose there exist m linear functions {9: 1m that by replacing N MLP modules with g: 's, simulates f. Given c > O, there exists {(2,, f fx;))] C
D S X so that A trained on {Ix; fax;)) by GD with squared loss learns f with [l fll <
Our hypothesis builds on the algorithmic alignment framework of (Xu et al., 2020), which states that 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 to encode log-and-exp transforms 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 target function that our model can extrapolate well. We can obtain via specialized features or feature transforms using domain knowledge (Lample & Charton, 2020; Webb al., 2020), or via representation learning (e.g., BERT) with unlabeled out-of-distribution data in X' \ 'D (Chen et al., 2020; Devlin al., 2019; Hu al., 2020; Mikolov al., 2013b; Peters et al., 2018). This brings new perspective of how representations help extrapolation in various application areas. For example, in natural language processing, pretrained representations (Mikolov 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 finance, identifying the right "factors" features is crucial for deep leaming models as the financial markets may frequently be in extrapolation regimes (Banz, 1981; Fama & French, 1993; Ross, 1976).
g h into a feature embedding h and "simpler"
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, provide theoretical and empirical support for our hypothesis.
4.2 THEORETICAL AND EMPIRICAL SUPPORT
We validate hypothesis three DP tasks: 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 simple yet fundamental task: learning the max degree of graph, a special case of DP with one iteration. As a corollary of Theorem I, the commonly used sum-based GNN (Eqn. 1) cannot extrapolate well (proof in Appendix B.4).
7

---

Published as a conference paper at ICLR 2021
sum pooling
maximin pooling
original features
improved features
70.6
11.0
43.8
0.0
6.1
0.0
1.1
0.7
extrapolate interpolate max degree
extrapolate Interpolate shortast path
extrapolate dist
extrapolate mass n-body problem
interpolate
(a) Importance of architecture
(b) Importance of representation
Figure 6: Extrapolation for algorithmic tasks. Each column indicates the task and mean average percentage error (MAPE). Encoding appropriate non- mi-liitarity in the architecture or representation less helpful for interpolation, but significantly improves extrapolation. Left: max degree and shortest path, GNNs that appropriately encode max/min extrapolate well, but GNNs with sum- sum-pooling do not. Right: With improved input representation, GNNs extrapolate better for the n-body problem.
33.6
complete
0.1
cycle
13,9
expander
ladder
11.6
4regular
ladder
generrll
4regular
12.5
expander
2.4
path
94.5
complete
8.9
(a) Max degree with GNNs that encode
(b) Shortest path with GNNs that encode min
Figure 7: Importance of the training graph structure. Rows indicate the graph structure covered by the training set and the extrapolation error (MAPE). In max degree, GNNs with readout extrapolate well the max/min degrees of the training graphs are not restricted (Theorem 3). In shortest path, the extrapolation errors of min GNNs follow a U-shape in the sparsity of the training graphs. More results may be found in Appendix D.2.
Corollary 1. GNNs with sum sum-aggreen agggagggenntion and sum- nea-reddout do not extrapolate well in Max Degree.
To achieve linear algorithmic aligmment, we can encode the only non-linearity, the max function, in the readout. Theorem 3 confirms that a GNN with max max-rtadout can extrapolate well in this task.
Theorem 3. (Extrapolation with GNNs). Assume all nodes have the same feature. Let and g' be the max/min degree function, respectively. Let g(.)) be the training set. If Al((G)) g'.),,
-Nuaxx, "ymaa,,'((C (G.) N ) Jizz spans IR, where Nmax and Nmin are the number of nodes that have max/min degree G,, then one-layer max-readout GNNs trained on (G with squared loss in the NTK regime learn
Theorem 3 does not follow immediately from Theorem 2, because MLP modules in GNNs only receive indirect supervision. We analyze the Graph NTK (Du al., 2019b) 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 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 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 mex-readout are better than GNNs with sum-readout (Figure 6a), confirming Corollary 1 and Theorem 3. For shortest path, GNNs with min-readout and aggregagggtto are better than GNNs with sum-readout (Figure 6a).
Experiments confirm the importance of training graphs structure (Figure 7). Interestingly, the two tasks favor different graph structure. For max degree, as Theorem predicts, GNNs extrapolate well when trained trees, complete graphs, expanders, and general graphs, and extrapolation errors
8

---

Published as a conference paper at ICLR 2021
higher when trained 4-regular, eycles, 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 dense graphs likely learn degenerative solutions.
representations that help extrapolation. Finally, we show a good imput representa- tion helps extrapolation. We study the n-body problem (Battaglia al., 2016; Watters al., 2017) (Appendix C.7), that is, predicting the time evolution of n objects in a gravitational system. Following previous work, the input 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 's mass mus position _ and velocity v!!' time t. The edge features are set
Experiments:
We train GNNs predict the velocity of each object time + 1. The true velocity f(G; u) for object u is approximately
f(G:u) u) vi + a', dt, a' c. Σ
me
x
(5)
effu
where C' is a constant. To learn f, the MLP modules need to learn nonlinear function. Therefore, GNNs do not extrapolate well to unseen masses or distances ("original features" in Figure 6b). We instead use improved representation h(G) encode non-linearity. At time t. we transform the (t)
edge features of (u, v) from zero to w (a.w) my (a! ( { /|| xx ( 113. 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).
CONNECTIONS TO OTHER OUT-OF-DISTRIBUTION SETTINGS
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 the training data range.
Domain adaptation studies generalization to a specific target domain (Ben-David et al., 2010; Blitzer al., 2008; Mansour 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- non-lincarities and may mitigate extrapolation by matching the target and source distributions, though the correctness of the learned mapping depends on the label distribution (Zhao al., 2019).
Self-supervised learning on a large amount of unlabeled data can learn useful non non- lineliiesssiee 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 outto--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 specific 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 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 al., 2020; Sinha et al., 2018; Staib & Jegelka, 2019). We instead look at more global perturbations. Still, one would expect that modifications that help extrapolation in general also improve robustness local perturbations.
CONCLUSION
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 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.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGMENTS
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 Chevron-Mi MI'T Energy Fellowship. This research was also supported by JST ERATO JPMJER 1201 and JSPS Kakenhi JP1SHO5291. MZ. was supported by ODNI, IARPA, via the BETTER Program contract 2-19-188 1111999111111 19951 6000055. The views, opinions, and/or findings contained in this article are those of the author and should not interpreted as representing the official 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.
REFERENCES
Kartik Ahuja, Jun Wang, Amit Dhurandhar, Karthikeyan Shanmugam, and Kush R. Varshney. Empirical or invariant risk minimization? a sample complexity perspective. In International Conference Learning Representations, 2021.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems, pp. 6155-6166, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via parameterization. In International Conference on Machine Learning, pp. 242-252, 2019b.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprini arXiv: 1907,0090000000 0.08999, 2019.
Raman Arora, Amitabh Basu, Poorya Mianjy, and Anirbit Mukherjee. Understanding deep neural networks with rectified linear units. In International Conference on Learning Representations, 2018.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu. Zhiyuan Li, Russ Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, pp- 8139-8141 -88888888888118 2019b.
Sanjeev Arora, Simon S. Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets small-data tasks. In International Conference on Learning Representations, 2020.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Generalization of two-layer neural networks: An asymptotic viewpoint. In International Conference on Learning Representations, 2020.
Rolf W Banz. The relationship between return and market value of common stocks. Journal of financial economics, 9()::3118, 1981.
Etienne Barnard and LFA Wessels. Extrapolation and interpolation in neural network classifiers. IEEE Control Systems Magazine, 12(5):50-53, 1992.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. In Advances in Neural Information Processing Systems, pp. 4502-4510, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 1 2018.
10

---

Published as a conference paper at ICLR 2021
Richard Bellman. On routing problem. Quarterly of applied mathematics, 16(1):87-90, 1958.
Richard Bellman. Dynamic programming. Science, 153(3731):34-377 1966.
Shai Ben-David. John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79(1-2)1151-175, 2010.
Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. In Advances in Neural Information Processing Systems, pp. 11873-12884, 2019.
John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman. Learning bounds for domain adaptation. In Advances in neural information processing systems, pp. 129-136, 2008.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, pp. 10835-10845, 2019.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference Machine Learning, 2020.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in Neural Information Processing Systems, 31 13:3366-- 3046, 2018.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems, pp. 2933-2943, 2019.
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Veličković. Principal neighbourhood aggregation for graph nets. Advances Neural Information Processing Systems, 2020.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of for the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171-4186, 2019.
Simon Du Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima of deep neural networks. In International Conference on Machine Leurning, pp. 1675- 1685, 2019a.
Simon S. Du and Jason D. Lee. On the power of over parametrization in neural networks with quadratic activation. In International Conference Machine Learning, 2018
Simon Du, Kangcheng Hou, Russ R Salakhutdinov, Barnabas Poczos, Ruosong Wang, and Keyulu
Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. In Advances in Neural Information Processing Systems, pp. 5724-5734 2019b.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes peeetetrametr neural networks, In International Conference Learning Representations, 2019c.
Eugene F Fama and Kenneth R French. Common risk factors in the returns on stocks and bonds. Journal of financial economics, 33(3):3-56, 1993.
Ken-Ichi Funahashi. On the approximate realization of continuous mappings by neural networks. Neural networks, 2(3):183-192, 1989.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain. Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-ad adversarill training of neural networks. The Journal of Machine Learning Research, (701):2096-200 2030, 2016.
11

---

Published as a conference paper at ICLR 2021
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. arXiv preprint arXiv:1904.121D90 2019.
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference Machine Learning, pp. 1273-1272, 2017.
Joel Goh and Melvyn Sim. Distributionally robust optimization and its tractable approximations. Operations research, 58:902-9113 2010.
Pamela Haley and DONALD Soloway. Extrapolation limitations of multilayer feedforward neural networks. In International Joint Conference on Neural Networks, volume 4, pp. 25-30. IEEE, 1992.
Boris Hanin and David Rolnick. Complexity of linear regions in deep networks. In International Conference Machine Learning, pp. 2596-2604, 2019.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning, pp. 1225-1234, 2016.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 9729-9738, 2020.
Matthias Hein, Maksym Andriushchenko, and Julian Bitterwolf. Why relu networks yield high- confidence predictions far away from the training data and how to mitigate the problem. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 41-50, 2019.
Dan Hendrycks, Xiaoyuan Liu, Eric Wallace, Adam Dziedzic, Rishabh Krishnan, and Dawn Song. Pretrained transformers improve out-of-distribttion robustness. In Association for Computational Linguistics, 2020.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359-366, 1989.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitmik, Percy Liang, Vijay Pande, and Jure Leskovee. Strategies for pre-training graph neural networks. In International Conference Learning Representations, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances neural information processing systems, pp. 8571-8580, 2018.
Michael Janner, Sergey Levine, William T. Freeman, Joshua B. Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-centric models. International Conference on Learning Representations, 2019.
Justin Johnson, Bharath Hariharan, Laurens der Maaten, Judy Hoffman, Li Fei-Fei, C Lawrence Zitmick, and Ross Girshick. Inferring and executing programs for visual reason- ing. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1989929998, 2017.
Vera Kurkova. Kolmogorov's theorem and multilayer neural networks. Neural networks, 5(3): 501-506, 1992.
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
Guillaume Lample and Franois Charton. Deep learning for symbolic mathematies. In International Conference on Learning Representations, 2020.
Alan Lapedes and Robert Farber. Nonlinear signal processing using neural networks: Prediction and system modelling. Technical report, 1987.
12

---

Published as a conference paper at ICLR 2021
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient
descent on structured data. In Advances in Neural Information Processing Systems, pp. 1157-8166, 2018.
Roi Livni, Shai Shalev-Shwartz, and Ohad Shamir. On the computational efficiency of training neural networks. In Advances in neural information processing systems, pp. 555-863, 2014.
Andreas Madsen and Alexander Rosenberg Johansen. Neural arithmetic units. In International Conference on Learning Representations, 2020.
Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly. Gradient Descent Quantizes ReLU Network Features. arXiv e-prints, art. arXiv:1803. met::1883.0833 003608 March 2018.
Yishay Mansour, Mehryar Mohri, and Afshin Rostamizadeh. Domain adaptation: Learning bounds and algorithms. In Conference on Learning Theory, 2009.
Jiayuan Mao, Chuang Gan, Pushmeet Kohli, Joshua B. Tenenbaum, and Jiajun Wu. The neuro- symbolic concept leamner Interpreting scenes, words, and sentences from natural supervision. In International Conference Learning Representations, 2019.
David B McCaughan. On the properties of periodic perceptrons. In International Conference on Neural Networks, 1997.
Tomas Mikolov, Quoe V Le, and Ilya Sutskever. Exploiting similarities among languages for machine translation. arXiv preprint arXiv:
2013a.
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. In Advances in Neural Information Processing Systems, 2013b.
Atsushi Nitanda and Taiji Suzuki. Optimal rates for averaged stochastic gradient descent under neural tangent kernel regime. In International Conference on Learning Representations, 2021.
Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations, 2020.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 78(5):477-0012, 2016.
Matthew Peters, Mark Neumann, Mohit lyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations. In Proceedings of the 2018 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long Papers), pp. 2227-22333333 2018.
Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. The Journal of Machine Learning Research, 19(1:1300-1132, 2018.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk mini- mization. In International Conference on Learning Representations, 2021.
Stephen A Ross. The arbitrage theory of capital asset pricing. Journal of Economic Theory, 13(3): 341-360, 1976.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
Adam Santoro, Felix Hill, David Barrett, Ari Morcos, and Timothy Lillicrap. Measuring abstract reasoning in neural networks. In International Conference on Machine Learning, pp. 4477-4486,
2018.
Pedro Savarese, Itay Evron, Daniel Soudry, and Nathan Srebro. How do infinite width bounded norm networks look in function space? In Conference Learning Theory, 2019.
13

---

Published as a conference paper at ICLR 2021
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models. In International Conference on Learning Representations. 2019.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions Neural Networks, 20(1):61-80, 2009.
Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness with principled adversarial training. In International Conference on Learning Representations, 2018.
Mei Song, Andrea Montanari, and P Nguyen. A mean field view of the landscape of two-layers neural networks. Proceedings of the National Academy of Sciences, 115:E76655 2018.
JM Sopena and R Alquezar. Improvement of learning in recurrent networks by substituting the sigmoid activation function. In International Conference on Artificial Neural Networks, pp. 417-420. Springer, 1994.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018
Matthew Staib and Stefanie Jegelka. Distributionally robust optimization and generalization in kernel methods. In Advances Neural Information Processing Systems, pp. 9134-9144, 2019.
Andrew Trask, Felix Hill, Scott E. Reed, Jack Rae, Chris Dyer, and Phil Blunsom. Neural arithmetic logic units. Advances in Neural Information Processing Systems, pp. 8035-8044, 2018.
Leslie G Valiant. A theory of the learnable. In Proceedings of the sixteenth annual ACM symposium on Theory of computing, pp. 436-445, ACM, 1984.
Vladimir Vapnik. The nature of statistical learning theory. Springer science & business media, 2013.
Petar Velickovic, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execution of graph algorithms. In International Conference Learning Representations, 2020.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning a physics simulator from video. In Advances in neural information processing systems, pp. 4539-4547, 2017.
Taylor Webb, Zachary Dulberg, Steven Frankland, Alexander Petrov, Randall OReilly, and Jonathan Cohen. Learning representations that support extrapolation. In International Conference on Machine Learning, pp. 10136-10146. PMLR, 2020.
Francis Williams, Matthew Trager, Daniele Panozzo, Claudio Silva, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. In Advances in Neural Information Processing Systems, pp. 8376-8385, 2019.
Shijie Wu and Mark Dredze. Beto, bentz, becas: The surprising cross-lingual effectiveness of bert. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 833-844 2019.
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning graphs with jumping knowledge networks. In International Conference on Machine Learning, pp. 5453-5462, 2018.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations, 2019.
Keyulu Xu, Jingling Li, Mozhi Zhang, Simon S. Du. Ken ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? In International Conference on Learning Representations, 2020. URL https: / wwwp/eeeee.e.ereeeeee. net forum?i www.://meem.......rer..........e
14

---

Published as a conference paper at ICLR 2021
Kexin Yi, Jiajun Wu, Chuang Gan, Antonio Torralba, Pushmeet Kohli, and Josh Tenenbaum. Neural- symbolic vqa: Disentangling reasoning from vision and language understanding. Advances in Neural Information Processing Systems, pp. 1031-1042. 2018.
Michelle Yuan, Mozhi Zhang, Benjamin Van Durme, Leah Findlater, and Jordan Boyd-Graber. Interactive refinement of cross-lingual word embeddings. In Proceedings of Empirical Methods in Natural Language Processing, 2020.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
Mozhi Zhang, Keyulu Xu, Ken-ichi Kawarabayashi, Stefanie Jegelka, and Jordan Boyd-Graber. Are girls neko or shōjo? cross-lingual alignment of non-isomorphic embeddings with iterative normalization. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 3 3180-3189, 2019.
Han Zhao, Shanghang Zhang, Guanhang Wu, José MF Moura, Joao P Costeira, and Geoffrey J Gordon. Adversarial multiple source domain adaptation. In Advances in neural information processing systems, pp. 8559-8570, 2018.
Han Zhao, Remi Tachet Des Combes, Kun Zhang, and Geoffrey Gordon. On learning invariant representations for domain adaptation. In International Conference Machine Learning, pp-
2019.
Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In International Conference Learning Representations, 2021.
15

---

Published as a conference paper at ICLR 2021
A THEORETICAL BACKGROUND
In this section, we introduce theoretical background on neural tangent kernel (NTK), which draws an equivalence between the training dynamics of infinitely-widd (or ultra-wide) neural networks and that of kernel regression with respect to the neural tangent kernel.
Consider general neural network f(0, 2:): X' RR where 0 IR' is the parameters in the network and & E X' is the input. Suppose we train the neural network by minimizing the squared loss over training data, (()) 2 Σ (O, mx)) y; by gradient descent with infinitesimally small learning rate, i.e., d01!) = (t)((tt), Let u(t) ([((t)).:));= be the network outputs. u(t) follows the dynamics
du(t) - (H())((((()) - y),
(6)
dt
where (() is an x matrix whose (i. entry is
H(t @f(((tt) xx)) af(O(t). xx,)
00
00
line of works show that for sufficiently wide networks, (t) stays almost constant during training.
ie., (t) (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 infinity, (()) converges a fixed matrix, the neural tangent kernel (NTK) (Jacot et al., 2018):
( (a)(0(1,)) x) 05(0(t), x') d-W 00
2TK( = E
(8)
00
where W is Gaussian.
Therefore, the learning dynamics of sufficiently 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 neural network at convergence on any specific training set, denoted by fNTK(Z), (x), can be precisely characterized, and is equivalent to the following kernel regression solution
fNTK(W) (z) = (NTK(x,
NTK (z, xx)) -TK Irain Y.
(9)
where NrK is the n x n kernel for training data, NTK ax:) is the kernel value between test data x and training data æ;, 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 al., 2018 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 more precisely characterize the function learned by a sufficiently-wide neural network given any specific 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 non-trivial compute or analyze the functional form of what neural network learns using Eqn. 9, because the kernel regression solution using neural tangent kernel only gives point- woint-wise evaluation. Thus, we instead analyze the function learned by network in the NTK's induced feature space, because representations in the feature space would give a functional form.
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 fits training data. Here the min-norm refers the RKHS norm.
Lemma 2
Lemma 2.
Let <(x) be a feature map induced by a neural tangent kernel, for any x E Rd. The solution to kernel regression Eqn. 9 is equivalent fxru(z) 6(x) BNIK. where BWIK is
min ||
((x:) B Vi, for 1,..., n.
16

---

Published as a conference paper at ICLR 2021
We prove Lemma 2 in Appendix B.6. To analyze the learned functions the min- 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 ReL.U activation. It follows easily from the kernel formula described in Jacot et al. (2018); Arora al. (2019b); Bietti & Mairal (2019).
Lemma 3. An infinite-dimensionn dimensional feature map Φ(x) induced by the neural tangent kernel of a two-layer multi-layer perceptron with ReLU activation function is
& (xx =c c(x (k) x =20)
0) (h""" ((( I ==))...).
(10)
where (R) NN, I). with k going to infinity. e is a constant, and is the indicator function
We prove Lemma 3 in Appendix B.7. The feature maps for other architectures, c.g., Graph Neural Networks (GNNs) can be derived similarly. We analyze the Graph Neural Tangent Kernel (GNTK) for simple GNN architecture in Theorem
We then use Lemma 2 and to characterize the properties of functions learned by an over- parameterized 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 fitting the training data.
However, there still remains many technical challenges. For example, provable extrapolation (exact or asymptotic) often not achieved with most training data distribution. Understanding the desirable condition requires significant insights into the geometry properties of training data distribution, and how they interact with the solution learned by neural networks. Our insights and refined analysis shows in Rd 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 infinitely wide neural networks, the feature space of infinite dimension. The analysis of infinite dimensional spaces poses non-trivial technical challenges too.
Since different theorems have their respective challenges and insightsttechniques, we refer the interested readers to the respective proofs for details. In Lemma (proof in Appendix B.2), Theorem 2 (proof in Appendix B.3), and Theorem 1 (proof in Appendix B.1) we analyze over wer-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
PROOFS
B.
PROOF OF THEOREM 1
To show neural network outputs f(x) converge to linear function along all directions we will analyze the function learned by a neural network on the training set {<,;, (<.),)))) yi)) by studying the functional representation in the network's neural tangent kernel RKHS space.
Recall from Section that in the NTK regime, i.e., networks are infinitely wide, randomly initialized, and trained by gradient descent with infinitesimally 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 a E Rdd the network output is given by
f(x) = ((0(z), ((x1)), (o(z). o(xx))) NTK Y. train
where NTK is the data x and training data æ,, and Y is training labels. By Lemma 2, the kernel regression solution is also equivalent to the min-norm solution the NTK RKHS space that fits all training data
kernel for training data, (db(r)) o (a::) is the kernel value between test
f(x) = d(x) BNTK.
(1I)
where the representation coefficient 3NTK is
min || 2
<(x;) B Vis for i 1, 1,....
17

---

Published as a conference paper at ICLR 2021
The feature map ((x: for two-layer MLP with ReLU activation is given by Lemma 3
& (x)) = c' x I 1 (K)T 200 w (k)" II (x) x ==20....)
(12)
where w((a) ~ (O, D) with l going to infinity. c' is a constant, and I is the indicator function. Without loss of generality, we assume the bias term be 1. For simplicity of notations, we denote each data æ plus bias term by, i.e., & = [a]]] (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 #i to
and æ = Zo + hv = (1 + A)zo. where we introduce the notation of æ and À for convenience, are given by Eqn. 11 and Eqn. 12
f( = NTK
2D ( I (N) Bu 20) (h) Zo (w) zo 20)...)
[(a) = NTK @ I (k) 2 20 (T) a. (x) # 2 0) :..):
where we have 2ó = [20l [1] and &t= = [(1 + \) [1]. It follows that
f(0) - ((xo) = TATK (t T (b) a こ 0) z I w (A)T 2i 2 0)
(13)
(k) 2. (k)) Φ こ 0) - (k)
zo (w
2o 2 .....
(14)
By re-arranging the terms, we get the following equivalent form of the entries:
2 [w T a 2 0) zo I I(w"" 20 2 0)
(15)
a (Iw a 2 0) - 1 (w 20 2 0) +1 [w] T do 2 0)) 20 11 (w T 2o 0)
(16)
(I(w a 2 0)
(w 2i 2 0)) + (a no) II (w T 2 2 0)
(17)
= [a|1] (I (w a 2 0) - -1 1(w T 2 2 0)) + [bv00]] 11 (w 2o 2 0)
(18)
Similarly, we have
w a 1l (w T é 2 0) - w 2o 1 [( T 0o 20)
T
(19)
w T w (I(w" (w" T a 2 0) I [ db 2 0) +1 (w xo 2 0)) w 2o I1 [I T æi 2 0)
(20)
w T (I (w à 0) - 1 ( [ 2 2 0)) + w (a 2 1 (w 20 2 0)
(21)
= w [=]1] (T [w a 2 0) - T (w zi 2 0)) + w [hv 00) T (w T En 2 0)
(22)
Again, let us denote the part of GNTK corresponding to each ur by B... Moreover, let us denote the part corresponding to Eqn. 18 by By and the part corresponding to Eqn. 22 by B Then we have
f()) f(20)
(23)
h
/ B [x/A|1/h] (I (w T a 2 0)
-I (I(w" T Lo 2 0)) dPww)
(24)
/ B [0/0] I1 [l(w 2 2 0) dP(w)
(25)
/ Ben TT [a/h|/1/ /h (T (w T á 2 0) - T Tuu 2o 2 0)) dPP(w)
Taio
(26)
f B2. T. [000]] T uw ai > 0) dP((w)
(27)
Note that all Bw are finite constants that depend on the training data. Next, we show that as f + oc. each of the terms above converges in O(L/e) to some constant coefficient Br that depend on the training data and the direction v. Let us first consider Eqn. 25. We have
f 11 (w) Lo 2 0) dP(w)) =
f I (w [ [1] 2 0) dP(ww
(28)
I (w [ ////t > 0) dP
(29)
/ (w T [o/0] 2 0) dIP (w)
as t + 00
(30)
18

---

Published as a conference paper at ICLR 2021
Because By are finite constants, it follows that
/ b [0|0] I (w Bi 2 0) dP(w) + f Bu [ 0] (Iw 0] 0) dP(w)),
(31)
where the right hand side is constant that depends training data and direction v. Next, we show the convergence for Eqn. 31. Given error >> 0, because B [w|0] are finite constants, need to bound the following by C' e for some constant C²,
/ w T æi 2 0) - (w T [b|0]] 2 0) dIP(w)]
(32)
/ I (w ar
| > 0) II (w [Zo [0] 0) dP((w)
(33)
Observe that the two terms in Eqn. 33 represent the volume of half-(balls) that are orthogonal to vectors [xy 1] and 2zp(0.. Hence, Eqn. 33 is the volume of the non-overlapping part of the two (half)balls, which is created by rotating angle 0 along the last coordinate. By symmetry, Eqn. 33 is linear in 0. Moreover, the angle = arctan( for some constant C. Hence, it follows that
if Il (w [0o|1]] 2 0) - 1Iww" [20 2 0) dP(ww) = C₁
(34)
C.
(35)
= 0(L/t)
(36)
In the last inequality, we used the fact that arctan x < x for 0. Hence, 0(1/t) < implies I = O(1/e) as desired. Next, we consider Eqn. 24.
f [</h1//h] (I (I a -I 2o 2 0)) dF(w)
(37)
Let us first analyze the convergence of the following:
I (w a 2 0) - [(w T Lo 2 0) dT((w)
(38)
/ 1l (w [(1 A) 11] N 0) - 11 (w [ 2 0) dP((w)d((w))
(39)
/
w [ool
T
0 (w")
[ [1] N 0) dPP ddP(() +
(40)
The convergence to 0 follows from Eqn. 34. Now we consider the convergence rate. The angle is at most T+X times of that Eqn. 34. Hence, the is as follows
( (1-
T A_ o -
h/t T+ T O
( =0
=0
h
o
(41)
(h +1tt t)t
Now we get back to Eqn. 24, which simplifies as the following.
f 1
1 h
1
(I a 2) [ww Ey 2 0) dP
(42)
h
We compare the rate of growth of left hand side and the rate of decrease of right hand side (indicators).
t
1
0 as t + x0
(43)
h (h + 1)t
h+
as 1 x
(44)
(h + l)b
(h + t)t
Hence, the indicators decrease faster, and it follows that Eqn. 24 converges to 0 with rate 0(
Moreover, we can bound w with standard concentration techniques. Then the proofs for Eqn. 26 and Eqn. 27 follow similarly. This completes the proof.
19

---

Published as a conference paper at ICLR 2021
B.2 PROOF OF LEMMA
Overview of proof. To prove exact extrapolation given the conditions on training data, we analyze the function learned by the neural network in functional form. The network's learned function can be precisely characterized by solution in the network's s neural tangent kernel feature space which has minimum RKHS norm among functions that can fit all training data, i.c., it corresponds the optimum of constrained optimization problem. We show that the global optimum of this constrained optimization problem, given the conditions on training data, precisely the same function as the underlying true function.
Setup and preparation. Let X = ( Fnf and Y = (y.....,, Vo denote the training set input features and their labels. Let Bg E R.a denote the true parameters/weights for the underlying linear function g. i.c.,
g(x) = By for all x R.
Recall from Section A that in the NTK regime, where networks are infinitely wide, randomly initialized, and trained by gradient descent with infinitesimally small learning rate, the learning dynamics of neural network equivalent to that of 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 's feature space. That is, the function learned by the neural network (in the ntk regime) can be precisely characterized as
f(x) 6(x) BNIK.
where the representation coefficient ANTK is
min ||3||2
(45)
s.t. ((x;) 3= = ya for i I,...,
(46)
An dinnti-dimeeiionall feature map (xx) for two-layer ReLU network is described Lemma 3
② (x) = (x w (k) 2 0) (M) x ((I((( (R) 2 ........
where (k) NO, I). with is going to infinity. c' is constant, and the indicator function. That is, there are infinitely many directions ur 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 weight of the neural network's learned function in the neural tangent kernel feature space, ANIK, consists of weight vectors for each x (A)) x 2 0) E Rd and (*) x (k) x 22 R. For simplicity of notation, we will does not matter for our purposes. For any w E Rd, we denote by 3ur = (B0 :..,
refer to particular w, without considering the index (k), which
.... A(di) E Rd
the weight vectors corresponding to a (w x 2 0), and denote by Gr E Rd the weight for w (w
0).
Observe that for any No I) Rdd any other vectors in the same direction will activate the same set of æ; E Rd, That is, if w æ; 2 0 for any w E IR then (k w) æ; 2 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 Bn and f'o to denote the combined effect of all weights k .... a(d)) E Rd and Biuu E IR for all kw with is > in the same direction of w. That is, for each Uni(unit sphere) E RRd, we define 3(5 as the total effect of weights in the same direction
( = 1 1 -A0LIT
w Tuu
=1)
dPP(u,, for [d]
(47)
[|o|||
||2]]
20

---

Published as a conference paper at ICLR 2021
where -V I). Note that to ensure the Bw is well-defined number, here work with the polar representation and integrate with respect to an angle. Then 3w is well-defined. But for simplicity of exposition, use the plain notation of integral. Similarly, we define (3n as reloading the notation of
3 = f /uTl
f
w
u
1)
| 2
dl(
(48)
||w|||
Here, in Eqn. 48 we have extra term of ||u compared to Equ. 47 because the NTK features that Eqn. 48 corresponds to, w T x I (w x: 2 0), has an extra w T term. So we need to take into account the scaling. This abstraction enables us to make claims on the high-level parameters Bwr and Br, only, which we will show to be sufficient to determine the learned function.
[w]]
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
/ (Bw) B
2 + (B
2 +... (B
min
(Bww) aP(ww)
(49)
s.t.
B B'. w xz dl((w) B% xs Vi [m],
(50)
where the density of w is now uniform on the unit sphere of Rd. Observe that since w from a uniform distribution, the probability density funetion P(w) is constant. This means every æ; is activated by half of the w on the unit sphere, which implies we can now write the right hand side of Eqn. 50 in the form of left hand side, i.e., integral form. This allows further simplify Eqn. 50 as
(Bw + B w -2 -.) x, dF vi E [n].
(51)
up 2:20
where Eqn. 51 follows from the following steps of simplification
/
3 a + +.. x + Bw w
(1)
wTan (ddP(w) = =3 + x.. wie
2,20
1 x
(1)
+ + ( (d) Bu w Ta, x, dP(w)
dP(w) (B
xt (1) + + + 200 (0)
vi e [m].
dP(ww)
22 20
w TT.220 20
3 a
(1) +... + 3ld) x (d) Brm w T (dd((w
w 22,2, 20
(3)) (1)
(1) +... + ((d) x aP(w) Vi [n].
2
T 10:20
(BT + B w 22 x; dl(w) Vi e [m].
10 20
Claim 1. Without loss of generality, assume the scaling factor e in NTK feature map Φ(x) is 1. Then the global optimum the constraint optimization problem Eqn. 49 subject Eqn. 51, i.e.
/ (BLL)
+ (3 + +... +...+ (B(8)
2 + (B', ) dP(w)
min B
(52)
s.f.
(BTH. ~ + B'n w - 2. 2.35) dP(w) = 0 vi [n].
(53)
2,
satisfies Bu + Byg ur -2 for all w.
21

---

Published as a conference paper at ICLR 2021
This claim implies the exact extrapolation want prove, i.e. (xT g(x). This is because, our claim holds, then for any a E R
[NTK(1))
Fw w dP(ww)
w LuTxz x20
2 [ dP w)
w x20
dIP 23,
T220
2B g(x)
2
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 convex. Moreover, the constraint Eqn. 53 is affine. Hence. by KKT, solution that satisfies the Lagrange condition will be the global optimum. We compute the Lagrange multiplier as
((3, = (8!)) + (B80) +...+ + (1!00 + (Bmm)) dF(ww
(54)
(GT + B'u w
-2. x; dP(w)
(55)
Setting the partial derivative of C(B, A) with respect to each variable to zero gives
OC
2,3( P(w)
+
X a
(k) I (w0 a; 2 0) = 0
(56)
(1((h)
1=1
OC
8 P(w)
+ Σ
As w x; II [ww x; 2 0) =
(57)
i=1
0C
/
(0TH + By w 2. 2. x; dP(w)
(58)
ax;
w 0220
It is clear that the solution in Claim 1 immediately satisfies Eqn. 58. Hence, it remains to show there exist a of A; for i E [n] that satisfies Eqn. 56 and Eqn. 57. We can simplify Eqn. 56 as
(k) A a
II [ww x; 2 0),
(59)
1=1
where e is a constant. Similarly, we can simplify Eqn. 57 as
B w x; (w 0)
(60)
Observe that combining Eqn. 59 and Eqn. 60 implies that the constraint Eqn. 60 can be further simplified as
Be w T Bw
(61)
It remains to show that given the condition training data, there exists set of X, so that Eqn. 59 and Eqn. 61 are satisfied.
Global optimum via the geometry of training data. Recall that assume our training data 1d {( p=) satisfies for any w E R4, there exist d linearly independent (œ Jiz1 C X, where X = {x;})_ so that w x. 2 0 and xx e X for 1..d, e.g., an orthogonal basis of I and their opposite vectors. We will show that under this data regime, we have
22

---

Published as a conference paper at ICLR 2021
(a) for any particular w, there indeed exist set of X that can satisfy the constraints Eqn. 59 and Eqn. 61 for this particular w.
(b) For any wi and w2 that activate the exact same set of {æ,), the same set of A can satisfy the constraints Eqn. 59 and Eqn. 61 of both w1 and W2.
(c) Whenever rotate wi to u2 that the set of x; being activated changed, we still find X that satisfy constraint of both wi and w2.
Combining (a), (b) and (c) implies there exists set of À that satisfy the constraints for all w. Hence, remains to show these three claims.
We first prove Claim (a). For each w. must find set of X that the following hold.
1/h =c c.
A
a
(k)
II (w x. 2 2),
i=1
B" = w Bw
Bw + B w 29g
Here. 3, and are fixed, and w vector on the unit sphere. easy to see that Bw is then determined by Bg and w, and there indeed exists a solution (solving consistent linear system). Hence we are left with a linear system with d linear equations
3 C Σ 1=1
x (k)
I(w T x; 2 0) 4K E [d]
to solve with free variables being ); so that ur activates F;, i.e., w x. 0. Because the training data {( }=) satisfies for any w, there exist at least linearly independent æ; that activate w. This guarantees for any w we must have at least d free variables. It follows that there must exist solutions X; to the linear system. This proves Claim (a).
Next, we show that (b) for any w and w2 that activate the exact same set of {x;}. the same set of X can satisfy the constraints Eqn. 59 and Eqn. 61 of both w and W2. Because wi and wy are activated by the same set of æ,, this implics
Bwn C. X x, II [w) x. 2 0) =c c. Σ
A
x. A.,..
(w) x. 2 20) Bwn
Since A already satisfy constraint Eqn. 59 for w1, they also satisfy that for w2- Thus, it remains to show that Bw, + Bw, w Bw, + Bww w assuming Bwo Buo Bw: w Bw, and B'wz mn Bwg. This indeed holds because
Bwn B w1 Bay + Bw W22
B I Bw: -
w Bw, w T =w, Bw: W
Wi w1 13 = w Ww 32 2m
1 B 1. BT
Bwi Bun
Here, we used the fact that wi and wz are vectors on the unit sphere. This proves Claim (b).
Finally, we show (c) that Whenever we rotate a w, to W2 that the set of æ; being activated changed, we can still find X that satisfy constraint of both wi and w2. Suppose we rotate wi to W2 so that wy lost activation with æτ, æ2, ..., Ep which in the set of linearly independent æ; being activated by wn and their opposite vectors D; are also the training set (without loss of generality). Then u2 must now also get activated by E1, 222 ... Ey. This is because if wo x; O, must have w. (-x;)= 0.
Recall that in the proof of Claim (a), we only needed the X from linearly independent x; that we used solve the linear systems, and their opposite the free variables to solve the linear system of
23

---

Published as a conference paper at ICLR 2021
d equations. Hence, we can set À to 0 for the other æ, while still satisfying the linear system. Then, suppose there exists X that satisfy
B
X a (k)
where the 2; are the linearly independent vectors that activate wi with opposite vectors in the training set, which we have proved in (a). Then can satisfy the constraint for Bw, below
B = c. Σ E (-xz)
A x
(k)
==1
1=++1
by setting d, = -Ae for = 1..pp Indeed, this gives
B
Σ (->:) (-ar,)
Σ
A x; (k)
==1
x=p+1
d
c. X Z (k)
Thus, we can also find X that satisfy the constraint for Bwz. Here, we do not consider the case where W2 is parallel with an æ; because such W2 has measure zero. Note that we can apply this argument iteratively because the flipping the sign always works and will not create any inconsistency.
Moreover, we can show that the constraint for fnm is satisfied by similar argument as in proof of Claim (b). This follows from the fact that our construction makes Bw; Bwn- Then we can follow the same argument as in (b) to show that Bw, + f'uy = Bwn + B w1- This completes the proof of Claim (c).
In summary, combining Claim (a), (b) and (c) gives that Claim I holds. That is, given our training data, the global optimum to the constrained optimization problem finding the min-norm solution among functions that fit the training data satisfies B.. + B' w 20g. We also showed that this claim implies exact extrapolation, i.e., the network's 's learned function f(x)) is equal to the true underlying function g(x) for all Rd. This completes the proof.
B.3 PROOF OF THEOREM 2
Proof of the asymptotic convergence to extrapolation builds upon our proof of exact extrapolation, ie., Lemma The proof idea is that if the training data distribution has support at all directions, when the number of samples n -> xo, asymptotically the training set will converge to some imaginary training set that satisfies the condition for exact extrapolation. Since 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 first show the training data converge to data set that satisfies the exact extrapolation condition in Lemma 1. Suppose training data ( [mr)"= are sampled from distribution whose support contains connected set S that intersects all directions, i.e., for so that fww E S.
any non-zero w E Rdd, there exists
Let us denote by S the set of datasets that satisfy the condition in Lemma 1. In fact, we will use relaxed condition in the proof of Lemma (Lemma in the main text uses a stricter condition for simplicity of exposition). Given general dataset X and dataset S S of the same size n, let (X, S) denote a matching of their data points, i.e., a outputs a sequence of pairs
o(X, S); (zx st)
for i E [n]
s.t. X [x, )
S = {s: |
Let l: Rd x Rd + R be the /2 distance that takes in a pair of points. We then define the distance between the datasets d( (( S) as the minimum sum of 12 distances of their data points over all
24

---

Published as a conference paper at ICLR 2021
possible matching.
min Σ 8 (X, 1=1
X ||| = n
d( d(x S)
X * |||
We can then define "closest distance to perfect dataset" function D X - R which maps dataset X to the minimum distance of X to any dataset in S
D D* (X) = min (X,
SES
It is casy to see that for any dataset X = fa D (X) can be bounded by the minimum of the closest distance to perfect dataset DD sub-datasets X of size 2d.
[n/2d]
D" ((x;) S
min D" (1x)))-(k-1)224 k=
(62)
This is because for any S E S, and any S S', we must have S' because dataset satisfies exact extrapolation condition long contains some key points. Thus, adding data will not hurt, i.e., for any X1 S X2, we always have
SX (X2)
Now let us denote by X, random dataset of size n where each æ; E X, is sampled from the training distribution. Recall that our training data {æ; ) are sampled from distribution whose support contains a connected set S* that intersects all directions, i.e., for any non-zero w E R there exists so that kw S* It follows that for random dataset X2a of size 2d, the probability that DD* (Xza)
happens is less than 1 for any >> 0.
of size 2d, e.g.. orthogonal basis and their opposite vectors. Observe that if we scale any by & > O, the resulting dataset is still in S by the definition of S. We denote the set of datasets where we are allowed to scale elements So by Sg. follows that
First there must exist S. = {s; 2
IP (D" (X2a) > E) = P
min d 2d, S) ><) SES
< P
min
(X 2d, S) >
SES,
= P
min min SES
l (a (Xed, 5)₂) >e)
1
min min SESp
(a (2dd 22d:SS) S)₂) Se
< 11
min min máx E (a (X2d, S) Sc
SESo
#=1
<8<1
where we denote the bound of P (D" (X2a) > E) by 8 1, and the last step follows from
min min max (a X 2d S) S 0 SEE 1=1
which further follows from the fact that for any 8; E Sg. by the assumption on training distribution, we can always find k > 0 so that ks; E S", a connected set in the support of training distribution. By
the connectivity of support S ks, cannot be an isolated point in S so for any e > 0, we must have
fx(x)dx
k ke-ka, Se.
25

---

Published as a conference paper at ICLR 2021
Hence, now apply Eqn. 62 to bound D D* (Xn). Given any >> 0, we have
PP (x,) > c) = PP (D" (X,) ()
[n/2d]
$ 1 - P
min D' (=k Sc)
k=1
[n/2d]
$ $1- 1
P D (1z))]=(k1))221 2()
ko2d
[n/2d]
P P(D' (1z)) ()))))))))22+) ))) >c)
8Ln/2d]
Here 5 < 1. This implies D' (X,,,)
_ i...
lim n-00
(D*( c) = Ve > 0
(63)
Eqn. 63 says as the number of training samples n - oc, our training set will converge in probability to dataset that satisfies the requirement for exact extrapolation.
Asymptotic convergence predictions. Let NTK(x, x') Rd Rd - R denote the neural tangent kernel for a two-layer ReLU MLP. It is easy to see that if x -> x", then VTK((,.) kernel matrix for training data.
(TK (Arora et al. (2019h)). Let NTK denote the
We have shown that our training set converges to perfect data set that satisfies conditions of exact extrapolation. Moreover, note that our training set will only have a finite number of (not increase with n) v; that are not precisely the same as those in perfect dataset. This is because perfect data only contains a finite number of key points and the other points can be replaced by any other points while still being perfect data set. Thus, we have NTKuain N*, where No is the some perfect data set.
NTK matrix for
Because neural tangent kernel is positive definite, we have NTK Rd, the prediction of NTK is
+ N Recall that for any
urin
fNTK = (NTK(2, 21), (T, NTK
train
where NTKuain is the x n kernel for training data, NTK 2;) is the kernel value between test data x and training data æ;, and is training labels.
Similarly, we have (NTK(æ, æ1), ..., NT,,,,TKK xn)) (NTK(æ, xf), "i,, x;)), where 2 is a perfect data set that our training set converges to. Combining this with NTK - N gives
fNIK E fÑIK 9.
where /NTK is the function learned using our training set, and [ÑTK is that learned using a perfect data set, which is equal to the true underlying function g. This completes the proof.
B.4 PROOF OF COROLLARY
In order for GNN with linear aggregations
(k) a
Σ MLP(() (1,0), (
(EN(()
hgg = (LP((+1)
Σ K9
WEG
to extrapolate in the maximum degree task, it must be able to simulate the underlying function
hG max
Σ
uEG eEA
26

---

Published as a conference paper at ICLR 2021
Because the max function cannot be decomposed as the composition of piece-wise linear functions, the MLP(K+1) module in GNN must learn function that is not piece-wise linear over domains outside the training data range. Since Theorem proves for two-layer overparameterized MLPs, here we also assume HLP(K+1) is a two-layer overparameterized MLP, although the result can be extended more layers. It then follows from Theorem that for any input and label (and thus gradient), MLP(K+1) will converge linear functions along directions from the origin. Hence, there are always domains where the GNN cannot learn a correct target function.
B.5 PROOF OF THEOREM 3
Our proof applies the similar proof techniques for Lemma and to Graph Neural Networks (GNNs). This is essentially analysis of Graph Neural Tangent Kernel (GNTK), i.e., neural tangent kernel of GNNs.
We first define the simple GNN architecture will be analyzing, and then present the GNTK for this architecture. Suppose G = (V, E) is an input graph without edge feature, and Vu E R is the node feature of any node u V. Let us consider the simple one-layer GNN whose input G and output is hg
hg = W
max MEG (EN(U))
Σ W x.
(64)
Note that our analysis can be extended 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 graph neural network f0, G) 9 + where E Ih' is the parameters in the network and G E G is the input graph. Then the neural tangent kernel is
H @ G.) 05(0,G,) 00
00
where O are the give natural feature map. Let us denote, for any node u, the degree of by
parameters, Hence, the gradients with respect to all parameters
hu
E.
(65)
vEN((M
It then follows from simple computation of derivative that the following is a feature map of the GNTK for Eqn. 64
(GG) - c.
max wEG
(A) hw)
Σ1 arg max (k) h -lwu...)
(66)
DEG
BEG
where (R) (O, I), with going infinity. is constant, and is the indicator function.
Next, given training data {G,,, y 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 Lemma 1 and 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 æ. = 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 coefficients for w in the same direction as in Lemma I and 2.
Similarly, for any w. let us define B. E as the linear coefficient corresponding to
Σ I = arg max Tn hh, hy in RKHS space, and denote by A R the weight for uEG
wEG
max (w h..). Similarly, can combine the effect of all 3 in the same direction in Lemma
27

---

Published as a conference paper at ICLR 2021
and 2. We define the combined effect with Sw and f This allows to about with two directions, + and
Recall that the underlying reasoning function, maximum degree, is
g(G) max hw..
uEG
We formulate the constrained optimization problem, i.c., min-norm solution in GNTK feature space that fits all training data, as
/ 8 d(w)
min B,8
s.t.
Σ uEG,
1 u== arg max w ne Bw hu + max (w h.) (w-hh))BddP(u = max hy Vi E [n], wEG.
vEG
uEG;
where G; the i-th training graph and w ~ (O, 1). By combining the effect of 3, and taking the derivative of the Lagrange for the constrained optimization problem and setting to zero, get the global optimum solution satisfy the following constraints.
G
X
Σ h. II arg max h. eEG,
(67)
1=1
MEG.
B_
A Σ h II
u arg min h. vEG
(68)
1=1
MEG.
3
X max h. wEG,
(69)
3 C.
A min hy NEG,
(70)
i=
max hn uEG;
B Σ I
arg max lr vEG:
-h + 34 wEG,
-nax hu
(71)
uEG,
+ B
u arg min h.p
hr, + 3
min h.
Vi E [n]
(72)
DEG;
NEG;
aEG;
where is some constant, A; are the Lagrange parameters. Note that here we used the fact that there are two directions +1 and 11. This enables the simplification of Lagrange derivative. For a similar step-by-step derivation of Lagrange, refer to the proof of Lemma 1.
Let us consider the solution 3 1 and B4 = 3 = 3 0. It is clear that this solution can fit the training data, and thus satisfies Eqn. 71. Moreover, this solution is equivalent the underlying reasoning function, maximum degree, g gG) max.eg h...
Hence, remains to show that, given our training data, there exist X; so that the remaining four constraints are satisfies for this solution. Let rewrite these constraints linear systems where the variables are A
Enec, hu .1
u arg max h... vEG.
3 B
Evec, hu
u = arg min h
(73)
3
VEG,
3
i=
max h weG:
min ht
wEG;
28

---

Published as a conference paper at ICLR 2021
By standard theory of linear systems, there exist X to solve Eqn. 73 if there are at least four training data G2 whose following vectors linear independent
EuEa, has II arg max he vEG,
max h. -N wEG: min ha Nuuin
Euec, hu -1l
arg min he DEG,
(74)
max ha uEG, min ha uEG;
max h. wEG: min h wEG.
Here, Nimax denotes the number of nodes that achieve the maximum degree in the graph G2, and Nitin denotes the number of nodes that achieve the min degree in the graph G, By the assumption of our training data that there at least four G; 9 with linearly independent Eqn. 74, Hence, simple GNN learns the underlying function as desired.
This completes the proof.
B.6 PROOF LEMMA
Let W denote the span of the feature maps of training data 2;, i.e.
W = span (ø (x+1) (222) ...., (xn)).
Then decompose the coordinates of fNIK in the RKHS space, BNTK. into vector Po for the component of fNIK in the span of training data features W, and a vector B1 for the component the orthogonal complement W i.e.,
ANTK = 3o + B1.
First, note that since fNIK must be able to the training data (NTK is universal kernel as we will discuss next), i.e.,
((x.) GNTK = Vi- Thus, have d(xx) T y:- Then, Go is uniquely determined by the kernel regression solution with respect to the neural tangent kernel
fNTX
((z)) o(I.))
NTK Y.
Irain
where NTK is the n x n kernel for training data, (Ø(x), (x;;)) is the kernel between test data æ and training data æ; and Y is training labels.
The kernel regression solution /NTK is uniquely determined because the neural tangent kemne NTK imain positive definite assuming term (Du et al., 2019c). In any case, the solution is a min-norm by pseudo-inverse.
training data parallet, which be enforced with bias
Moreover, a unique kernel regression solution fNTK that spans the training data features corresponds to unique representation in the RKHS space Bo-
Since Bo and (1 orthogonal, also have the following
112 = ||3 + 3:112 = || + 11
This implies the norm of BNTK is at least as large the norm of any B such that o(x;) BNTX y Moreover, observe that the solution to kernel regression Eqn. 9 the feature span of training data, given the kernel matrix for training data is full rank.
SNTX (w) = ((0(#), ((x ()..)).... ( (o(x), ((x))) NTK Y. Imin
Since 31 is for the component of fNIK the orthogonal complement of training data feature span, we must have G1 O. It follows that ANTK is equivalent to
min || B
s.t. ((x,) T B = Vis for = 1,..., n.
desired.
29

---

Published as a conference paper at ICLR 2021
B.7 PROOF OF LEMMA 3
We first compute the neural tangent kernel (T æ²) for two-layer multi-layer perceptron (MLP) with ReLU activation function, and then show that it can be induced by the feature space ((x)
specified in the lemma that NTK(æ, x²) (6(x),,(zz))
Recall that Jacot et al. (2018) have derived the general framework for computing the neural tangent kernel of 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. (2018) and substituting the general activation function with ReLU gives the kernel formula for two-layer MIP 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 htlayee multi-layer perceptron can be recursively defined via a dynamic programming process. Here, [ Rd x Rd - R for = O.../ is the covariance for the i-th layer.
x (x, x²) = x: T Taa,"
(2, 2))
E(- x¹)
A (xx, x') = M(i-1) (z',z)
)x- xx)
(ax x') c
E
(o(u)a(v)).
The derivative covariance is defined similarly:
y (x, x') c
E (0,
[d())((v)
a
Then the neural tangent kernel for an h-layer network is defined as
NTK (2, x²)
I xi (x, x²)
(o,)))
Ammi
where we let (x, x²) = 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.c. o(a)
max(0, a) = a I[(a 2 0) and o(a)= = 1(aa 2 0).
2
h
(T (x. x') = Σ
(z. (2, x²)
II k=
(K)) (1)) (x,
i=1
= = (x), x') Σ(11 (xx, + E (')
So we can get the NTK via (x) and (101 x'), x (x, x'). Precisely.
y 100 (xxx') x Ta,
x T
x x² x -
x x
x ) x x' )
(1) (x,x²) =
x. x
(2, c E [u I(u 2 0) v I[v 2 0)].
To sample from ( we let L be a decomposition of A (1), such that A()) = LLT Here, we can see that L = (z, x') x²_ Thus, sampling from (O, A (1) is equivalent to first sampling w NO I). and output
Lw w (x. xx) ).
Then we have the equivalent sampling (u, v) = (w x, w x'). It follows that
(x) x²) =
E w~.)((.))
(w T x II (w x 2 0) w x' (w T x' 2 0)]
30

---

Published as a conference paper at ICLR 2021
It follows from the same reasoning that
- (z, (xx
(w 2 0)
[( x' 0)
(o.))
The neural tangent kernel for two-layer MLP with ReLU activation is then
NTK (x. 50 (ar, a²)
501) [z, x.) x (x, (xr
x x'' II 2 0) II w x' 2 0)]
=c. w~N(
TTe x ww x 0) w x' 11 (w x' 2 0)]
C (O,)
Next, we use the kernel formula to compute a feature map for two-layer MLP with ReLU activation function. Recall that by definition valid feature map must satisfy the following condition
NTK (1) x') = (Ø(æ), o(x')
It easy that the way represent NTK formula makes easy find such decomposition. The following infinitediimensional feature map would satisfy the requirement because the inner product of ((x) and ((x') for any æ, x' would be equivalent the expected value in NTK, after we integrate with respect to the density function of w.
(x. ( (x x 1220) 2) (k)) x (x) (R)
① (x) = =d(2
x x2 .)......
where (k(k) NO, I), with k going to infinity. e' is constant, and I is the indicator function. Note that here the density of features of oxx) is determined by the density of w. i.c. Gaussian.
C EXPERIMENTAL DETAILS
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 simple nonlinear functions in various dimen- sions and training/test distributions: quadratic, cosine, square root, and 11 norm functions, with MLPs with wide range of hyper-parameetrs.
This validates our implications on MLPs generally cannot extrapolate in tasks with nonlinear target functions, unless the nonlinear function is directionally linear out-of-distribuue In the latter case, the extrapolation error is more sensitive to the
2) Computation of the R-Squared of MLP's learned functions along (thousands of) randomly sampled directions in out-of-dise of-iittribttion domain.
This validates Theorem and shows the convergence rate is very high in practice, and often happens immediately out training range.
3) Learning tasks where the target functions are linear functions with MLPs. These validate Theorem 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 and cover all directions, and hence, MLPs extrapolate.
b) Training data distribution is restricted in some directions, e.g., restricted to be posi- tive/negaivee/constan some feature dimensions. This shows when training distribu- tion is restrictive in directions, MLPs may fail to extrapolate.
Exact extrapolation with infinitely-wide neural networks. i.c. exact computation with neural tangent kernel (NTK) on the data regime in Lemma 1. This is mainly for theoretical understanding.
4) ML.Ps with cosine, quadratic, and tanh activation functions.
31

---

Published as a conference paper at ICLR 2021
5) Learning maximum degree of graphs with Graph Neural Networks. Extrapolation graph structure, number of nodes, and node features. To show the role of architecture for extrapolation. study the following GNN architecture regimes.
a) GNN with graph-level max-poo max-pooling and neighbor-level sum sum- num-pooling By Theorem 3, this GNN architecture extrapolates in max degree with appropriate training data. b) GNN with graph-level and neighbor-level sum-poole aum-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, learning tasks only consider graph structure. We consider training sets sampled from various graph structure, and find only those satisfy conditions Theorem 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 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 U-shape curve of extrapolation error with respect to the sparsity of training graphs.
a) Specific graph structure: path, cycle, tree, expander, ladder, complete graphs, general graphs, 4-regular graphs.
b) Random graphs with a range of probability p an edge between any two nodes. Smaller p samples sparse graphs and large 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--odd problem with GNNs, the edge features are simply set to O.
b) Improved features. We show although our edge features do not bring in new information, it helps extrapolation.
C.1 LEARNING SIMPLE NON-LINEAR FUNCTIONS
Dataset details. We consider four tasks where the underlying functions are simple non-linear functions g Rdd + R. Given an input x Rd, the label computed by y g(x) for all æ. We consider the following four families of simple functions g.
a) Quadratic functions g(x=) = x: T Ax. In each dataset, we randomly sample In the simplest case where A = I, g(x)) >i= x²
a) Cosine functions g(x)) = Σ cos (2tr n.)).
Square root functions g(x) Σd VT. Here, the domain X' of is restricted to the space in Rdd with non-negative value in each dimension.
32

---

Published as a conference paper at ICLR 2021
d) L1 norm functions g(z) |a|1 Ei= z
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 uniformly sampled from hyper-cube. Training and validation data are sampled from I-a, aj with a E {0.5, 1.0), i... each dimension of x Rd is uniformly sampled from a, a). Test data are sampled from [-a, a] with f2.0, 5.0, 10.0).
ii) Training and validation data are uniformly sampled from sphere, where every point has L2 distance from the origin. We sample from r {0.5, 1.0). Then, we sample random Gaussian vector in R We obtain the training or validation data = q | r. This corresponds uniform sampling from the sphere.
Test data are sampled (non-uniformly) from a hyper-ball. We first sample r uniformly from [0.0, 2.0], [0.0, 5.0], and (0.0, 10.0]. Then, we sample random Gaussian vector in R. We obtain the test data æ //ll This corresponds to (non-ume mon-uniform) sampling from a hyper-ball Rd
b) We sample 20, 000 training data, 1, 000 validation data, and 20, 000 test data.
c) We sample input dimension from {1, 2, 8).
d) For quadratic functions, sample the entries of A uniformly from [-1, 1].
Model and hyperparameter settings. We consider the multi-layer perceptron (MLP) architecture.
MLP(x) w (wid-2)) (...o (w())zz))
We search the following hapereparameters for MLPs
a) Number of layers d from (2,4)
b) Width of each (A) from (64, 128, 512).
c) Initialization schemes.
i) The default initialization in Py PyTorh.
ii) The initialization scheme in neural tangent kernel theory, i.c., we sample entries of W from (O, 1) and scale the output after each (k)) by / -.. where d₃ is the output dimension of (k).
d) Activation function is ReLU.
We train the MLP with the mean squared error (MSE) loss, and Adam and SGD optimizer. We consider the following hyper- pppee-paraetter for training
a) Initial learning rate from {5e 2, le 2,5c 3, le -3). Learning rate decays 0.5 for every 50 epochs
b) Batch size from {32, 64, 128).
Weight decay is set to 1e 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.c., 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 data from different ranges, the mean absolute percentage етrоr (MAPE) loss, which scales the error by the actual value, better measures the extrapolation performance
1
A
F
MAPE
A
33

---

Published as a conference paper at ICLR 2021
where A, is the actual value and F; is the predicted value. Hence, in experiments, also report the MAPE.
C.2 R-SQUARED FOR OUT-OF-DISTRIBUTTON DIRECTIONS
We perform linear regression fit the predictions of MLPs along randomly sampled directions in out-of-distribution regions, and compute the R-squared (or R²) for these directions. This experiment to validate Theorem and show that the convergence rate (to a linear function) very high in practice.
Definition. R-squared, also known as coefficient of determination, assesses how strong the linear relationship is between input and output variables. The closer R-squared is to the stronger the linear relationship is, with being perfectly linear.
Datasets and models. We perform the R-squared computation on over 2, 000 combinations of datasets, test/train distributions, and hyper-panc apper-parameters, e.g., learning rate, batch size, MLP layer, width, initialization. These are described in Appendix C.I.
Computation. For each combination of dataset and model syper-parameters described in Sec- tion C.I, we save the trained MLP model f R + IR. For each dataset and model combination, we then randomly sample 5, 000 directions via Gaussian vectors (o, I). For each of these directions w, we compute the intersection point Tw of direction w and the training data distribution support (specified by hyper-sphere or hyper-cube; see Section C. I for details).
We then collect 100 predictions of the trained MLPP f along direction (assume w is normalized) with
{ (
+
F (...)) 10
-w. ( (x +
r
w
100
(75)
10
k=
where r is the range of training data distribution support (see Section C.I). 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- non-linear functions, we confirm that more than 96% of the R-squared results are above 0.99. This empirically confirms Theorem and shows that the convergence rate is fact fast practice. Along most directions, MLP's learned function becomes linear immediately out of the training data support.
C.3 LEARNING LINEAR FUNCTIONS
Dataset details, We consider the tasks where the underlying functions are linear Rd + R. Given an imput æ e Rd the label is computed by y g(x) Ax for all æ. 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 from 2,, 32).
c) We sample entries of A uniformly from [-a, a). where we sample E (5.0, 10.03.
d) The shape and support of training, validation, and test data distributions.
i) Training, validation, and test data are uniformly sampled from hyper-cube. Train- ing and validation data are sampled from [-a, a] with a E (5.0, 10.0), i.e., each dimension of I is uniformly sampled from [-a, a). Test data are sampled from [-a, a]d with E {20.0, 50.0}.
ii) Training and validation data are uniformly sampled from sphere, where every point has L2 distance from the origin. We sample r from r E (5.0, 10.0). Then, we sample random Gaussian vector in Rd We obtain the training or validation data x = q | т. This corresponds to uniform sampling from the sphere.
34

---

Published as a conference paper at ICLR 2021
Test data sampled (non- nnon-uniformly) from hyper-ball. We first sample r uniformly from [0.0, 20.0] and [0.0, 50.0].. Then, we sample a random Gaussian vector q in IR We obtain the test data 9 2/l1/122 from a hyper-ball in IRd.
This corresponds to (non-uniform) sampling
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 first dimension of any training data æ, to a fixed number 0.1, and randomly sample the remaining dimensions according d).
ii) We restrict the first k dimensions of any training data æ, to be positive. For input dimension 32, we only consider the hyper-cube training distribution, where we sample the first & dimensions from [0, a) and sample the remaining dimensions from [-a, al. For input dimensions and 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 & from {I, 16.32].
iii We restrict the first & dimensions of any training data æ; to be negative. For input dimension 32, we only consider the hyper-cube training distribution, where we sample the first &s dimensions from I- a, 0] and sample the remaining dimensions from [-a, a]. For input dimensions and 2, we consider both hyper-cube and hyper-sphere training distribution by performing rejection sampling. For input dimension 2, we consider from {1.2}. For input dimension 32, we consider & from {1, 16, 32}.
Model and hyperparameter settings. For the regression task, we search the same set of hyper- parameters as those in simple non-linear functions (Section C.1).We report the test error with the same validation procedure as in Section C.I.
Exact computation with neural tangent kernel Our experiments with MLPs validate Theorem asymptotic extrapolation for neural networks trained in regular regimes. Here, we also validate Lemma 1, exact extrapolation with finite data regime, by training infinitely-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 2d training examples according to the conditions in Lemma Specifically, we first sample an orthogonal basis and their opposite vectors X = fe,, ) --1 We then randomly sample 100 orthogonal transform matrices Q the QR decomposition. Our training samples are QX, i.e multiply each point in X by Q. This gives 100 training sets with 2d data points satisfying the condition in Lemma 1.
We perform kernel regression on these training sets using two-layer neural tangent kernel (NTK). Our code for exact computation of NTK is adapted from Arora et al. (2020); Novak et al. (2020). We verify that the test losses are all precisely 0, up to machine precision. This empirically confirms Lemma 1.
Note that due to the difference of taaer--r settings in different implementations NTK, to reproduce our experiments and achieve zero test error, the implementation by Arora et al. (2020) is assumed.
C.4 MLPS WITH COSINE, QUADRATIC, AND TANH ACTIVATION
This section describes the experimental settings for extrapolation experiments for MLPs with cosine, quadratic, and tanh activation functions. We train MLPs to leam the following functions:
a) Quadratic function g(z) Ax, where a randomly sampled matrix.
1=1 cos(2n ax)) b) Cosine function g(z) = Σd=
C) Hyperbolic tangent function g(x) Ed tanh(w:).
d) Linear function g(x: W i: +
35

---

Published as a conference paper at ICLR 2021
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, and test data from [-5,5] 5] For cosine, sample input dimension d from {1. 2}, training and validation data from [-100, 100]", and test data from |-200, 200] For tanh, we sample input dimension from 8), training and validation data from 1-000, 100] and test data from [-200, 200) For linear, we use a subset of datasets from Appendix C.3: I and 8 input dimensions with hyper-cube training distributions.
Model and hyperparameter settings. We use the same hyperparameters from Appendix C.I, except we fix the batch size to 128, as the batch size has minimal impact activation is hard to optimize, so we only report models with training MAPE less than 1.
on models. ML.ss with cos
C.5
MAX DEGREF
Dataset details. We consider the task of finding 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 EveV( 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 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 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 sampled uniformly from [-5.0, 5.0)s, i.e., a three-dimensional vector where each dimension sampled from [-5.0, 5.0). There are two schemes for test sets, in the first 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 hG by first iteratively aggregating and transforming the neighbors of all node vectors (k))) (vector for node u in layer k), and perform max or sum-pooling over all node features h,, to obtain bg. Formally, we have
(k) = Σ MLP(() (h!--1 (((-"",,(-")) 1) hG = MLP(K++) graph-pooling (h( u E
(EN(u)
(76)
Here, (u)) denotes the neighbors of u, K is the number of GNN iterations, and graph-pooling is hyper-parameter with choices as max or sum. (0) 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 from max or sum.
c) Width of all ML.Ps are 256.
36

---

Published as a conference paper at ICLR 2021
d) The number of layers for MLP(() with k 1..K are set to 2. The number of layers for MLP((+1) is set to 1.
We train the GNNs with the mean squared error (MSE) loss, and Adam and SGD optimizer. We for training
search the following
a) Initial learning rate set to 0.01.
b) Batch size is set to 64.
Weight decay is set to le - 5.
d) Number of epochs is 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 byper-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.6 SHORTEST PATH
Dataset details. We consider the task of finding the length of the shortest path on a graph, from given source to target nodes. Given any graph G -V, E). the node features, besides regular node features, encode whether node is source and whether node is target t. The edge features are scalar representing the edge weight. For unweighted graphs, all edge weights are Then the label g()) is the length of the shortest path from to / 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 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 0.9 with equal probability).
c) The number of vertices of graphs [V| for training and validation sets sampled uniformly from 40)...40.. The number vertices 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, the first 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, 11(v I), f)] with h sampled from 5-5.0,
e) After sampling a graph and edge weights, we sample source and by randomly sampling 8, i and selecting the first pair s whose shortest path involves at most 3 hops. This enables to solve the task using GNNs with iterations.
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 random graphs with different p. We consider p = 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 input graph G, GNN learns the output he by first iteratively aggregating and
37

---

Published as a conference paper at ICLR 2021
transforming the neighbors of all node vectors la (vector for node u in layer k), and perform max or sum-pooling over all node features h,, to obtain ba. Formally, we have
(k) min MLP(A) (h (h(k-1), h/ ((k-1), ( hag = MLP(K+1) vEN(u)
min fu
(77)
VEG
Here, Muu) denotes the neighbors of u, K is the number of GNN iterations, and for neighbor aggregation both min and sum. ha is the input node feature of node ((,,o) 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.
c) The number of layers for MLP(() with k L..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 to 0.01.
b) Batch size is set to 64.
c) Weight decay set to 5.
d) Number of epochs is set to 250.
We perform the same model selection and validation as Section C.5.
C.7 N-BODY PROBLEM
Task description. The n-body problem asks neural network to predict how stars in 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, c.g., 0.001 seconds.
Mathematically, in an n-body system = {X.) such as solar systems, all n stars {X+ exert distance and mass- desendentn gravitational forces each other, there were n(n 1) relations or forces in the system. Suppose X; at time t is at position x; and has velocity vf. The overall forces a star X, receives from other stars is determined by physics laws the following
G.
m, mj
(x)) (=)-
(78)
where G is the gravitational constant, and is the mass of star X.. Then acceralation at is determined by the net force F. and the mass of star m,
a = =FFF/m
(79)
Suppose the velocity of star X; at time is vf. Then assuming the time steps dt, i.e., difference be- tween time frames, are sufficiently small, the velocity at the next time frame f+ 1 can be approximated by
++ of + at dt.
(80)
Given m. æt. and vi. our task asks the neural network to predict of for all stars X, In our task, we consider two extrapolation schemes
a) The distances between stars ||a=| x// la are out-of-distributioo for test set, i.e., different sampling ranges from the training set.
b) The masses of stars my are out-of--istriburi for test set, i.c., 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.
38

---

Published as a conference paper at ICLR 2021
Dataset details. We first describe the simulation and sampling of our training We sample 100 videos of n-body system evolution, each with 500 rollout, i.c., time steps. We consider the orbit situation: there exists huge center star and several other stars. We sample the initial states, i..., position, velocity, masses, acceleration etc according to the following parameters.
a) The mass the center star 100kg.
b) The masses of other stars are sampled from [0.02, 9.0
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द of other objects randomly sampled from all angles, with distance [10.0, 100.0)m.
c) The velocity of the center star is
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 least 30m. 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 m...20Jm, but have in-distribution masses.
b) We sample the mass OOD test set follows
i) The mass of the center star is 200kg, i.e., twice of that in the training set. The masses of other stars are sampled from [0.04, 18.0]kg, compared to [0.02, 9.0 in the training set.
iii) The distances are in-distribution, i.e., sampling process 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 our GNN runs complete graph with nodes being the stars X,. GNN learns the star (node) representations by aggregating and transforming the interactions (forces) of all other node vectors
O₂ = MLP(2)
Σ DES((M)
MLP(1) (he huy W(w.o)
(81)
Here, h, the input feature of node v, including mass, position and velocity
ha (m,, Ba, v.)
W(a.)) is the input edge feature of edge (u, v). The loss is computed and backpropagated via the MSE loss of
[o:. ..., mn [ans1, ans 1ll
where o, denotes the output of GNN for node and ans; denotes the label for node in the next frame.
We search the following hyper-parameten for GNNs
a) Number of GNN iterations set to 1.
39

---

Published as a conference paper at ICLR 2021
b) Width of all MLPs are set to 128.
The number of layers for MLP(!) is 4. The number of layers for MLP(2) is
d) We consider nwo representations of edge/relations (((.))-
i) The first one is simply 0.
ii) The better representation, which makes the underlying target function more linear, is
my
(ar) æ!)
(6.5)
| 2 11
We train the GNN with the mean squared error (MSB) loss, and Adam optimizer. We search the following hyper-parameters for training
a) Initial learning rate set to 0.005. learning rate decays 0.5 for every 50 epochs
b) Batch size is set to 32.
Weight decay set to le 5.
d) Number of epochs is set to 2, 2000
D VISUALIZATION AND ADDITIONAL EXPERIMENTAL RESULTS
D.1 VISUALIZATION RESULTS
In this section, we show additional visualization results the MLP's learned function out training distribution (in black color) 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, show additional experimental results for graph-based reasoning tasks.
40

---

Published as a conference paper at ICLR 2021
Figure 8: (Quadratic function). Both panels show the learned v.s. true y x + x... In each figure, color OOD predictions by MLPs in black, underlying function in grey, and in-distribution predictions in blue. The support of training distribution a square (cube) for the top panel, and is circle (sphere) for the bottom panel.
41

---

Published as a conference paper at ICLR 2021
Figure 9: (Cos function). Both panels show the learned v.s. true y = cos(2t 211) + cos(2 22). In each figure, we color OOD predictions by MLPs in black, underlying function in grey, and in- distribution predictions in blue. The support of training distribution a square (cube) for both top and bottom panels, but with different ranges.
42

---

Published as a conference paper at ICLR 2021
Figure 10: (Cos function). Top panel shows the learned v.s. true cos(2 x1) ) + cos(2n x2) where the support of training distribution is circle (sphere). Bottom panel shows results for cosine in 1D, i.e. y = cos(2 r). In each figure, we color OOD predictions by MI.Ps in black, underlying function in grey, and in in-distribution predictions in blue.
43

---

Published as a conference paper at ICLR 2021
Figure 11: (Sqrt function). Top panel shows the learned v.s. true / + 22 where the support of training distribution is square (cube). Bottom panel shows the results for the square root function in 1D, i.e. = vr. In each figure, we color OOD predictions by MIPs in black, underlying function in grey, and in-distribution predictions in blue.
44

---

Published as a conference paper at ICLR 2021
Figure 12: (L1 function). Both panels show the learned v.s. true y |a].. In the top panel, the MLP successfully learns to extrapolate the absolute function. In the bottom panel, an MLP with different hyper-parameters fails extrapolate. In each figure, color OOD predictions by MLPs in black. underlying function in grey, and in-distribution predictions in blue.
45

---

Published as a conference paper at ICLR 2021
Figure 13: (L1 function). Both panels show the learned v.s. true y = |x|| + [122] In the top panel, the MLP successfully learns to extrapolate the Il norm function. In the bottom panel, an MLP with different hyper-parameters fails to extrapolate. each figure, we color OOD predictions by MLPs in black, underlying function in grey, and in-distribution predictions in blue.
46

---

Published as a conference paper at ICLR 2021
Figure 14: (Linear function). Both panels show the learned v.s. true y= + F2, with the support of training distributions being square (cube) for top panel, and circle (sphere) for bottom panel. MLPs successfully extrapolate the linear function with both training distributions. This is explained by Theorem 2: both sphere and cube intersect all directions. In each figure, we color OOD predictions by MLPs in black, underlying function in grey, and in an-distribution predictions in blue.
47

---

Published as a conference paper at ICLR 2021
D.2 EXTRA EXPERIMENTAL RESULTS
In this section, we show additional experimental results.
48

---

Published as a conference paper at ICLR 2021
10 3
10 1
MAPE
10
10 3
all
fix1
pos1
pos16 pos32
Figure 15: Density plot of the test errors in MAPE. The underlying functions linear, but train MLPs on different distributions, whose support potentially miss some directions. The training support for "all" are hyper-cubes that intersect all directions. In "fix1", we the first dimension of training data to a fixed number. In "posX", we restrict the first dimensions of training data to be positive. We can see that MLPs trained on "all" extrapolate the underlying linear functions, but MLPs trained on datasets with missing directions, i.e., "fixl" and "posX", often cannot extrapolate well.
49

---

Published as a conference paper at ICLR 2021
86.6
75.7
94.2
94.2
0.1
0.3
0.1
path
4-regular
ladder
cycle
expander
complete
tree
general
Figure 16: Maximum degree:
continuous and "spurious" node features. Here, each node has node feature in R$ that shall not contribute to the answer of maximum degree. GNNs with graph-level max- nax-pooling extrapolate to graphs with OOD node features and graph structure, graph sizes, if trained on graphs that satisfy the condition in Theorem 3.
50

---

Published as a conference paper at ICLR 2021
3099.6
94.5
97.2
99.2
91.1
107.2
78.5
11.0
6.4
0.1 expender
0.1 complete
0.0
0.0
path
4-regular
ladder
cycle
tree
general
Figure 17: Maximum degree:
max-pooling v.s. sum-pooling. In each sub-figure, left column shows test emrors for GNNs with graph-level max-pooling; right column shows test errors for GNNs with graph-level sum-pooling. x-axis shows the graph structure covered in training set. GNNs with sum-pooling fail extrapolate, validating Corollary 1. GNNs with max-pooling encodes appropriate non-linear operations, and thus extrapolates under appropriate training sets (Theorem 3).
51

---

Published as a conference paper at ICLR 2021
10
MAPE
p=1.0
Figure 18: Shortest path: random graphs. We train GNNs with neighbor and graph-level min on random graphs with probability p of edge between any two vertices. x-axis denotes the p for the training set, and y-axis denotes the test/cxtrapolation error on unseen graphs. The test errors follow U-shape: errors are high if the training graphs are very sparse (small p) or dense (large p). The same pattern is obtained if we train on specific graph structure.
52


Reviews:
[{"review_id": "sP3shlpXqlP", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 9, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary\n\nThe paper studies how neural networks extrapolate. The authors theoretically\nexamine two-layer ReLU MLPs with mean squared loss in the NTK regime and,\nbuilding on these results, GNNs. They find that the MLPs quickly converge to\nlinear functions along any direction from the origin, but can provably learn a\nlinear target function where the training distribution is sufficiently diverse.\nFor GNNs, they propose a hypothesis that the success of extrapolating\nalgorithmic tasks to new data relies on encoding task-specific non-linearities\nin the architecture or features. The theoretical results are supported by\nempirical results which sometimes go beyond the specific conditions of the\ntheorems (like increasing the number of layers in the MLP to 4 in Appendix\nC.1.).\n\n## Pros\n\n- The paper provides both theoretical and practical insight into the\n  extrapolation capabilities of neural networks, especially GNNs.\n- I especially liked the part about GNNs and the hypothesis that if we can\n  encode the non-linearities outside the MLPs so the MLPs only have to learn\n  linear functions, GNNs will extrapolate well.\n- Overall I found the paper very interesting and fun to read.\n\n## Concerns\n\n- The theoretical MLP results are very specific. Sometimes this is not apparent\n  from either the abstract or the discussion of the results. Some of the\n  constraints:\n  + The MLPs have two layers, which I find the most limiting constraint as most\n    practical MLPs have more layers.\n  + The mean squared loss is used throughout the paper. I think this is not\n    emphasized enough (it is mentioned only a single time in the paper). As far\n    as I understand the proofs also rely on the loss, so the loss should be\n    included in the conditions of the theorems.\n  + We are under the NTK regime, which is of course evident from the techniques\n    used. Nevertheless, this is not mentioned in the abstract.\n  + The MLPs are ReLU MLPs which is emphasized sufficiently in the paper. The\n    authors include preliminary empirical results for other activations\n    functions in the Appendix (sin, quadratic, and tanh).\n\n## Questions\n\n- Could the proofs of Theorem 3 and Theorem 5 be generalized to MLPs with more\nlayers?\n- Can we gain some insight into extrapolation with other loss functions like\nsoftmax based on these results?\n\n## Reasons for ranking\n\nI found the paper really interesting and gained much insight from it. Some of\nthe constraints of the MLPs are not emphasized enough and the writing is more\ngeneral at parts than the results warrant. Even with the constraints I believe\nthat this is an important step and sheds light on the extrapolation capabilities\nof neural networks. If the constraints can be made clearer I'm willing to\nimprove my score further.\n\n## Minor comments\n\n- Second to last paragraph on page 5: \"for Theoreom 5\" should be \"for Theorem 5\".\n- Caption of Figure 1: outisde => outside\n- In 4.2., \"Experiments: architectures that help extrapolation\": \"GNNs with\n  max-readout are better than GNNs with sum-readout (Fig 6a)\" should be Fig 5a.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "qmzcnHg2Pu", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 9, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper analyzes the extrapolate ability of MLPs and GNNs. In contrast to the existing theoretical works that focus on generalizability  and capacity of these models, this paper emphasizes the behavior of training algorithm using gradient descent. It takes analogy of kernel regression via the neural tangent kernel as an example to study the bias induced by the gradient descent algorithm. The presentation of this paper is clear and well-organized with the most significant result shown in the first section, raising interest of the readers, as opposed to leaving them behind a massive amount of proofs. The contribution of this paper is significant as well since it draws attention of the researcher to theoretical analysis on the bias induced from the implementations of the algorithms as compared to the theoretical analysis on the model structure itself. Model extrapolation is also closely connected to topics such as meta-learning, multi-task learning, domain adaptation and semi-supervised learning since the ability of model extrapolation will limit its performance when applied to other tasks. \n\nPros:\n1. This paper has shown some interesting results: for instance, MLP with ReLU trained by GD will converge to linear functions along any direction from origin outside the support of the training data. This coincide with the idea that MLP are piecewise linear in different regions. The proof is complicated though and requires the analogy to the kernel regression as basis.  This result seems to suggest that the learning of MLP on data manifold supported by training data is also local linear and without support of training data, the induction follows the inertia of linearity.  It is curious to see if this is due to the piecewise linearity of ReLU function.  Maybe we will have better nonlinear extrapolation for MLP using tanh and other sigmoid functions. \n2.  Comparison between GNN and Dynamic programming algorithm is very intuitive and inspiring. It suggests that max/min aggregate as opposed to more commonly used sum-aggregate in GNN is more suitable for extrapolation and the similarity between max/min aggregate GNN and DP is also very convincing. In general, this paper built up a good intuition before diving into the proof, which is well-appreciated. \n3.   The suggestion to improve extrapolation is to put the nonlinearity into the architecture of the GNN or into the input representation is useful. For instance, replacing sum-aggregate to min/max aggregate helps to achieve good extrapolation. It also explains why the pre-trained embeddings such as BERT can be used in other tasks and still extrapolate well.\n\n\nSuggestions:\n1. Limitations of the study scope. This paper only discuss results of neural network using ReLU and GD. Although GD is widely used, the ReLU as the activation function plays a critical role in the study of extrapolation. It is necessary to provide analysis on the use of other common activation function to understand if the extrapolation ability is expanded. \n2.  It is interesting to see more connection with domain adaptation and semi-supervised learning as well.", "extra_scores": {}, "extra_text": {}}, {"review_id": "ErscdGWPVMS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 9, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper investigates the extrapolation power of MLPs and GNNs (trained by gradient descent with mean squared loss) from a theoretical perspective. The authors show results of extensive experiments that back up their theoretical findings.  \n\nIn particular, the authors study the question of what these neural networks learn outside the training distribution, and identify conditions when they extrapolate well. Their findings suggest that ReLU MLPs extrapolate well in linear tasks, with a fast convergence rate (O(1/\\epsilon). GNNs (having MLP modules) extrapolate well when the non-linear operations are encoded in either network architecture or data representation, so as the inner MLP modules are aligned with only linear functions.\n\nThe paper is well written, ideas and definitions clearly explained and experiments laid out in detail. The theoretical contributions of the work are important as they enhance our understanding of how these networks learn and how well they generalize. These findings help us design GNNs based on the data and problem at hand. As such, this work addresses a fundamental question in GNN understanding and must be published. \n\nSome comments/questions to the authors:\n- In Section 3.2, “diversity” of a distribution is informally defined in terms of training support and direction. A more thorough definition would be helpful. \n- The title of the paper is somewhat misleading: “from feedforward to GNN” insinuates that there are other network types that are discussed in the paper.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "16iWXj6OMC_", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper tackles the challenging question of how deep networks might learn to extrapolate knowledge outside the support of their training distribution. The paper contributes both with novel theoretical arguments as well as with empirical evidence collected on targeted cases. Differently from other recent approaches to the problem, the theoretical analyses presented here are non-asymptotic and provide precise information about the kind of functions that MLPs can learn in the proximity of the training region. Moreover, the authors provide compelling arguments about the need to explicitly encoding (task-specific) non-linearities in the input representation and/or in the model architecture in order to promote successful extrapolation.\nOverall, the paper addresses important issues and can be considered at the frontier of deep learning research. The paper is well-written and the recent literature is properly reviewed. In light of this, I think the paper would be of interest to the ICLR community. However, I would like to explicitly mention that I was not able to carefully review all the details and proofs reported in the Appendix, which is of an unusual length (almost 40 pages) for an ICLR paper.\n\nComments for possible improvements:\n- The analyses reported in Appendix D.3 / C.4 regarding the extrapolation capability of MLPs with different activation functions (sin, tanh, quadratic) are relevant and should be emphasized. They could also be expanded, for example by considering some data generation tasks analyzed in the main text.\n- It would be very interesting to extend this analysis to other simple problems, where MLPs cannot extrapolate appropriately. I am specifically referring to the simple counting and arithmetic tasks discussed in [1], where generalization outside the training distribution was achieved by adding ad-hoc gate units to the network. I think this domain is particularly relevant here, given that arithmetics is mentioned by the authors in the opening sentence of the paper.\n\n[1]\tA. Trask, F. Hill, S. Reed, J. Rae, C. Dyer, and P. Blunsom, “Neural Arithmetic Logic Units,” in arXiv:1808.00508, 2018.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies how (two layer) neural nets extrapolates. The paper is beautifully written and the authors very successfully answered all the questions. They managed to update the paper, clarify the assumptions and add additional experiments. 
Decision: Accept (Oral)