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

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

This validates our implications on MLPs generally cannot extrapolate in tasks with nonlinear target functions, unless the nonlinear function is directionally linear out-of-distribution. In the latter case, the extrapolation error is more sensitive to the hyper-parameters.

2) Computation of the R-Squared of MLP’s learned functions along (thousands of) randomly sampled directions in out-of-distribution domain. This validates Theorem 1 and shows the convergence rate is very high in practice, and often happens immediately out of training range.

3) Learning tasks where the target functions are linear functions with MLPs. These validate Theorem 2 and Lemma 1 , i.e., MLPs can extrapolate if the underlying function is linear under conditions on training distribution. This section includes four ablation studies: a) Training distribution satisfy the conditions in Theorem 2 and cover all directions, and hence, MLPs extrapolate.

b) Training data distribution is restricted in some directions , e.g., restricted to be positive/negative/constant in some feature dimensions. This shows when training distribution is restrictive in directions, MLPs may fail to extrapolate.

c) Exact extrapolation with infinitely-wide neural networks , i.e., exact computation with neural tangent kernel (NTK) on the data regime in Lemma 1 . This is mainly for theoretical understanding.

4) MLPs with cosine, quadratic, and tanh activation functions.

5) Learning maximum degree of graphs with Graph Neural Networks. Extrapolation on graph structure, number of nodes, and node features. To show the role of architecture for extrapolation, we study the following GNN architecture regimes. a) GNN with graph-level max-pooling and neighbor-level sum-pooling. By Theorem 3 , this GNN architecture extrapolates in max degree with appropriate training data.

b) GNN with graph-level and neighbor-level sum-pooling. By Corollary 1 , this default GNN architecture cannot extrapolate in max degree.

To show the importance of training distribution, i.e., graph structure in training set, we study the following training data regimes.

a) Node features are identical , e.g., 1 1 . In such regimes, our learning tasks only consider graph structure. We consider training sets sampled from various graph structure, and find only those satisfy conditions in Theorem 3 enables GNNs with graph-level max-pooling to extrapolate.

b) Node features are spurious and continuous. This also requires extrapolation on OOD node features. GNNs with graph-level max-pooling with appropriate training sets also extrapolate to OOD spurious node features.

6) Learning the length of the shortest path between given source and target nodes, with Graph Neural Networks. Extrapolation on graph structure, number of nodes, and edge weights. We study the following regimes. a) Continuous features. Edge and node features are real values. This regime requires extrapolating to graphs with edge weights out of training range.

Test graphs are all sampled from the “general graphs” family with a diverse range of structure. Regarding the type of training graph structure, we consider two schemes. Both schemes show a U-shape curve of extrapolation error with respect to the sparsity of training graphs.

a) Specific graph structure: path, cycle, tree, expander, ladder, complete graphs, general graphs, 4-regular graphs.

b) Random graphs with a range of probability p p of an edge between any two nodes. Smaller p p samples sparse graphs and large p p samples dense graphs.

7) Physical reasoning of the n 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 n -body problem with GNNs, the edge features are simply set to 0 0 .

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 : ℝ d → ℝ g:\mathbb{R}^{d}\rightarrow\mathbb{R} . Given an input 𝒙 ∈ ℝ d \bm{x}\in\mathbb{R}^{d} , the label is computed by y = g ⁡ ( 𝒙 ) y=g(\bm{x}) for all 𝒙 \bm{x} . We consider the following four families of simple functions g g . a) Quadratic functions g ⁡ ( 𝒙 ) = 𝒙 ⊤ ​ A ​ 𝒙 g(\bm{x})=\bm{x}^{\top}A\bm{x} . In each dataset, we randomly sample A A . In the simplest case where A = I A=I , g ⁡ ( 𝒙 ) = ∑ i = 1 d x i 2 g(\bm{x})=\sum_{i=1}^{d}x_{i}^{2} .

a) Cosine functions g ⁡ ( 𝒙 ) = ∑ i = 1 d cos ⁡ ( 2 ​ π ⋅ 𝒙 i ) g(\bm{x})=\sum_{i=1}^{d}\cos\left(2\pi\cdot\bm{x}_{i}\right) .

c) Square root functions g ⁡ ( 𝒙 ) = ∑ i = 1 d 𝒙 i g(\bm{x})=\sum_{i=1}^{d}\sqrt{\bm{x}_{i}} . Here, the domain 𝒳 \mathcal{X} of 𝒙 \bm{x} is restricted to the space in ℝ d \mathbb{R}^{d} with non-negative value in each dimension.

d) L1 norm functions g ⁡ ( 𝒙 ) = | 𝒙 | 1 = ∑ i = 1 d | 𝒙 i | g(\bm{x})=|\bm{x}|_{1}=\sum_{i=1}^{d}|\bm{x}_{i}| .

We sample each dataset of a task by considering the following parameters a) The shape and support of training, validation, and test data distributions. i) Training, validation, and test data are uniformly sampled from a hyper-cube. Training and validation data are sampled from [ − a , a ] d [-a,a]^{d} with a ∈ { 0.5 , 1.0 } a\in\{0.5,1.0\} , i.e., each dimension of 𝒙 ∈ ℝ d \bm{x}\in\mathbb{R}^{d} is uniformly sampled from [ − a , a ] [-a,a] . Test data are sampled from [ − a , a ] d [-a,a]^{d} with a ∈ { 2.0 , 5.0 , 10.0 } a\in\{2.0,5.0,10.0\} .

ii) Training and validation data are uniformly sampled from a sphere, where every point has L ​ 2 L2 distance r r from the origin. We sample r r from r ∈ { 0.5 , 1.0 } r\in\{0.5,1.0\} . Then, we sample a random Gaussian vector 𝒒 \bm{q} in ℝ d \mathbb{R}^{d} . We obtain the training or validation data 𝒙 = 𝒒 / ‖ 𝒒 ‖ 2 ⋅ r \bm{x}=\bm{q}/\|\bm{q}\|_{2}\cdot r . This corresponds to uniform sampling from the sphere. Test data are sampled (non-uniformly) from a hyper-ball. We first sample r r uniformly from [ 0.0 , 2.0 ] , [ 0.0 , 5.0 ] , [0.0,2.0],[0.0,5.0], and [ 0.0 , 10.0 ] [0.0,10.0] . Then, we sample a random Gaussian vector 𝒒 \bm{q} in ℝ d \mathbb{R}^{d} . We obtain the test data 𝒙 = 𝒒 / ‖ 𝒒 ‖ 2 ⋅ r \bm{x}=\bm{q}/\|\bm{q}\|_{2}\cdot r . This corresponds to (non-uniform) sampling from a hyper-ball in ℝ d \mathbb{R}^{d} .

b) We sample 20,000 20,000 training data, 1,000 1,000 validation data, and 20,000 20,000 test data.

c) We sample input dimension d d from { 1 , 2 , 8 } \{1,2,8\} .

d) For quadratic functions, we sample the entries of A A uniformly from [ − 1 , 1 ] [-1,1] .

##### Model and hyperparameter settings.

We consider the multi-layer perceptron (MLP) architecture. MLP ​ ( 𝒙 ) = 𝑾 ( d ) ⋅ σ ⁡ ( 𝑾 ( d − 1 ) ​ σ ​ ( … ​ σ ​ ( 𝑾 ( 1 ) ​ 𝒙 ) ) ) \displaystyle\text{MLP}(\bm{x})=\bm{W}^{(d)}\cdot\sigma\left(\bm{W}^{(d-1)}\sigma\left(...\sigma\left(\bm{W}^{(1)}\bm{x}\right)\right)\right) We search the following hyper-parameters for MLPs a) Number of layers d d from { 2 , 4 } \{2,4\} .

b) Width of each 𝑾 ( k ) \bm{W}^{(k)} from { 64,128,512 } \{64,128,512\} .

c) Initialization schemes. i) The default initialization in PyTorch.

ii) The initialization scheme in neural tangent kernel theory, i.e., we sample entries of 𝑾 k \bm{W}^{k} from 𝒩 ⁡ ( 0 , 1 ) \mathcal{N}(0,1) and scale the output after each 𝑾 ( k ) \bm{W}^{(k)} by 2 d k \sqrt{\frac{2}{d_{k}}} , where d k d_{k} is the output dimension of 𝑾 ( k ) \bm{W}^{(k)} .

