How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks

1 INTRODUCTION

• The analysis connects over-parameterized neural-network training dynamics to kernel regression through the neural tangent kernel.This explicitly incorporates the implicit bias induced by gradient-descent training despite neural networks’ high capacity and universal approximation.Source: §1 · pp. 1, 3
• ReLU MLPs trained by gradient descent converge to linear functions along directions from the origin, limiting extrapolation for most nonlinear tasks.The convergence rate for two-layer networks is O(1/t), and convergence often occurs close to the training data.Source: §1 · pp. 2–3
• MLPs can extrapolate linear targets when the training distribution has sufficiently diverse geometry.These results are non-asymptotic and quantify behavior near the training distribution rather than merely relying on finitely many ReLU linear regions.Source: §1 · p. 2, Figure 1 · p. 2
• GNNs can extrapolate dynamic-programming tasks when appropriate task-specific nonlinearities are encoded in the architecture and input representation.The hypothesis is proved in a simplified Graph NTK setting and supported empirically on max degree, shortest paths, and n-body tasks, including unseen graph sizes, structures, edge weights, and node features.Source: §1 · p. 2, Figure 2 · p. 2

1.1 RELATED WORK

• Prior work reported poor MLP extrapolation on simple polynomials, whereas this paper gives general ReLU MLP conditions and a non-asymptotic extrapolation analysis.It also presents the first theoretical study of GNN extrapolation and extends extrapolation to unseen features and structures.Source: §1.1 · p. 3

Setting

• The analysis studies feedforward and graph neural networks trained by gradient descent or its variants using squared loss.The underlying task is to learn a function g over domains such as vectors or graphs from labeled training data.Source: §2 · p. 3
• Extrapolation evaluates a model on domain points outside the training support, with error defined as expected loss between predictions and the underlying function.In-distribution generalization instead assumes identical training and test supports.Source: §2 · p. 3

Graph Neural Networks

• GNNs operate on graphs with node and edge features, recursively updating node representations through neighbor aggregation.They may additionally aggregate final node representations into a graph representation through a readout operation, depending on the task.Source: §2 · p. 3
• The preliminaries note that replacing commonly used sum aggregation and sum readout may improve extrapolation.This motivates later analysis of GNN design choices in Section 4.Source: §2 · p. 3

3.1 LINEAR EXTRAPOLATION BEHAVIOR OF RELU MLPS

• Outside the training range, two-layer ReLU MLP predictions converge along every direction from the origin to a linear function at rate O(1/t).This behavior is formalized in the NTK regime and is empirically observed across depths, widths, learning rates, and batch sizes.Source: §3.1 · pp. 4–5
• Because of this directional linearity, ReLU MLPs fail on most nonlinear targets but extrapolate linear functions and suitable directionally linear functions such as the L1 norm.They do not extrapolate quadratic or square-root targets correctly, while suitable hyperparameters enable correct L1-norm extrapolation.Source: §3.1 · p. 5

3.2 WHEN RELU MLPS PROVABLY EXTRAPOLATE WELL

• ReLU MLPs learn linear targets exactly from carefully chosen 2d examples, while broader success requires training support that contains a connected subset reaching every direction.Experiments show small extrapolation error when all directions are represented and much larger error when examples cover only some directions; missing feature combinations from spurious correlations can therefore cause incorrect extrapolation.Source: §3.2 · p. 5

3.3 MLPS WITH OTHER ACTIVATION FUNCTIONS

• MLPs with non-ReLU activations extrapolate best when activation and target match, but performance depends on training-range width, network depth, and input dimensionality.Tanh requires a sufficiently wide training range, quadratic activation extrapolates well only with two-layer networks, and cosine activation is difficult to optimize in high dimensions.Source: §3.3 · p. 6

4.1 HYPOTHESIS: LINEAR ALGORITHMIC ALIGNMENT HELPS EXTRAPOLATION

