PaperlayerView exact revision on arXiv

Source-linked AI summary

What Can Neural Networks Reason About?

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

arXiv:1905.13211v4cs.LGcs.AIcs.CVcs.NEstat.ML

TL;DR

The paper asks why some neural-network structures generalize better than others on reasoning tasks despite comparable expressive power. It introduces algorithmic alignment between network computation and reasoning algorithms, formalizes its effect on sample complexity, and finds that GNNs align with dynamic programming across several tasks while also revealing limits on harder problems.

  • Problem

    Limited understanding of how network structure relates to generalization in reasoning makes it difficult to characterize learnable tasks, explain existing models, and design architectures for new tasks.

  • Method

    The paper develops a theoretical framework that measures alignment between a neural network’s computation structure and the algorithmic structure of an exact reasoning process.

  • Results

    Across several reasoning-task categories, the framework explains why GNNs learn dynamic-programming tasks and predicts that better algorithmic alignment reduces sample complexity.

  • Takeaways & Limitations

    Algorithmic alignment provides a basis for analyzing existing reasoning architectures and guiding network design for new tasks.

  • Takeaways & Limitations

    The analysis focuses on exact reasoning processes with clear algorithmic structure and leaves approximation algorithms and unknown structures for future work.

Abstract

from arXiv · show

Neural networks have succeeded in many reasoning tasks. Empirically, these tasks require specialized network structures, e.g., Graph Neural Networks (GNNs) perform well on many such tasks, but less structured networks fail. Theoretically, there is limited understanding of why and when a network structure generalizes better than others, although they have equal expressive power. In this paper, we develop a framework to characterize which reasoning tasks a network can learn well, by studying how well its computation structure aligns with the algorithmic structure of the relevant reasoning process. We formally define this algorithmic alignment and derive a sample complexity bound that decreases with better alignment. This framework offers an explanation for the empirical success of popular reasoning models, and suggests their limitations. As an example, we unify seemingly different reasoning tasks, such as intuitive physics, visual question answering, and shortest paths, via the lens of a powerful algorithmic paradigm, dynamic programming (DP). We show that GNNs align with DP and thus are expected to solve these tasks. On several reasoning tasks, our theory is supported by empirical results.

1 INTRODUCTION

The paper addresses why structured neural networks generalize better on reasoning tasks by formalizing alignment between network computation and reasoning algorithms. It connects this alignment to sample efficiency, explains GNN success on dynamic-programming tasks, and identifies limits on harder problems.

  • Reasoning models often succeed when their architectures explicitly encode task-relevant structure, especially pairwise relations and recursive updates in GNNs.
  • The paper introduces algorithmic alignment, measuring how closely a network’s computation graph matches the algorithmic structure of a reasoning process.When alignment is high, the network needs to learn simpler algorithmic steps rather than an entire reasoning procedure.
  • A sample complexity bound decreases with better algorithmic alignment under simplifying assumptions.
  • Dynamic programming unifies intuitive physics, visual question answering, and shortest paths, while multi-iteration GNNs align with DP and are expected to learn these tasks sample-efficiently.The paper reports that these expectations are supported by empirical results.
  • The framework distinguishes architectures across increasingly complex tasks, including summary statistics, relational argmax, dynamic programming, and NP-hard problems.It predicts that Deep Sets and GNNs can learn summary statistics, GNNs can learn relational argmax and DP, and GNNs cannot sample-efficiently learn subset sum.
  • The formalization extends earlier claims about relational inductive biases by quantifying network–task structure and deriving implications for learning.

2 PRELIMINARIES

The preliminaries describe neural-network structures for reasoning over object sets and relate them to algorithmic computation. They contrast permutation-invariant set processing with GNN message passing, which additionally models pairwise relations and can align with Bellman–Ford.

  • The framework represents each universe as a set of objects, each encoded by a vector, and learns to answer questions about unseen universes.
  • A GNN’s induced computation structure can match Bellman–Ford for shortest paths, leaving its modules to learn a simple update equation.
  • For multiple objects, a plain MLP applied to concatenated representations often fails to generalize.
  • Deep Sets impose permutation invariance so outputs remain unchanged under input reordering.
  • GNNs recursively update each node by aggregating neighboring-node representations through message passing.
  • Unlike Deep Sets, GNNs can focus on pairwise relations while retaining permutation invariance.

3 THEORETICAL FRAMEWORK: ALGORITHMIC ALIGNMENT