d) Activation function σ \sigma is set to ReLU.

We train the MLP with the mean squared error (MSE) loss, and Adam and SGD optimizer. We consider the following hyper-parameters for training a) Initial learning rate from { 5 ​ e − 2 , 1 ​ e − 2 , 5 ​ e − 3 , 1 ​ e − 3 } \{5e-2,1e-2,5e-3,1e-3\} . Learning rate decays 0.5 0.5 for every 50 50 epochs

b) Batch size from { 32 , 64 , 128 } \{32,64,128\} .

c) Weight decay is set to 1 ​ e − 5 1e-5 .

d) Number of epochs is set to 250 250 .

##### Test error and model selection.

For each dataset and architecture, training hyper-parameter setting, we perform model selection via validation set, i.e., we report the test error by selecting the epoch where the model achieves the best validation error. Note that our validation sets always have the same distribution as the training sets.

We train our models with the MSE loss. Because we sample test data from different ranges, the mean absolute percentage error (MAPE) loss, which scales the error by the actual value, better measures the extrapolation performance MAPE = 1 n ​ | A i − F i A i | , \displaystyle\text{MAPE}=\frac{1}{n}\left|\frac{A_{i}-F_{i}}{A_{i}}\right|, where A i A_{i} is the actual value and F i F_{i} is the predicted value. Hence, in our experiments, we also report the MAPE.

### C.2 R-squared for Out-of-distribution Directions

We perform linear regression to fit the predictions of MLPs along randomly sampled directions in out-of-distribution regions, and compute the R-squared (or R 2 R^{2} ) for these directions. This experiment is to validate Theorem 1 and show that the convergence rate (to a linear function) is very high in practice.

##### 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 1 1 , the stronger the linear relationship is, with 1 1 being perfectly linear.

##### Datasets and models.

We perform the R-squared computation on over 2,000 2,000 combinations of datasets, test/train distributions, and hyper-parameters, e.g., learning rate, batch size, MLP layer, width, initialization. These are described in Appendix C.1 .

##### Computation.

For each combination of dataset and model hyper-parameters as described in Section C.1 , we save the trained MLP model f : ℝ d → ℝ f:\mathbb{R}^{d}\rightarrow\mathbb{R} . For each dataset and model combination, we then randomly sample 5,000 5,000 directions via Gaussian vectors 𝒩 ⁡ ( 𝟎 , 𝑰 ) \mathcal{N}(\bm{0},\bm{I}) . For each of these directions 𝒘 \bm{w} , we compute the intersection point 𝒙 𝒘 \bm{x}_{\bm{w}} of direction 𝒘 \bm{w} and the training data distribution support (specified by a hyper-sphere or hyper-cube; see Section C.1 for details).

We then collect 100 100 predictions of the trained MLP f f along direction 𝒘 \bm{w} (assume 𝒘 \bm{w} is normalized) with { ( 𝒙 𝒘 + k ⋅ r 10 ⋅ 𝒘 ) , f ⁡ ( 𝒙 𝒘 + k ⋅ r 10 ⋅ 𝒘 ) } k = 0 100 , \displaystyle\left\{\left(\bm{x}_{\bm{w}}+k\cdot\frac{r}{10}\cdot\bm{w}\right),f\left(\bm{x}_{\bm{w}}+k\cdot\frac{r}{10}\cdot\bm{w}\right)\right\}_{k=0}^{100}, (75) where r r is the range of training data distribution support (see Section C.1 ). We perform linear regression on these predictions in Eqn. 75 , and obtain the R-squared.

##### Results.

We obtain the R-squared for each combination of dataset, model and training setting, and randomly sampled direction. For the tasks of learning the simple non-linear functions, we confirm that more than 96 % 96\% of the R-squared results are above 0.99 0.99 . This empirically confirms Theorem 1 and shows that the convergence rate is in fact fast in practice. Along most directions, MLP’s learned function becomes linear immediately out of the training data support.

### C.3 Learning Linear Functions

