G RAPH N EURAL N ETWORKS ARE D YNAMIC P ROGRAMMERS
A BSTRACT
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.
Speciﬁcally, 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 quantiﬁed?
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 ﬁndings in the literature, and hope it will serve as a foundation for building stronger algorithmically aligned GNNs.
I NTRODUCTION
One of the principal pillars of neural algorithmic reasoning (Veli ˇ ckovi ´ c & Blundell, 2021) is training neural networks that execute algorithmic computation in a high-dimensional latent space.
A fundamental question in this space is: which architecture should be used to learn a particular algorithm?
We seek architectures with low sample complexity : generalising better with fewer training examples.
The key theoretical advance towards this aim was made by Xu et al.
Therein, the notion of algorithmic alignment is formalised, favouring architectures that align better to the algorithm, in the sense that we can separate them into modules, which correspond to the computations of the target algorithm’s subroutines.
The theory predicts that graph neural networks (GNNs) algorithmically align with dynamic programming (DP).
However, it quickly became apparent that it is not enough to just train any GNN—many recent works propose specialised GNNs that align with linearithmic (Freivalds et al., 2019), iterative (Tang et al., 2020) or structured algorithms (Veli ˇ ckovi
´
c et al., 2020).
We believe that the fundamental reason why so many isolated efforts needed to look into learning speciﬁc classes of algorithms is the fact the GNN-DP connection has not been sufﬁciently explored .
Indeed, the original work of Xu et al.
(2019) merely mentions in passing that the formulation of
DP algorithms seems to align with GNNs, and demonstrates one example (Bellman-Ford).
Our thorough investigation of the literature yielded no concrete follow-up to this initial claim.
Here, we make a ﬁrst step towards a framework that could allow us to identify GNNs that could align particularly well with certain classes of DP, rather than assuming a “one-size-ﬁts-all” GNN.
We interpret the operations of both DP and GNNs from the lens of category theory and abstract algebra.
In doing so, several previously shown results will naturally arise as corollaries, and we hope it opens up the door to a broader uniﬁcation between algorithmic reasoning and the geometric deep learning blueprint (Bronstein et al., 2021).
GNN S , DYNAMIC PROGRAMMING , AND THE CATEGORICAL CONNECTION
We will use the deﬁnition of GNNs based on Bronstein et al.
Let a graph be a tuple of nodes and edges , G =
( V, E ) , with one-hop neighbourhoods deﬁned as N u
= { v ∈ V | ( v, u ) ∈ E } .
Further, a node feature matrix X ∈ R | V |× k gives the features of node u as x u ; we omit edge- and graph-level features for clarity.
A (message passing)
GNN over this graph is then executed as:
[Equation]
where ψ : R k × R k
→ R k is a message function , φ : R k
→ R k is a readout function , and (cid:76) is a permutation-invariant aggregation function (such as (cid:80) or max ).
Both ψ and φ can be realised as
MLPs, but many special cases exist, giving rise to, e.g., attentional GNNs (Veliˇckovi´c et al., 2017).
Dynamic programming is deﬁned as a process that solves problems in a divide et impera fashion: imagine that we want to solve a problem instance x .
DP proceeds to identify a set of subproblems , η ( x ) , such that solving them ﬁrst, and recombining the answers, can directly lead to the solution for x : f ( x )
= ρ ( { f ( y ) |
y ∈ η ( x ) } ) .
Eventually, we decompose the problem enough until we arrive at a 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 a bottom-up fashion.
This rule is often expressed programmatically:
[Equation]
To initiate our discussion on why DP can be connected with GNNs, it is a worthwhile exercise to show how Equation 2 induces a graph structure.
To see this, we leverage a categorical analysis of dynamic programming ﬁrst proposed by De Moor (1994).
Therein, dynamic programming algorithms are reasoned about as a composition of three components (presented here on a high level): dp =
[Equation]
(cid:124)(cid:123)(cid:122)(cid:125) recombine ◦ σ (cid:124)(cid:123)(cid:122)(cid:125) score ◦
[Equation]
(cid:124)(cid:123)(cid:122)(cid:125) expand (3)
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).
Let V be the space of all subproblems, and R an appropriate value space (e.g. the real numbers).
Then, expansion is deﬁned as η : V → P ( V ) , giving the set of all subproblems relevant for a given problem.
Note that this also induces a set of edges between subproblems, E ;
namely, ( x, y ) ∈
E if x ∈ η ( y ) .
Each subproblem is scored by using a function σ : P ( V ) → P ( R ) .
Finally, the individual scores are recombined using the recombination function, ρ : P ( R ) → R .
The ﬁnal dynamic programming primitive therefore computes a function dp : V → R in each of the subproblems of interest.
T HE DIFFICULTY OF CONNECTING GNN S 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-ﬁnding problems, typically takes place over “tropical” objects like ( N ∪ {∞} , min , +) 1 , 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 deﬁne 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 S → R , where S is a ﬁnite 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 ﬁnite sets, and “ R -valued quantities” on it.
By “category” here we mean a collection of objects (all ﬁnite sets) together with a notion of composable arrows (functions between ﬁnite sets).
To draw our GNN-DP connection, we need to devise an abstract object which can capture both the GNN’s message passing/aggregation stages (Equation 1) and the DP’s scoring/recombination stages (Equation 2).
It may seem quite intuitive that these two concepts can and should be relatable, and category theory is a very attractive tool for “making the obvious even more obvious” (Fong & Spivak, 2019).
Indeed, recently concepts from category theory have enabled the construction of powerful GNN architectures beyond permutation equivariance (de Haan et al., 2020).
We will construct the integral transform by composing transformations over our input features in a way that will depend minimally on the speciﬁc choice of R .
In doing so, we will build a computational diagram that will be applicable for both GNNs and DP (and their own choices of R ), and hence allowing for focusing on making components of those diagrams as aligned as possible.
T HE F OUR A RROWS OF THE I NTEGRAL T RANSFORM
[Figure]
The pullback then “pulls” the node features f ( u ) along the span, into edge features g ( e vu ) =
f ( u ) .
Once all sender features are pulled back to their edges, the pushforward then “collects” all of the edge features that send to a particular receiver, by pushing them along the span.
A directed graph ( V, E ) is deﬁned by two functions s, t : E → V , deﬁning the sender and receiver nodes for each edge.
We express this in the following diagram:
[Equation]
The general question is: given some data f on
V assigning features f ( v ) to each v ∈ V , how to transform it via this diagram into some other data on V ?
If we are able to do this, we will be able to characterise both the processes of sending messages between nodes in GNNs and scoring subproblems in DP.
Initially, we will focus on the cases where the “message” sent along edge e ∈ E depends only on the sender and not the receiver.
This is the case in many path-ﬁnding algorithms, such as Bellman-Ford.
We provide details on how to generalise the messages to be receiverconditioned in Appendix B.
We start by deﬁning a kernel transformation k :
[ E, R ] → [ E, R ] that performs some computation over edge features, e.g. in the case of Bellman-Ford, adding the sender distance d s ( e ) (previously pulled back to the edge via s ∗ ) to the edge weight w s ( e ) → t ( e ) .
Here we use [ E, R ] as a shorthand notation for the set of functions E
→ R .
Using the kernel, we can complete the following diagram:
[Equation]
Here N [ R ] is the collection of of multisets in R , or “bags of R -values”.
This object is explained more rigorously in Appendix A.
These four arrows, taken together, form our integral transform: starting with node features, performing computations over edges, ﬁnally aggregating the edge messages in the receivers.
In doing so, we have updated our initial node features (which are elements of [ V, R ] ).
We now explain each of the four arrows in turn.
s ∗ is the pullback arrow, which takes data deﬁned on V and returns data deﬁned on E .
To do this, we precompose the node features with the source function.
That is, ( s ∗ f )( v ) :
= f ( s ( v )) .
Then, the kernel is applied to the resulting edge features, integrating the sender’s features with any provided edge features (e.g. edge weights).
After applying our kernel, we have edge messages m : E → R as a result.
We now need to send these messages to the receiver nodes, for which we employ the pushforward .
We deﬁne
( t ∗ m )( v ) =
(cid:80) e ∈ t − 1 ( v ) m ( e ) , which we interpret as a formal sum in N [ R ] .
Intuitively, ( t ∗ m )( v ) is the “bag of incoming values” at v .
The reason why we require multisets is because there may be several edges in t ’s preimage that store the same value; this also aligns nicely with related work that studies GNN expressive power using multisets (Xu et al., 2018).
Finally, we apply an aggregator (cid:76) , pointwise for every receiver, to compute updated features for every node.
To have such an aggregator, we need one property from R : that it have the structure of a commutative monoid.
(see A for details)
Note the generality of our construction: nearly all of the operations at play here rely on the properties of the sets E and V .
The only arrow which is constraining the choice of R is the aggregator (cid:76) ; it necessitates that R be a commutative monoid, which is a sensible choice for both pathﬁnding algorithms and GNNs.
Hence, we have deﬁned a generic blueprint where arbitrary R s can be inserted, hence amplifying the connection between GNNs (where R often includes real-valued vectors) and
DP (where R may be, e.g., the “tropical” numbers).
I MMEDIATE COROLLARIES OF THE INTEGRAL TRANSFORM VIEW
Since we designed the aforementioned integral transform in a way that it can describe both GNNs (Equation 1) and DP (Equation 2), there exist some “embarrassingly obvious” ways of making the architecture of the GNN more algorithmically aligned to the target DP algorithm.
Perhaps the clearest one is the ﬁnal arrow, the aggregator ( (cid:76) ).
If we make the GNN’s choice of aggregator function match the one used by the target algorithm, this should lead to immediate gains to sample complexity and generalisation.
Indeed, this aligns well with one of the earliest lines of research in algorithmic reasoning: deploying GNNs with aggregators that align to the problem.
Already in Veliˇckovi´c et al.
(2019), it was demonstrated that using the max aggregation in GNNs yields superior performance on executing pathﬁnding algorithms, especially when compared to the more popular choice of (cid:80) .
Similar ﬁndings were observed by Tang et al.
Richter & Wattenhofer (2020), who also clearly indicated environments where max is the superior aggregator.
Lastly, theoretical works such as Xu et al.
Li et al.
(2021) have illustrated the additional useful corollaries to extrapolation and noisylabel training when using aligned aggregators.
The fact that already such a vast collection of research works can be uniﬁed as analysing one arrow in a categorical diagram is a statement to the potential of our proposed analysis.
Any invested effort in analysing GNNs and DP in this manner could provide strong returns when organising existing research and proposing future work.
I DEAS FOR FUTURE WORK
What other arrows can be “inspected” in this way?
The most natural candidate is the kernel arrow, which corresponds to both the GNN’s message function ψ (Equation 1) to the DP scoring function (Equation 2).
Aligning the two more closely may require looking deeper into the kernel arrow to formalise entirely, and we leave such analyses for future work to maintain a focused contribution.
While it might seem that the pullback and pushforward arrows are fairly static, this is primarily because both the graph structure (in the case of GNNs) and the expansions (in the case of DP) are assumed ﬁxed known upfront.
This is a simplifying assumption that need not always hold (for example, if our algorithm relies on any kind of data structure ).
Another promising avenue for future work could dive deeper into these two arrows as well, and speciﬁcally to develop GNNs that automatically propose subproblems within their computations.
Lastly, it is not at all unlikely that analyses similar to ours have already been used to describe other ﬁelds of science—beyond algorithmic reasoners.
The principal ideas of span and integral transform are central to deﬁning Fourier series (Willerton, 2020), and appear in the analysis of Yang-Mills equations in particle physics (Eastwood et al., 1981).
Properly understanding the common ground behind all of these deﬁnitions may, in the very least, lead to interesting connections, and a shared understanding between the various ﬁelds they span.
A T HE MULTISET MONAD , COMMUTATIVE MONOIDS AND AGGREGATORS
Given a set S , we deﬁne N [ S ] := { p : S → N | # { p ( r ) (cid:54) = 0 } < ∞} , the natural-valued functions of ﬁnite support on S .
This has a clear correspondence to multisets over S : p sends each element of
S to the amount of times it appears in the multiset.
We can write its elements formally as (cid:80) s ∈ S n s s , where all but ﬁnitely many of the n s are nonzero.
Given a function f : S → T between sets, we can deﬁne a function N [ f ] : N [ S ] → N [ T ] , as follows: N [ f ]( (cid:80) s ∈ S n s s ) :
= (cid:80) s ∈ S n s f ( s ) , which we can write as (cid:80) t ∈ T m t t , where m t = (cid:80)
f
( s )=
t n s .
For each S , we can also deﬁne two special functions.
The ﬁrst is unit : S → N [ S ] , sending each element to its indicator function (i.e. an element x ∈ S to the multiset {{ x }} ).
The second is join : N [ N [ S ]]
→ N [ S ] , which interprets a nested sum as a single sum.
These facts tell us that N [ − ] is a monad , a special kind of self-transformation of the category of sets.
Monads are very general tools for computation, used heavily in functional programming languages (e.g. Haskell) to model the semantics of wrapped or enriched types.
Monads provide a clean way for abstracting control ﬂow, as well as gracefully handling functions with side effects (Wadler, 1990).
It is well-known that the algebras for the monad N [ − ] are the commutative monoids , sets equipped with a commutative and associative binary operation and a unit element.
Concretely, a commutative monoid structure on a set R is equivalent to deﬁning an aggregator function (cid:76) : N [ R ]
→ R compatible with the unit and monad composition.
Here, compatibility implies it should correctly handle sums of singletons and sums of sums, in the sense that the following two diagrams commute ; that is, they yield the same result regardless of which path is taken:
[Equation]
The ﬁrst diagram explains that the outcome of aggregating a singleton multiset (i.e. the one produced by applying unit ) with (cid:76) is equivalent to the original value placed in the singleton.
The second diagram indicates that the (cid:76) operator yields the same results over a nested multiset, regardless of whether we choose to directly apply it twice (once on each level of nesting), or ﬁrst perform the join function to collapse the nested multiset, then aggregate the collapsed multiset with (cid:76) .
Conveniently, such a (cid:76) is exactly what is required to appropriately connect the outcomes of our previously derived pullback ( s ∗ ) and pushforward ( t ∗ ) operations into a commutative diagram:
[Equation]
As such, so long as R is endowed with a commutative monoid structure, we can use this construction to deﬁne a suitable (cid:76) to serve as our aggregation function.
B R ECEIVER - DEPENDENT KERNELS
The construction in Section 4 does not immediately support receiver-dependent features, as any feature attached to a non-sender node is lost in the initial pullback step.
To get around this issue without losing the spirit of the integral transform, we assume that each e ∈ E has an opposite e ∗ ∈ E (s.t. s ( e )
=
t ( e ∗ ) , s ( e ∗ ) =
t ( e ) ).
While such a construction is natural for undirected graphs, it is possible to make it work while having directionality constraints.
Speciﬁcally, the kernel can mask these extra edges as an implementation detail (e.g. attaching an edge weight of “ + ∞ ” in the context of pathﬁnding), so this construction does not lose generality.
Now, we can express receiver dependence by factorising the kernel arrow as follows:
[Equation]
The ﬁrst arrow sends g ( e ) to g ( e, e ∗ ) .
This gives our k (cid:48) access to both sender and receiver features.
Again, we can perform any necessary receiver masking within k (cid:48) to avoid contributions from “virtual” edges.
In general, more complex kernels may require additionally enlarging E by a constant factor.
For example, if E is given an additional symmetry E → E we could use this to express ternary kernels, and so on.