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

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

Figure 4(a) 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 4(a) shows that MLPs can extrapolate well when the target function is linear. However, this is not always true. In 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 1 states that over-parameterized MLPs can learn a linear target function with only 2 ​ d 2d examples.

###### Lemma 1 .

Let g ⁡ ( 𝐱 ) = 𝛃 ⊤ ​ 𝐱 g(\bm{x})=\bm{\beta}^{\top}\bm{x} be the target function for 𝛃 ∈ ℝ d \bm{\beta}\in\mathbb{R}^{d} . Suppose { 𝐱 i } i = 1 n \{\bm{x}_{i}\}_{i=1}^{n} contains an orthogonal basis { 𝐱 ^ i } i = 1 d \{\hat{\bm{x}}_{i}\}_{i=1}^{d} and { − 𝐱 ^ i } i = 1 d \{-\hat{\bm{x}}_{i}\}_{i=1}^{d} . If we train a two-layer ReLU MLP f f on { ( 𝐱 i , y i ) } i = 1 n \{(\bm{x}_{i},y_{i})\}_{i=1}^{n} with squared loss in the NTK regime, then f ⁡ ( 𝐱 ) = 𝛃 ⊤ ​ 𝐱 f(\bm{x})=\bm{\beta}^{\top}\bm{x} for all 𝐱 ∈ ℝ d \bm{x}\in\mathbb{R}^{d} .

Lemma 1 is mainly of theoretical interest, as the 2 ​ d 2d examples need to be carefully chosen. Theorem 2 builds on Lemma 1 and identifies a more practical condition for successful extrapolation: if the support of the training distribution covers all directions (e.g., a hypercube that covers the origin), the MLP converges to a linear target function with sufficient training data.

###### Theorem 2 .

(Conditions for extrapolation). Let g ⁡ ( 𝐱 ) = 𝛃 ⊤ ​ 𝐱 g(\bm{x})=\bm{\beta}^{\top}\bm{x} be the target function for 𝛃 ∈ ℝ d \bm{\beta}\in\mathbb{R}^{d} . Suppose { 𝐱 i } i = 1 n \{\bm{x}_{i}\}_{i=1}^{n} is sampled from a distribution whose support 𝒟 \mathcal{D} contains a connected subset 𝒮 \mathcal{S} , where for any non-zero 𝐰 ∈ ℝ d \bm{w}\in\mathbb{R}^{d} , there exists k > 0 k>0 so that k ​ 𝐰 ∈ 𝒮 k\bm{w}\in\mathcal{S} . If we train a two-layer ReLU MLP f : ℝ d → ℝ f:\mathbb{R}^{d}\rightarrow\mathbb{R} on { ( 𝐱 i , y i ) } i = 1 n \{(\bm{x}_{i},y_{i})\}_{i=1}^{n} with squared loss in the NTK regime, f ⁡ ( 𝐱 ) ⟶ p 𝛃 ⊤ ​ 𝐱 f(\bm{x})\stackrel{{\scriptstyle p}}{{\longrightarrow}}\bm{\beta}^{\top}\bm{x} as n → ∞ n\rightarrow\infty .

Experiments: 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., 𝒟 \mathcal{D} includes all directions. Empirically, the extrapolation error is indeed small when the condition of Theorem 2 is satisfied (“all” in Figure 4(b) ). In contrast, the extrapolation error is much larger when the training examples are restricted to only some directions (Figure 4(b) and Figure 3 ).

Relating to previous works, Theorem 2 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 an image collection. Therefore, the condition for Theorem 2 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 a linear function if the training data has full (feature) rank. MLPs are more expressive, so identifying the linear target function requires additional constraints.

