How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks — part 3/11

Cleaned-text code points [31514, 46232). Bibliography omitted here; source.md includes it.

Experiments: representations that help extrapolation. Finally, we show a good input representation helps extrapolation. We study the n n -body problem ( Battaglia et al., 2016 ; Watters et al., 2017 ) (Appendix C.7 ), that is, predicting the time evolution of n n objects in a gravitational system. Following previous work, the input is a complete graph where the nodes are the objects ( Battaglia et al., 2016 ) . The node feature for u u is the concatenation of the object’s mass m u m_{u} , position 𝒙 u ( t ) \bm{x}_{u}^{(t)} , and velocity 𝒗 u ( t ) \bm{v}_{u}^{(t)} at time t t . The edge features are set to zero. We train GNNs to predict the velocity of each object u u at time t + 1 t+1 . The true velocity f ⁡ ( G , u ) f(G;u) for object u u is approximately f ⁡ ( G , u ) ≈ 𝒗 u t + 𝒂 u t ⋅ d ​ t , 𝒂 u t = C ⋅ ∑ v ≠ u m v ‖ 𝒙 u t − 𝒙 v t ‖ 2 3 ⋅ ( 𝒙 v t − 𝒙 u t ) , \displaystyle f(G;u)\approx\bm{v}_{u}^{t}+\bm{a}_{u}^{t}\cdot dt,\quad\bm{a}_{u}^{t}=C\cdot\sum\limits_{v\neq u}\frac{m_{v}}{\|\bm{x}_{u}^{t}-\bm{x}_{v}^{t}\|_{2}^{3}}\cdot\Big(\bm{x}_{v}^{t}-\bm{x}_{u}^{t}\Big), (5) where C C is a constant. To learn f f , the MLP modules need to learn a nonlinear function. Therefore, GNNs do not extrapolate well to unseen masses or distances (“original features” in Figure 6(b) ). We instead use an improved representation h ⁡ ( G ) h(G) to encode non-linearity. At time t t , we transform the edge features of ( u , v ) (u,v) from zero to 𝒘 ( u , v ) ( t ) = m v ⋅ ( 𝒙 v ( t ) − 𝒙 u ( t ) ) / ‖ 𝒙 u ( t ) − 𝒙 v ( t ) ‖ 2 3 \bm{w}_{(u,v)}^{(t)}=m_{v}\cdot\big(\bm{x}_{v}^{(t)}-\bm{x}_{u}^{(t)}\big)/\|\bm{x}_{u}^{(t)}-\bm{x}_{v}^{(t)}\|_{2}^{3} . 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 6(b) ).

## 5 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 to the training data range.

Domain adaptation studies generalization to a specific target domain ( Ben-David et al., 2010 ; Blitzer et al., 2008 ; Mansour et al., 2009 ) . Typical strategies adjust the training process: for instance, use unlabeled samples from the target domain to align the target and source distributions ( Ganin et al., 2016 ; Zhao et al., 2018 ) . Using target domain data during training may induce useful non-linearities and may mitigate extrapolation by matching the target and source distributions, though the correctness of the learned mapping depends on the label distribution ( Zhao et al., 2019 ) .

Self-supervised learning on a large amount of unlabeled data can learn useful non-linearities beyond the labeled training data range ( Chen et al., 2020 ; Devlin et al., 2019 ; He et al., 2020 ; Peters et al., 2018 ) . Hence, our results suggest an explanation why pre-trained representations such as BERT improve out-of-distribution robustness ( Hendrycks et al., 2020 ) . In addition, self-supervised learning could map semantically similar data to similar representations, so some out-of-domain examples might fall inside the training distribution after the mapping.

Invariant models aim to learn features that respect 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 are confounders or anti-causal variables ( Ahuja et al., 2021 ; Rosenfeld et al., 2021 ) , this may essentially increase the training data range, since variations in the invariant features may be ignored by the model.

