PaperlayerView exact revision on arXiv

Source-linked AI summary

How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks

Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S. Du, Ken-ichi Kawarabayashi, Stefanie Jegelka

arXiv:2009.11848v5cs.LGcs.AIcs.CVstat.ML

TL;DR

The paper asks why neural networks extrapolate inconsistently beyond their training support. Using gradient-descent and neural-tangent-kernel analyses of MLPs and GNNs, it finds that linearization and training-data diversity govern MLP extrapolation, while task-specific nonlinearities support GNN extrapolation.

  • Problem

    Prior work reports mixed extrapolation results: MLPs fail on some simple tasks, whereas GNNs succeed on some complex algorithmic tasks.

  • Method

    The paper analyzes gradient-descent-trained ReLU MLPs and GNNs through over-parameterized-network and neural tangent kernel connections.

  • Results

    ReLU MLPs linearize along directions from the origin and extrapolate linear targets when training data are diverse; GNN extrapolation is supported when architecture and features encode appropriate nonlinearities.

  • Takeaways & Limitations

    Extrapolation can depend on aligning the target computation with linear behavior learned by the network, while encoding task-specific nonlinearities helps preserve nonlinear algorithmic structure.

  • Takeaways & Limitations

    The analysis assumes ReLU activations for its main feedforward results and proves the GNN hypothesis only for a simplified case.

Abstract

from arXiv · show

We study how neural networks trained by gradient descent extrapolate, i.e., what they learn outside the support of the training distribution. Previous works report mixed empirical results when extrapolating with neural networks: while feedforward neural networks, a.k.a. multilayer perceptrons (MLPs), do not extrapolate well in certain simple tasks, Graph Neural Networks (GNNs) -- structured networks with MLP modules -- have shown some success in more complex tasks. Working towards a theoretical explanation, we identify conditions under which MLPs and GNNs extrapolate well. First, we quantify the observation that ReLU MLPs quickly converge to linear functions along any direction from the origin, which implies that ReLU MLPs do not extrapolate most nonlinear functions. But, they can provably learn a linear target function when the training distribution is sufficiently "diverse". Second, in connection to analyzing the successes and limitations of GNNs, these results suggest a hypothesis for which we provide theoretical and empirical evidence: the success of GNNs in extrapolating algorithmic tasks to new data (e.g., larger graphs or edge weights) relies on encoding task-specific non-linearities in the architecture or features. Our theoretical analysis builds on a connection of over-parameterized networks to the neural tangent kernel. Empirically, our theory holds across different training settings.

1 INTRODUCTION

The paper studies why neural networks show mixed extrapolation behavior, establishing linearization and diversity conditions for MLPs and extending these insights to GNNs.

  • Prior studies report that MLPs fail on some simple polynomial extrapolation tasks, while GNNs generalize to larger graphs in challenging algorithmic tasks.
  • The paper defines extrapolation through gradient-descent training and analyzes how implicit bias constrains over-parameterized neural networks outside the training distribution.
  • MLPs can extrapolate linear targets when the training distribution is sufficiently diverse, especially when it provides support across directions.
  • For dynamic-programming tasks, the paper hypothesizes and supports that GNN extrapolation depends on encoding appropriate nonlinearities in architecture and input representation.
  • The paper connects over-parameterized network training to neural tangent kernels and provides theoretical and empirical support for its extrapolation hypothesis.
  • ReLU MLP predictions converge to linear functions along directions from the origin, with convergence rate O(1/t), limiting extrapolation of most nonlinear functions.
  • The work positions itself as an initial formal analysis of GNN extrapolation, including unseen graph sizes, features, and structures.

2 PRELIMINARIES

The preliminaries define extrapolation relative to the training-distribution support and introduce the gradient-descent setting, MLPs, and graph neural networks.

  • Extrapolation concerns predictions on a domain larger than the training-distribution support, rather than the in-distribution case where training and test domains coincide.
  • The paper defines extrapolation error as the expected loss between model and underlying-function outputs over a distribution outside the training support.
  • The analysis focuses on neural networks trained by gradient descent or variants with squared loss, covering feedforward and graph neural network architectures.
  • GNNs operate on graphs whose nodes and edges carry feature vectors, recursively computing node representations through neighbor aggregation.
  • Graph-level representations or final node representations provide the network output after message-passing iterations.
  • Replacing commonly used sum aggregation and sum readout may help extrapolation according to the paper’s results.