• In shortest path, min-aggregation aligns GNN updates with Bellman-Ford’s linear distance-plus-edge-weight step, whereas sum-aggregation requires learning a nonlinear function that MLPs generally cannot extrapolate.The same intuition extends to dynamic-programming tasks with recursive structures similar to GNN updates.Source: §4.1 · pp. 6–7
• Linear algorithmic alignment hypothesizes that neural networks extrapolate well when architecture or representations encode nonlinearities and MLP modules only learn nearly linear steps.The framework requires modules to align with linear functions for extrapolation, rather than merely with easy-to-learn functions for interpolation.Source: §4.1 · p. 7

4.2 THEORETICAL AND EMPIRICAL SUPPORT

• Encoding max/min operations improves extrapolation: max-readout GNNs outperform sum-readout on max degree, while min-readout and min-aggregation outperform sum-readout on shortest path.For the n-body problem, improved edge features also improve extrapolation to unseen masses or distances by making the learned functions linear.Source: Figure 6 · p. 8, §4.2 · pp. 8–9
• For max degree, sum-aggregation and sum-readout do not extrapolate well, while encoding max in the readout enables max-readout GNNs to learn the target in the NTK regime.The guarantee requires diverse training graph structures; it fails when all training graphs share the same max or min degrees.Source: §4.2 · p. 8
• Training-graph structure materially affects extrapolation: max-degree models perform well on trees, complete graphs, expanders, and general graphs, while shortest-path errors vary U-shaped with training-graph sparsity.Encoding nonlinear operations or transformations can also support extrapolation beyond GNNs, including log-and-exp features for multiplication and symbolic-operation libraries.Source: §4.1 · p. 7, §4.2 · pp. 8–9

Overview

• Related methods may improve extrapolation by learning useful nonlinearities beyond the training range or mapping relevant test data into the training range.These two mechanisms organize the connections discussed in this section.Source: §5 · p. 9

Domain adaptation

• Domain adaptation uses target-domain information to align source and target distributions, potentially inducing useful nonlinearities, though correctness depends on the label distribution.Typical strategies incorporate unlabeled target samples during training.Source: §5 · p. 9

Self-supervised learning

• Self-supervised learning can learn nonlinearities beyond labeled data and map semantically similar out-of-domain examples to representations inside the training distribution.This provides a proposed explanation for why pre-trained representations such as BERT improve out-of-distribution robustness.Source: §5 · p. 9

Invariant models

• Invariant models can effectively increase the training range by learning features whose invariant variations are ignored, when those invariances are correctly identified.The passage notes this can occur in the linear case and with confounders or anti-causal variables.Source: §5 · p. 9

Distributional robustness

• Distributional robustness targets local adversarial perturbations, whereas this work studies more global perturbations, although extrapolation improvements may also enhance local robustness.The connection is presented as an expectation rather than a demonstrated result.Source: §5 · p. 9

6 CONCLUSION

• The conclusion identifies conditions for desired MLP and GNN extrapolation and attributes GNN success on complex algorithmic tasks to task-specific nonlinearities in architectures or features.The theoretical results and the paper’s empirical findings agree with prior empirical results.Source: §6 · p. 9

A THEORETICAL BACKGROUND

• In the NTK regime, sufficiently wide neural networks trained by gradient descent are equivalent to kernel gradient descent and converge to an NTK kernel-regression solution.For suitable scaling and random initialization, the neural tangent kernel matrix remains effectively constant during training.Source: A THEORETICAL BACKGROUND · p. 16, §B.1 · p. 17
• The analysis characterizes extrapolation error through the difference between the NTK-learned function and the true function over the input domain.The framework includes explicit feature maps for two-layer ReLU MLPs and analogous constructions for GNNs.Source: A THEORETICAL BACKGROUND · pp. 16–17
• The NTK solution is the minimum-RKHS-norm function in the induced feature space that fits all training data.This feature-space formulation supplies a functional representation for analyzing extrapolation beyond pointwise kernel evaluations.Source: A THEORETICAL BACKGROUND · p. 16
• The paper characterizes over-parameterized networks by solving a constrained minimum-norm problem in NTK feature space subject to fitting the training data.This characterization underlies the analysis of learned-function properties in the NTK regime.Source: A THEORETICAL BACKGROUND · p. 17
• Provable extrapolation generally requires geometric conditions on training data directions and, for graphs, their structure, while infinite-dimensional NTK spaces create additional technical challenges.These conditions determine how the training distribution interacts with the network’s learned solution.Source: A THEORETICAL BACKGROUND · p. 17