The framework explains differing generalization through algorithmic alignment: how easily a network’s computational structure simulates a reasoning algorithm. It formalizes this relationship in PAC learning and derives sample-complexity implications, while noting theoretical and practical limitations.

  • Networks can have equal expressive power yet generalize differently because their structures may represent reasoning algorithms with different ease.The framework studies generalization rather than representation, since MLPs, Deep Sets, and GNNs can represent broad function classes.
  • 3.1 FORMALIZATION OF ALGORITHMIC ALIGNMENT: The PAC framework treats sample complexity as the number of training samples needed to achieve low test error with high probability.The paper defines learnability using error and failure-probability parameters and defines sample complexity as the minimum qualifying sample count.
  • 3.1 FORMALIZATION OF ALGORITHMIC ALIGNMENT: Algorithmic alignment measures whether network modules can simulate an algorithm and whether those module functions have low sample complexity.Alignment requires that module functions generate the reasoning function and that their learning costs satisfy a bounded aggregate criterion.
  • 3.1 FORMALIZATION OF ALGORITHMIC ALIGNMENT: Theorem 3.5 links sample-efficient MLP learning to polynomially simple functions, while algorithm steps spanning many objects can require more samples.The paper identifies for-loops as an example of complex algorithm steps that can increase the relevant polynomial complexity.
  • 3.2 BETTER ALGORITHMIC ALIGNMENT IMPLIES BETTER GENERALIZATION: Under simplifying assumptions, the sample-complexity bound increases with alignment value M, so better alignment yields a lower bound.The result concerns sequential module training with auxiliary labels and uses learned functions satisfying a Lipschitz condition.
  • 3.2 BETTER ALGORITHMIC ALIGNMENT IMPLIES BETTER GENERALIZATION: For pairwise-relation summary statistics, the sample-complexity bound for MLP is O(ℓ^2) times larger than for GNN in the theorem’s setting.This illustrative corollary reflects better alignment between GNN computation and the target relation than between MLP computation and that target.

4 PREDICTING WHAT NEURAL NETWORKS CAN REASON ABOUT

Algorithmic alignment predicts which architectures learn reasoning tasks sample-efficiently: networks generalize better when their computation structure matches the task’s reasoning algorithm. Experiments support these predictions across summary statistics, relational argmax, dynamic programming, and subset sum.

  • 4.1 SUMMARY STATISTICS: Deep Sets align with summary-statistics algorithms, whereas MLPs must learn the aggregation loop and therefore require more samples.The framework predicts better sample complexity for Deep Sets than MLPs on sum and max tasks.
  • 4.1 SUMMARY STATISTICS: 96% test accuracy for Deep Sets versus 9% for MLP on maximum value difference; sorting makes MLP achieve perfect accuracy.Sorting reduces the reasoning algorithm to a simple subtraction, which even MLPs can learn with low sample complexity.
  • 4.2 RELATIONAL ARGMAX: GNNs generalize better than Deep Sets on relational argmax because one-iteration message passing directly aggregates pairwise relations.On the furthest-pair task, Deep Sets reach 21% test accuracy while GNNs exceed 90%.
  • 4.3 DYNAMIC PROGRAMMING: GNNs align with dynamic programming by mapping node representations to subproblem answers and aggregation to the DP update.With enough iterations, GNNs can sample-efficiently learn DP algorithms with simple updates such as sum, min, or max.
  • 4.3 DYNAMIC PROGRAMMING: On the monster-trainer shortest-path task, GNNs with four iterations achieve 94% accuracy, nearly matching seven-iteration GNNs at 95%.Models with worse algorithmic alignment improve more slowly as training data increases; GNN4 gains 23% from 40,000 to 80,000 samples versus 0.2% for Deep Sets.
  • 4.3 DYNAMIC PROGRAMMING: The framework formulates visual question answering and intuitive physics as dynamic programming, explaining why GNNs are effective for these task families.For tasks outside DP, such as NP-hard problems, the framework predicts that GNNs cannot learn them sample-efficiently.
  • 4.4 DESIGNING NEURAL NETWORKS WITH ALGORITHMIC ALIGNMENT: Neural Exhaustive Search mirrors exhaustive subset enumeration and reaches 98% test accuracy on subset sum, where GNNs are predicted not to generalize.NES uses an LSTM and MLP on each subset, then aggregates subset results with max pooling and an MLP.

5 CONCLUSION