3 HOW FEEDFORWARD NEURAL NETWORKS EXTRAPOLATE

The paper analyzes how gradient-descent-trained ReLU MLPs extrapolate and finds that their predictions become linear along rays outside the training support. MLPs therefore extrapolate nonlinear targets poorly in general, but can extrapolate linear targets when training data covers sufficiently diverse directions.

  • Linear extrapolation behavior: ReLU MLP predictions converge to linear functions along any direction from the origin outside the training data range.Theorem 1 establishes this behavior for two-layer networks in the NTK regime, while experiments observe it across depths, widths, learning rates, and batch sizes.
  • Linear extrapolation behavior: The linearization occurs at rate O(1/t) for two-layer ReLU MLPs in the NTK regime.The convergence coefficients and constants depend on the training data and the direction of extrapolation.
  • Linear extrapolation behavior: ReLU MLPs generally cannot extrapolate targets that are nonlinear along directions away from the origin.Theorem 1 implies that targets must be almost-linear along these directions for the MLP to match them outside the training data.
  • When ReLU MLPs extrapolate well: MLPs extrapolate linear target functions well when the training distribution is sufficiently diverse across directions.Theorem 2 gives a support condition containing a connected subset that reaches every nonzero direction, while experiments find lower extrapolation error when this condition holds.
  • When ReLU MLPs extrapolate well: Spurious correlations can harm extrapolation because they leave combinations of features and directions absent from the training data.The paper uses camels appearing only in deserts as an example of missing feature combinations that violate the diversity condition.
  • MLPs with other activation functions: With non-ReLU activations, extrapolation is strongest when the activation resembles the target, but depth and optimization impose additional limitations.Quadratic activation extrapolates quadratic targets well for two-layer networks but not four-layer networks, while cosine activations are difficult to optimize in high dimensions.

4 HOW GRAPH NEURAL NETWORKS EXTRAPOLATE

The paper explains GNN extrapolation through linear algorithmic alignment: architecture and representations encode task-specific nonlinearities so MLP modules learn nearly linear steps. Theory and experiments support this account across dynamic-programming tasks, while graph structure and feature design determine success.

  • 4.1 HYPOTHESIS: LINEAR ALGORITHMIC ALIGNMENT HELPS EXTRAPOLATION: Min-aggregation GNNs can align with Bellman-Ford because their MLP modules learn the linear update d[k−1][v]+w(v, u), whereas sum-aggregation requires a nonlinear update.The contrast explains why min-aggregation can extrapolate in shortest path while sum-aggregation is not expected to extrapolate well.
  • 4.1 HYPOTHESIS: LINEAR ALGORITHMIC ALIGNMENT HELPS EXTRAPOLATION: GNNs can extrapolate when architecture and input representations encode appropriate nonlinearities, leaving MLP modules to learn nearly linear steps.This is the paper’s linear algorithmic alignment hypothesis for neural networks solving dynamic-programming tasks.
  • 4.2 THEORETICAL AND EMPIRICAL SUPPORT: Max-degree extrapolation improves when training graphs are not restricted in max/min degrees; errors are higher for 4-regular, cycle, and ladder graphs.The theorem’s condition treats graph-structure diversity as analogous to directional diversity for MLPs.
  • 4.2 THEORETICAL AND EMPIRICAL SUPPORT: GNNs with sum-aggregation and sum-readout do not extrapolate well in Max Degree, while encoding max in the readout enables extrapolation.Theorem 3 establishes that one-layer max-readout GNNs learn max degree under its stated training and NTK conditions.
  • 4.2 THEORETICAL AND EMPIRICAL SUPPORT: Shortest-path extrapolation errors follow a U-shape as training-graph sparsity changes, with sparse or dense training likely producing degenerative solutions.The reported behavior indicates that the graph structures covered during training matter differently across tasks.
  • 4.2 THEORETICAL AND EMPIRICAL SUPPORT: For the n-body problem, improved edge representations encode nonlinearity so MLP modules learn linear functions, improving extrapolation to unseen masses or distances.The new edge features add no information; they change the learning problem faced by the MLP modules.

5 CONNECTIONS TO OTHER OUT-OF-DISTRIBUTION SETTINGS

The paper relates extrapolation to other out-of-distribution settings, distinguishing global extrapolation from local robustness and describing possible links through representations, target-domain data, and invariances.

  • 5 CONNECTIONS TO OTHER OUT-OF-DISTRIBUTION SETTINGS: The paper frames related methods as helping extrapolation by learning useful nonlinearities beyond the training range or mapping relevant test data into that range.This perspective is applied to domain adaptation, self-supervised learning, and invariant models.
  • 5 CONNECTIONS TO OTHER OUT-OF-DISTRIBUTION SETTINGS: Domain adaptation may mitigate extrapolation by using target-domain data to align target and source distributions.The passage describes this as a possible connection rather than an established result of the paper.
  • 5 CONNECTIONS TO OTHER OUT-OF-DISTRIBUTION SETTINGS: Self-supervised learning may learn nonlinearities beyond labeled training ranges and map semantically similar out-of-domain examples to similar representations.The paper connects these mechanisms to the reported out-of-distribution robustness of pretrained representations.
  • 5 CONNECTIONS TO OTHER OUT-OF-DISTRIBUTION SETTINGS: Invariant models may effectively increase the training range when variations in invariant features are ignored.This connection depends on the model learning the relevant invariances.
  • 5 CONNECTIONS TO OTHER OUT-OF-DISTRIBUTION SETTINGS: Unlike distributional robustness, which considers small adversarial perturbations, the paper studies more global distributional perturbations.The paper expects modifications that help extrapolation generally may also improve local robustness.

6 CONCLUSION

The paper formally characterizes extrapolation by gradient-trained neural networks through NTK-induced feature spaces, identifying conditions for successful extrapolation. It reports linearization of ReLU MLPs, exact or asymptotic recovery under suitable data geometry, and a GNN hypothesis centered on encoded nonlinearities.

  • Main conclusions: The paper identifies data-geometry conditions under which MLPs extrapolate linear targets exactly or asymptotically.The analysis requires suitable directional coverage and connects exact extrapolation to the global optimum in NTK feature space.
  • Main conclusions: The proposed explanation for GNN extrapolation is that appropriate task-specific nonlinearities must be encoded in architectures or input features.The paper proves the hypothesis in a simplified case and provides empirical support in more general settings.
  • Theoretical framework: The analysis uses the equivalence between sufficiently wide, gradient-trained networks and NTK kernel regression.The learned function is analyzed as a minimum-RKHS-norm interpolant in the NTK-induced feature space.
  • Theoretical framework: NTK feature-space analysis characterizes the learned function by solving a minimum-norm constrained optimization problem that fits the training data.This functional representation makes extrapolation properties analyzable beyond point-wise kernel predictions.

B.2 PROOF OF LEMMA 1

The proof establishes exact extrapolation by expressing the NTK solution as a minimum-norm function and showing that, under the specified training-data geometry, its optimum equals the underlying linear target.

  • B.2 PROOF OF LEMMA 1: The proof reformulates the learned network as a minimum-norm solution in the NTK feature space subject to fitting every training example.The global optimum is analyzed through a convex constrained optimization problem.
  • B.2 PROOF OF LEMMA 1: For two-layer ReLU networks, the infinite-dimensional feature map supplies direction-indexed feature weights used to represent the learned function.The proof groups weights associated with the same direction and uses the resulting coefficients to simplify the optimization.
  • B.2 PROOF OF LEMMA 1: The training-data condition provides at least d linearly independent activated points for every direction w, ensuring solvability of the relevant linear system.These activated points supply enough free variables for constructing the required Lagrange multipliers.
  • B.2 PROOF OF LEMMA 1: Claims (a), (b), and (c) extend constraint satisfaction from individual directions to all directions on the unit sphere.The proof handles a fixed direction, directions with identical activation sets, and rotations that change the activated training points.
  • B.2 PROOF OF LEMMA 1: The resulting optimum satisfies βw + β′w · w = 2βg and yields exact extrapolation, fNTK(x) = g(x), for every x ∈ R^d.Thus the learned function equals the underlying linear target across the entire domain.

B.3 PROOF OF THEOREM 2