B.1 PROOF OF THEOREM 1

• Theorem 1’s proof studies ReLU MLP outputs in the NTK RKHS and shows convergence toward a linear function along any direction from the origin.For the analyzed terms, the proof obtains an O(1/ϵ) convergence rate as the distance t along the direction grows.Source: §B.1 · pp. 17, 19

Constrained optimization in NTK feature space

• In the NTK regime, gradient-descent learning is equivalent to kernel regression, whose learned function is the minimum-RKHS-norm interpolant in the NTK feature space.This functional characterization reduces the proof to a constrained optimization problem over NTK feature-space weights.Source: §B.2 · pp. 20–21

Global optimum via the geometry of training data

• The constrained NTK optimization has a global optimum equal to the true function, yielding exact extrapolation f(x) = g(x) for all x ∈ R^d.The objective is convex and the constraints are affine, so KKT conditions establish global optimality; the training-data geometry supplies the required multipliers.Source: §B.2 · pp. 22, 24
• The proof uses training data containing an orthogonal basis and opposite vectors to construct Lagrange multipliers satisfying the constraints for every activation direction.The construction covers fixed activation sets, directions sharing an activation pattern, and rotations that change which training points activate.Source: §B.2 · pp. 22–23
• Because every direction activates at least d linearly independent training vectors, the resulting linear systems have enough free variables to admit solutions.This establishes the first part of the multiplier construction for arbitrary unit-sphere directions.Source: §B.2 · p. 23

Asymptotic convergence of data sets

• If the support contains a connected set intersecting every direction, the random training set converges in probability to a dataset satisfying the exact-extrapolation condition.The proof bounds the distance to a perfect dataset using finite subsets and shows this distance vanishes as n →∞.Source: §B.3 · pp. 24–26
• The argument relies on exact-extrapolation datasets containing finitely many key points, so adding further training data does not destroy the condition.This permits the random dataset to approach a perfect dataset while only finitely many points differ from its key points.Source: §B.3 · pp. 25–26

Asymptotic convergence of predictions

• Continuity of the two-layer ReLU MLP neural tangent kernel makes the training kernel matrix converge to the kernel matrix of a perfect dataset.The proof combines dataset convergence with NTK continuity and convergence of the corresponding kernel vectors.Source: §B.3 · p. 26
• Consequently, the NTK prediction function learned from the random training set converges to the function learned from the perfect dataset, which equals the true underlying function.Positive definiteness of the neural tangent kernel ensures the inverse-kernel prediction expression is well-defined.Source: §B.3 · p. 26

• To extrapolate the maximum-degree task, a GNN with linear aggregations must simulate the underlying maximum function.Source: §B.4 · p. 26
• Because max cannot be composed from piecewise-linear functions, MLP(K+1) must learn a non-piecewise-linear function outside the training range.Source: §B.4 · p. 27
• Under the two-layer overparameterized assumption, Theorem 1 implies MLP(K+1) becomes linear along every direction from the origin, so some domains cannot be learned correctly.Source: §B.4 · p. 27

Architecture and GNTK

• The proof studies a simple one-layer GNN on graphs without edge features and derives its corresponding GNTK feature map.The analysis is presented for simplicity but can extend to edge features, ReLU activations, alternative neighbor aggregation, and graph-level pooling.Source: §B.5 · p. 27

Minimum-norm solution

• Assuming uniform node features, the proof reduces Gaussian parameter directions to the two cases +1 and −1 and formulates the GNTK minimum-norm solution fitting all training graphs.The two-direction reduction simplifies the Lagrange optimization used to characterize the learned function.Source: §B.5 · pp. 27–28

Recovered function

• The resulting solution fits the training data and is equivalent to the underlying reasoning function g(G) = maxu∈G hu.The proof identifies this function as the graph’s maximum degree.Source: §B.5 · p. 28

Training-data condition

• With at least four training graphs whose associated maximum- and minimum-degree vectors are linearly independent, the GNN learns the underlying maximum-degree function.This condition guarantees Lagrange parameters satisfying the remaining constraints.Source: §B.5 · pp. 28–29

B.6 PROOF OF LEMMA 2