The paper frames algorithmic alignment as a way to understand which reasoning tasks neural networks learn well and why GNNs generalize on many popular tasks. It presents preliminary sample-complexity results while leaving reasoning paradigms beyond dynamic programming and representation–reasoning integration for future work.

  • 5 CONCLUSION: The paper introduces algorithmic alignment to formalize how neural-network structure interacts with a task’s reasoning process.It uses this framework to study generalization when learning underlying reasoning processes.
  • 5 CONCLUSION: GNNs generalize on many popular reasoning tasks because their underlying reasoning processes resemble dynamic programming.This connection is presented as an explanation of GNN success and a suggestion of current architectural limits.
  • 5 CONCLUSION: Future work includes designing networks for reasoning paradigms beyond dynamic programming and exploring neural architectures organized around algorithmic structures.The paper describes this as an interesting direction rather than a completed result.
  • 5 CONCLUSION: A broader open question is how to disentangle and eventually integrate representation and reasoning.The paper notes that reasoning assumes good representations of concepts and objects in the world.

A PROOF OF PROPOSITION 3.1

This proof establishes that GNNs can represent the permutation-invariant continuous functions covered by Deep Sets. It does so by embedding Deep Sets within one-layer GNNs and extending the construction across multiple message-passing layers.

  • A PROOF OF PROPOSITION 3.1: The proof uses Deep Sets’ universal approximation result for permutation-invariant continuous functions.The cited result assumes elements from a compact subset of R^d and set size bounded by N.
  • A PROOF OF PROPOSITION 3.1: Any Deep Sets model can be expressed by a GNN with one message-passing iteration.The construction makes the message function ignore the target input, yielding the required per-element transformation and aggregation.
  • A PROOF OF PROPOSITION 3.1: Multiple GNN layers preserve this construction because a function ρ can be represented as a composition of functions implemented by individual GNN layers.This extends the one-iteration argument to GNNs with multiple message-passing iterations.
  • A PROOF OF PROPOSITION 3.1: Therefore, GNNs are universal approximators for permutation-invariant continuous functions.This conclusion follows by combining the Deep Sets approximation result with the GNN simulation construction.

B PROOF OF PROPOSITION 3.2

This proof constructs MLPs that simulate GNN computations by representing message functions, aggregations, and repeated message-passing iterations through stacked layers. Variable-sized inputs are handled with separate size-specific MLPs and a selecting meta-layer.

  • B PROOF OF PROPOSITION 3.2: For a fixed input-set size, an MLP can simulate a GNN by receiving the concatenated node representations in an arbitrary ordering.The construction assumes the set size is bounded by M.
  • B PROOF OF PROPOSITION 3.2: Each GNN message-passing iteration is reproduced by MLP layers that compute pairwise transformations and encode summation through weights.The resulting hidden vector contains the concatenated updated node representations.
  • B PROOF OF PROPOSITION 3.2: Stacking the corresponding MLP layers simulates all K message-passing iterations and the final output function.The construction then reproduces the GNN’s computation using a single MLP.
  • B PROOF OF PROPOSITION 3.2: Variable-sized inputs are handled by constructing one MLP for each size from 1 through M and using a meta-layer to select the matching output.The selector uses the input-set size information through indicator functions.

C PROOF OF THEOREM 3.5

The proof extends a scalar-output sample-complexity theorem to vector-valued functions. It treats output components independently and combines their guarantees with a union bound.

  • C PROOF OF THEOREM 3.5: Theorem 3.5 generalizes an earlier scalar-case theorem under its stated assumptions.The paper points to Arora et al. (2019a) for the complete list of assumptions.
  • C PROOF OF THEOREM 3.5: The proof applies the scalar theorem independently to each component of the output vector.This converts the vector-valued problem into a collection of scalar-valued learning problems.
  • C PROOF OF THEOREM 3.5: A union bound combines per-component error and failure guarantees into an overall vector-output guarantee.The component-level rates are chosen so the total error and failure probability are bounded by the target ε and δ.
  • C PROOF OF THEOREM 3.5: The sample-complexity expression can be further simplified when m is treated as a constant.The section states this simplification after deriving the vector-valued bound.

D PROOF OF THEOREM 3.6

