[Paper Review] How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks

Keyulu Xu, Mozhi Zhang+4 more|arXiv (Cornell University)|Sep 24, 2020

Domain Adaptation and Few-Shot Learning90 references108 citations

TL;DR

The paper analyzes how gradient-descent trained neural networks extrapolate outside training data, showing ReLU MLPs converge to linear functions along directions from the origin, and that GNNs can extrapolate when task-specific non-linearities are encoded in architecture or features. It provides theoretical results (NTK-based) and empirical validation across DP tasks.

ABSTRACT

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.

Motivation & Objective

• Quantify how gradient-descent trained neural networks extrapolate outside the training support.
• Explain why MLPs struggle with nonlinear extrapolation while GNNs can succeed in DP-like tasks.
• Identify conditions under which MLPs and GNNs extrapolate well.
• Relate feedforward extrapolation insights to GNN architectures and representations.

Proposed method

• Analyze over-parameterized networks in the neural tangent kernel (NTK) regime to connect training dynamics with kernel regression.
• Prove linear extrapolation behavior for two-layer ReLU MLPs along directions from the origin with rate O(1/t).
• Show that MLPs extrapolate linear targets well when the training distribution is sufficiently diverse (Theorem 2).
• Propose and test the hypothesis that GNNs extrapolate well when task-specific non-linearities are encoded in architecture or input representations (Theorem 3 and accompanying experiments).
• Use Graph NTK to analyze extrapolation in simplified GNN cases and validate on max degree, shortest path, and n-body tasks.
• Discuss the role of architecture (e.g., max/min readouts) and input representations in enabling extrapolation.

Experimental results

Research questions

• RQ1 When do ReLU MLPs trained by gradient descent extrapolate well beyond the training distribution?
• RQ2 Under what conditions can GNNs extrapolate nonlinear tasks, and how do architecture and input representations influence this?
• RQ3 Can encoding task-specific non-linearities in GNNs or representations enable extrapolation to unseen graph sizes, structures, or edge weights?
• RQ4 How does the geometry of the training data affect extrapolation for linear targets in MLPs and DP-like tasks in GNNs?

Key findings

• ReLU MLPs extrapolate to linear functions along directions from the origin with rate O(1/t).
• MLPs can extrapolate linear target functions if the training distribution covers enough directions (diverse geometry).
• GNNs can extrapolate well in DP-like tasks if appropriate non-linearities are encoded in architecture or features, supported by theory (Graph NTK) and experiments.
• Replacing sum-aggregation with architecture that mirrors DP updates (e.g., min/max reads) enables better extrapolation in tasks like max degree and shortest path.
• Improved input representations can enable GNNs to extrapolate nonlinear dynamics (e.g., n-body) by shifting nonlinearity into the representation rather than the MLPs.