The proof of asymptotic extrapolation shows that directional support in the training distribution makes sampled datasets approach a perfect extrapolation dataset, while NTK continuity transfers this convergence to predictions.

  • B.3 PROOF OF THEOREM 2: If the training distribution contains a connected support set intersecting every direction, increasingly large random datasets approach datasets satisfying the exact-extrapolation condition.The proof uses a finite set of key points and convergence in probability of the sampled training set.
  • B.3 PROOF OF THEOREM 2: The proof measures dataset proximity through matching points and the minimum distance to any dataset satisfying the exact-extrapolation condition.This distance is reduced to suitable sub-datasets of size 2d.
  • B.3 PROOF OF THEOREM 2: Adding data does not hurt the exact-extrapolation condition because datasets containing the required key points remain valid after augmentation.This monotonicity allows bounds based on small sub-datasets to control larger training sets.
  • B.3 PROOF OF THEOREM 2: The limiting perfect-dataset predictor equals the true function, so the learned predictions converge to the underlying function.The conclusion follows by combining dataset convergence, kernel convergence, and exact extrapolation for the perfect dataset.
  • B.3 PROOF OF THEOREM 2: Continuity of the ReLU MLP neural tangent kernel makes the training kernel matrix and test-kernel evaluations converge as the sampled dataset approaches a perfect dataset.The proof compares the actual NTK predictor with the predictor trained on the limiting perfect dataset.

B.4 PROOF OF COROLLARY 1

The corollary applies ReLU MLP linearization to a GNN maximum-degree task, showing that correct extrapolation requires a non-piecewise-linear module outside the training range.

  • B.4 PROOF OF COROLLARY 1: In the maximum-degree task, a GNN with linear aggregations must simulate the underlying maximum function to extrapolate correctly.The required computation is assigned to the MLP(K+1) module.
  • B.4 PROOF OF COROLLARY 1: Because the maximum function cannot be decomposed into piecewise-linear compositions, the relevant GNN module must learn a function that is not piecewise linear outside training data.This requirement concerns domains beyond the training data range.
  • B.4 PROOF OF COROLLARY 1: ReLU MLP linearization implies that an over-parameterized two-layer MLP(K+1) converges to linear functions along directions from the origin.Consequently, there are domains where the GNN cannot learn the correct target function.

B.5 PROOF OF THEOREM 3

The proof analyzes a simple GNN through its Graph Neural Tangent Kernel and shows that sufficiently diverse graph training data let it recover the maximum-degree function.

  • Architecture: The analyzed one-layer GNN operates on graphs without edge features and produces a graph-level output from node features.The architecture and its GNTK are defined before analyzing the learned function.
  • Feature-space reduction: The proof treats uniform node features as a one-dimensional problem with positive and negative directions.This reduces the analysis to combining coefficients associated with the two directions.
  • GNTK analysis: The learned predictor is characterized as the min-norm solution in GNTK feature space that fits all training graphs.The optimization is expressed using Lagrange parameters and constraints derived from the GNTK feature representation.
  • Conclusion: The resulting solution fits the training data and equals the underlying maximum-degree function g(G) = maxu∈G hu.Thus, under the stated training-data condition, the simple GNN learns the desired function.
  • Diversity condition: With at least four training graphs whose relevant vectors are linearly independent, suitable Lagrange parameters exist for the target solution.The proof rewrites the remaining constraints as a linear system and invokes linear independence.

B.6 PROOF OF LEMMA 2

This proof connects over-parameterized MLP training to kernel regression by representing the neural tangent kernel in an RKHS feature space and identifying its minimum-norm interpolant.

  • Feature-space decomposition: The training-data feature span is W = span(φ(x1), φ(x2), ..., φ(xn)).The NTK solution is decomposed into components inside and orthogonal to this span.
  • Kernel regression: The NTK predictor must fit the training labels, so its in-span component is determined by the kernel-regression solution.The constraints φ(xi)^TβNTK = yi imply φ(xi)^Tβ0 = yi.
  • Uniqueness: Positive definiteness of the training NTK makes the kernel-regression solution unique when no two training examples are parallel, or via a pseudoinverse otherwise.A bias term can enforce the non-parallel condition.
  • Minimum norm: The orthogonal component vanishes in the minimum-norm solution because ||β0 + β1||_2^2 = ||β0||_2^2 + ||β1||_2^2.Any nonzero orthogonal component would increase the RKHS norm without changing training predictions.
  • ReLU NTK construction: For ReLU MLPs, the paper derives the NTK using the general recursive kernel framework and an infinite-dimensional Gaussian feature map.The ReLU and derivative are substituted into the recursive covariance construction.

C EXPERIMENTAL DETAILS