To summarize, we 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 σ ⁡ ( x ) = tanh ⁡ ( x ) \sigma(x)=\tanh(x) , cosine σ ⁡ ( x ) = cos ⁡ ( x ) \sigma(x)=\cos(x) ( Lapedes & Farber, 1987 ; McCaughan, 1997 ; Sopena & Alquezar, 1994 ) , and quadratic σ ⁡ ( x ) = x 2 \sigma(x)=x^{2} ( Du & Lee, 2018 ; Livni et 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, but 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 as 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.

## 4 How Graph Neural Networks Extrapolate

Above, we saw 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 et 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 an example: training GNNs to solve the shortest path problem. For this task, prior works observe that a 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 ) : 𝒉 u ( k ) = min v ∈ 𝒩 ⁡ ( u ) ⁡ MLP ( k ) ​ ( 𝒉 u ( k − 1 ) , 𝒉 v ( k − 1 ) , 𝒘 ( v , u ) ) . \displaystyle\bm{h}_{u}^{(k)}=\min_{v\in\mathcal{N}(u)}\text{MLP}^{(k)}\big(\bm{h}_{u}^{(k-1)},\bm{h}_{v}^{(k-1)},\bm{w}_{(v,u)}\big). (2) We first provide an intuitive explanation (Figure 2(a) ). Shortest path can be solved by the Bellman-Ford (BF) algorithm ( Bellman, 1958 ) with the following update: d ⁡ [ k ] ​ [ u ] = min v ∈ 𝒩 ⁡ ( u ) ⁡ d ⁡ [ k − 1 ] ​ [ v ] + 𝒘 ⁡ ( v , u ) , d[k][u]=\min_{v\in\mathcal{N}(u)}d[k-1][v]+\bm{w}(v,u), (3) where 𝒘 ⁡ ( v , u ) \bm{w}(v,u) is the weight of edge ( v , u ) (v,u) , and d ​ [ k ] ​ [ u ] d[k][u] is the shortest distance to node u u within k k steps. The two equations can be easily aligned: GNNs simulate the BF algorithm if its MLP modules learn a linear function d ​ [ k − 1 ] ​ [ v ] + 𝒘 ​ ( v , u ) d[k-1][v]+\bm{w}(v,u) . Since MLPs can extrapolate linear tasks, this “alignment” may explain why min-aggregation GNNs can extrapolate well in this task.

For comparison, we can reason why we would not expect GNNs with the more commonly used sum-aggregation (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 1 suggests that they will not extrapolate most nonlinear functions outside the training support.

We can generalize the above intuition to 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. 1 ) ( Xu et al., 2020 ) .

###### Definition 3 .

Dynamic programming (DP) is a recursive procedure with updates Answer ​ [ k ] ​ [ s ] = DP-Update ​ ( { Answer ​ [ k − 1 ] ​ [ s ′ ] } , s ′ = 1 ​ … ​ n ) , \displaystyle\text{Answer}[k][s]=\text{DP-Update}(\left\{\text{Answer}[k-1][s^{\prime}]\right\},s^{\prime}=1...n), (4) where Answer [ k ] ​ [ s ] [k][s] is the solution to a sub-problem indexed by iteration k k and state s s , and 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-linearities into the model architecture and input representations so 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 : 𝒳 → ℝ f:\mathcal{X}\rightarrow\mathbb{R} be the underlying function and 𝒩 \mathcal{N} a neural network with m m MLP modules. Suppose there exist m m linear functions { g i } i = 1 m \{g_{i}\}_{i=1}^{m} so that by replacing 𝒩 \mathcal{N} ’s MLP modules with g i g_{i} ’s, 𝒩 \mathcal{N} simulates f f . Given ϵ > 0 \epsilon>0 , there exists { ( x i , f ⁡ ( x i ) ) } i = 1 n ⊂ 𝒟 ⊊ 𝒳 \{(x_{i},f(x_{i}))\}_{i=1}^{n}\subset\mathcal{D}\subsetneq\mathcal{X} so that 𝒩 \mathcal{N} trained on { ( x i , f ⁡ ( x i ) ) } i = 1 n \{(x_{i},f(x_{i}))\}_{i=1}^{n} by GD with squared loss learns f ^ \hat{f} with ‖ f ^ − f ‖ < ϵ \|\hat{f}-f\|<\epsilon .