Distributional robustness considers small adversarial perturbations of the data distribution, and ensures that the model performs well under these ( Goh & Sim, 2010 ; Sagawa et al., 2020 ; Sinha et al., 2018 ; Staib & Jegelka, 2019 ) . We instead look at more global perturbations. Still, one would expect that modifications that help extrapolation in general also improve robustness to local perturbations.

## 6 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 an explanation how GNNs have been able to extrapolate well in complex algorithmic tasks: encoding appropriate non-linearity in architecture and features can help extrapolation. Our results and hypothesis agree with empirical results, in this paper and in the literature.

#### 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 a Chevron-MIT Energy Fellowship. This research was also supported by JST ERATO JPMJER1201 and JSPS Kakenhi JP18H05291. MZ was supported by ODNI, IARPA, via the BETTER Program contract 2019-19051600005. The views, opinions, and/or findings contained in this article are those of the author and should not be 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.

## Appendix 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-wide (or ultra-wide) neural networks and that of kernel regression with respect to the neural tangent kernel.

Consider a general neural network f ⁡ ( 𝜽 , 𝒙 ) : 𝒳 → ℝ f(\bm{\theta},\bm{x}):\mathcal{X}\rightarrow\mathbb{R} where 𝜽 ∈ ℝ m \bm{\theta}\in\mathbb{R}^{m} is the parameters in the network and 𝒙 ∈ 𝒳 \bm{x}\in\mathcal{X} is the input. Suppose we train the neural network by minimizing the squared loss over training data, ℓ ⁡ ( 𝜽 ) = 1 2 ​ ∑ i = 1 n ( f ⁡ ( 𝜽 , 𝒙 i ) − y i ) 2 \ell(\bm{\theta})=\frac{1}{2}\sum_{i=1}^{n}(f(\bm{\theta},\bm{x}_{i})-y_{i})^{2} , by gradient descent with infinitesimally small learning rate, i.e., d ​ 𝜽 ​ ( t ) d ​ t = − ∇ ℓ ​ ( 𝜽 ​ ( t ) ) \frac{d\bm{\theta}(t)}{dt}=-\nabla\ell(\bm{\theta}(t)) . Let 𝒖 ⁡ ( t ) = ( f ⁡ ( 𝜽 ⁡ ( t ) , 𝒙 i ) ) i = 1 n \bm{u}(t)=(f(\bm{\theta}(t),\bm{x}_{i}))_{i=1}^{n} be the network outputs. 𝒖 ⁡ ( t ) \bm{u}(t) follows the dynamics d ​ 𝒖 ​ ( t ) d ​ t = − 𝑯 ⁡ ( t ) ​ ( 𝒖 ⁡ ( t ) − y ) , \displaystyle\frac{d\bm{u}(t)}{dt}=-\bm{H}(t)(\bm{u}(t)-y), (6) where 𝑯 ⁡ ( t ) \bm{H}(t) is an n × n n\times n matrix whose ( i , j ) (i,j) -th entry is 𝑯 ( t ) i ​ j = ⟨ ∂ f ⁡ ( 𝜽 ⁡ ( t ) , 𝒙 i ) ∂ 𝜽 , ∂ f ⁡ ( 𝜽 ⁡ ( t ) , 𝒙 j ) ∂ 𝜽 ⟩ . \displaystyle\bm{H}(t)_{ij}=\biggl<\frac{\partial f(\bm{\theta}(t),\bm{x}_{i})}{\partial\bm{\theta}},\frac{\partial f(\bm{\theta}(t),\bm{x}_{j})}{\partial\bm{\theta}}\biggr>. (7) A line of works show that for sufficiently wide networks, 𝑯 ⁡ ( t ) \bm{H}(t) stays almost constant during training, i.e., 𝑯 ⁡ ( t ) = 𝑯 ⁡ ( 0 ) \bm{H}(t)=\bm{H}(0) in the limit ( Arora et al., 2019a ; Arora et al., 2019b ; Allen-Zhu et al., 2019a ; Du et al., 2019c ; Du et al., 2019a ; Li & Liang, 2018 ; Jacot et al., 2018 ) . Suppose network parameters are randomly initialized with certain scaling, as network width goes to infinity, 𝑯 ⁡ ( 0 ) \bm{H}(0) converges to a fixed matrix, the neural tangent kernel (NTK) ( Jacot et al., 2018 ) : NTK ( 𝒙 , 𝒙 ′ ) = 𝔼 𝜽 ∼ 𝒲 ⟨ ∂ f ⁡ ( 𝜽 ⁡ ( t ) , 𝒙 ) ∂ 𝜽 , ∂ f ⁡ ( 𝜽 ⁡ ( t ) , 𝒙 ′ ) ∂ 𝜽 ⟩ , \displaystyle\text{NTK}(\bm{x},\bm{x}^{\prime})=\mathop{\mathbb{E}}_{\bm{\theta}\sim\mathcal{W}}\biggl<\frac{\partial f(\bm{\theta}(t),\bm{x})}{\partial\bm{\theta}},\frac{\partial f(\bm{\theta}(t),\bm{x}^{\prime})}{\partial\bm{\theta}}\biggr>, (8) where 𝒲 \mathcal{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 a neural network at convergence on any specific training set, denoted by f NTK ​ ( 𝒙 ) f_{\text{NTK}}(\bm{x}) , can be precisely characterized, and is equivalent to the following kernel regression solution f NTK ​ ( 𝒙 ) = ( NTK ​ ( 𝒙 , 𝒙 1 ) , … , NTK ​ ( 𝒙 , 𝒙 n ) ) ⋅ NTK train − 1 ​ 𝒀 , \displaystyle f_{\text{NTK}}(\bm{x})=(\text{NTK}(\bm{x},\bm{x}_{1}),...,\text{NTK}(\bm{x},\bm{x}_{n}))\cdot\text{NTK}_{\text{train}}^{-1}\bm{Y}, (9) where NTK train \text{NTK}_{\text{train}} is the n × n n\times n kernel for training data, NTK ​ ( 𝒙 , 𝒙 i ) \text{NTK}(\bm{x},\bm{x}_{i}) is the kernel value between test data 𝒙 \bm{x} and training data 𝒙 i \bm{x}_{i} , and 𝒀 \bm{Y} is the training labels.

We can in fact exactly calculate the neural tangent kernel matrix for certain architectures and activation functions. The exact formula of NTK with ReLU activation has been derived for feedforward neural networks ( Jacot et al., 2018 ) , convolutional neural networks ( Arora et al., 2019b ) , and Graph Neural Networks ( Du et al., 2019b ) .

Our theory builds upon this equivalence of network learning and kernel regression to more precisely characterize the function learned by a sufficiently-wide neural network given any specific training set. In particular, the difference between the learned function and true function over the domain of 𝒳 \mathcal{X} determines the extrapolation error.

However, in general it is non-trivial to compute or analyze the functional form of what a neural network learns using Eqn. 9 , because the kernel regression solution using neural tangent kernel only gives point-wise evaluation. Thus, we instead analyze the function learned by a network in the NTK’s induced feature space , because representations in the feature space would give a functional form.

Lemma 2 makes this connection more precise: the solution to the kernel regression using neural tangent kernel, which also equals over-parameterized network learning, is equivalent to a min-norm solution among functions in the NTK’s induced feature space that fits all training data. Here the min-norm refers to the RKHS norm.

###### Lemma 2 .

Let ϕ ⁡ ( 𝐱 ) \phi(\bm{x}) be a feature map induced by a neural tangent kernel, for any 𝐱 ∈ ℝ d \bm{x}\in\mathbb{R}^{d} . The solution to kernel regression Eqn. 9 is equivalent to f NTK ​ ( 𝐱 ) = ϕ ​ ( 𝐱 ) ⊤ ​ 𝛃 NTK f_{\text{NTK}}(\bm{x})=\phi(\bm{x})^{\top}\bm{\beta}_{\text{NTK}} , where 𝛃 NTK \bm{\beta}_{\text{NTK}} is min 𝜷 ⁡ ‖ 𝜷 ‖ 2 \displaystyle\min_{\bm{\beta}}\|\bm{\beta}\|_{2} s.t. ϕ ( 𝒙 i ) ⊤ 𝜷 = y i , for i = 1 , … , n . \displaystyle\phi(\bm{x}_{i})^{\top}\bm{\beta}=y_{i},\;\;\;\text{for }i=1,...,n.

We prove Lemma 2 in Appendix B.6 . To analyze the learned functions as the min-norm solution in feature space, we also need the explicit formula of an induced feature map of the corresponding neural tangent kernel. The following lemma gives a NTK feature space for two-layer MLPs with ReLU activation. It follows easily from the kernel formula described in Jacot et al. (2018) ; Arora et al. (2019b) ; Bietti & Mairal (2019) .

###### Lemma 3 .

An infinite-dimensional feature map ϕ ⁡ ( 𝐱 ) \phi(\bm{x}) induced by the neural tangent kernel of a two-layer multi-layer perceptron with ReLU activation function is ϕ ⁡ ( 𝒙 ) = c ⁡ ( 𝒙 ⋅ 𝕀 ⁡ ( 𝒘 ( k ) ⊤ ​ 𝒙 ≥ 0 ) , 𝒘 ( k ) ⊤ ​ 𝒙 ⋅ 𝕀 ⁡ ( 𝒘 ( k ) ⊤ ​ 𝒙 ≥ 0 ) , … ) , \displaystyle\phi\left(\bm{x}\right)=c\left(\bm{x}\cdot\mathbb{I}\left(\bm{w}^{(k)^{\top}}\bm{x}\geq 0\right),\bm{w}^{(k)^{\top}}\bm{x}\cdot\mathbb{I}\left(\bm{w}^{(k)^{\top}}\bm{x}\geq 0\right),...\right), (10) where 𝐰 ( k ) ∼ 𝒩 ⁡ ( 𝟎 , 𝐈 ) \bm{w}^{(k)}\sim\mathcal{N}(\bm{0},\bm{I}) , with k k going to infinity. c c is a constant, and 𝕀 \mathbb{I} is the indicator function.

We prove Lemma 3 in Appendix B.7 . The feature maps for other architectures, e.g., Graph Neural Networks (GNNs) can be derived similarly. We analyze the Graph Neural Tangent Kernel (GNTK) for a simple GNN architecture in Theorem 3 .

We then use Lemma 2 and 3 to characterize the properties of functions learned by an 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) is 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 ℝ d \mathbb{R}^{d} space, we need to consider the directions of training data. In graphs, we need to consider, in addition, the graph structure of training data. We refer readers to detailed proofs for the intuition of data conditions. Moreover, since NTK corresponds to infinitely wide neural networks, the feature space is of infinite dimension. The analysis of infinite dimensional spaces poses non-trivial technical challenges too.

Since different theorems have their respective challenges and insights/techniques, we refer the interested readers to the respective proofs for details. In Lemma 1 (proof in Appendix B.2 ), Theorem 2 (proof in Appendix B.3 ), and Theorem 1 (proof in Appendix B.1 ) we analyze over-parameterized MLPs. The proof of Corollary 1 is in Appendix B.4 . In Theorem 3 we analyze Graph Neural Networks (proof in Appendix B.5 ).

## Appendix B Proofs

### B.1 Proof of Theorem 1

To show neural network outputs f ⁡ ( 𝒙 ) f(\bm{x}) converge to a linear function along all directions 𝒗 {\bm{v}} , we will analyze the function learned by a neural network on the training set { ( 𝒙 i , y i ) } i = 1 n \{(\bm{x}_{i},y_{i})\}_{i=1}^{n} , by studying the functional representation in the network’s neural tangent kernel RKHS space.

Recall from Section A that in the NTK regime, i.e., networks are 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.