The experiments test the theory across nonlinear MLP tasks, GNN algorithmic tasks, and varied architectures, representations, training distributions, and extrapolation regimes.

  • MLP tasks: The experiments evaluate nonlinear functions including quadratic, cosine, square root, and l1 norm targets across dimensions and distribution shapes.MLP experiments vary training and test supports, including hyper-cubes, spheres, and hyper-balls.
  • GNN tasks: GNN experiments study maximum degree and shortest-path tasks under changes in graph structure, node count, edge weights, and node features.Training graphs include specific structures and random graphs spanning different sparsities.
  • Representation ablation: For n-body reasoning, the experiments compare original edge features with improved representations designed to make the target function more linear.The original representation fails on unseen masses and distances, whereas the improved representation enables extrapolation in the reported setting.
  • MLP results: More than 96% of R-squared results exceed 0.99, showing that learned MLP functions become linear along most directions immediately outside training support.This empirically confirms the predicted fast convergence to directional linearity.
  • NTK validation: Exact two-layer NTK regression achieves test losses of precisely 0 up to machine precision on the Lemma 1 data regime.This experiment directly checks the theoretical extrapolation result for the NTK setting.

C.7 N-BODY PROBLEM

The n-body experiments test GNN prediction under out-of-distribution distances and masses, comparing an original representation with an improved edge-feature representation.

  • Task: The task predicts each star’s next-frame state in an orbiting n-body system governed by distance- and mass-dependent gravitational interactions.The GNN predicts node outputs corresponding to the true labels in the next frame.
  • Training distribution: The training data contain three stars, a 100kg center star, other-star masses in [0.02, 9.0]kg, and distances in [10.0, 100.0]m.The orbit simulations use 100 videos with 500 rollout steps and provide 10,000 training examples.
  • OOD regimes: Test extrapolation changes either pairwise distances or masses beyond training while keeping the other variable in distribution.The mass-OOD set doubles the center mass to 200kg and expands other-star masses to [0.04, 18.0]kg.
  • GNN model: The model is a one-iteration Interaction Network on a complete graph, aggregating interactions among all stars.Node features include mass, position, and velocity, while edge features encode pairwise relations.
  • Representations: The improved edge representation makes the underlying target function more linear than the original zero edge representation.This comparison is the experimental test of whether representation choice supports extrapolation.

D.1 VISUALIZATION RESULTS

These visualizations compare MLP predictions with the underlying functions inside and outside training support. Outside distribution, learned functions become linear along directions from the origin, although linear targets can extrapolate successfully.

  • General visualization pattern: MLP predictions agree with the underlying functions within the training range, while OOD predictions become linear along directions from the origin.This directional linearity is identified as the behavior explained by Theorem 1.
  • Other nonlinear targets: Cosine and square-root visualizations compare black MLP OOD predictions with grey underlying functions and blue in-distribution predictions.The cosine plots include one-dimensional and two-dimensional cases, while the square-root plots use analogous dimensions.
  • L1 targets: For the absolute function, one MLP successfully extrapolates while another with different hyperparameters fails.The same success-versus-failure pattern appears for both |x| and |x1| + |x2|.
  • Linear targets: MLPs can extrapolate linear functions when the training distribution is sufficiently diverse.Both square and circle supports intersect all directions, explaining successful extrapolation for y = x1 + x2.

D.2 EXTRA EXPERIMENTAL RESULTS

Additional experiments examine linear-function extrapolation under training supports that may miss directions and GNN extrapolation across graph structures. They show that support geometry and pooling choice affect extrapolation, while shortest-path errors vary non-monotonically with training sparsity.

  • MLP training distributions: MLPs trained on hypercubes intersecting all directions extrapolate underlying linear functions, unlike supports that restrict or fix dimensions.The density plot compares “all”, “fix1”, and “posX” training distributions using test errors in MAPE.
  • Maximum degree: GNNs with graph-level max-pooling extrapolate across OOD node features, graph structures, and graph sizes under Theorem 3’s training condition.The node features are continuous and intended not to contribute to maximum degree.
  • Maximum degree: GNNs with sum-pooling fail to extrapolate, whereas max-pooling can extrapolate under appropriate training sets.The comparison is presented as validating Corollary 1 and Theorem 3.
  • Shortest path: Shortest-path extrapolation errors follow a U-shaped curve as training-graph edge probability changes.Errors are high for very sparse or very dense training graphs, and the same pattern appears for specific graph structures.
Loading 2009.11848v5…