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

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

We sample the same test set as in our experiments with MLPs. For training set, we sample 2 ​ d 2d training examples according to the conditions in Lemma 1 . Specifically, we first sample an orthogonal basis and their opposite vectors 𝑿 = { 𝒆 i , − 𝒆 i } i = 1 d \bm{X}=\{\bm{e}_{i},-\bm{e}_{i}\}_{i=1}^{d} . We then randomly sample 100 100 orthogonal transform matrices Q Q via the QR decomposition. Our training samples are Q ​ 𝑿 Q\bm{X} , i.e., multiply each point in 𝑿 \bm{X} by Q Q . This gives 100 100 training sets with 2 ​ d 2d data points satisfying the condition in Lemma 1 .

We perform kernel regression on these training sets using a two-layer neural tangent kernel (NTK). Our code for exact computation of NTK is adapted from Arora et al. (2020) ; Novak et al. (2020) . We verify that the test losses are all precisely 0 0 , up to machine precision. This empirically confirms Lemma 1 .

Note that due to the difference of hyper-parameter settings in different implementations of NTK, to reproduce our experiments and achieve zero test error, the implementation by Arora et al. (2020) is assumed.

### C.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 learn the following functions: a) Quadratic function g ⁡ ( 𝒙 ) = 𝒙 ⊤ A ​ 𝒙 g(\bm{x})=\bm{x}\top A\bm{x} , where A A is a randomly sampled matrix.

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

c) Hyperbolic tangent function g ⁡ ( 𝒙 ) = ∑ i = 1 d tanh ⁡ ( 𝒙 i ) g(\bm{x})=\sum_{i=1}^{d}\tanh({\bm{x}_{i}}) .

d) Linear function g ⁡ ( 𝒙 ) = W ​ 𝒙 + b g(\bm{x})=W\bm{x}+b .

##### Dataset details.

We use 20,000 training, 1,000 validation, and 20,000 test data. For quadratic, we sample input dimension d d from { 1 , 8 } \{1,8\} , training and validation data from [ − 1 , 1 ] d [-1,1]^{d} , and test data from [ − 5 , 5 ] d [-5,5]^{d} . For cosine, we sample input dimension d d from { 1 , 2 } \{1,2\} , training and validation data from [ − 100,100 ] d [-100,100]^{d} , and test data from [ − 200,200 ] d [-200,200]^{d} . For tanh, we sample input dimension d d from { 1 , 8 } \{1,8\} , training and validation data from [ − 100,100 ] d [-100,100]^{d} , and test data from [ − 200,200 ] d [-200,200]^{d} . For linear, we use a subset of datasets from Appendix C.3 : 1 and 8 input dimensions with hyper-cube training distributions.

##### Model and hyperparameter settings.

We use the same hyperparameters from Appendix C.1 , except we fix the batch size to 128, as the batch size has minimal impact on models. MLPs with cos activation is hard to optimize, so we only report models with training MAPE less than 1.

### C.5 Max Degree

Dataset details. We consider the task of finding the maximum degree on a graph. Given any input graph G = ( V , E ) G=(V,E) , the label is computed by the underlying function y = g ⁡ ( G ) = max ⁡ ∑ v ∈ 𝒩 ⁡ ( u ) u ∈ G ⁡ 1 y=g(G)=\max\limits_{u\in G}\sum_{v\in\mathcal{N}(u)}1 . For each dataset, we sample the graphs and node features with the following parameters a) Graph structure for training and validation sets. For each dataset, we consider one of the following graph structure: path graphs, cycles, ladder graphs, 4-regular random graphs, complete graphs, random trees, expanders (here we use random graphs with p = 0.8 p=0.8 as they are expanders with high probability), and general graphs (random graphs with p = 0.1 p=0.1 to 0.9 0.9 with equal probability for a broad range of graph structure). We use the networkx library for sampling graphs.

b) Graph structure for test set. We consider the general graphs (random graphs with p = 0.1 p=0.1 to 0.9 0.9 with equal probability).