The proof inductively controls representation errors across sequentially learned modules. Lipschitzness, algorithmic learnability, stability, and union bounds preserve O(ϵ) error through all layers, yielding learnability of g.

  • Inductive learnability: The proof bounds the test error of the final learned function by controlling intermediate representation errors at each layer.The induction compares learned and correct module outputs and propagates these bounds forward.
  • Base case: For the first module, learnability gives error below ϵ with probability at least 1 −δ, and aggregation yields O(ϵ) input error for the next module.A union bound controls the constant number of aggregation terms.
  • Inductive step: Assuming O(ϵ) error for inputs at layer k, the induction step bounds the corresponding error at layer k + 1.The learned and correct inputs are denoted ˆz and z, respectively.
  • Inductive step: Lipschitzness bounds the output change caused by perturbed inputs as L1∥ˆz −z∥.This controls the first term in the decomposition of the module error.
  • Inductive step: Algorithmic learnability and stability make the learned function on perturbed samples close to the correct function on correct samples.Stability applies when training inputs are only slightly perturbed.
  • Conclusion: With three bad events each having probability at most O(δ), the total module error is O(ϵ) with probability at least 1 −O(δ), completing the proof.The bound is ∥ˆf(ˆz) −f(z)∥≤L1O(ϵ) + L0O(ϵ) + ϵ = O(ϵ).

E PROOF OF COROLLARY 3.7

The corollary compares sequential GNN-style learning with a giant MLP for pairwise squared differences. Encoding all pairwise polynomial terms directly creates an O(ℓ2) extra sample-complexity cost for the MLP.

  • Polynomial representation: A giant MLP learns the pairwise function (Xi −Xj)2 for ℓ2 terms and encodes those terms in its weights.The number of polynomial terms is therefore of order ℓ2.
  • Polynomial representation: The pairwise function f(x, y) = (x −y)2 is represented as a polynomial, enabling application of the sample-complexity theorem.The proof begins by expressing the function in polynomial form.
  • Aggregation: A one-layer GNN applies an MLP to each pair and sums the outputs, producing aggregation error O(ℓ2 · ϵ).The same aggregation error arises when an MLP directly learns g, so it does not affect the comparison.
  • Direct MLP comparison: The target function g can be expressed as a sum of polynomials regardless of the concatenation order of the objects.The proof applies the polynomial sample-complexity theorem to this representation.
  • Conclusion: Comparing the direct-MLP and sequential-learning bounds yields an O(ℓ2) sample-complexity difference.The direct MLP requires O((ℓ2 + log(1/ˆδ))/ˆϵ2), whereas the sequential comparison uses O(log(1/ˆδ)/ˆϵ2).

F PROOF OF CLAIM 4.1

The claim is proved by contradiction: no single-object function summed over two objects can represent the pairwise function g(x, y) when unequal inputs must produce a nonzero value.

  • Contradiction argument: If f(x) + f(y) = g(x, y) held for all x and y, setting x = y would force f(x) = 0 for every x.This follows because g(x, x) = 0.
  • Contradiction argument: For unequal x and y, the resulting sum would remain zero while g(x, y) is nonzero, producing a contradiction.Therefore, no such single-object decomposition exists.

G EXPERIMENTS: DATA AND TRAINING DETAILS

The experiments define synthetic reasoning tasks, training protocols, and model configurations, including a shortest-path task built around a dynamic-programming algorithm. The shortest-path procedure reaches paths with up to 2k stops using half the Bellman-Ford iterations.

  • Treasure tasks: The treasure datasets use 25 objects with sampled locations, values, and colors, and report test accuracy at the hyperparameters with best validation accuracy.The datasets include 50,000/5,000/5,000 and 60,000/6,000/6,000 train/validation/test splits.
  • Training setup: Training uses Adam, cross-entropy loss, scheduled learning-rate decay, 150 epochs, and tuned batch sizes, with architecture-specific hidden dimensions and depths.Dropout at rate 0.5 is applied before the last two hidden layers of MLP1.
  • Treasure tasks: The treasure-label task asks for the colors of the two most distant treasures and encodes answers into 21 classes.Labels are ordered by the pair of colors and represented one-hot.
  • Monster trainer: For shortest paths with at most seven stops, the algorithm runs forward and reverse Bellman-Ford-style updates for four iterations.The two recurrences compute distances from the source and to the target.
  • Monster trainer: After k iterations, enumerating a midpoint combines the forward and reverse distances to obtain a shortest path with at most 2k stops.The combined quantity is distances[k][u] + distancet[k][u].
  • Monster trainer: The shortest-path algorithm needs half the iterations of Bellman-Ford, while its dataset uses 200,000 training and 6,000 validation and test examples.Rejection sampling balances shortest-path lengths from 3 through 7 and removes trivial questions.
Loading 1905.13211v4…