• B.6 proves that the NTK coefficient vector decomposes as β_NTK = β_0 + β_1 over the training-feature span and its orthogonal complement.The training-feature span is W = span(φ(x_1), ..., φ(x_n)).Source: §B.6 · p. 29
• B.6 uses interpolation and positive definiteness to establish a unique kernel-regression solution, with the minimum-norm interpretation available by pseudoinverse.Positive definiteness holds assuming no two training data are parallel, which can be enforced with a bias term.Source: §B.6 · p. 29
• B.6 shows the NTK solution has no orthogonal component, because the training kernel-regression solution lies in the feature span and orthogonality gives ||β_NTK||_2^2 = ||β_0||_2^2 + ||β_1||_2^2.Thus β_1 = 0, and the NTK solution has no larger norm than any interpolating coefficient vector.Source: §B.6 · p. 29

B.7 PROOF OF LEMMA 3

• B.7 derives the neural tangent kernel for an h-layer MLP through recursively defined covariance and derivative-covariance quantities.The proof specializes the general framework to ReLU activation for the two-layer MLP case.Source: §B.7 · p. 30
• B.7 computes the two-layer ReLU MLP NTK by substituting σ(a) = max(0, a) and σ̇(a) = I(a ≥ 0) into the recursive framework.The resulting calculation uses Gaussian sampling of the pair (w^T x, w^T x′).Source: §B.7 · pp. 30–31
• B.7 constructs an infinite-dimensional feature map whose inner product equals the NTK expectation, with feature density determined by Gaussian weights.The map uses weights w^(k) sampled from N(0, I), and its validity follows by matching the kernel’s expected inner product.Source: §B.7 · p. 31

C.1 LEARNING SIMPLE NON-LINEAR FUNCTIONS

• Nonlinear-function experiments study quadratic, cosine, square-root, and l1-norm targets across various dimensions and training/test distributions to assess MLP extrapolation claims.Source: C EXPERIMENTAL DETAILS · p. 31, §C.1 · p. 32

C.3 LEARNING LINEAR FUNCTIONS

• Linear-target experiments test whether MLPs extrapolate under diverse supports and fail when training data omit directions through fixed, positive, or negative coordinate restrictions.The study uses linear targets g(x) = Ax and compares unrestricted distributions with supports that restrict selected dimensions.Source: C EXPERIMENTAL DETAILS · p. 31, §C.3 · p. 35

C.4 MLPS WITH COSINE, QUADRATIC, AND TANH ACTIVATION

• Additional experiments evaluate whether cosine, quadratic, and tanh activations alter extrapolation behavior relative to the ReLU-based settings.The evaluated targets include quadratic and linear functions, with the section focused on MLPs using cosine, quadratic, and tanh activation functions.Source: C EXPERIMENTAL DETAILS · p. 31, §C.4 · p. 35

C.2 R-SQUARED FOR OUT-OF-DISTRIBUTION DIRECTIONS

• For simple non-linear functions, more than 96% of R-squared results exceed 0.99, and along most directions, MLP learned functions become linear immediately outside training-data support.The study computes R-squared over more than 2,000 combinations of datasets, test/train distributions, and hyperparameters.Source: §C.2 · p. 34

C.3 LEARNING LINEAR FUNCTIONS

• NTK kernel regression achieves test losses precisely 0 up to machine precision on training sets satisfying Lemma 1, empirically confirming exact extrapolation.The experiment uses 2d examples formed from orthogonal basis vectors and their opposites, transformed by 100 random orthogonal matrices.Source: §C.3 · p. 35

C.4 MLPS WITH COSINE, QUADRATIC, AND TANH ACTIVATION

• C.4 evaluates cosine, quadratic, tanh, and linear targets using separate training and test ranges, with 20,000 training and test examples and batch size fixed at 128.Cosine-activation models are reported only when training MAPE is below 1 because they are difficult to optimize.Source: §C.4 · p. 36

C.6 SHORTEST PATH

• C.6 SHORTEST PATH uses GNNs to predict source-target shortest-path length, with graph sizes, edge weights, and test distributions designed to evaluate structural and edge-weight extrapolation.Training and validation graphs have 20–40 vertices and edge weights in [1.0, 5.0], while test graphs have 50–70 vertices and may use weights in [1.0, 10.0]; sampled source-target pairs have paths of at most 3 hops.Source: §C.6 · p. 37