c) The number of vertices of graphs | V | |V| for training and validation sets are sampled uniformly from [ 20 ​ … ​ 30 ] [20...30] . The number of vertices of graphs | V | |V| for test set is sampled uniformly from [ 50..100 ] [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 1 .

ii) Spurious (continuous) features. Node features in training and validation sets are sampled uniformly from [ − 5.0 , 5.0 ] 3 [-5.0,5.0]^{3} , i.e., a three-dimensional vector where each dimension is sampled from [ − 5.0 , 5.0 ] [-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 [-5.0,5.0]^{3} . In the second case we extrapolate node features, we sample node features uniformly from [ − 10.0 , 10.0 ] 3 [-10.0,10.0]^{3} .

e) We sample 5,000 5,000 graphs for training, 1,000 1,000 graphs for validation, and 2,500 2,500 graphs for testing.

##### Model and hyperparameter settings.

We consider the following Graph Neural Network (GNN) architecture. Given an input graph G G , GNN learns the output h G h_{G} by first iteratively aggregating and transforming the neighbors of all node vectors h u ( k ) h_{u}^{(k)} (vector for node u u in layer k k ), and perform a max or sum-pooling over all node features h u h_{u} to obtain h G h_{G} . Formally, we have h u ( k ) = ∑ v ∈ 𝒩 ⁡ ( u ) MLP ( k ) ​ ( h v ( k − 1 ) , h u ( k − 1 ) ) , h G = MLP ( K + 1 ) ​ ( graph-pooling ​ { h u ( K ) : u ∈ G } ) . \displaystyle h_{u}^{(k)}=\sum\limits_{v\in\mathcal{N}(u)}\text{MLP}^{(k)}\left(h_{v}^{(k-1)},h_{u}^{(k-1)}\right),\quad h_{G}=\text{MLP}^{(K+1)}\left(\text{graph-pooling}\{h_{u}^{(K)}:{u\in G}\}\right). (76) Here, 𝒩 ⁡ ( u ) \mathcal{N}(u) denotes the neighbors of u u , K K is the number of GNN iterations, and graph-pooling is a hyper-parameter with choices as max or sum. h u ( 0 ) h_{u}^{(0)} is the input node feature of node u u . We search the following hyper-parameters for GNNs a) Number of GNN iterations K K is 1 1 .

b) Graph pooling is from max or sum.

c) Width of all MLPs are set to 256 256 .

d) The number of layers for MLP ( k ) \text{MLP}^{(k)} with k = 1 . . K k=1..K are set to 2 2 . The number of layers for MLP ( K + 1 ) \text{MLP}^{(K+1)} is set to 1 1 .

We train the GNNs with the mean squared error (MSE) loss, and Adam and SGD optimizer. We search the following hyper-parameters for training a) Initial learning rate is set to 0.01 0.01 .

b) Batch size is set to 64 64 .

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

d) Number of epochs is set to 300 300 for graphs with continuous node features, and 100 100 for graphs with uniform node features.

##### Test error and model selection.

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

### C.6 Shortest Path

Dataset details. We consider the task of finding the length of the shortest path on a graph, from a given source to target nodes. Given any graph G = ( V , E ) G=(V,E) , the node features, besides regular node features, encode whether a node is source s s , and whether a node is target t t . The edge features are a scalar representing the edge weight. For unweighted graphs, all edge weights are 1 1 . Then the label y = g ⁡ ( G ) y=g(G) is the length of the shortest path from s s to t t on G 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 p=0.6 which are expanders with high probability), and general graphs (random graphs with p = 0.1 p=0.1 to 0.9 0.9 with equal probability for a broad range of graph structure). We use the networkx library for sampling graphs.

b) Graph structure for test set. We consider the general graphs (random graphs with p = 0.1 p=0.1 to 0.9 0.9 with equal probability).

