What Can Neural Networks Reason About?
Authors: Keyulu Xu
Jingling Li
Mozhi Zhang
Simon S. Du
Ken-ichi Kawarabayashi
Stefanie Jegelka
Venue:  (2020)
Keywords: reasoning
deep learning theory
algorithmic alignment
graph neural networks

Abstract: 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.

Published as a conference paper at ICLR 2020
WHAT CAN NEURAL NETWORKS REASON ABOUT?
Keyulu Xu', Jingling Lit, Mozhi Zhang', Simon S. Du', Ken-ichi Kawarabayashi",
Stefanie Jegelka?
I "aassacusettts Institute of Technology (MIT)
+University of Maryland
IIstitute for Advanced Study (IAS)
"National Institute of Informatics (NII)
(keyulu, stefje) @mit edu
ABSTRACT
Neural networks have succeeded in many reasoning tasks. Empirically, these tasks require specialized network structures, c.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 network structure generalizes better than others, although they have equal expressive power. In this paper, we develop 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 aligmment and derive a sample complexity bound that decreases with better alignment. This framework offers explanation for the empirical success of popular reasoning models, and suggests their limitations. As an example, we unify seemingly different reasoning tasks, such 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 solve these tasks. On several reasoning tasks, our theory is supported by empirical results.
INTRODUCTION
Recently, there have been many advances in building neural networks that can learn to reason. Reasoning spans variety of tasks, for instance, visual and text-based question answering (Johnson et al., 2017a; Weston et al., 2015; Hu et al., 2017; Fleuret et al., 2011; Antol et al., 2015), intuitive physics, i.e., predicting the time evolution of physical objects (Battaglia et al., 2016; Watters et al., 2017; Fragkiadaki et al., 2016; Chang et al., 2017), mathematical reasoning (Saxton et al., 2019; Chang et al., 2019) and visual IQ tests (Santoro al., 2018; Zhang al., 2019).
Curiously, neural networks that perform well in reasoning tasks usually possess Specific struc- tures (Santoro al., 2017). Many successful models follow the Graph Neural Network (GNN) framework (Battaglia et al., 2018; 2016; Palm et al., 2018; Mrowca et al., 2018; Sanchez-Gonzan al., 2018; Janner et al., 2019). These networks explicitly model pairwise relations and recursively update each object's 's representation by aggregating its relations with other objects. Other computa- tional structures, e.g., neural symbolic programs (Yi et al., 2018; Mao al., 2019; Johnson al., 2017b) and Deep Sets (Zaheer et al., 2017), are effective on specific tasks.
However, there is limited understanding of the relation between the generalization ability and network structure for reasoning. What tasks can neural network (sample efficiently) learn to reason about? Answering this question is crucial for understanding the empirical success and limitations of existing models, and for designing better models for new reasoning tasks.
This paper is an initial work towards answering this fundamental question, by developing a theoretical framework to characterize what tasks a neural network can reason about. We build observation that reasoning processes resemble algorithms. Hence, we study how well a reasoning algorithm aligns with the computation graph of the network. Intuitively, if they align well, the network only needs to learn simple algorithm steps to simulate the reasoning process, which leads to better sample efficiency. We formalize this intuition with numeric measure of algorithmic alignment,
simple

---