C.7 N-BODY PROBLEM

• The task predicts each star’s next-frame state from an n-body system governed by distance- and mass-dependent gravitational forces.The network predicts velocities for all stars at the next time step, using sufficiently small time increments to approximate the dynamics.Source: §C.7 · p. 38
• The experiment tests two extrapolation schemes: out-of-distribution pairwise distances and out-of-distribution star masses.The distance and mass ranges differ between the test and training distributions.Source: §C.7 · p. 38
• Training uses 10,000 simulated examples from 100 orbit-system videos, with three stars, a 100kg center star, and other-star masses sampled from [0.02, 9.0]kg.Other stars begin at randomly oriented distances in [10.0, 100.0]m, and frames are rejection-sampled so all pairwise distances are at least 30m.Source: §C.7 · p. 39
• The two 5,000-example OOD test sets vary either distances to [1..20]m with in-distribution masses or masses while keeping distances in-distribution.For mass OOD, the center star is 200kg and other-star masses are sampled from [0.04, 18.0]kg, compared with [0.02, 9.0]kg in training.Source: §C.7 · p. 39
• The model is a one-iteration GNN on a complete star graph that aggregates transformed interactions, using node features containing mass, position, and velocity.Training minimizes node-level next-frame mean squared error.Source: §C.7 · p. 39
• The GNN uses width-128 MLP modules, with four layers in MLP(1) and two in MLP(2), and an initial learning rate of 0.005.The learning rate decays 0.5 every 50 epochs; batch size is 32, and weight decay is 1e −5.Source: §C.7 · p. 40
• The study compares zero edge representations with a relation representation that makes the underlying target function more linear.The latter is identified as the better edge/relationship representation in the supplied setup.Source: §C.7 · p. 40

General visualization findings

• MLP predictions agree with the true functions within the training range but become linear along directions from the origin out of distribution.The in-distribution behavior follows generalization arguments, while the out-of-distribution directional linearity is explained by Theorem 1.Source: §D.1 · p. 40

Graph-based reasoning tasks

• Additional experiments report visualization results for graph-based reasoning tasks.Source: §D.1 · p. 40

Additional function visualizations

• Additional visualizations examine quadratic, cosine, and square-root functions across one- and two-dimensional settings with cube, sphere, or varied-range training supports.The visualizations compare learned functions, true functions, and in-distribution predictions using black, grey, and blue, respectively.Source: Figure 8 · p. 41, §D.1 · p. 41, Figure 9 · p. 42

L1 function

• MLPs can successfully extrapolate y = |x| and y = |x1| + |x2| under some settings, but different hyper-parameters can cause extrapolation failure.The figures show both successful and unsuccessful extrapolation cases for the absolute-value and l1-norm functions.Source: Figure 12 · p. 45, Figure 13 · p. 46

Linear function

• MLPs successfully extrapolate the linear function y = x1 + x2 for both cube and sphere training supports.The result is explained by Theorem 2 because both supports intersect all directions.Source: Figure 14 · p. 47

MLP extrapolation

• MLPs extrapolate underlying linear functions when training supports intersect all directions, but often fail when directions are missing.The “fix1” and “posX” distributions omit directions and produce poorer extrapolation than “all.”Source: Figure 15 · p. 49

Maximum degree

• GNNs with graph-level max-pooling extrapolate maximum degree to OOD node features, graph structures, and graph sizes when training graphs satisfy Theorem 3.The node features are in R3 but should not affect the maximum-degree answer.Source: Figure 16 · p. 50
• For maximum degree, sum-pooling fails to extrapolate, whereas max-pooling succeeds under appropriate training sets because it encodes the required nonlinear operations.This validates Corollary 1 for sum-pooling and supports Theorem 3 for max-pooling.Source: Figure 17 · p. 51

Shortest path

• Shortest-path extrapolation errors follow a U-shape, becoming high when training graphs are very sparse or very dense.The same pattern appears when training uses specific graph structures rather than random graphs.Source: Figure 18 · p. 52