Our hypothesis builds on the algorithmic alignment framework of ( Xu et al., 2020 ) , which states that a 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 is to encode log-and-exp transforms in 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 2(b) ). Sometimes, we can decompose the target function f f as f = g ∘ h f=g\circ h into a feature embedding h h and a “simpler” target function g g that our model can extrapolate well. We can obtain h h via specialized features or feature transforms using domain knowledge ( Lample & Charton, 2020 ; Webb et al., 2020 ) , or via representation learning (e.g., BERT) with unlabeled out-of-distribution data in 𝒳 ∖ 𝒟 \mathcal{X}\setminus\mathcal{D} ( Chen et al., 2020 ; Devlin et al., 2019 ; Hu et al., 2020 ; Mikolov et al., 2013b ; Peters et al., 2018 ) . This brings a new perspective of how representations help extrapolation in various application areas. For example, in natural language processing, pretrained representations ( Mikolov et 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” or features is crucial for deep learning models as the financial markets may frequently be in extrapolation regimes ( Banz, 1981 ; Fama & French, 1993 ; Ross, 1976 ) .

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, we provide theoretical and empirical support for our hypothesis.

### 4.2 Theoretical and Empirical Support

We validate our hypothesis on three DP tasks: max degree, shortest path, and n n -body problem, and prove the hypothesis for max degree. We highlight the role of graph structures in extrapolation.

Theoretical analysis. We start with a simple yet fundamental task: learning the max degree of a graph, a special case of DP with one iteration. As a corollary of Theorem 1 , the commonly used sum-based GNN (Eqn. 1 ) cannot extrapolate well (proof in Appendix B.4 ).

###### Corollary 1 .

GNNs with sum-aggregation and sum-readout do not extrapolate well in Max Degree.

To achieve linear algorithmic alignment, we can encode the only non-linearity, the max function, in the readout. Theorem 3 confirms that a GNN with max-readout can extrapolate well in this task.

###### Theorem 3 .

(Extrapolation with GNNs). Assume all nodes have the same feature. Let g g and g ′ g^{\prime} be the max/min degree function, respectively. Let { ( G i , g ( G i ) } i = 1 n \{(G_{i},g(G_{i})\}_{i=1}^{n} be the training set. If { ( g ⁡ ( G i ) , g ′ ​ ( G i ) , g ⁡ ( G i ) ⋅ N i max , g ′ ​ ( G i ) ⋅ N i min ) } i = 1 n \{(g(G_{i}),g^{\prime}(G_{i}),g(G_{i})\cdot N^{\max}_{i},g^{\prime}(G_{i})\cdot N^{\min}_{i})\}_{i=1}^{n} spans ℝ 4 \mathbb{R}^{4} , where N i max N^{\max}_{i} and N i min N^{\min}_{i} are the number of nodes that have max/min degree on G i G_{i} , then one-layer max-readout GNNs trained on { ( G i , g ⁡ ( G i ) ) } i = 1 n \{(G_{i},g(G_{i}))\}_{i=1}^{n} with squared loss in the NTK regime learn g g .

Theorem 3 does not follow immediately from Theorem 2 , because MLP modules in GNNs only receive indirect supervision. We analyze the Graph NTK ( Du et al., 2019b ) to 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 3 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 C -regular graphs (regular graphs with degree C 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 et 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 max-readout are better than GNNs with sum-readout (Figure 6(a) ), confirming Corollary 1 and Theorem 3 . For shortest path, GNNs with min-readout and min-aggregation are better than GNNs with sum-readout (Figure 6(a) ).

Experiments confirm the importance of training graphs structure (Figure 7 ). Interestingly, the two tasks favor different graph structure. For max degree, as Theorem 3 predicts, GNNs extrapolate well when trained on trees, complete graphs, expanders, and general graphs, and extrapolation errors are higher when trained on 4-regular, cycles, or ladder graphs. For shortest path, extrapolation errors follow a U-shaped curve as we change the sparsity of training graphs (Figure 7(b) and Figure 18 in Appendix). Intuitively, models trained on sparse or dense graphs likely learn degenerative solutions.