Dataset details. We consider the tasks where the underlying functions are linear g : ℝ d → ℝ g:\mathbb{R}^{d}\rightarrow\mathbb{R} . Given an input 𝒙 ∈ ℝ d \bm{x}\in\mathbb{R}^{d} , the label is computed by y = g ⁡ ( 𝒙 ) = A ​ 𝒙 y=g(\bm{x})=A\bm{x} for all 𝒙 \bm{x} . For each dataset, we sample the following parameters a) We sample 10,000 10,000 training data, 1,000 1,000 validation data, and 2,000 2,000 test data.

b) We sample input dimension d d from { 1 , 2 , 32 } \{1,2,32\} .

c) We sample entries of A A uniformly from [ − a , a ] [-a,a] , where we sample a ∈ { 5.0 , 10.0 } a\in\{5.0,10.0\} .

d) The shape and support of training, validation, and test data distributions. i) Training, validation, and test data are uniformly sampled from a hyper-cube. Training and validation data are sampled from [ − a , a ] d [-a,a]^{d} with a ∈ { 5.0 , 10.0 } a\in\{5.0,10.0\} , i.e., each dimension of 𝒙 ∈ ℝ d \bm{x}\in\mathbb{R}^{d} is uniformly sampled from [ − a , a ] [-a,a] . Test data are sampled from [ − a , a ] d [-a,a]^{d} with a ∈ { 20.0 , 50.0 } a\in\{20.0,50.0\} .

ii) Training and validation data are uniformly sampled from a sphere, where every point has L ​ 2 L2 distance r r from the origin. We sample r r from r ∈ { 5.0 , 10.0 } r\in\{5.0,10.0\} . Then, we sample a random Gaussian vector 𝒒 \bm{q} in ℝ d \mathbb{R}^{d} . We obtain the training or validation data 𝒙 = 𝒒 / ‖ 𝒒 ‖ 2 ⋅ r \bm{x}=\bm{q}/\|\bm{q}\|_{2}\cdot r . This corresponds to uniform sampling from the sphere. Test data are sampled (non-uniformly) from a hyper-ball. We first sample r r uniformly from [ 0.0 , 20.0 ] [0.0,20.0] and [ 0.0 , 50.0 ] , [0.0,50.0], . Then, we sample a random Gaussian vector 𝒒 \bm{q} in ℝ d \mathbb{R}^{d} . We obtain the test data 𝒙 = 𝒒 / ‖ 𝒒 ‖ 2 ⋅ r \bm{x}=\bm{q}/\|\bm{q}\|_{2}\cdot r . This corresponds to (non-uniform) sampling from a hyper-ball in ℝ d \mathbb{R}^{d} .

e) We perform ablation study on how the training distribution support misses directions. The test distributions remain the same as in d). i) We restrict the first dimension of any training data 𝒙 i \bm{x}_{i} to a fixed number 0.1 0.1 , and randomly sample the remaining dimensions according to d).

ii) We restrict the first k k dimensions of any training data 𝒙 i \bm{x}_{i} to be positive. For input dimension 32 32 , we only consider the hyper-cube training distribution, where we sample the first k k dimensions from [ 0 , a ] [0,a] and sample the remaining dimensions from [ − a , a ] [-a,a] . For input dimensions 1 1 and 2 2 , we consider both hyper-cube and hyper-sphere training distribution by performing rejection sampling. For input dimension 2 2 , we consider k k from { 1 , 2 } \{1,2\} . For input dimension 32 32 , we consider k k from { 1 , 16 , 32 } \{1,16,32\} .

iii) We restrict the first k k dimensions of any training data 𝒙 i \bm{x}_{i} to be negative. For input dimension 32 32 , we only consider the hyper-cube training distribution, where we sample the first k k dimensions from [ − a , 0 ] [-a,0] and sample the remaining dimensions from [ − a , a ] [-a,a] . For input dimensions 1 1 and 2 2 , we consider both hyper-cube and hyper-sphere training distribution by performing rejection sampling. For input dimension 2 2 , we consider k k from { 1 , 2 } \{1,2\} . For input dimension 32 32 , we consider k k from { 1 , 16 , 32 } \{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.1 .

##### Exact computation with neural tangent kernel

Our experiments with MLPs validate Theorem 2 asymptotic extrapolation for neural networks trained in regular regimes. Here, we also validate Lemma 1 , exact extrapolation with finite data regime, by training an 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.