c) The number of vertices of graphs | V | |V| for training and validation sets are sampled uniformly from [ 20 ​ … ​ 40 ] [20...40] . The number of vertices of graphs | V | |V| for test set is sampled uniformly from [ 50..70 ] [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 ] [1.0,5.0] . There are two schemes for test sets, in the first case we do not extrapolate edge weights, so we sample edge weights uniformly from [ 1.0 , 5.0 ] [1.0,5.0] . In the second case we extrapolate edge weights, we sample edge weights uniformly from [ 1.0 , 10.0 ] [1.0,10.0] . All node features are [ h , 𝕀 ⁡ ( v = s ) , 𝕀 ⁡ ( v = t ) ] [h,\mathbb{I}(v=s),\mathbb{I}(v=t)] with h h sampled from [ − 5.0 , 5.0 ] [-5.0,5.0] .

e) After sampling a graph and edge weights, we sample source s s and t t by randomly sampling s s , t t and selecting the first pair s s , s s whose shortest path involves at most 3 3 hops. This enables us to solve the task using GNNs with 3 3 iterations.

f) We sample 10,000 10,000 graphs for training, 1,000 1,000 graphs for validation, and 2,500 2,500 graphs for testing.

We also consider the ablation study of training on random graphs with different p p . We consider p = 0.05..1.0 p=0.05..1.0 and report the test error curve. The other parameters are the same as described above.

##### Model and hyperparameter settings.

We consider the following Graph Neural Network (GNN) architecture. Given an input graph G G , GNN learns the output h G h_{G} by first iteratively aggregating and transforming the neighbors of all node vectors h u ( k ) h_{u}^{(k)} (vector for node u u in layer k k ), and perform a max or sum-pooling over all node features h u h_{u} to obtain h G h_{G} . Formally, we have h u ( k ) = min v ∈ 𝒩 ⁡ ( u ) ⁡ MLP ( k ) ​ ( h v ( k − 1 ) , h u ( k − 1 ) , w ( u , v ) ) , h G = MLP ( K + 1 ) ​ ( min u ∈ G ⁡ h u ) . \displaystyle h_{u}^{(k)}=\min\limits_{v\in\mathcal{N}(u)}\text{MLP}^{(k)}\left(h_{v}^{(k-1)},h_{u}^{(k-1)},w_{(u,v)}\right),\quad h_{G}=\text{MLP}^{(K+1)}\left(\min_{u\in G}h_{u}\right). (77) Here, 𝒩 ⁡ ( u ) \mathcal{N}(u) denotes the neighbors of u u , K K is the number of GNN iterations, and for neighbor aggregation we run both min and sum. h u ( 0 ) h_{u}^{(0)} is the input node feature of node u u . w ( u , v ) w_{(u,v)} is the input edge feature of edge ( u , v ) (u,v) . We search the following hyper-parameters for GNNs a) Number of GNN iterations K K is set to 3 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 256 .

e) The number of layers for MLP ( k ) \text{MLP}^{(k)} with k = 1 . . K k=1..K are set to 2 2 . The number of layers for MLP ( K + 1 ) \text{MLP}^{(K+1)} is set to 1 1 .

We train the GNNs with the mean squared error (MSE) loss, and Adam and SGD optimizer. We consider the following hyper-parameters for training a) Initial learning rate is set to 0.01 0.01 .

b) Batch size is set to 64 64 .

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

d) Number of epochs is set to 250 250 .

We perform the same model selection and validation as in Section C.5 .

### C.7 N-Body Problem

Task description. The n-body problem asks a neural network to predict how n stars in a physical system evolves according to physics laws. That is, we train neural networks to predict properties of future states of each star in terms of next frames, e.g., 0.001 0.001 seconds.

