How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks

1 INTRODUCTION . Humans extrapolate well in many tasks . For example , we can apply arithmetics to arbitrarily large numbers . One may wonder whether a neural network can do the same and generalize to examples arbitrarily far from the training data ( Lake et al. , 2017 ) . Curiously , previous works report mixed extrapolation results with neural networks . Early works demonstrate feedforward neural networks , a.k.a . multilayer perceptrons ( MLPs ) , fail to extrapolate well when learning simple polynomial functions ( Barnard & Wessels , 1992 ; Haley & Soloway , 1992 ) . However , recent works show Graph Neural Networks ( GNNs ) ( Scarselli et al. , 2009 ) , a class of structured networks with MLP building blocks , can generalize to graphs much larger than training graphs in challenging algorithmic tasks , such as predicting the time evolution of physical systems ( Battaglia et al. , 2016 ) , learning graph algorithms ( Velickovic et al. , 2020 ) , and solving mathematical equations ( Lample & Charton , 2020 ) . To explain this puzzle , we formally study how neural networks trained by gradient descent ( GD ) extrapolate , i.e. , what they learn outside the support of training distribution . We say a neural network extrapolates well if it learns a task outside the training distribution . At 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 et al. , 2016 ; Soudry et al. , 2018 ) . In our analysis , we explicitly consider such implicit bias through the analogy of the training dynamics of over-parameterized neural networks and kernel regression via the neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ) . Starting with feedforward networks , the simplest neural networks and building blocks of more complex architectures such as GNNs , we establish that the predictions of over-parameterized MLPs with ReLU activation trained by GD converge to linear functions along any direction from the origin . We prove a convergence rate for two-layer networks and empirically observe that convergence often occurs close to the training data ( Figure 1 ) , which suggests ReLU MLPs can not 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 to the training distribution . In contrast , our results are non-asymptotic and quantify what kind of functions MLPs will learn close to the training distribution . Second , we identify a condition when MLPs extrapolate well : the task is linear and the geometry of the training distribution is sufficiently “ diverse ” . To our knowledge , our results are the first extrapolation results of this kind for feedforward neural networks . We then relate our insights into feedforward neural networks to GNNs , to explain why GNNs extrapolate well in some algorithmic tasks . Prior works report successful extrapolation for tasks that can be solved by dynamic programming ( DP ) ( Bellman , 1966 ) , which has a computation structure aligned with GNNs ( Xu et al. , 2020 ) . DP updates can often be decomposed into nonlinear and linear steps . Hence , we hypothesize that GNNs trained by GD can extrapolate well in a DP task , if we encode appropriate non-linearities in the architecture and input representation ( Figure 2 ) . Importantly , encoding non-linearities may be unnecessary for GNNs to interpolate , because the MLP modules can easily learn many nonlinear functions inside the training distribution ( Cybenko , 1989 ; Hornik et al. , 1989 ; Xu et al. , 2020 ) , but it is crucial for GNNs to extrapolate correctly . We prove this hypothesis for a simplified case using Graph NTK ( Du et al. , 2019b ) . Empirically , we validate the hypothesis on three DP tasks : max degree , shortest paths , and n-body problem . We show GNNs with appropriate architecture , input representation , and training distribution can predict well on graphs with unseen sizes , structures , edge weights , and node features . Our theory explains the empirical success in previous works and suggests their limitations : successful extrapolation relies on encoding task-specific non-linearities , which requires domain knowledge or extensive model search . From a broader standpoint , our insights go beyond GNNs and apply broadly to other neural networks . To summarize , we study how neural networks extrapolate . First , ReLU MLPs trained by GD converge to linear functions along directions from the origin with a rate of O ( 1/t ) . Second , to explain why GNNs extrapolate well in some algorithmic tasks , we prove that ReLU MLPs can extrapolate well in linear tasks , leading to a hypothesis : a neural network can extrapolate well when appropriate nonlinearities are encoded into the architecture and features . We prove this hypothesis for a 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 polynomials ( Barnard & Wessels , 1992 ; Haley & Soloway , 1992 ) . We instead show a general pattern of how ReLU MLPs extrapolate and identify conditions for MLPs to extrapolate well . More recent works study the implicit biases induced on MLPs by gradient descent , for both the NTK and mean 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-dimensional settings of the NTK regime , MLP is asymptotically at most a linear predictor in certain scaling limits ( Ba et al. , 2020 ; Ghorbani et al. , 2019 ) . We study a different setting ( extrapolation ) , and our analysis is non-asymptotic in nature and does not rely on random matrix theory . Prior works explore GNN extrapolation by testing on larger graphs ( Battaglia et al. , 2018 ; Santoro et al. , 2018 ; Saxton et al. , 2019 ; Velickovic et al. , 2020 ) . We are the 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 → R with a training set { ( xi , yi ) } ni=1 ⊂ D , where yi = g ( xi ) and D is the support of training distribution . Previous works have extensively studied in-distribution generalization where the training and the test distributions are identical ( Valiant , 1984 ; Vapnik , 2013 ) ; i.e. , D = X . In contrast , extrapolation addresses predictions on a domain X that is larger than the support of the training distribution D. We will say a model extrapolates well if it has a small extrapolation error . Definition 1 . ( Extrapolation error ) . Let f : X → R be a model trained on { ( xi , yi ) } ni=1 ⊂ D with underlying function g : X → R. Let P be a distribution over X \ D and let ` : R× R→ R be a loss function . We define the extrapolation error of f as Ex∼P [ ` ( f ( x ) , g ( x ) ) ] . We focus on neural networks trained by gradient descent ( GD ) or its variants with squared loss . We study two network architectures : feedforward and graph neural networks . Graph Neural Networks . GNNs are structured networks operating on graphs with MLP modules ( Battaglia et al. , 2018 ; Xu et al. , 2019 ) . Let G = ( V , E ) be a graph . Each node u ∈ V has a feature vector xu , and each edge ( u , v ) ∈ E has a feature vector w ( u , v ) . GNNs recursively compute node representations h ( k ) u at iteration k ( Gilmer et al. , 2017 ; Xu et al. , 2018 ) . Initially , h ( 0 ) u = xu . For k = 1 .. K , GNNs update h ( k ) u by aggregating the neighbor representations . We can optionally compute a graph representation hG by aggregating the final node representations . That is , h ( k ) u = ∑ v∈N ( u ) MLP ( k ) ( h ( k−1 ) u , h ( k−1 ) v , w ( v , u ) ) , hG = MLP ( K+1 ) ( ∑ u∈G h ( K ) u ) . ( 1 ) The final output is the graph representation hG or final node representations h ( K ) u depending on the task . We refer to the neighbor aggregation step for h ( k ) u as aggregation and the pooling step in hG as readout . Previous works typically use sum-aggregation and sum-readout ( Battaglia et al. , 2018 ) . Our results indicate why replacing them may help extrapolation ( Section 4 ) . 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 , we 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 1 illustrates examples of how ReLU MLPs extrapolate when trained by GD on various nonlinear functions . These examples suggest that outside the training support , the predictions quickly become linear along directions from the origin . We systematically verify this pattern by linear regression on MLPs ’ predictions : the coefficient of determination ( R2 ) is always greater than 0.99 ( Appendix C.2 ) . That is , ReLU MLPs “ linearize ” almost immediately outside the training data range . We formalize this observation using the implicit biases of neural networks trained by GD via the neural tangent kernel ( NTK ) : optimization trajectories of over-parameterized networks trained by GD are equivalent to those of kernel regression with a 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 A . Definition 2 . ( Informal ) A neural network trained in 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 over-parameterized neural networks via NTK ( Allen-Zhu et al. , 2019a ; b ; Arora et al. , 2019a ; b ; Cao & Gu , 2019 ; Du et al. , 2019c ; a ; Li & Liang , 2018 ; Nitanda & Suzuki , 2021 ) . We instead analyze extrapolation . Theorem 1 formalizes our observation from Figure 1 : outside the training data range , along any direction tv from the origin , the prediction of a two-layer ReLU MLP quickly converges to a linear function with rate O ( 1t ) . The linear coefficients βv and the constant terms in the convergence rate depend on the training data and direction v. The proof is in Appendix B.1 . Theorem 1 . ( Linear extrapolation ) . Suppose we train a two-layer ReLU MLP f : Rd → R with squared loss in the NTK regime . For any direction v ∈ Rd , let x0 = tv . As t → ∞ , f ( x0 + hv ) − f ( x0 ) → βv · h for any h > 0 , where βv is a constant linear coefficient . Moreover , given > 0 , for t = O ( 1 ) , we have | f ( x0+hv ) −f ( x0 ) h − βv| < . ReLU networks have finitely many linear regions ( Arora et al. , 2018 ; Hanin & Rolnick , 2019 ) , hence their predictions eventually become linear . In contrast , Theorem 1 is a more fine-grained analysis of how MLPs extrapolate and provides a convergence rate . While Theorem 1 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.1 and C.2 ) . Our proof technique potentially also extends to deeper networks . Theorem 1 implies which target functions a ReLU MLP may be able to match outside the training data : only functions that are almost-linear along the directions away from the origin . Indeed , Figure 4a shows ReLU MLPs do not extrapolate target functions such as x > Ax ( quadratic ) , ∑d i=1 cos ( 2π ·x ( i ) ) ( cos ) , and ∑d i=1 √ x ( i ) ( sqrt ) , where x ( i ) is the i-th dimension of x . With suitable hyperparameters , MLPs extrapolate the L1 norm correctly , which satisfies the directional linearity condition . Figure 4a provides one more positive result : MLPs extrapolate linear target functions well , across many different hyperparameters . While learning linear functions may seem very limited at 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 .

Summary: 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.