Published as a conference paper at ICLR 2020
Summary statisties What is the maximum value difference among treusures?
Relational argmax What are the colors of the furthest pair of objects?
Dynomic programming What is the cost to defeat monster X by following the optimal puth?
NP-hard problem Subset sum: Is there subset that sums O?
Figure 1: Overview of reasoning tasks with increasingly complex structure. Each task category shows an example task on which we perform experiments in Section 4. Algorithmic alignment suggests that (a) Deep Sets and GNNs, but not MLP, can sample efficiently learn summary statistics, (b) GNNs, but not Deep Sets, can learn relational argmax, (c) GNNs can learn dynamic programming. an algorithmic paradigm that we show to unify many reasoning tasks, (d) GNNs cannot learn subset sum (NP-hard), but NES, network we design based on exhaustive search, can generalize. Our theory agrees with empirical results (Fig. 3).
and show initial support for our hypothesis that algorithmic alignment facilitates learning: Under simplifying assumptions, show sample complexity bound that decreases with better alignment.
Our framework explains the empirical success of popular reasoning models and suggests their limitations. As concrete examples, we study four categories of increasingly complex reasoning tasks: summary statistics, relational argmax (asking about properties of the result of comparing multiple relations), dynamic programming, and NP-hard problems (Fig. 1). Using alignment, we characterize which architectures are expected to learn each task well: Networks inducing permutation invariance, such as Deep Sets (Zaheer et al., 2017), can learn summary statisties, and one- me-iteration GNNs can learn relational argmax. Many other more complex tasks, such as intuitive physics, visual question answering, and shortest paths despite seeming different - can all solved via powerful algorithmic paradigm: dynamic programming (DP) (Bellman, 1966). Multi-iteration GNNs algorithmically align with DP and hence are expected to sample-efficiently learn these tasks. Indeed, they do. Our results offer explanation for the popularity of GNNs in the relational reasoning literature, and also suggest limitations for tasks with even more complex structure. As an example of such task, consider subset sum, NP-hard problem where GNNs indeed fail. Overall, empirical results (Fig. 3) agree with our theoretical analysis based on algorithmic alignment (Fig. 1). These findings also suggest how take into account task structure when designing new architectures.
The perspective that structure in networks helps is not new. For example, in a well-known position paper, Battaglia al. (2018) argue that GNNs are suitable for relational reasoning because they have relational inductive biases, but without formalizations. Here, we take such ideas one step further, by introducing a formal definition (algorithmic alignment) for quantifying the relation between network and task structure, and by formally deriving implications for learning. These theoretical ideas are the basis for characterizing what reasoning tasks network can learn well. Our algorithmic structural condition also differs from structural assumptions common in learning theory (Vapnik, 2013; Bartlett & Mendelson, 2002; Bartlett et al., 2017; Neyshabur et al., 2015; Golowich et al., 2018) and specifically aligns with reasoning.
In summary, we introduce algorithmic alignment to analyze learning for reasoning. Our initial theo- retical results suggest that algorithmic alignment is desirable for generalization. On four categories of reasoning tasks with increasingly complex structure, we apply our framework to analyze which tasks some popular networks can learn well. GNNs algorithmically align with dynamic program- ming, which solves a broad range of reasoning tasks. Finally, our framework implies guidelines for designing networks for new reasoning tasks. Experimental results confirm our theory.
2 PRELIMINARIES
We begin by introducing notations and summarizing common neural networks for reasoning tasks. Let S denote the universe, i.e., a configuration/set of objects to reason about. Each object E S is represented by vector X. This vector could be state descriptions (Battaglia et al., 2016; Santoro
2

---

Published as a conference paper at ICLR 2020
Graph Neural Network
Bellman-Ford algorithm
... GNN iter:
for k ISI
for u in S:
need to lear for-loops
for u S:
hu = Σv MLP(h,/k-1, hu
d[k][u] = min, [[k-1][v] + cost (v, u)
Learns a simple reasoning step
Figure 2: Our framework suggests that better algorithmic alignment improves generalization. As example, our framework explains why GNN generalizes when learning answer shortest paths. A correct reasoning process for the shortest paths task is the Bellman-Ford algorithm. The computation structure of GNN (left) aligns well with Bellman-Ford (right): the GNN can simulate Bellman-Ford by merely learning simple reasoning step. i.e., the relaxation step in the last line (a sum, and min over neighhoring nodes v) via aggregation operation. In contrast, giant MLP or Deep Set must learn the structure of the entire for-loop. Thus, the GNN is expected to generalize better when learning shortest paths, as confirmed experiments (Section 4.3).
al., 2017) features learned from data such images (Santoro al., 2017). Information about the specific question can also be included in the object representations. Given a set of universes {S1,-.., SM} and answer labels (y), ...: yM3 S V, aim to learn function g that answer questions about unseen universes, = (S).
Multi-layer perceptron (MLP). For single-object universe, applying an MLP on the object representation usually works well. But when there are multiple objects, simply applying an MLP to the concatenated object representations often does not generalize (Santoro et al., 2017).
Deep Sets. As the input to the reasoning function is an unordered set, the function should be permutation-invariin i.c., the output is the same for all input orderings. To induce permutation invariance in neural network, Zaheer et al. (2017) propose Deep Sets, of the form
= MLP2 Σ MLP1 (
(2.1)
SES
Graph Neural Networks (GNNs). GNNs are originally proposed for learning on graphs (Scarselli al., 2009b). Their structures follow message passing scheme (Gilmer et al., 2017; Xu al., 2018; 2019), where the representation (k) of each node s (in iteration k)) is recursively updated by aggregating the representation of neighboring nodes. GNNs can be adopted for reasoning by considering objects as nodes and assuming all objects pairs are connected, i.e., complete graph (Battaglia et al., 2018):
(x(h) Σ tE MIP(A) (N(
1. 1 hs MLP2 E SES mn())
(2.2)
where fg is the answer/output and K is the number of GNN layers. Each object's representation initialized as ( X.. Although other aggregation functions are proposed, we use sum in our experiments. Similar to Deep Sets, GNNs are also permutation invariant. While Deep Sets focus on individual objects, GNNs can also focus on pairwise relations.
The GNN framework includes many reasoning models. Relation Networks (Santoro et al., 2017) and Interaction Networks (Battaglia et al., 2016) resemble one-layer GNNs, Recurrent Relational Networks (Palm et al., 2018) apply LSTMs (Hochreiter & Schmidhuber, 1997) after aggregation.
THEORETICAL FRAMEWORK: ALGORRITHMIC ALIGNMENT
Next, study how the network structure and task may interact, and possible implications for generalization. Empirically, different network structures have different degrees of success in learning reasoning tasks, e.g., GNNs learn relations well, but Deep Sets often fail (Fig. 3). However, all these networks are universal approximators (Propositions 3.1 and 3.2). Thus, their differences in test accuracy must come from generalization.

---

Published as a conference paper at ICLR 2020
We observe that the answer many reasoning tasks may be computed via reasoning algorithm; we further illustrate the algorithms for some reasoning tasks in Section 4. Many neural networks represent algorithms (Pérez et al., 2019). For example, Deep Sets universally represent permutation-invall
set functions (Zaheer et al., 2017; Wagstaff et al., 2019). This also holds for GNNs and MIPss we show in Propositions 3.1 and 3.2 (our setting differs from Scarselli et al. (2009a) and Xu et al. (2019), who study functions on graphs):
Proposition 3.1. Let f: _R + R. be any continuous function over sets S of bounded cardinality [S] S N. If f is permutation-invaariee R. then f can be approximated arbitrarily closely by GNN (of any depth).
to the elements in S, and the elements are in compact set in
Proposition 3.2. For any GNN N. there is an MLP that can represent all functions N can represent.
But, empirically, not all network structures work well when learning these algorithms, i.c., they generalize differently. Intuitively, network may generalize better if it can represent function "more easily". We formalize this idea by algorithmic alignment, formally defined in Definition 3.4. Indeed, not only the reasoning process has algorithmic structure: the neural network's ss architecture induces a computational structure on the function it computes. This corresponds to an algorithm that prescribes how the network combines computations from modules. Fig. illustrates this idea for GNN, where the modules are its MLPs applied to pairs of objects. In the shortest paths problem, the GNN matches the structure of the Bellman-Forr Ford algorithm: to simulate the Bellman- Ford with a GNN, the GNN's MLP modules only need to learn simple update equation (Fig. 2). In contrast, if we want to represent the Bellman-Ford algorithm with a single MI.,, it needs to simulate an entire for-loop, which is much more complex than one update step. Therefore, we expect the GNN to have better sample complexity than MLP when learning to solve shortest path problems.
This perspective suggests that neural network which better aligns with correct reasoning process (algorithmic solution) can more easily learn a reasoning task than a neural network that does not align well. If we look more broadly at reasoning, there may also exist solutions which only solve task approximately, or whose structure is obtuse. In this paper, we focus on reasoning tasks whose underlying reasoning process is exact and has clear algorithmic structure. We leave the study of approximation algorithms and unknown structures for future work.
3.1 FORMALIZATION OF ALGORITHMIC ALIGNMENT
We formalize the above intuition in a PAC learning framework (Valiant, 1984). PAC learnability formalizes simplicity as sample complexity, i.e., the number of samples needed ensure low test error with high probability. It refers to learning algorithm A that, given training samples (x,, yi y outputs function f = A((z,, yı )) The learning algorithm here is the neural network and its training method, e.g., gradient descent. A function is simple if it has low sample complexity.
Definition 3.3. (PAC learning and sample complexity). Fix an error parameter c > 0 and failure probability 8 E 1). Suppose {2,, 3 are i.i.d. samples from some distribution D, and the data satisfies yı1 = g((x;) for some underlying function g. Let f f= A({x,, MM be the function generated by learning algorithm A. Then is (M, e, 8-learmable with A if
PAND [llf(=) ا(=)) s 2
(3.1)
The sample complexity CA (g, e, is the minimum M so that g is (M, 8) learnable with A.
With the PAC learning framework, define numeric measure of algorithmic alignment (Defini- tion 3.4), and under simplifying assumptions, we show that the sample complexity decreases with better algorithmic alignment (Theorem 3.6).
Formally, a neural network aligns with an algorithm if it can simulate the algorithm via limited number of modules, and each module is simple, i.e., has low sample complexity.
Definition 3.4. (Algorithmic alignment). Let g be a reasoning function and A a neural network with n modules N₁. The module functions fi...., fn generate g for A if, by replacing M, with f., the network N simulates g. Then (M, e, )-aagooithmmiclly aligns with g if (1) fi.... fa generate g and (2) there are learning algorithms A, for the NE's such that max; CA. (Si, e, 8) S M.
4

---

Published as a conference paper at ICLR 2020
Good algorithmic alignment, i..., small M, implies that all algorithm steps fü simulate the algorithm 9 are easy to leamn Therefore, the algorithm steps should not simulate complex programming constructs such as for-loops, whose sample complexity is large (Theorem 3.5).
Next, we show how to compute the algorithmic alignment value M. Algorithmic alignment resembles Kolmogorov complexity (Kolmogorov, 1998) for neural networks. Thus, it is generally non-trivial to obtain the optimal alignment between a neural network and an algorithm. However, one important difference to Kolmogorov complexity is that any algorithmic alignment that yields decent sample complexity is good enough (unless we want the tightest bound). In Section 4, we will see several examples where finding a good alignment is not hard. Then, we can compute the value of an alignment by summing the sample complexity of the algorithm steps with respect to the modules, e.g. MLPs. For ilustration, we show an example of how one may compute sample complexity of MLP modules.
A line of works show one can analyze the optimization and generalization behavior of overparameter- ized neural networks via neural tangent kernel (NTK) (Allen-Zih
al., 2019; Arora et al., 2019a;b;
2020; Du et al., 2019c;a; Jacot et al., 201 8; Li & Liang, 2018). Building upon Arora et al. (2019a), Du et al. (2019b) show that infinitely-wide GNNs trained with gradient descent can provably learn certain smooth functions. The current work studies a broader class of functions, e.g., algorithms, compared to those studied in Du et al. (2019b), but with more simplifying assumptions.
Here, Theorem 3.5, proved in the Appendix, summarizes and extends Theorem 6.1 of Arora et al. (2019a) for over- per-parameterized MLP modules to vector-valued functions. Our framework can be used with other sample complexity bounds for other types of modules, too.
Theorem 3.5. (Sample complexity for overparameterized MLP modules). Let A be an overpa- rameterized and randomly initialized two-layer MLP trained with gradient descent for a sufficient
(
number of iterations. Suppose : .Rd + IRm with components g(z) = Σ, az (0)
(3 x)p3
where 8 E BR a ERd
and p3 = 1 or Pj (1) = (1 (-). The sample complexity CAlg, 8) is
C)( =0 max, ΣK Ej= Pj 1a.1 1105" 11 1e2 + log lloggmm) (m/8)
(3.2)
(e/m)²
Theorem 3.5 suggests that functions that are "simple" when expressed as a polynomial, e.g.. via a Taylor expansion, are sample efficiently learnable by an MLP module. Thus, algorithm steps that perform computation over many objects may require many samples for an MLP module to learn, since the number K of polynomials of such complex algorithm steps.
|| can increase in Eqn. (3.2). "For loop" is example
3.2 BETTER ALGORITHMIC ALIGNMENT IMPLIES BETTER GENERALIZATION
We show an initial result demonstrating that algorithmic alignment is desirable for generalization. Theorem 3.6 states that, in simplifying setting where we sequentially train modules of a network with auxiliary labels, the sample complexity bound increases with algorithmic alignment value M.
While we do not have auxiliary labels in practice, we observe the same pattern for end-to-end learning in experiments (Section 4). We leave sample complexity analysis for and-to-end-learnigg to future work. We prove Theorem 3.6 in Appendix D.
Theorem 3.6. (Algorithmic alignment improves sample complexity). Fix and 8. Suppose | D. where [S,] < N. and y g g(S)) for some g. Suppose .,,,.N are network N's align via functions
(S,, ys MLP modules in sequential order. Suppose N and (M, e, fi, fn. Under the following assumptions, g is (M, Occ), -eaaraall by N.
a) Algorithm stability, Let A be the learning algorithm for the N, 's. Suppose f = A(fx; Vi MM and f = =A(f,, ys M.) For any x, l[f(x) f())l S Lo max; ||z, 2 1l for some Lo. b) Sequential learning. We train Ni's sequentially: Ni has input samples (a ff with a( obtained from St. For j > 1, the input a;(²) for forN, are the outputs from the previous modules, but labels are generated by the correct functions fz-1, f1 on a(").
c) Lipschitzness. The learned functions fq satisfy (1f)())) f())l s Ta 1|: for some Ta
5

---

Published as a conference paper at ICLR 2020
100%
96%
96%
100%
96%
92%
9%
21%
9%
GNN3
GNN1
Deep Sets
MLP
Sorted MLP
GNN3
GNN1
Deep Sets
MLP
(a) Maximum value difference.
(b) Furthest pair.
96%
94%
98%
62%
72%
69%
61%
80%
27%
11%
GNN7
GNN4
GNN3
GNN2
GNN1
Deep Sets
MLP
NES
GNNB
GNN1
Deep Sets
MLP
(c) Monster trainer.
(d) Subset sum. Random guessing yields 50%.
Figure 3: Test accuracies on reasoning tasks with increasingly complex structure. Fig. shows an overview of the tasks. GNNK is GNN with k iterations. (a) Summary statistics. All models except MLP generalize. (b) Relational argmax. Deep Sets fail. (c) Dynamic programming. Only GNNs with sufficient iterations generalize. (d) NP-hard problem. Even GNNs fail, but NES generalizes.
In our analysis, the Lipschitz constants and the universe size are constants going into O(c) and O()). As an illustrative example, we use Theorem 3.6 and 3.5 to show that GNN has polynomial improvement sample complexity over MIPP when learning simple relations. Indeed, GNN aligns better with summary statistics of pairwise relations than MLP does (Section 4.1).
Corollary 3.7. Suppose universe S has & objects X1,..., Xt, and g -X, In the setting of Theorem 3.6, the sample complexity bound for MLP is ( times larger than for GNN.
- xx)²
4 PREDICTING WHAT NEURAL NETWORKS CAN REASON ABOUT
Next, we apply our framework analyze the neural networks for reasoning from Section 2: MLP, Deep Sets, and GNNs. Using algorithmic alignment, we predict whether each model can generalize on four categories of increasingly complex reasoning tasks: summary statistics, relational argmax, dynamic programming, and an NP-hard problem (Fig. 3). Our theoretical analysis is confirmed with experiments (Dataset and training details are in Appendix G). To empirically compare sample complexity of different models, make sure all models perfectly fit training sets through extensive hyperparameter tuning. Therefore, the test accuracy reflects how well model generalizes.
The examples in this section, together with our framework, suggest explanation why GNNs are widely successful across reasoning tasks: Popular reasoning tasks such as visual question answering and intuitive physics can be solved by DP. GNNs align well with DP, and hence are expected to learn sample efficiently.
4.1 SUMMARY STATISTICS
As discussed in Section 2, we assume each object X has a state representation X [/1, .... ha]. where each h; Rd, is feature vector. An MLP can learn simple polynomial functions of the state representation (Theorem 3.5). In this section, we show how Deep Sets use MLP as building blocks to learn summary statistics.
Questions about summary statistics are common in reasoning tasks. One example from CLEVR (John- son et al., 2017a) is "How many objects are either small cylinders red things?" Deep Sets (Eqn. 2.1) align well with algorithms that compute summary statistics over individual objects. Suppose we want to compute the sum of a feature over all objects. To simulate the reasoning algorithm, we can use the first MLP in Deep Sets extract the desired feature and aggregate them using the pooling layer. Under this alignment, each MLP only needs to lear simple steps, which leads to good sample complexity. Similarly, Deep Sets learn compute max or min of feature by using smooth approximations like the softmax max,es X, z log(Z.sx exp(X,)). In contrast, if we train an MI.P perform sum max, the MLP must learn complex for-loop and therefore needs more

---

Published as a conference paper at ICLR 2020
100
GNN7
Test accuracy (%)
GNN4
GNN3
GNN2
GNN1
Deep Sts
MUP
40
80
120
160
Number of training samples (thousand)
Figure 4: Test accuracy vs. training set size for models trained on sub-sampled training sets and evaluated on the same test set of monster trainer (DP task). Test accuracies increase faster when neural network aligns well with an algorithmic solution of the task. For example, the test accuracy of GNN4 increases by 23% when the number of training samples increases from 40, 000 to 80, 000, which is much higher than that of Deep Sets (0.2%)
samples. Therefore, our framework predicts that Deep Sets have better sample complexity than MLP when learning summary statistics.
Maximum value difference. We confirm our predictions by training models to compute the maximun value difference task. Each object in this task is treasure X = [hn, f22 hs] with location h1, value f2, and color h3. We train models to predict the difference in value between the most and the least valuable treasure, y(S) max,es (X,) minges h2()...
The test accuracy follows our prediction (Fig. 3a). MLP does not generalize and only has 9% test accuracy, while Deep Sets 96%. Interestingly, if we sort the treasures by value (Sorted MLP in Fig. 3a), MLP achieves perfect test accuracy. This observation can be explained with our theory-when the treasures are sorted, the reasoning algorithm is reduced to a simple subtraction: (x()),, which has low sample complexity for even MLPs (Theorem 3.5). GNNs also have high test accuracies. This is because summary statistics are a special case of relational argmax, which GNNs can learn as shown next.
ySS h2((11s)
4.2 RELATIONAL ARGMAX
Next, we study relational argmax: tasks where we need to compare pairwise relations and answer a question about that result. For example, a question from Sort-of-CLEVR (Santoro et al., 2017) asks "What is the shape of the object that is farthest from the gray object?", which requires comparing the distance between object pairs.
One-iteration GNN aligns well with relational argmax, as it sums over all pairs of objects, and thus can compare, c.g. via softmax, pairwise information without learning the "for loops". In contrast, Deep Sets require many samples to learn this, because most pairwise relations cannot be encoded as a sum of individual objects:
Claim 4.1. Suppose g(z, y) if and only if a y. There is no such that g(x. y) = f(xx) + f(y)-
Therefore, we train a Deep Set to compare pairwise relations, one of the MLP modules has to learn complex "for loop", which leads to poor sample complexity. Our experiment confirms that GNNs generalize better than Deep Sets when learning relational argmax.
Furthest pair. As an example of relational argmax, we train models to identify the furthest pair use the same object settings as the maximum value difference task. We train models to find the colors of the two treasures with the largest distance. The answer is pair of colors, encoded as an integer category:
among a set of objects.
SS) = (ha(X.,) hy(X,.)) s.t. [X,,, = arg max , see (ll,, - h1 (Xex)ll,
Distance as a pairwise function satisfies the condition in Claim 4.1. As predicted by our framework, Deep Sets has only 21 % test accuracy, while GNNs have more than 90% accuracy.
7

---

Published as a conference paper at ICLR 2020
4.3 DYNAMIC PROGRAMMING
We observe that a broad class of relational reasoning tasks can be unified by the powerful algorithmic paradigm dynamic programming (DP) (Bellman, 1966). DP recursively breaks down problem into simpler sub-problems. It has the following general form:
Answer k = DP- Uppated { [Answer[k: - 1[b1 .j = 1....)
(4.1)
where Answer(k|]d is the solution to the sub-problem indexed by iteration and state i, and DP- Update is an task-specific update function that computes Answer[k][]] from Answerik - 1][5]'s.
GNNs algorithmically align with a class of DP algorithms. We can interpret GNN as DP algorithm.
where node representations h, (A) are [is]er]k][[]] and the GNN aggregation step is the DP-Update. Therefore, Theorem 3.6 suggests that GNN with enough iterations can sample efficiently learn any DP algorithm with a simple up-upe function, e.g. sum/min/maxx
Shortest paths. As an example, we experiment with GNN on Shortest paths, a standard DP problem. Shortest paths can be solved by the Bellman-Ford algorithm (Bellman, 1958), which recursively updates the minimum distance between each object u and the source s:
distance[1] cost(s, u), distance [e] [u] min, [dittance k- 1][u] + cost(u, u)}, (4.2)
As discussed above, GNN aligns well with this DP algorithm. Therefore, our framework predicts that GNN has good sample complexity when learning to find shortest paths. To verify this, we test different models on monster trainer game, which is a shortest path variant with unkown cost functions that need to be learned by the models. Appendix G.3 describes the task in details.
In Fig. 3c, only GNNs with at least four iterations generalize well. The empirical result confirms our theory: neural network can sample efficiently learn task if aligns with a correct algorithm. Interestingly, GNN does not need many iterations as Bellman-Ford. While Dellman-Ford needs N = iterations, GNNs with four iterations have almost identical test accuracy as GNNs with seven iterations (94% vs 95%). This can also be explained through algorithmic alignment, as GNN aligns with an optimized version of Bellman-Fort Ford, which we explain in Appendix G.3.
Fig. 4 shows how the test accuracies of different models vary with the number of sampsed training points. Indeed, the test accuracy increases more slowly for models that align worse with the task, which implies they need more training samples to achieve similar generalization performance. Again, this confirms our theory.
After verifying that GNNs can sample-effciinnlly learn DP, we show that two popular families of reasoning tasks, visual question answering and intuitive physics, can be formulated as DP. Therefore, our framework explains why GNNs are effective in these tasks.
Visual question answering. The Pretty-CLEVE dataset (Palm et al., 2018) is an extension of Sort-of-CLEVR (Santoro et al., 2017) and CLEVR (Johnson et al., 2017a). GNNs work well on these datasets. Each question in Pretty-CLEVE has state representations and asks "Starting at object X, if each time we jump to the closest object, which object K jumps away?". This problem can be solved by DP, which computes the answers for ie jumps from the answers for (I- - jumps.
closest( [1]
[[]] arg miny d(i,j), closest [:] [+]
= closest(k- 1] [[ooses((1]]) for k > 1, (4.3)
where closest [k]]t] is the answer for jumping & times from object and d(i,j) j) is the distance between the i-th and the j-th object.
Intuitive physics. Battaglia al. (2016) and Watters et al. (2017) train neural networks to predict object dynamics in rigid body scenes and n- n-body systems. Chang et al. (2017) and Janner et al. (2019) study other rigid body scenes. If the force acting on physical object stays constant, we compute the object's 's trajectory with simple functions (physics laws) based on its initial position and force. Physical interactions, however, make the force change, which means the function to compute the object's dynamics has to change too. Thus, DP algorithm would recursively compute the next force changes in the system and update DP states (velocity, momentum, position etc of objects) according to the (learned) forces and physics laws (Thijssen, 2007).
for & = 1..K : time = min.,, sarre-chaang-ttine time(state(k 1, a). state(k- [I: 1,5)),
(4.4)
for i I..N: state (ite/k]p k|6] Update- seree-----oorosstttteeee ]][r], time), j = I..N,
(4.5)
8

---

Published as a conference paper at ICLR 2020
foree-change-Bee computes the time which the force between object and will change. Update- by-forces updates the state of each object at the next force change time. In rigid body systems, force changes only at collision. In datasets where object collides more than once between time frames, one-iteration algorithm/GNN can work (Battaglia et al., 2016). More iterations are needed multiple collisions occur between two consecutive frames (Li & Liang, 2018). In n-body systems, forces change continuously but smoothly. Thus, finite-iteration DP/GNN can be viewed as a form of Runge-Kutta method (DeVries & Hamill, 1995).
4.4 DESIGNING NEURAL NETWORKS WITH ALGORITHMIC ALIGNMENT
While DP solves many reasoning tasks, it has limitations. For example. NP-hard problems cannot be solved by DP. follows that GNN also cannot sample-effiiientl learn these hard problems. Our framework, however, goes beyond GNNs. If we know the structure of a suitable underlying reasoning algorithm, we can design a network with a similar structure to learn If we have no prior knowledge about the structure, then neural architecture search over algorithmic structures will be needed.
Subset Sum. As an example, we design a new architecture that can learn to solve the subset sum problem: Given a set of numbers, does there exist subset that sums to 0? Subset sum is NP-hard (Karp, 1972) and cannot be solved by DP. Therefore, our framework predicts that GNN cannot generalize this task. One subset sum algorithm is exhaustive search, where enumerate all 2 possible subsets and check whether has zero-sum. Following this algorithm, we design a similarly structured neural network which we call Neural Exhaustive Search (NES). Given a universe, NES enumerates all subsets of objects and passes each subset through an LSTM followed by a MLP. The results are aggregated with a max- max- max-pooling layer and MLP:
MLP22mmax,5 MLP, o LSTMMX
X1,..., X|=| E T)).
(4.6)
This architecture aligns well with subset-sum, since the first MLP and LSTM only need learn a simple step, checking whether subset has zero sum. Therefore, expect NES to generalize well in this task. Indeed, NES has 98% test accuracy, while other models perform much worse (Fig. 3d).
CONCLUSION
This paper is an initial step towards formally understanding how neural networks can learn to reason. In particular, we answer what tasks a neural network can learn to reason about well, by studying the generalization ability of learning the underlying reasoning processes for a task. To this end, we introduce an algorithmic alignment framework to formalize the interaction between the structure of a neural network and a reasoning process, and provide preliminary results on sample complexity. Our results explain the success and suggest the limits of current neural architectures: Graph Neural Networks generalize many popular reasoning tasks because the underlying reasoning processes for those tasks resemble dynamic programming.
Our algorithmic alignment perspective may inspire neural network design and opens up theoretical avenues. An interesting direction for future work is to design, e.g. via algorithmic alignment, neural networks that can learn other reasoning paradigms beyond dynamic programming, and to explore the neural architecture search space of algorithmic structures.
From a broader standpoint, reasoning assumes a good representation of the concepts and objects in the world. To complete the picture, it would also be interesting to understand how better disentangle and eventually integrate "representation" and "reasoning".
ACKNOWLEDGMENTS
We thank Zi Wang and Jiajun Wu for insightful discussions. This research was supported by NSF CAREER award 1553284, DARPA DSO's Lagrange program under grant FA86501 PA86501827834 and Chevron-MIT Energy Fellowship. This research was also supported by JST ERATO JPMJER 201 and JSPS Kakenhi JP18H05291. MZ was supported by DARPA award 1-10011-15---111 13 under subcontract to Raytheon BBN Technologies. The views, opinions, and/or findings contained in this article are those of the author and should not be interpreted representing the official views or policies, either expressed or implied, of the Defense Advanced Research Projects Agency or the Department of Defense.

---

Published as a conference paper at ICLR 2020
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in Neural Information Processing Systems, pp. 6155-6166, 2019.
Stanislaw Antol, Aishwarya Agrawal, Jiasen Lu, Margaret Mitchell, Dhruv Batra, Lawrence Zitnick, and Devi Parikh. Vqa: Visual question answering. In Proceedings of the IEEE international conference on computer vision, pp. 11425-2433 2015.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grainet analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference Machine Learning, pp. 322-332, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances Neural Information Processing Systems, pp. $139-8148 2019b.
Sanjeev Arora, Simon S. Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets small-data tasks. In International Conference on Learning Representations, 2020.
Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463-482 2002.
Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. neural networks. In Advances in Newral Information Processing Systems, pp. 6240-6249, 2017.
normallzed margin bounds for
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. In Advances in Neural Information Processing Systems, pp. 4502-4510, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 180.00161 2018.
Richard Bellman. On routing problem. Quarterly of applied mathematics, 16(1):87-90, 1958.
Richard Bellman. Dynamic programming. Science, 153(3731):34-77, 1966.
Michael Chang, Abhishek Gupta, Sergey Levine, and Thomas L. Griffiths. Automatically composing representation transformations as a means for generalization. In International Conference on Learning Representationss, 2019.
Michael B Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach to learning physical dynamics. In International Conference on Learning Representations, 2017.
Paul L Veieies and Patrick Hamill. A first course in computational physics, 1995.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pp. 1675- 1685, 2019a.
Simon Du, Kangcheng Hou, Russ R Salakhutdinov, Barnabas Poczos, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. In Advances in Neural Information Processing Systems, pp. 5724-5734, 2019b.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over- parameterized neural networks. In International Conference on Learning Representations, 2019c.
10

---

Published as a conference paper at ICLR 2020
François Fleuret, Ting Li, Charles Dubout, Emma K Wampler, Steven Yantis, and Donald Geman. Comparing machines and humans on visual categorization test. Proceedings of the National Academy Sciences, 108(43):16221-77233 2011.
Katerina Fragkiadaki, Pulkit Agrawal, Sergey Levine, and Jitendra Malik. Learning visual predictive models of physics for playing billiards. In International Conference Leurning Representations, 2016.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference Machine Learning, pp. 1273-1272, 2017.
Noah Golowich, Alexander Rakhlin, and Ohad Shamir. Size-independent sample complexity of neural networks. In Conference On Learning Theory, pp. 297-299, 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1997.
Ronghang Hu, Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Kate Saenko. Learning to reason: End-to-end module networks for visual question answering. In Proceedings of the IEEE International Conference on Computer Vision, pp. 804-813, 2017.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances neural information processing systems, pp. 8571-8580, 2018
Michael Janner, Sergey Levine, William T. Freeman, Joshua B. Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-centric models. In International Conference on Learning Representations, 2019.
Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Li Fei-Fei, C Lawrence Zitnick, and Ross Girshick. Clevr: A diagnostic dataset for compositional language and elementary visual reasoning. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2901-2910, 2017a.
Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Judy Hoffman, Li Fei-Fei, C Lawrence Zitnick, and Ross Girshick. Inferring and executing programs for visual reason- ing. In Proceedings of the IEEE International Conference on Computer Vision, pp. 2989-2998, 2017b.
Richard M Karp. Reducibility among combinatorial problems. In Complexity of computer computa- tions, pp. 85- Springer, 1972.
Andrei N Kolmogorov. On tables of random numbers. Theoretical Computer Science, 207(2): 387-395, 1998.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances Neural Information Processing Systems, pp. 8157-8166. 2018.
Jiayuan Mao, Chuang Gan, Pushmeet Kohli, Joshua B. Tenenbaum, and Jiajun Wu. The neuro- symbolic concept leamner Interpreting scenes, words, and sentences from natural supervision. In International Conference on Learning Representations, 2019.
Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li F Fei-Fei, Josh Tenenbaum, and Daniel L Yamins. Flexible neural representation for physics prediction. In Advances in Neural Information Processing Systems, pp. 8799-88 881088800, 10, 2018.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. Nasedbased capacity control in neural networks. In Conference on Learning Theory, pp. 1376-1401, 2015.
Rasmus Palm, Ulrich Paquet, and Ole Winther. Recurrent relational networks. Advances in Neural Information Processing Systems, pp. 3368-3378 2018.

---

Published as a conference paper at ICLR 2020
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. In International Conference on Learning Representations, 2019.
Alvaro Sanchez-Gonzalez, Nicolas Heess, Jost Tobias Springenberg. Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks learnable physics engines for inference and control. In International Conference on Machine Learning, pp. 4467-4476, 2018.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. simple neural network module for relational reasoning. In Advances in neural information processing systems, pp. 4967-4976, 2017.
Adam Santoro, Felix Hill, David Barrett, Ari Morcos, and Timothy Lillicrap. Measuring abstract reasoning in neural networks. In International Conference on Machine Learning, pp. 4477-4486, 2018.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models. In International Conference on Learning Representations, 2019.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. Computational capabilities of graph neural networks. IEEE Transactions Neural Networks, 20 (1):81-102, 2009a.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions Neural Networks, 20(1):61-80, 2009b.
Jos Thijssen. Computational physies. Cambridge university press, 2007.
Leslie G Valiant. A theory the learnable. In Proceedings of the sixteenth annual ACM symposium on Theory of computing, pp. 436-445, ACM, 1984.
Vladimir Vapnik. The nature of statistical learning theory. Springer science & business media, 2013.
Edward Wagstaff, Fabian B Fuchs, Martin Engelcke, Ingmar Posner, and Michael Osborne. On the limitations of representing functions on sets. In International Conference on Machine Learning. 2019.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning a physics simulator from video. In Advances in neural information procexsing systems, pp. 4539-4547, 2017.
Jason Weston, Antoine Bordes, Sumit Chopra, Alexander Rush, Bart van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks. arXiv preprint arXiv: 00t:::502.05550 006905698, 2015.
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning on graphs with jumping knowledge networks. In International Conference on Machine Learning, pp. 5453-5462, 2018.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations, 2019.
Kexin Yi, Jiajun Wu, Chuang Gan, Antonio Torralba, Pushmeet Kohli, and Josh Tenenbaum. Neural- symbolic vqa: Disentangling reasoning from vision and language understanding. In Advances in Neural Information Processing Systems, pp. 1031-1042, 2018.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh. Barnabas Poczos, Ruslan R Salakhutdinov, and Alexander Smola. Deep sets. In Advances in Neural Information Processing Systems, pp- 3391-3401, 2017.
Chi Zhang, Feng Gao, Baoxiong Jia, Yixin Zhu, and Long-Chun Zhu. Raven: A dataset for relational and analogical visual reasoning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5317-5327, 2019.
12

---

Published as a conference paper at ICLR 2020
A PROOF OF PROPOSITION 3.1
We will prove the universal approximation of GNNs by showing that GNNs have at least the same expressive power Deep Sets, and then apply the universal approximation of Deep Sets for permutation invariant continuous functions.
Zaheer al. (2017) prove the universal approximation of Deep Sets under the restriction that the set size is fixed and the hidden dimension is equal to the set size plus one. Wagstaff et al. (2019) extend the universal approximation result for Deep Sets by showing that the set size does not have to be fixed and the hidden dimension is only required to be at least as large as the set size. The results for our purposes can summarized as follows.
Universal approximation of Deep Sets. Assume the elements are from compact set in R Any continuous function on a set of size bounded by N. i.e. f R - R, that is permutation invariant to the elements in S can be approximated arbitrarily close by some Deep Sets model with sufficiently large width and output dimension for its MLPss
Next we show any Deep Sets can be expressed by some GNN with one message passing iteration. The computation structure one-layer GNNs is shown below.
hg Σ a X1), hg = g 1ES
Σ hs
(A.1)
.9ES
where & and g are parameterized by MLPs. If ò is function that ignores X, so that di (X,; X₁) p(X.) for some p. e.g., by letting part of the weight matricies in á be 0, then we essentially get Deep Sets in the following form.
h, = pXx,) hs =
Σ he L
(A.2)
For any such p. we can get the corresponding e via the construction above. Hence for any Deep Sets, we can express it with an one-layer GNN. The same result applies to GNNs with multiple layers (message passing iterations), because multiple (kk's, ss which that GNNs are universal approximators for permutation invariant continuous functions.
express function p(X,) by the composition of express with GNN layer via our construction above. It then follows
B PROOF OF PROPOSITION 3.2
For any GNN N, we construct an MLP that able to do the exact same computation as N. will then follow that the MLP can represent any function A can represent. Suppose the computation structure of N is the following.
(k) (10 m(--)) hs 9
> (((
(B.1)
where f and g are parameterized by MLPs. Suppose the set size is bounded by M (the expressive power of GNNs also depend on M Wagstaff et al. (2019)). We first show the result for fixed size input, i.e., MLPs approach to deal with variable sized input.
simulate GNNs if the input set has a fixed size, and then apply ensemble
Let the input to the MLP be vector concatenated by 1 (0)s,, "s, in some arbitrary ordering. For each message passing iteration of N, any (()) can be represented by an MLP. Thus, for each pair of (h1 (h(6-1)) l1g (1(--)),, we can set weights in the MLP that the the concatenation of all ((k,kk1) (k-1) 11 (k-11)
become the hidden vector after some layers of the MLP. With the vector of (h
(k-1) hss (k-1)) as
input, in the next few layers of the MLP we can construct weights so that we have the concatenation
of h Eies f(k) (1g
encode (k-1) 1) (k- (k-1) as the result of the hidden dimension, because
summation with weights in MLPs. So far, we can simulate an iteration of GNN A with layers of MLP.
We repeat the process for K times by stacking the similar layers. Finally, with concatenation of
13

---

Published as a conference paper at ICLR 2020
hi as our hidden dimension in the MLP, similarly, we can simulate hg = 9 his h(²))) layers of MLP. Stacking all layers together, have obtained MLP that simulate ..
with
To deal with variable sized inputs, we construct M MLPs that can simulate the GNN for each imput set size 1,... M. Then construct meta-layer, whose weights represent (universally approximate) the summation of the output of M MLPs multiplied by an indicator function of whether each MLPs has the same size the input (these need be input information). The meta layer weights on top can then essentially select the output from of MLP that has the same size as the set input and then exactly simulate the GNN. Note that the MLP we construct here has the requirement for how we input the data and the information of set sizes etc. In practice, we can have M MLPs and decide which MIP to use depending on the input set size.
C PROOF OF THEOREM 3.5
Theorem 3.5 is generalization Theorem 6.1 in (Arora et al.. 2019a), which addresses the scalar case. See (Arora et al., 2019a) for complete list of assumptions.
Theorem C.1. (Arora et al., 20/9a) Suppose we have R. R, E, o, where 3, Rd. a E and pj = 1 or P, 21 (1 E N+). Let A be an overparameterized two-layer MLP that is randomly initialized and truined with gradient descent for sufficient number of iterations.
The sample complexity C,(g,e,5) is O Σ, p,lo,)/1118 +log(1/6)
To extend the sample complexity bound to vector-valued functions, we view each entry/component of the output vector as an independent scalar-valued output. We can then apply a union bound to bound the error rate and failure probability for the output vector, and thus, bound the overall sample complexity.
Let and 8 be the given error rate and failure probability. Moreover, suppose we choose some error rate eo and failure probability 8, for the output/function each entry. Applying Theorem C.I to each component
g(x)(1) Eaj0 (3)TT))
Ds g:
(C.1)
yields sample complexity bound of
Co(gn.c, 60) =o (
(1) (1)
10
12
+ log (1/5o)
(C.2)
for each g: (2:). Now let us bound the overall each entry. The probability that we fail to learn each of the is at most on- Hence, by a union bound, the probability that we fail learn any of the Ji is at most m dg. Thus, with probability at least 1 - -mmg we successfully learn all g: for = 1, m,, so the eтrоr for every entry is bounded by eg- The error for the vector output then at most 5 ED mEg.
rate and failure probability given eo and dg for
Setting mddg = å and meg = gives us 6p mp function for each output entry independently with error e/m and failure rate 5/m, we can successfully learn the entire vector-valued function with rate and 8. This yields the following overall sample complexity bound:
and ég <. Thus, if we can successfully learn m
(2) (4) ca(o,c,8)=0 = max; E, Dj la
113
12
+ log (m/5)
(C.3)
(e/m)²
Regarding m as constant, we can further simplify the sample complexity to
ca(g,c,0)=0 =0 max; E, p, 1/ (4)
-11550 + log (1/6)
(C.4)
2
14

---

Published as a conference paper at ICLR 2020
D PROOF OF THEOREM 3.6
We will show the learnability result by an inductive argument. Specifically, we will show that under our setting and assumptions, the error between the learned function and correct function the test set will not blow up after the transform of another learned function fy. assuming learnability previous f1, ... f- 1 by induction. Thus, we can essentially provably learn at all layers/iterations and eventually learn g.
Suppose we have performed the sequential learning. Let us consider what happens at the test time. Let fy be the correct correct functions as defined in the algorithmic aligmment. Let fy be the functions learned by algorithm Aj and MLP N- We have input S ~ D, and our goal is to bound ([)(S) - [(S) with high probability. To show this, we bound the error of the intermediate representation vectors, Le., the output of fj and fj. and thus, the input to fy+1 and fJ+1-
Let us first consider what happens for the first module Mi. fi and fi have the same input distribution I ~ DD where x are obtained from S, e.g. the pairwise object representations as in Eqn. 2.2. Hence, by the learnability assumption on A1. (x) - [x))] < e with probability at least - 5. The error for the input of N2 is then Occ) with failure probability ( because there are constant number of terms of aggregation of fi 's output, and we can apply union bound to upper bound the failure probability.
Next, we proceed by induction. Let us fix a k. Let = denote the input for Su, which are generated by the previous 's, and let & denote the input for ft. which are generated by the previous fy's. Assume ||2 2| I O(c) with failure probability at most (()). We aim to show that this holds for + 1. For the simplicity of notation, let f denote the correct function fa and let denote the learned function ft. Since there are a constant number of terms for aggregation, our goal is then to bound [lf(2) - f2)]l. By triangle inequality, we have
[[())) - [())] = (1 f(=) + f(=) s(=]ll
(D.1)
S (1 - f(=))l + +1[f(=)- - s(=)l
(D.2)
We can bound the first term with the Lipschitzness assumption of ff as the following.
[I(()) f(=)) S - =||
(D.3)
To bound the second term, our key insight is that f is a learnale correct function, so by the learnability coefficients in algorithmic alignment, is close to the function f learned by the learning algorithm A on the correct samples, i.e., f is close to f = A ([2,, y= Moreover, f is generated by the learning algorithm A on the perturbed samples, i.e., f = AA (f&, }: By the algorithm stability assumption, if and f should be close if the input samples are only slightly perturbed. It then follows that
(If(2) f(=)]1 [f(()) f(=) f(z) f(=)]l
(D.4)
< - f()) + 11f(2) f(x)l
(D.5)
Lo max || 2.l +e w.p. 2 5
(D.6)
where and &; are the training samples at the same layer kt Here, we apply the same induction condition as what we had for z and 2: || - lll < O(e) with failure probability at most 0()). We can then apply union bound to bound the probability of any bad event happening. Here, we have 3 bad events each happening with probability at most 0(5). Thus, with probability at least O(8), we have
lIf(=) [(=)]| < (,(()) + IoO(c) + +(=0(e)
(D.7)
This completes the proof.
E PROOF OF COROLLARY 3.7
Our main insight is that giant MLP learns the same function (X;-Xx them in the weights. This leads to the ( extra sample complexity through Theorem 3.5, because the number of polynomial terms (X, X2) is of order &²
- X,)² for 62 times and encode
15

---

Published as a conference paper at ICLR 2020
First of all, the function f(r v) (x y)² can be expressed the following polynomial.
(x y) (11 - 1] [x vl)
(E.1)
We have 3 =1 -1-11111 11. [1. so ||31| 4. Hence, by Theorem 3.5, it takes O ( samples for an MLP to learn f(r (a- y)². Under the sequential training setting, an one-layer GNN applies an MLP lear f, and then sums up the outcome of f(x,, X2) for all pairs X,, X, Here, we essentially get the aggregation error O(E² <) from & pairs. However, we will that applying an MLP to learn will also incur the same aggregation error. Hence, we do not need to consider the aggregation error effect when we compare the sample complexities.
Now we consider using MLP to learn the function g. No matter in what order the objects X, are concatenated, we can express g with the sum of polynomials the following.
g(S) = _00
(E.2)
where Pig has at the i-th entry, -1 at the j-th entry and 0 elsewhere. Hence || 4. It then follows from Theorem 3.5 and union bound that it takes O(( + log(1/8))/22 to learn g. where é Ee and 8 8 Here, we have discussed above, the same aggregation error occurs in the aggregation process of f. so we can simply consider é for both. Thus, comparing O[log(L/0)/E* and ((((2 + { gives us the O(e2) difference.
F PROOF OF CLAIM 4.1
We prove the claim by contradiction. Suppose there exists f such that f(x) + f(y) g(x, v) for any a and y. This implies that for any x, we have f(x) + f(x) = g(x, x) = 0. follows that f(xx = for any x. Now consider some x and so that x: * y. We must have f(x) + f(y) = 0 0. However, g(x, 3) # because x # y. Hence, there exists x and so that f(x) + f(y) + g(z,g). We have reached contradiction.
G EXPERIMENTS: DATA AND TRAINING DETAILS
G.1 FANTASTIC TREASURE: MAXIMUM VALUE DIFFERENCE
Dataset generation. In the dataset, sample 50, 000 training data, 5, 000 validation data, and 5, 000 test data. For each model, we report the test accuracy with the hyperparameter setting that achieves the best validation accuracy. In each training sample, the input universe consists of 25 treasures XI...,,22: For each treasure X,, we have X; = [/,,, f22 ha], where the location h; is sampled uniformly from [0..20] the value h2 sample uniformly form [0...00], and the color hg is sampled uniformly from [1..6]. The task is to answer what the difference is in value between the most and least valuable treasure. We generate the answer label for universe
follows: we find the the maximum difference in value among all treasures and set it to Then we make the label into one-hot encoding with 100 + 101 classes.
Hyperparameter setting. We train all models with the Adam optimizer, with learning rate from le - 3, 5e 4, and 1e-4 4, and we decay the learning rate by 0.5 every 50 steps. We use cross-entropy loss. We train all models for 150 epochs. We tune batch size of 128 and 64.
For GNNs and HRN, we choose the hidden dimension of MLP modules from 128 and 256. For DeepSet and MLP, we choose the hidden dimension of MLP modules from 128, 256, 2500, 5000. For the MLP and DeepSet model, we choose the number of of hidden layers for MLP moduels from 4 and 8, 16. For GNN and HRN, set the number of hidden layers of the MLP modules to 3, 4. Moreover, dropout with rate 0.5 is applied before the last two hidden layers of MLP;, iee. the last MLP module in all models.
G.2 FANTASTIC TREASURE: FURTHEST PAIR
Dataset generation. the dataset, we sample 60, 000 training data, 6, 000 validation data, and 6, 000 test data. For each model, report the test accuracy with the hyperparameter setting that
16

---

Published as a conference paper at ICLR 2020
achieves the best validation accuracy. In each training sample, the imput universe consists of 25 treasures X1...., X25. For each treasure Xs, we have X; = [hi, h2, hal, where the location hı is sampled uniformly from [0..20]$ S the value h2 is sample uniformly form [0.. and the color hg is sampled uniformly from [1..6]. The task is to answer what are the colors of the two treasure that are the most distant from each other. We generate the answer label y for universe S as follows: we find the pair of treasures that are the most distant from each other, say (X,, X2). Then we order the pair (X,), (xx))))
obtain an ordered pair (a. b) with S (aka. a min h (/a((i), ha ha(X,)) and (b = max (a.hhh((,...)).))))( (xx(X,)) H) where (xx(X;) denotes the color of X;. Then we compute the label from (a, b) by counting how many valid pairs of colors are smaller than (a, b) (a pair (k. 1) is smaller than (a, b) iff i). k < or ii). k- = and I < b). The label is one-hot encoding of the minimum cost with 6 (6- 1)/2 + 6 = 21 classes.
Hyperparameter setting. We train all models with the Adam optimizer, with learning rate from le 3, 5e - 4, and le 4, and we decay the learning rate by 0.5 every 50 steps. We use cross-entropy loss. We train all models for 150 epochs. We tune batch size of 128 and 64.
For the MLP and DeepSet model, we choose the number of of hidden layers of MLP modules from 4 and 8, 16. For GNN and HRN models, we set the number of hidden layers of the MLP modules from 3 and 4. For DeepSet and MLP models, we choose the hidden dimension of MLP modules from 128, 256, 2500, 5000. For GNNs and HRN, we choose the hidden dimension of MLP modules from 128 and 256. Moreover, dropout with rate 0.5 is applied before the last two hidden layers of MLP1, iee. the last MLP module in all models.
G.3 MONSTER TRAINER
Task description. We are a monster trainer who lives in a world S with 10 monsters. Each monster X = [h,, h2] has a location h1 E [0..10]² and a unique combat level h2 E [1..10]. In each game, the trainer starts at random location with level zero, Xuminer [Po. 0], and receives a quest to defeat the level-k monster. At each time step, the trainer can challenge any more powerful monster X, with cost equal to the product of the travel distance and the level difference c( c(Xuainer; X) |l/ (X - h1 (X)(ll,, x (h2(X) - h2 (Xaminer)). After defeating monster X, the trainer's level upgrades to (X),, and the trainer moves to hı (X). We ask the minimum cost of completing the quest, i.e., defeating the level-k monster. The range of cost (number of classes for prediction) is 200. To make games even more challenging, we sample games whose optimal solution involves defeating three to seven non-quest monsters.
A DP algorithm for shortest paths that needs half of the iterations of Bellman-Ford. We provide a DP algorithm as the following. To compute a shortest-path from a source object to a target object t with at most seven stops, we run the following updates for four iterations:
distance, [1]]]] = cost(s, k,,
distance, [R][u]] = min, {distance, [k- - [][(] + cost(r, u)}
(G.1)
distance, [1][u] = cost(u, t), distance, [K][u] = min (distance, - [][0] + cost(u, })}
(G.2)
Update Eqn. G.1 identical to the Bellman-Food algorithm Equ. 4.2, and distance, is the shortest distance from to u with at most h stops. Update Eqn. G.2 is a reverse Bellman- algorithm, and distance, [k][[]] the shortest distance from u to with at most h stops. After running Eqn. G. and Eqn. G.2 for k iterations, we compute shortest path with at most 2 stops by enumerating a mid-point and aggregating the results of the two Bellman-Fow Ford algorithms:
min.. { ddistance,(M][u + distance: [K][u]]
(G.3)
Thus, this algorithm needs half of the iterations of Bellman-Fon Ford.
Dataset generation. In the dataset, we sample 200, 000 training data, 6, 000 validation data, and 6, 000 test data. For each model, we report the test accuracy with the hyperparameter setting that achieves the best validation accuracy. In each training sample, the input universe consists of the trainer and 10 monsters Xo..., X10, and the request level k. i.e., we need to challenge monster k. We have X; [4, A2], where h1 i indicates the combat level, and the location h2 10.. is sampled uniformly from [0.. We generate the answer label y for a universe S as follows. We implement shortest path algorithm compute the minimum cost from the trainer monster k
17

---

Published as a conference paper at ICLR 2020
where the cost is defined in task description. Then the label y is one-hot encoding of minimum cost with 200 classes. Moreover, when we sample the data, we apply rejection sampling to ensure that the minimum cost's shortest path is of length 3, 4, 5, 6, 7 with equal probability. That is, eliminate the trivial questions.
Hyperparameter setting. We train all models with the Adam optimizer, with learning rate from 2e 4 and 5e 4, and we decay the learning rate by 0.5 every 50 steps. We use cross-entropy loss. We train all models for 300 epochs. We tune batch size of 128 and 64.
For the MLP model, we choose the number of layers from and 8, 16. For other models, we choose the number of hidden layers of MI.P modules from 3 and 4. For GNN models, we choose the hidden dimension of MLP modules from 128 and 256. For DeepSet and MLP models, we choose the hidden dimension of MLP modules from 128, 256, 2500. Moreover, dropout with rate 0.5 is applied before the last two hidden layers of MLP1, Le., the last MLP module in all models.
G.4 SUBSET SUM
Dataset generation. In the dataset, we sample 40, 000 training data, 4, 000 validation data, and 4, 0000 test data. For each model, report the test accuracy with the hyperparameter setting that achieves the best validation accuracy. In each training sample, the input universe S consists 6 numbers X1;-.. ..., Xe. where each X; is uniformly sampled from [-200..200]. The goal is to decide if there exists a subset that sums up to O. In the data generation, we carefully decrease the number of questions that have trivial answers: 1)we control the number of samples where 0 (XI,..., Xe) to be around 1% of the total training data; 2) we further control the number of samples where X1 + ... + X6
j E [1..6] so that X; -X, to be around 1.5% of the total training data. In addition, we apply rejection sampling to make sure that the questions with answer yes (aka, such subset exists) and answer no (aka, no such subset exists) are balanced (i.e., 20,000 samples for each class in the training data).
Hyperparameter setting. We train all models with the Adam optimizer, with learning rate from le - 3, 5e - 4, and le 4, and we decay the learning rate by 0.5 every 50 steps. We use cross-entropy loss. We train all models for 300 epochs. The batch size we use for all models is 64.
For DeepSets and MIP models, we choose the number of of hidden layers of the MLP modules from
4, 8, 16. For GNN and HRN models, we set the number of hidden layers of the last MLP modules to
4. For DeepSets and MLP, choose the hidden dimension of MLP modules from 128, 256, 2500, 5000. For GNN and HRN models, we choose the hidden dimension of MLP modules from 128 and
256. Moreover, dropout with rate 0.5 is applied before the last two hidden layers of MLP i.e. the last MLP module in all models.
The model Neural Exhaustive Search (NES) enumerates all possible non-empty subsets of S, and passes the numbers of to an MI.,, in a random order, to obtain the hidden feature. The hidden feature then passed single-direction layer LSTM of hidden dimension 128. Afterwards, NES applies an aggregation function to these 26 hidden states obtained by the LSTM to obtain the final output. For NES, we set the number of hidden layers of the last MLP, i.e. MLP2, to 4, the number of hidden layers of the MLPs prior to the last MLP, iee. MLP1, to 3, and we choose the hidden dimension of all MLP modules from 128 and 256.
18


Reviews:
[{"review_id": "rJeOFDKjFr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work seeks theoretical and empirical proof of the reasoning capacity of neural networks. The authors build on a body of research that demonstrates the usefulness of different neural network architectures for different reasoning problems. For example, Deep Sets have been proposed to answer questions about sets (e.g., a summary statistic), and GNNs about graph related problems, such as shortest path.\n\nI anticipate that readers would be very satisfied with the intuition behind the main result: neural networks that “align” with known algorithmic solutions are better able to learn the solutions. Many architectures have been proposed over the years, often with a high-level justification for the architecture’s form. For example, Relation Networks noted the difficulty with learning n^2 relations using an MLP, which is an observation reflected in this work’s explanation of the difficulty with learning a for loop. \n\nProvided here is a justification for these high-level design decisions. The authors provide some theory and experimental results to demonstrate their proposed notion of alignment, and show that NNs that align with known algorithmic solution do well, while those that do not align do not do well. In particular, I appreciate both the positive and negative evidence, since demonstrating lack of alignment (and poor performance) is a necessary condition to show alongside alignment (and good performance).\n\nI’d like to caution the authors regarding their main conclusion, which is stated a few times in the paper:\n\n“This perspective suggests that whether a neural network can learn a reasoning task depends on whether there exists an algorithmic solution that the network aligns with”.\n\nI think this logic is not precisely correct, and I would modify this to:\n\n“If the structure of a neural network aligns with a known algorithmic solution, then it can more easily learn a reasoning task than a neural network does not align”. \n\nThis is a subtle but important difference. In particular, the original logic does not capture situations where an algorithmic solution is not known, but a neural network can otherwise still learn a solution (consider object classification). I think even the corrected logic as I’ve spelled it out above might not be quite right either, since it does not consider situations where the algorithmic solution exists, but it obtuse. Would a neural network easily learn such a task? \n\nOverall I think the paper is clearly written, and the experiments are adequate. Unfortunately I am not well-versed in the theoretical literature on this topic, so my assessment of the proofs is limited, and I will need to defer to the other reviewers on these matters. My surface level assessment of them is that the logic seems generally sound, but I cannot make any strong statements placing them in the context of previous work, nor can I properly evaluate the nuances. Nonetheless, as a whole, I think this is a strong contribution and a nicely put together piece of work.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "Hylr6tgjKB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a measure of classes of algorithmic alignment that measure how \"close\" neural networks are to known algorithms, e.g. dynamic programming (DP). The measure is based on the number of samples needed such that the expected generalization error is less than epsilon with 1-delta probability, where epsilon and delta are free parameters.\n\nThe paper proves the link between several classes of known algorithms and neural network architectures by showing how their sample complexity varies. For instance the paper shows that Graph Neural Network (GNN), can approximate any DP algorithm in a sample efficient manner, whereas MLP and deep sets (permutation invariant NN) can't. The paper empirically verifies their claims on 4 toy datasets, each representing an increasingly complex algorithm needed to solve the problem. \n\nI recommend this paper be accepted, since I think it's an important direction of research, and it formalizes a lot of intuition about neural network architectures.\n\nIt would be very interesting if the authors could actually compute the number of samples, M, for different NN architectures on the toy datasets, and show how it matches empirical findings. This could be a powerful tool if it could be made easy to use for the common practitioner.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "HJe7pdvCYr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a framework, dubbed algorithmic alignment, based on PAC learning and sample complexity, with the aim to explain generalization on reasoning tasks for different neural architectures. The framework roughly states that in order for the model to be able to learn and successfully generalize on a reasoning task, it needs to be able to easily learn (to approximate) steps of the reasoning tasks. The authors use this framework to propose an increasingly difficult set of tasks, designed to showcase the type of models that would be fit or unfit to solve them. The resulting experiments corroborate the theory, showing the limits of MLPs, Deep Sets, and consequently Graph Neural Networks. The final claim that an NP-hard task needs an enumerative architecture, and then experimental validation of that claim is nice and fits into the theory.\nThe added benefit of the paper is that the authors show as a side-effect that visual question answering and intuitive physics\n\nOverall, the paper presents a meaningful contribution to the theory of learning, formalizing the means of quantifying the capabilities of architectures to solve tasks of certain complexity. The paper, though dense, is well well written, and carries an interesting conclusion that better algorithmic alignment brings the sample complexity down, i.e. models with better algorithmic alignment to the task (function they want to approximate) should generalize better.\nThe formalization presented in the paper, though remarkably intuitive, might be difficult to practically use for more elaborate models and it is not clear whether it can be numerically computed. The paper (i.e. the reader) would certainly benefit from more examples of algorithmic alignment comparison of different models, such as one done in Corollary 3.7.\n\nQuestion:\n- difference to Kolmogorov complexity is that any algorithmic alignment that yields decent sample complexity is good enough - how do you define decent?\n- You state: “in Section 4, we will show that we can usually derive a near-optimal alignment by avoiding as many “for loops” in algorithm steps as possible.” yet I did not see that there. Was that effectively shown in Corollary 3.7?\n\nSlightly related work: On the Turing Completeness of Modern Neural Network Architectures", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a framework which qualifies how well given neural architectures can perform on reasoning tasks. From this, they show a number of interesting empirical results, including the ability of graph neural network architectures for learn dynamic programming.

This substantial theoretical and empirical study impressed the reviewers, who strongly lean towards acceptance. My view is that this is exactly the sort of work we should be show-casing at the conference, both in terms of focus, and of quality. I am happy to recommend this for acceptance.
Decision: Accept (Spotlight)