Mathematically, in an n-body system S = { X i } i = 1 n S=\{X_{i}\}_{i=1}^{n} , such as solar systems, all n stars { X i } i = 1 n \{X_{i}\}_{i=1}^{n} exert distance and mass-dependent gravitational forces on each other, so there were n ⁡ ( n − 1 ) n(n-1) relations or forces in the system. Suppose X i X_{i} at time t t is at position 𝒙 i t \bm{x}_{i}^{t} and has velocity 𝒗 i t \bm{v}_{i}^{t} . The overall forces a star X i X_{i} receives from other stars is determined by physics laws as the following 𝑭 i t = G ⋅ ∑ j ≠ i m i × m j ‖ 𝒙 i t − 𝒙 j t ‖ 2 3 ⋅ ( 𝒙 j t − 𝒙 i t ) , \displaystyle\bm{F}_{i}^{t}=G\cdot\sum\limits_{j\neq i}\frac{m_{i}\times m_{j}}{\|\bm{x}_{i}^{t}-\bm{x}_{j}^{t}\|_{2}^{3}}\cdot\left(\bm{x}_{j}^{t}-\bm{x}_{i}^{t}\right), (78) where G G is the gravitational constant, and m i m_{i} is the mass of star X i X_{i} . Then acceralation 𝒂 i t \bm{a}_{i}^{t} is determined by the net force 𝑭 i t \bm{F}_{i}^{t} and the mass of star m i m_{i} 𝒂 i t = 𝑭 i t / m i \displaystyle\bm{a}_{i}^{t}=\bm{F}_{i}^{t}/m_{i} (79) Suppose the velocity of star X i X_{i} at time t t is 𝒗 i t \bm{v}_{i}^{t} . Then assuming the time steps d ​ t dt , i.e., difference between time frames, are sufficiently small, the velocity at the next time frame t + 1 t+1 can be approximated by 𝒗 i t + 1 = 𝒗 i t + 𝒂 i t ⋅ d ​ t . \displaystyle\bm{v}_{i}^{t+1}=\bm{v}_{i}^{t}+\bm{a}_{i}^{t}\cdot dt. (80) Given m i m_{i} , 𝒙 i t \bm{x}_{i}^{t} , and 𝒗 i t \bm{v}_{i}^{t} , our task asks the neural network to predict 𝒗 i t + 1 \bm{v}_{i}^{t+1} for all stars X i X_{i} . In our task, we consider two extrapolation schemes a) The distances between stars ‖ 𝒙 i t − 𝒙 j t ‖ 2 \|\bm{x}_{i}^{t}-\bm{x}_{j}^{t}\|_{2} are out-of-distribution for test set, i.e., different sampling ranges from the training set.

b) The masses of stars m i m_{i} are out-of-distribution for test set, i.e., different sampling ranges from the training set.

Here, we use a physics engine that we code in Python to simulate and sample the inputs and labels. We describe the dataset details next.

##### Dataset details.

We first describe the simulation and sampling of our training set. We sample 100 100 videos of n-body system evolution, each with 500 500 rollout, i.e., time steps. We consider the orbit situation: there exists a huge center star and several other stars. We sample the initial states, i.e., position, velocity, masses, acceleration etc according to the following parameters. a) The mass of the center star is 100 ​ k ​ g 100kg .

b) The masses of other stars are sampled from [ 0.02 , 9.0 ] ​ k ​ g [0.02,9.0]kg .

c) The number of stars is 3 3 .

d) The initial position of the center star is ( 0.0 , 0.0 ) (0.0,0.0) .

d) The initial positions 𝒙 i t \bm{x}_{i}^{t} of other objects are randomly sampled from all angles, with a distance in [ 10.0,100.0 ] ​ m [10.0,100.0]m .

e) The velocity of the center star is 𝟎 \bm{0} .

f) The velocities of other stars are perpendicular to the gravitational force between the center star and itself. The scale is precisely determined by physics laws to ensure the initial state is an orbit system.

For each video, after we get the initial states, we continue to rollout the next frames according the physics engine described above. We perform rejection sampling of the frames to ensure that all pairwise distances of stars in a frame are at least 30 ​ m 30m . We guarantee that there are 10,000 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 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 5,000 data points for each dataset. a) We sample the distance OOD test set to ensure all pairwise distances of stars in a frame are from [ 1..20 ] ​ m [1..20]m , but have in-distribution masses.

b) We sample the mass OOD test set as follows i) The mass of the center star is 200 ​ k ​ g 200kg , i.e., twice of that in the training set.