Graph Neural Networks are Dynamic Programmers
Authors: Andrew Joseph Dudzik
Petar Veličković
Venue: NeurIPS 2022 Accept (2022)
Keywords: algorithmic reasoning
graph neural networks
category theory
polynomial functors
bellman-ford
integral transform
pullback
pushforward
monads
message passing
dynamic programming

Abstract: Recent advances in neural algorithmic reasoning with graph neural networks (GNNs) are propped up by the notion of algorithmic alignment. Broadly, a neural network will be better at learning to execute a reasoning task (in terms of sample complexity) if its individual components align well with the target algorithm. Specifically, GNNs are claimed to align with dynamic programming (DP), a general problem-solving strategy which expresses many polynomial-time algorithms. However, has this alignment truly been demonstrated and theoretically quantified? Here we show, using methods from category theory and abstract algebra, that there exists an intricate connection between GNNs and DP, going well beyond the initial observations over individual algorithms such as Bellman-Ford. Exposing this connection, we easily verify several prior findings in the literature, produce better-grounded GNN architectures for edge-centric tasks, and demonstrate empirical results on the CLRS algorithmic reasoning benchmark. We hope our exposition will serve as a foundation for building stronger algorithmically aligned GNNs.

Graph Neural Networks are Dynamic Programmers
Andrew Dudzik" DeepMind adudzik@deepmind.comm com
Petar Veličković DeepMind petarvodeepmindccom com
Abstract
Recent advances in neural algorithmic reasoning with graph neural networks (GNNs) are propped up by the notion of algorithmic alignment. Broadly, a neural network will be better at learning to execute reasoning task (in terms of sam- ple complexity) its individual components align well with the target algorithm. Specifically, GNNs claimed to align with dynamic programming (DP), gen- eral problem-solving strategy which expresses many polynomial-time algorithms. However, has this alignment truly been demonstrated and theoretically quantified? Here we show, using methods from category theory and abstract algebra, that there exists an intricate connection between GNNs and DP, going well beyond the initial observations over individual algorithms such as Bellman-Ford. Exposing this connection, we easily verify several prior findings in the literature, produce better-gree grounded GNN architectures for edge-centric tasks, and demonstrate empiri- cal results on the CLRS algorithmic reasoning benchmark. We hope our exposition will serve foundation for building stronger algorithmically aligned GNNs.
Introduction
One of the principal pillars of neural algorithmic reasoning 27 is training neural networks that execute algorithmic computation in nigh-dimensional latent space. While this process in itself insightful, and can lead to stronger combinatorial optimisation systems [21]. it is valuable in terms of expanding the applicability of classical algorithms. Evidence of this value are emerging, with pre-trained algorithmic reasoners utilised in implicit planning [11] and self- learning [28]
fundamental question in this space is: which architecture should be used learn a particular algorithm (or collection of algorithms [36))? Naturally, we seek architectures that have low sample complexity, as they will allow
create models that generalise better with fewer training examples.
The key theoretical advance towards achieving this aim has been made by [37 Therein, the authors formalise the notion of algorithmic alignment, which states that should favour architectures that align better to the algorithm, in the sense that we can separate them into modules, which individually correspond the computations of the target algorithm's 's subroutines. It can be proved that architectures with higher algorithmic alignment will have lower sample complexity in the NTK regime (20). Further, the theory of [37] predicts that graph neural networks (GNNs) algorithmically align with dynamic programming [3 DPJ. The authors demonstrate this by forming an analogy to the Bellman-Ford algorithm (2)
Since DP is very general class of problem-solving techniques that can be used to express many classical algorithms, this finding has placed GNNs as the central methodology for neural algorithmic exccution 77). However, it quickly became apparent that it is not enough to just train any GNN-for many algorithmic tasks, careful attention is required. Several papers illustrated special cases of GNNs that align with sequential algorithms [31], linearithmic sequence processing [16], physics simulations
Equal contribution.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

23 iterative algorihtms [26,, data structures [29] or auxiliary memory (24) Some explanations for this lack of easy generalisation have arisen- we now have both geometric and causal views into how better generalisation can be achieved.
We believe that the fundamental reason why so many isolated efforts needed to look into learning specific classes of algorithms is the fact the GNN-DP connection has not been sufficiently explored. Indeed, the original work of [37 merely mentions in passing that the formulation of DP algorithms seems to align with GNNs, and demonstrates one example (Bellman-Ford), Our thorough investiga- tion of the literature yielded no concrete follow-up to this initial claim. But DP algorithms are very rich and diverse, often requiring broad spectrum of computations. Hence what we really need is a framework that could allow us identify GNNs that could align particularly well with certain classes of DP, rather than assuming a mine-size-iitsalll GNN architecture will exist.
As first step towards this, in this paper interpret the operations of both DP and GNNs from the lens of category theory and abstract algebra. We elucidate the GNN-DP connection by observing diagrammatic abstraction of their computations, recasting algorithmic alignment aligning the diagrams of (G)NNs to ones of the target algorithm class. In doing so, several previously shown results will naturally arise as corollaries, and we propose novel GNN variants that empirically align better to edge-centric algorithms. We hope our work opens up the door to broader unification between algorithmic reasoning and the geometric deep learning blueprint
2 GNNs, dynamic programming, and the categorical connection
Before diving into the theory behind our connection, we provide a quick recap on the methods being connected: graph neural networks and dynamic programming. Further, we cite related work to outline why it is sufficient to interpret DP from the lens of graph algorithms.
We will use the definition of GNNs based 5 Let graph be a tuple of nodes and edges, G = (V, E), with one-hop neighbourhoods defined as N₁ = {u E V (v, u) E E}. Further, a node feature matrix X E v|*k gives the features of node u as xu; we omit edge- and graph-level features for clarity, A (message passing) GNN over this graph is then executed as:
h Φ
Xus
v (b.:
vENN
where v/ : RR* R + IR* is message function, Φ RR x R - RR is readout function, and Q is a permutation-in -nniariiinnnnnaiaanntt uggregation function (such as Σ or max). Both vi and Φ can realised as MLPs, but many special cases exist, giving rise to, e.g., attentional GNNs 30]
divide et impera fashion: imagine that we want to solve a problem instance r. DP proceeds to identify a set of subproblems, ((x), such that solving them first, and recombining the answers, can directly lead to the solution for a: f(x)= p(( y 7)=)))) Eventually, we decompose the problem enough until we arrive at an instance for which the solution is trivially given (i.e. f(y which is known upfront). From these "base cases", we can gradually build up the solution for the problem instance we initially care for in bottom-up fashion. This rule is often expressed programmatically:
Dynamic programming is defined as a process that solves problems
dpíx] +
recombine (ercombinesccor(((e (escore(dpp)) yy), dp(x)) for y in expand(x))
To initiate our discussion on why DP can be connected with GNNs, it is worthwhile exercise to induces a graph structure. To see this, we leverage a categorical analysis of dynamic programming first proposed by [10]. Therein, dynamic programming algorithms are reasoned about as a composition of three components (presented here on high level):
show how Equation 2
dp
(3)
recombine
expand
Expansion selects the relevant subproblems; scoring computes the quality of each individual subprob- lem's solution w.r.t. the current problem, and recombining combines these solutions into a solution for the original problem (e.g. by taking the max, or average).
Therefore, we can actually identify every subproblem as node in a graph. Let V be the space of all subproblems, and R an appropriate value space (c.g. the real numbers). Then, expansion is defined

---

as 7:: V + P(V). giving the set of all subproblems relevant for given problem. Note that this also induces set of edges between subproblems, E; namely, (z, y) E i E nyy). Each subproblem is scored by using function (P(V) + PR).. Finally, the individual scores are recombined using the recombination function, / ((R) - R. The final dynamic programming primitive therefore computes function dp: V + R in each of the subproblems of interest.
Therefore. dynamic programming algorithms can be seen as performing computations over a graph of subproblems, which can usually be precomputed for the task at hand (since the outputs of are assumed known upfront for every subproblem). One specific popular example the Bellman- algorithm [2]. which computes single-soure source shortest paths from given source node, in graph G (V, E). In this case, the set of subproblems is exactly the set of nodes, V. and the expansion m(u) is exactly the set of one-hop neighbours of u in the graph. The algorithm maintains distances of every node to the source, du. The rule for iteratively recombining these distances as follows:
du f min duu min dy Wp-ra HEN=
(4)
where wi-+u is the distance between nodes and u. The algorithm's 's base cases are d, 0 for the source node, d, = too otherwise. Note that more general forms of Bellman-Foo pathfinding exist, for appropriate definitions of + and min (in general known as a semiring). Several recent research papers such as NBFNet (39 explicitly call on this alignment in their motivation.
3 The difficulty of connecting GNNs and DP
The basic technical obstacle to establishing a rigorous correspondence between neural networks and DP is the vastly different character of the computations they perform. Neural networks are built from linear algebra over the familiar real numbers, while DP, which is often a generalisation of path-finding problems, typically takes place over "tropical" objects like (N U {oo}, min, + + which are usually studied in mathematics as "degenerations" of Euclidean space. The two worlds cannot clearly be reconciled, directly, with simple equations.
However, if we define an arbitrary "latent space" R and make as few assumptions as possible, we can observe that many of the behaviors we care about, for both GNNs and DP, arise from looking at functions >> R. where is a finite set. R can be seen as the set of real-valued vectors in the case of GNNs, and the tropical numbers in the case of DP.
So our principal object of study is the category of finite sets, and "R-valued quantities" on it. By "category" here we mean a collection of objects (all finite sets) together with a notion of composable arrows (functions between finite sets).
To draw our GNN-DP connection, we need to devise an abstract object which can capture both the GNN's 's message passing/aggregatinn (Equation 2) It may seem quite intuitive that these two concepts can and should be relatable, and category theory is very attractive tool for "making the obvious even more obvious" " [15). Indeed, recently concepts from category theory have enabled the construction of powerful GNN architectures beyond permutation equivariance (.. Here, propose integral transforms as such object.
stages (Equation [ and the DP's sooring/recombination stages
We will construct the integral transform by composing transformations over our input features in a way that will depend minimally on the specific choice of Rt. In doing so, we will build a computational diagram that will be applicable for both GNNs and DP (and their own choices of Rt,, and hence allowing for focusing making components of those diagrams aligned possible.
4 The integral transform
An integral transform can be encoded in a diagram of this form, which we call a polynomial span:
2 Here we require the addition of special "oo" placeholder object to denote the vertices the DP expansion basn't 't reached so far.
3

---

X
Y
W
where W, X, Y and Z are finite sets. The arrows p, stand, respectively, for "input", "process", and "output". In context, the sets will have the following informal meaning: W represents the set over which we define our inputs. Z the set over which we define outputs. X and Y are, respectively, carrier sets for the arguments, and the message - we will clarify their meaning shortly.
Before proceeding, it is worthy to note the special case of X Y E, with p being the identity map. Such a diagram is commonly known as a span. A span that additionally has W = Z = V is equivalent to representation of directed graph with vertex set Z and edge set Y (V <- E - V); in this case i(e) and ((e) are the functions identifying the source and target nodes of each edge.
The key question is: given input data f on W. assigning features (ww) to each w W, how to transform it, via the polynomial span, into data on Z? we can do this, we will be able to characterise both the process of sending messages between nodes in GNNs and scoring subproblems in DP.
For us, data on a carrier set S consists of element of [S, R] := {f S + R}. where R is "set of possible values". For now, we will think of Rt as an arbitrary (usually infinite) set, though we will see later that should possess some algebraic structure; it should be a semiring.
The transform proceeds in three steps, following the edges of the polynomial span:
X, R
Y Y.R
T
0g
1 [Z, R]
(5)
[W, R]
We call the three arrows i" Psc. Og the pullback, the argument pushfoward, and the message puskfor- ward. Taken together, they form an integral transform-and we conjecture that this transform can be described as polynomial functor, where p& and og correspond to the dependent product and dependent sum from type theory (cf. Appendix for details).
The pullback i is the casiest to define. Since we have function i X + W (part of the polynomial span) and function f W R (our input data), we can produce data on X, that function in X + R, by composition. We hence define i" f o
Unfortunately, the other two arrows of the polynomial span point in the wrong direction for naive composition. For the moment, we will focus on how to define 0g and leave Ps. for later.
We start with message data Y + R. It may be attractive invert the output arrow in order to define a composition with -1, as was done in the case of the pullback. However, unless is bijective, the preimage o 0-1 2 P(Y) takes values in the power set of Y. There is an additional technicality: if the composition m o takes values in P(R). it will fail to detect multiplicities, we unable to tell from subset of R whether multiple messages had the same value.
So instead, our pushforward takes values in bag(R), the set of finite multisets (or bags) R, which we describe in more detail in appendix [ For the moment, it is enough to know that bag is equivalent formal sum, and we define an intermediate message pushforward (D) Seel-"(u)m()) E [2, bag [.)..
*For technical reasons, we ask that for each E Y, the preimage p '(v) = fa E X p(x) = y} should have total order. This is to properly support functions with non- commmuimu arguments, though this detail is unnecessary for our key examples, where arguments commute.

---

Clbb
a
gews)
glow)
message
pushforward g((te)
Cas ((
pushforward ((te))
g( (ha)
angument pushforward
Obd
glese)
glen))
((ew))
ese gf Che)
Ae((ewa 1))
d
f(d)
pullback
f(d)
pullback
b((esa.0) O))
Figure 1: The illustration of how pullback and pushforward combine to form the integral transform, for two specific cases. Left: Polynomial span V < E + E - V with trivial argument pushforward (identity). Each edge eur is connected to its sender and receiver nodes (u, v) via the span (black arrows). The pullback then "pulls" the node features [()) along the span, which the argument pushforward folds into edge features == = their edges, the message pushforward then "collects" all of the edge features that send to a particular receiver, by pushing them along the span. Right: Polynomial span V <- E' + E' + E' + V, situation more commonly found in GNNs. In this case, the pullback pulls sender and receiver node features into the argument function, h. The argument pushforward then computes, from these arguments, the edge messages, J. which are sent to receivers via the message pushforward, as before. See Appendix|e A for a visualisation of how these arrows translate into GNN code.
Once all sender features are pulled back to
All that is missing to complete our definition of og is an aggregator + bag(R) - R. As we will see later, specifying well-behaved aggregator the same as imposing a commutative monoid structure on R. With such an aggregator on R, we can define (u) θ (u)
We return Pg. which is constructed very similarly. The only difference that, while we deliberately regard the collection of messages as unordered, the collection of arguments used to compute message has an ordering we wish respect. So instead of the type bag(It), we use the type list(RI) of finite lists of elements of R, and our aggregator X list(R) - R is now akin to a fold operator.
We illustrate the use of these two aggregators in decomposed diagram:
Y list (R)
P10
[X R
Y R
1*
og
Z2 bag (1)
+
[W, R]
[2, R]
Note that any semiring (R. 8, 3) comes equipped with binary operators @. C that allow aggregators 8, Q to be defined inductively. In fact, the converse that every set with two such aggregators is semiring- is also true, if we assume some reasonable conditions on the aggregators, which we can explain in terms of one of the most utilised concepts in category theory and functional programming-mone
[3 Due to space constraints, we refer the interested reader to Appendices Bandd for full exposition of how we can use monads over lists and bags to constrain the latent space R to respect semiring structure.
5

---

For now, it's enough to know that our key examples of the real numbers (with multiplication and addition, for GNNs) and the tropical natural numbers (with addition and minimum, for DP) both allow for natural interpretations of X and + in the integral transform.
We now ready to show how the integral transform of algorithms and GNNs. We start with the Bellman-Ford algorithm [2] (Equation 4 that was traditionally used to demonstrate the concept of algorithmic alignment.
be used to instantiate popular examples
Bellman-Ford
Let It = (N U (oc),+ mi) be the "min-plus" semiring of extended natural numbers, with @ = and @ min. This the coefficient semiring over which the Bellman-Foo Ford algorithm takes place.
Let EV) be a weighted graph with source and target maps s,t : + and edge weights w E + R. For purely technical reasons, we also need to explicitly materialise bias function b V - R. which is, in practice, a constant-zero function (b(v) 0 for all E V) but will prove necessary for defining the argument pushforward
We interpret Bellman- Ford as the following polynomial span:
(V + E) + (V +
V+E
(6)
V + (V E)
V
Here "+" is the disjoint union of sets, defined as A+ B = {(a, 1)
| a E () b B}. Note
that [S + T, R = [S, R x [T. R), i.e. specifying data on a disjoint union is equivalent to specifying data on each component separately.
Initially, we describe each of the four sets of the polynomial span, making their role clear:
Input:
W = V + (V + E). Our input to Bellman- Ford includes: the current estimate of node distances (du; function in [V, RJ). edge weights (w; function in [F. R). and the previously discussed bias b, a function in [V. R). Hence our overall inputs are members of
RV FE, R) x [V, R = (V (V + F), R). justifying
choice of input space.
Arguments: X (V + F) (V + F). Here we collect the ingredients necessary to compute Bellman Dellman--Fod'ss subproblem solutions coming from neighbouring nodes. To do this, need to combine data in the nodes with data living edges- those are the arguments to the function. And since they meet in the edges, we "lift" our node distances [V, Rt] to edges they are sending from, giving us additional function in [E, R). Hence our argument carrier space is now (V + E) + (V + E) (the remaining three inputs remain unchanged).
Message: Y V E. Once the arguments are combined to compute messages, we are left with signal in each edge (containing the sum of corresponding d, and wav). and each node (containing just dy, for the purposes access to the previous optimal solution). Hence our messages are members of [V, R] [E, Rt.. justifyigg our choice of message space.
Output: Z = V. Lastly, the output of one step of Bellman-Ford are updated values d' which
interpret just (output) data living on V.
We now describe how to propagate data along each arrow of the diagram in turn, beginning with inputs (f. b, w) of node features f V + R. bias V - R. and edge weights w E + R:
Pullback, i": First, can note the input function i (V + +V - E) decomposes as the sum of two arrows. in V + F - V is the identity function on V and the source function F, and ig V + F + V F is just the identity. So we calculate the pullback i* (f, b, w) = (f.f o 8, b, w). giving us the arguments to compute messages.
Argument pushforward, Ps: Next, the process function simply identifies the two copies of V E, and sums their values. So the argument pushforward is p (f. f s, 6, w)

---

w)) This also allows us to interpret the bias function, b, as a "self-edge" the graph with weight
(f.fos] =f
Message pushforward, og: The output function o: V - V is the identity function on V and the target function E. So the message pushforward gives us (o= + b, (f (fos) w))(u) = (f( + b(u)) @ ① (())) (f (f) = min(f(u) - b(u), min f(u) + We-n).
Letting b(0) 0, we can see that this is exactly Equation 4
So we have produced the formula for the Follman-Ford algorithm directly from the polynomial span in Diagram
Note that Pe is aligned with using max aggregation in neural networks-diredde explaining several previous proposals, such as [3I]. But additionally, as defined, is aligned with concatenating all message arguments together and passing them through linear function, which is how such step implemented in GNNs' message functions. We now direct our polynomial span analysis at GNNs.
6 GNNs
We study the popular message passing neural nefwork (MPNN) model [9]] which can be interpreted using the following polynomial span diagram:
E+ ( +
E
1+V+E
Here the set refers to a singleton set-sometnnes also called () or unit-which is used as carrier for graph-level features. This implies the graph features will be specified as [1. RR] > R, as expected.
Given all these features, how would we compute messages? The natural way is to combine the features the sender and receiver node of each edge, features of said edge, and graph-level features these will form our arguments, and they need to all "meet" in the edges. This motivates our argument space F. + (F + F.) + F:: all of the above four, accordingly broadcast into their respective edge(s).
The input map, i, is then the unique map to the singleton, the sender and receiver functions on the two middle copies of E, and the identity on the last copy of E, i.c. i(a, c, d) {O, s(b), ((), d). The process map, collapses the four copies of E' into just one, to hold the computed message. Lastly, the output map, o, the target function, identifying the node to which the message will be delivered.
The actual computation performed by the network (over real values in R, which can support various semirings of interest) is exactly an integral transform, with an extra MLP processing step messages:
E +
E +
Ps
E. RR
MLP
00
1+V
[V. R
is useful to take a moment discuss what was just achieved: with single abstract template (the polynomial span), we have successfully explained both dynamic programming algorithm, and a GNN update rule-merely by choosing the correct support and latent space.
7 Improving GNNs with edge updates, with experimental evaluation
From now on, we will set V2, as our baseline GNNs will use fully connected graphs, and will accentuate the polynomial nature of our construction.
7

---

We now show how our polynomial span view can be used to directly propose better- aaiter- GNN architectures for certain algorithmic tasks. Since the MPNN diagram above oulputs only node features, to improve predictive performance on edge-centric algorithms, it is natural augmentation also update edge features, by adding edges to the output carrier (as done by, e.g., ①I:
V2 +
P
V
1 V V2 V2
V + V2
But notice that there is problem with the output arrow. Since we are using each message twice, is no longer function -it'd have to send each edge message two different objects! To resolve this, we need to appropriately augment the messages and the arguments. This is equivalent to specifying new polynomial span with output V22 which we can then recombine with Diagram
(8)
1 1+V V + V2
V2
Most edge-centric algorithms of interest (such as the paths [I4]), compute edge-level outputs reducing over choice of "intermmdiate" node. Hence, would be beneficial to produce messages with shape V which would then reduce to features over V2, There three possible ways broadcast both node and edge features into V we propose the following polynomial span, which materialises each of those arguments:
algorithm for all-pairs shortest
/3 + (V3 + V3 + V3) + + +3+
V3
(9)
1 1 V //2
1/2
Finally, inserting this into Diagram 7 gives us a corrected polynomial span with output V + V2,
4V2 + 7V3
22 + V3
(10)
1 V + /2
V + 1/2
Here have collapsed the copies of 2 and /3 in the argument position for compactness.
While Diagram doesn't make sense polynomial diagram of sets, we can clearly still implement it as architecture M.. since nothing stops us from sending the same tensor to two places. We want to investigate whether our proposed modification of Diagram 10] which materialises order- 3 messages, leads to improved algorithmic alignment on edge-centric algorithms. To support this evaluation, we initially use set six tasks from the recently proposed CLRS Algorithmic Reasoning Benchmark (32), which evaluates how well various (G)NNs align to classical algorithms, both in- and -itto---iitributionn We reuse exactly the data generation and base model implementations the publicly available code for the CLRS benchmark.

---

We implemented each of these options by making our GNN's message and update functions be two-layer MLPs with embedding dimension 24, and hidden layers of size 8 and 16. Our test results (out-of-distributtion) are summarised in Table [ For convenience, we also illustrate the in-distribution performance of our models via plots given in Appendix
Lastly, we scale up our experiments to 27 different tasks in CLRS, 96-dimensional embeddings, and using the PGN processor [29], which is the current state-of-the-art model on CLRS in terms of task win count 32
We summarise the performance improvement obtained by our V variant of PGN in Table aggregated across edge-centric tasks as well as ones that do not require explicit edge-level reasoning. For convenience, we provide the per-task test performance in Appendix F (Table 3).
We found that the V architecture was equivalent to, outperformed, the non non-polynomial (V²) one in all edge-centric algorithms (up to standard error). Additionally, this architecture appears to also provide some gains tasks without explicit edge-level reasoning requirements, albeit smaller on average and less consistently. Our result directly validates our theory's 's predictions, in the context of presenting better-aligned GNN for edge-centric algorithmic targets.
Table 1: Test (out-of- -ifriisiribution)) results of our models on all models on the six algorithms studied. while V corresponds to our proposal
2 corresponds to the baseline model offered by Diagram 7 in Diagram 10, which respects the polynomial span.
Algorithm
Varge
V llage
V2-small
V
Dijkstra
59.58% = 2.82
68.53% +: 2.40
56.10% + 3.25
60.32% +: 2.70
Find Maximum Subarray
8.33% + 0.50
9.06%
0.65
8.46% + 0.55
7.89% + 0.64
Floyd- Warshall
7.46% = 0.63
9.00% 0.81
6.66% = 0.62
8.23% = 0.62
Insertion Sort
13.39% = 1.27
24.67% - 2.44
14.69% = 1.32
20.23% = 2.21
Matrix Chain Order
67.64% 1.23
70.79% = 1.54
68.85% = 2.26
68.76% = 1,21
Optimal BST
53.03% = 2.80
54.56% + 4.34
46.65% + 3.82
51.94% + 4.60
Overall average
35.24%
39.43%
33.57%
36.23%
Table 2: Test (out-of-distribution) results across 27 tasks in CLRS, for the PGN processor network, averaged across edge-centric and other tasks. See Appendix F for the per-task test performances.
Algorithms
PGN
73-PGN
Average Improvement
Edge-centric algorithms
35.03%
39.03 39.08%
4.44% = 1.06
Other algorithms
35.37%
36.33%
1.01% = 0.11
Average of the two groups
35.20%
37.70%
2.73%
8 Conclusions
In this paper, we describe the use of category theory and abstract algebra to explicitly expand on the GNN-DP connection, which was previously largely handwaved on specific examples. We derived a generic diagram of an integral transform (based on standard categorical concepts like pullback, pushforward and commutative monoids), and argued why it is general enough to support both GNN and DP computations. With this diagram materialised, we were able to immediately unify large quantities of prior work as simply manipulating one arrow or element in the integral transform. We also provided empirical evidence of the utility of polynomial spans for analysing GNN architectures, especially in terms of algorithmic alignment. It is our hope that our findings inspire future research into aetter-aligned neural algorithmic reasoners, especially focusing on generalising or diving into several aspects of this diagram.
Lastly, it not at all unlikely that analyses similar to ours have already been used to describe other algorithmic reasoners. The principal ideas of span and integral transform
fields of science--eeeeee are central to defining Fourier series (35), and appear in the analysis of Yang-Mills equations in particle physics 13]. Properly understanding the common ground behind all of these definitions may, in the very least, lead to interesting connections, and shared understanding between the various fields they span.

---

Acknowledgments and Disclosure of Funding
We would like to thank Charles Blundell, Tai- Tai-Danae Bradley, Taco Cohen, Bruno Gavranović, Bogdan Georgiev, Razvan Pascanu, Karolis Spukas, Grzegorz Swirszez, and Vincent Wang- Maścianica for the very useful discussions and feedback on prior versions of this work.
Special thanks Tamara von Glehn for key comments helping formally connect integral transforms to polynomial functors.
This research was funded by DeepMind.
References
[1] Peter W Battaglia, Jessica 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 arXivv1806.01261, 2018.
[2] Richard Bellman. On routing problem. Quarterly of applied mathematics, 16(1):87-90, 1958.
[3] Richard Bellman. Dynamic programming. Science, 153(3731):34 37, 1966.
[4] Beatrice Bevilacqua, Yangze Zhou, and Bruno Ribeiro. Size-invariant graph representations for graph classification extrapolations. In International Conference Machine Learning, pages 837-851, PMLR, 2021.
[5] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104. 13478, 2021.
[6] Francesca Cagliari and Sandra Mantovani. Cartesianness: topological spaces, uniform spaces, and affine schemes. Topology and its Applications, 41:263-272 1991.
[7] Quentin Cappart, Didier Chételat, Elias Khalil, Andrea Lodi, Christopher Morris, and Petar Veličković. Combinatorial optimization and reasoning with graph neural networks. arXiv preprint arXiv:2102.09544, 2021.
[8] Eugenia Cheng. Iterated distributive laws. In Mathematical Proceedings of the Cambridge Philosophical Society, volume 150, pages 459-487. Cambridge University Press, 2011.
[9] Pim de Haan, Taco S Cohen, and Max Welling. Natural graph networks. Advances in Neural Information Processing Systems, 33:3636-3646, 2020.
[10] Oege De Moor. Categories, relations and dynamic programming. Mathematical Structures in Computer Science, 4(1):33 69, 1994.
[11] Andreea-loana Deac, Petar Veličković, Ognjen Milinkovic, Pierre-Lue Bacon, Jian Tang, and Mladen Nikolic. Neural algorithmic reasoners are implicit planners. Advances in Neural Information Processing Systems, 34, 2021.
[12] Andrew Dudzik. Quantales and byperstructures: Monads, mo' moopreolems. arXiv preprint arXiv: 1707.09277 2017.
[13] Michael Eastwood, Roger Penrose, and RO Wells. Cohomology and massless fields. Commun. Math. Phys, 7803)330--35, 1981.
[14] Robert Floyd. Algorithm shortest path. Communications of the ACM, 5(6):345, 1962.
[15] Brendan Fong and David I Spivak. An invitation to applied category theory; seven sketches in compositionnalty. Cambridge University Press, 2019.
[16] Karlis Freivalds, Emils Ozolinš, and Agris Šostaks. Neural shuffle-exchange networks-sequence processing in o (n log time. Advances in Neural Information Processing Systems, 32, 2019.
10

---

[17] Nicola Gambino and Joachim Kock. Polynomial functors and polynomial monads. In Math- ematical proceedings of the cambridge philosophical society. volume 154, pages 153-192. Cambridge University Press, 2013.
[18] Jeffrey Giansiracusa and Noah Giansiracusa, Equations of tropical varieties. Duke Mathematical Journal, 165(18):3379 3433,,33333330 2016.
[19] Justin Gilmer, Samuel S Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning. pages 2263-1272 PMLR, 2017.
[20] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[21] Vinod Nair, Sergey Bartunov, Felix Gimeno, Ingrid von Glehn, Pawel Lichocki, Ivan Lobov, Brendan O' D'Donoghue, Nicolas Sonnerat, Christian Tjandraatmadja, Pengming Wang, et al. Solving mixed integer programs using neural networks. arXiv preprint arXiv:2012. 13349, 2020.
[22] Susan Niefield. Cartesianness: topological spaces, uniform spaces, and affine schemes. Pure Appl. Alg., 23:147-163, 1982.
[23] Alvaro Sanchez-Gonzale Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning simulate complex physics with graph networks. In International Conference Machine Learning, pages 8459-8468. PMLR, 2020.
[24] Heiko Strathmann, Mohammadamin Barekatain, Charles Blundell, and Petar Veličkovíć. Persis- tent message passing. arXiv preprint arXiv:2103.01043, 2021.
[25] Daisuke Tambara. On multiplicative transfer. Communications Algebra, 21((): 1993.
[26] Hao Tang, Zhiao Huang, Jiayuan Gu, Bao-Liang Lu, and Hao Su. Towards graph-related problem solving by iterative homogeneous gnns. Advances in Neural Information Processing Systems, $3:15811-15822 2020.
[27] Petar Veličković and Charles Blundell. Neural algorithmic reasoning. Patterns, 2(7):100273, 2021.
[28] Petar Veličković, Matko Bošnjak, Thomas Kipf, Alexander Lerchner, Raia Hadsell, Raz- van Pascanu, and Charles Blundell. Reaslning--oduulaeed representations. arXiv preprint arXiv:2107.08881, 2021.
[29] Petar Veličković, Lars Buesing, Matthew Overlan, Razvan Pascanu, Oriol Vinyals, and Charles Blundell. Pointer graph networks. Advances in Neural Information Processing Systems, 33:2232-2244, 2020.
[30] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 1rXiv:1710.10903. 2017.
[31] Petar Veličković, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural exccution of graph algorithms. arXiv preprint arXiv: 10593, 2019.
[32] Petar Veličković, Adrià Puigdomênech Badia, David Budden, Razvan Pascanu, Andrea Ban- ino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The clrs algorithmic reasoning benchmark. arXiv preprint vrXiv:2205.15633 2022.
[33] Philip Wadler. Comprehending monads. Proceedings of the 1990 ACM Conference on LISP and Functional Programming, pages 61-78, 1990.
[34] Mark Weber. Polynomials in categories with pullbacks. Theory and Applications of Categories, 30(16):533-598 2015.
[35] Simon Willerton. Integral transforms and the pull-push perspective, i. The n- Category Café, 2020.
11

---

[36] Louis-Pascal Xhonneux, Andreea-loand Deac, Petar Veličković, and Jian Tang, How to transfer algorithmic reasoning knowledge learn new algorithms? Advances in Neural Information Processing Systems, 34, 2021.
[37] Keyulu Xu, Jingling Li, Mozhi Zhang, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? arXiv preprint wrXXvv/905..10 2019.
[38] Keyulu Xu, Mozhi Zhang, Jingling Li, Simon $ Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint arXiv:2009, 11848, 2020.
[39] Zhaocheng Zhu, Zuobai Zhang, Louis- Pascal Xhonneux, and Jian Tang. Neural bellman- ford networks: A general graph neural network framework for link prediction. Advances in Neural Information Processing Systems, 34, 2021.
Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] to [Yes], [No], or [N/A] . You are strongly encouraged to include justification to your answer, either by referencing the appropriate section of your paper or providing brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section ??.
Did you include the license the code and datasets? [No] The code and the data are proprietary.
Did you include the license to the code and datasets? [N/A]
Please do not modify the questions and only the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's s contributions and scope? [Yes] We propose novel approach to reinterpret both graph neural networks and dynamic programming, and show empirical gains from an architecture motivated by our blueprint.
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our work is of theoretical nature.
(d) Have you read the ethics review guidelines and ensured that your paper conforms
them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] Appropriate references to proofs are provided in all areas where proofs are missing.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or URL)? [No] We aim to release the code at a future point. The data is publicly available.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they chosen)? Yes]] The data generation and base model implementation is publicly available within the CLRS benchmark. We detail the model extensions we made.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes]
12

---

4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cite the CLRS benchmark.
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A] Our data abstract and algorithmically
generated.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "pp9_HrSgpVJ", "reviewer": "Reviewer_TCta", "summary": "This paper proposes an approach to interpret the operations of both DP and GNNs from the perspective of category theory and abstract algebra. They derive the GNN-DP connection by a diagrammatic abstraction in the computations and recast the algorithmic alignment to align the diagrams of GNNs to some specific algorithm classes. Some experiments on six edge-centric tasks from CLRS Algorithmic Reasoning Benchmark demonstrate the effectiveness of the proposed derive.", "questions": "- One major question is that in the evaluation part why use a two-layer MLP with different dimensional sizes, It seems GNNs are discussed in this paper, however, using MLP for evaluation is unclear to me. I am not very familiar with this topic, so please clarify it.\n- Furthermore, why choose in-domain and out-of-domain for evaluation, and what is the rationality behind it.", "limitations": "The evaluation part is not sufficient. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- This paper theoretically proposes the integral transform to connect GNNs and DP. Furthermore, they provide sufficient theoretical proof to prove the correctness. The theoretical part seems solid to me. \n- This paper is good at writing, the target problem is demonstrated clear and the proposed solution is well-discussed. \n\nWeakness:\n- The evaluation part is not sufficient and in this paper, only six algorithms are evaluated with different dimension lengths from in-domain and out-of-domain datasets.\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "9HDaZZkiBDT", "reviewer": "Reviewer_wgV4", "summary": "The authors connect GNN and DP by using integral transforms as a common form to represent the operations used in GNN and DP. The common form allows studying GNN and GP in unified framework and authors use the connection to design GNN architecture that better align with the underlying DP algorithm to achieve better empirical performance on benchmark.", "questions": "1. I find it difficult to fully grasp details for polynomial span diagram of GNNs. One way this paper can be strengthened is to provide more details on how Diagram (7) and (10) are implemented in practices and what does these abstract operations correspond to in GNN operations (perhaps a side-by-side figure with left being the integral transform diagram, and right being the things like pseudocode)\n2. The experiments part can be improved by adding negative examples where V^3 network does not improve over V^2 network for those non-edge-centric CLRS benchmarks, or as authors stated in Line 46-47, by showcasing examples that certain classes of GNN fit certain classes of DP, rather than showing V^3 outperforms V^2 throughout.\n3. Can authors explain more on the consistently low performance even for Dijkstra problem (~70%) even though it seems the GNN architecture well matches the DP problem structure here (also for results of other benchmarks with performance as low as ~10%). Can authors provide insights on in what directions performance can be boosted, e.g., for Find Maximum Subarray (currently with only ~10% performance)? The baseline I am mainly referring to are the empirical results in Xu et al. 2019 ICLR. It is hard to persuaded by the title “Graph Neural Networks are Dynamic Programmers” with the these underwhelming experiment performance. Is it because of the limited training data? Perhaps authors can consider downplaying their paper title.\n4. Information about experiments setup (including data sample size, hyperparameter tuning steps, etc. that can be replicate the experiments) should be added in line (rather than being put into checklist)", "limitations": "See above in \"Questions\" section", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**originality:** On the question being asked, this work is similar to Xu et al. 2019 ICLR but this work is original in that this study connect GNN and DP in a more formal way.\n\n**quality & clarity:** Introduction and background sections are clearly written and motivate this study well. And that dynamic programming (e.g., Bellman-Ford) and GNN be represented in the same framework is well demonstrated. I do not have the expertise on examining all the methods details in section 4 so i defer this to other reviewers. I find the quality of empirical results (section 7) can be improved (see below). \n\n**significance:** this study is important on progress to connect classical DP and GNN and may be used to better understand and propose new type of GNN. However, presentation on methods and empirical results can be improved to strengthen the impact of this paper to broader audience.", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "58DR_f043v1", "reviewer": "Reviewer_QbK2", "summary": "In this paper, the authors describe the connection between GNNs and dynamic programming via the category theory, specifically, the integral transform. It shows that both GNNs and DPs can be described as a polynomial span diagram and people can improve GNNs' edge updates from the perspective of the polynomial span diagram. \n\nWith the analysis from the polynomial span, the authors show that we can use order-3 edge updates to improve the performance of GNNs for edge-centric problems in the CLRS Algorithmic Reasoning Benchmark. \n", "questions": "Can the authors demonstrate the improvement of the order-3 updates with larger GNNs? For example, a 128-dimension embedding and 256-dimension hidden layer?  \n\nHow's the improved the GNNs compare to the state-of-the-art on the CLRS benchmark? Or, can you improve the state-of-the-art models via the order-3 updates?", "limitations": "Though it is an interesting theoretical analysis work for GNNs and DP, the experimental results/settings are not satisfactory enough. (As stated in the weakness.)", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThis is an interesting and intriguing paper for connecting GNNs with DP as well as elevating the understanding of GNNs. The paper builds the connection between GNNs and DP via the polynomial span diagram from category theory. It also demonstrates a bit the empirical usage of this analysis via the improved edge updates, tested on the CLRS Algorithmic Reasoning Benchmark.\n\nWeakness:\n\nI admit that I'm not an expert in category theory, so I expect a more clear/intuitive explanation of the \"integral transform\" and the \"polynomial span diagram\". \nThe experimental results on the CLRS algorithm reasoning benchmark are done via a very small-scale GNN (hidden size of 8/16), which could not fully demonstrate the practical impact of the order-3 updates. Hence, it doesn't surprise me that higher-order updates improve the small GNN models. ", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept