Graph Neural Networks are Dynamic Programmers
Anonymous Author(s)
AffiliationAddress email
Abstract
Recent advances in neural algorithmic reasoning with graph neural networks(GNNs) are propped up by the notion of algorithmic alignment.
Broadly, a neuralnetwork 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, thatthere exists an intricate connection between GNNs and DP, going well beyond theinitial observations over individual algorithms such as Bellman-Ford.
Exposingthis connection, we easily verify several prior findings in the literature, producebetter-grounded GNN architectures for edge-centric tasks, and demonstrate empirical results on the CLRS algorithmic reasoning benchmark.
We hope our expositionwill serve as a foundation for building stronger algorithmically aligned GNNs.
Introduction
One of the principal pillars of neural algorithmic reasoning [24] is training neural networks thatexecute algorithmic computation in a high-dimensional latent space.
While this process is in itselfinsightful, and can lead to stronger combinatorial optimisation systems [19], it is valuable in termsof expanding the applicability of classical algorithms.
Evidence of this value are emerging, withpre-trained algorithmic reasoners utilised in implicit planning [11] and self-supervised learning [25].
A fundamental question in this space is: which architecture should be used to learn a particularalgorithm (or collection of algorithms [33])?
Naturally, we seek architectures that have low samplecomplexity , as they will allow us to create models that generalise better with fewer training examples.
The key theoretical advance towards achieving this aim has been made by [34].
Therein, the authorsformalise the notion of algorithmic alignment , which states that we should favour architecturesthat align better to the algorithm, in the sense that we can separate them into modules, whichindividually correspond to the computations of the target algorithm’s subroutines.
It can be provedthat architectures with higher algorithmic alignment will have lower sample complexity in the NTKregime [18].
Further, the theory of [34] predicts that graph neural networks (GNNs) algorithmicallyalign with dynamic programming [3, DP].
The authors demonstrate this by forming an analogy to the
Bellman-Ford algorithm
Since DP is a very general class of problem-solving techniques that can be used to express manyclassical algorithms, this finding has placed GNNs as the central methodology for neural algorithmicexecution [7].
However, it quickly became apparent that it is not enough to just train any GNN—formany algorithmic tasks, careful attention is required.
Several papers illustrated special cases of GNNsthat align with sequential algorithms [28], linearithmic sequence processing [15], physics simulations[21], iterative algorihtms [23], data structures [26] or auxiliary memory [22].
Some explanations for this lack of easy generalisation have arisen—we now have both geometric [35] and causal [4] viewsinto how better generalisation can be achieved.
We believe that the fundamental reason why so many isolated efforts needed to look into learningspecific classes of algorithms is the fact the GNN-DP connection has not been sufficiently explored .
Indeed, the original work of [34] merely mentions in passing that the formulation of DP algorithmsseems 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.
But DP algorithms are veryrich and diverse, often requiring a broad spectrum of computations.
Hence what we really need is aframework that could allow us to identify GNNs that could align particularly well with certain classesof DP, rather than assuming a “one-size-fits-all” GNN architecture will exist.
As a first step towards this, in this paper we interpret the operations of both DP and GNNs from thelens of category theory and abstract algebra.
We elucidate the GNN-DP connection by observinga diagrammatic abstraction of their computations, recasting algorithmic alignment to aligning thediagrams of (G)NNs to ones of the target algorithm class.
In doing so, several previously shownresults will naturally arise as corollaries, and we propose novel GNN variants that empirically alignbetter to edge-centric algorithms.
We hope our work opens up the door to a broader unificationbetween algorithmic reasoning and the geometric deep learning blueprint [5].
GNNs, dynamic programming, and the categorical connection
Before diving into the theory behind our connection, we provide a quick recap on the methods beingconnected: graph neural networks and dynamic programming.
Further, we cite related work to outlinewhy it is sufficient to interpret DP from the lens of graph algorithms.
We will use the definition of GNNs based on [5].
Let a graph be a tuple of nodes and edges ,
G =
( V, E ) , with one-hop neighbourhoods defined as N u = { v ∈ V
| ( v, u ) ∈ E } .
Further, anode feature matrix X ∈ R | V |× k gives the features of node u as x u ; we omit edge- and graph-levelfeatures 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 apermutation-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 [27].
Dynamic programming is defined 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 first, and recombining the answers, can directly lead to the solution forx : f ( x ) =
ρ ( { f ( y ) |
y ∈ η ( x ) } ) .
Eventually, we decompose the problem enough until we arrive atan 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 abottom-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 toshow how Equation 2 induces a graph structure.
To see this, we leverage a categorical analysisof dynamic programming first proposed by [10].
Therein, dynamic programming algorithms arereasoned about as a composition of three components (presented here on a high level):
[Equation]
[Table] [Table]
Expansion selects the relevant subproblems; scoring computes the quality of each individual subproblem’s solution w.r.t.
the current problem, and recombining combines these solutions into a solutionfor the original problem (e.g. by taking the max, or average).
Therefore, we can actually identify every subproblem as a node in a graph .
Let V be the space of allsubproblems, and R an appropriate value space (e.g. the real numbers).
Then, expansion is defined as η : V → P ( V ) , giving the set of all subproblems relevant for a given problem.
Note that this alsoinduces a set of edges between subproblems, E ;
namely, ( x, y ) ∈
E if x ∈ η ( y ) .
Each subproblem isscored by using a function σ : P ( V ) → P ( R ) .
Finally, the individual scores are recombined usingthe recombination function, ρ : P ( R ) → R .
The final dynamic programming primitive thereforecomputes a function dp : V → R in each of the subproblems of interest.
Therefore, dynamic programming algorithms can be seen as performing computations over a graphof subproblems , which can usually be precomputed for the task at hand (since the outputs of η areassumed known upfront for every subproblem).
One specific popular example is the Bellman-Fordalgorithm [2], which computes single-source shortest paths from a given source node, s , in a graph
G =
( V, E ) .
In this case, the set of subproblems is exactly the set of nodes, V , and the expansionη ( u ) is exactly the set of one-hop neighbours of u in the graph.
The algorithm maintains distances ofevery node to the source, d u .
The rule for iteratively recombining these distances is as follows:
[Equation]
where w v → u is the distance between nodes v and u .
The algorithm’s base cases are d s = 0 for thesource node, d u = + ∞ otherwise.
Note that more general forms of Bellman-Ford pathfinding exist,for appropriate definitions of + and min (in general known as a semiring ).
Several recent researchpapers such as NBFNet [36] explicitly call on this alignment in their motivation.
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 fromlinear algebra over the familiar real numbers, while DP, which is often a generalisation of path-findingproblems, typically takes place over “tropical” objects like ( N ∪ {∞} , min , +) 1 , which are usuallystudied in mathematics as “degenerations” of Euclidean space.
The two worlds cannot clearly bereconciled, directly, with simple equations.
However, if we define an arbitrary “latent space” R and make as few assumptions as possible, wecan observe that many of the behaviors we care about, for both GNNs and DP , arise from looking atfunctions S → R , where S is a finite set.
R can be seen as the set of real-valued vectors in the caseof 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 composablearrows (functions between finite 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, andcategory theory is a very attractive tool for “making the obvious even more obvious” [14].
Indeed,recently concepts from category theory have enabled the construction of powerful GNN architecturesbeyond permutation equivariance [9].
Here, we propose integral transforms as such an object.
We will construct the integral transform by composing transformations over our input features in away that will depend minimally on the specific choice of R .
In doing so, we will build a computationaldiagram that will be applicable for both GNNs and DP (and their own choices of R ), and henceallowing for focusing on making components of those diagrams as aligned as possible.
The integral transform
An integral transform can be encoded in a diagram of this form, which we call a polynomial span :
[Equation]
where W, X, Y and Z are finite sets.
The arrows i, p, o stand, respectively, for “input”, “process”,and “output”.
In context, the sets will have the following informal meaning: W represents the setover 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 messages 2 —we will clarify their meaning shortly.
Before proceeding, it is worthy to note the special case of X
= Y =
E , with p being the identitymap.
Such a diagram is commonly known as a span .
A span that additionally has W = Z = V isequivalent to a representation of a directed graph with vertex set Z and edge set Y ( V ← E → V );in this case i ( e ) and o ( 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 f ( w ) to each w ∈ W , how totransform it, via the polynomial span, into data on Z ?
If we can do this, we will be able to characteriseboth the process of sending messages between nodes in GNNs and scoring subproblems in DP.
For us, data on a carrier set S consists of an element of [ S, R ] :=
{ f : S → R } , where R is a “set ofpossible values”.
For now, we will think of R as an arbitrary (usually infinite) set, though we will seelater that it should possess some algebraic structure; it should be a semiring.
The transform proceeds in three steps, following the edges of the polynomial span:
[Equation]
We call the three arrows i ∗ , p ⊗ , o ⊕ the pullback , the argument pushfoward , and the message pushforward .
Taken together, they form an integral transform —and we conjecture that this transform canbe described as a polynomial functor, where p ⊗ and o ⊕ correspond to the dependent product anddependent sum from type theory (cf.
Appendix C for details).
The pullback i ∗ is the easiest to define.
Since we have a function i : X → W (part of the polynomialspan) and a function f : W → R (our input data), we can produce data on X , that is, a function in
X → R , by composition.
We hence define i ∗ f = f ◦ i .
Unfortunately, the other two arrows of the polynomial span point in the wrong direction for naïvecomposition.
For the moment, we will focus on how to define o ⊕ and leave p ⊗ for later.
We start with message data m : Y → R .
It may be attractive to invert the output arrow o in orderto define a composition with o − 1 , as was done in the case of the pullback.
However, unless o isbijective, the preimage o − 1 : Z → P ( Y ) takes values in the power set of Y .
There is an additionaltechnicality: if the composition m ◦ o − 1 takes values in P ( R ) , it will fail to detect multiplicities; weare unable to tell from a 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) of R ,which we describe in more detail in the next section. For the moment, it is enough to knowthat a bag is equivalent to a formal sum, and we define an intermediate message pushforward( o ⊕ m )( u ) :
= σ e ∈ t − 1 ( u ) m ( e ) ∈ [ Z, bag ( R )] .
e ba e bc e bd e be e bb b a c d e f ( d ) g
( e bd ) g ( e ba )
g ( e bc ) g
( e be ) g
( e bb ) pushforward p u ll b a c k e ba e bc e bd e be e bb b a c d e f ( d ) g
( e bd ) g ( e ba )
g ( e bc ) g
( e be ) g
( e bb ) h (( e bd , 0))
h (( e bd , 1))
p u ll b a c k message pushforward argumentpushforward commonly
All that is missing to complete our definition of o ⊕ is an aggregator (cid:76) : bag ( R ) → R .
As wewill see later, specifying a well-behaved aggregator is the same as imposing a commutative monoidstructure on R .
With such an aggregator on R , we can define ( o ⊕ m )( u ) :
= (cid:76) ( o ⊕ m )( u ) .
We return to p ⊗ , which is constructed very similarly.
The only difference is that, while we deliberatelyregard the collection of messages as unordered, the collection of arguments used to compute a messagehas an ordering we wish to respect.
So instead of the type bag ( R ) , we use the type list ( R ) of finitelists of elements of R , and our aggregator (cid:78) : list ( R ) → R is now akin to a fold operator.
We illustrate the use of these two aggregators in a decomposed diagram:
[Equation]
Note that any semiring ( R, ⊗ , ⊕ ) comes equipped with binary operators ⊗ , ⊕ that allow aggregators 165 (cid:78) , (cid:76) to be defined inductively.
In fact, the converse—that every set with two such aggregatorsis a semiring—is also true, if we assume some reasonable conditions on the aggregators, whichwe can explain in terms of one of the most utilised concepts in category theory and functionalprogramming— monads [30].
Due to space constraints, we refer the interested reader to Appendices
A and B for a full exposition of how we can use monads over lists and bags to constrain the latentspace R to respect a semiring structure.
For now, it’s enough to know that our key examples of the real numbers (with multiplication andaddition, for GNNs) and the tropical natural numbers (with addition and minimum, for DP) bothallow for natural interpretations of (cid:78) and (cid:76) in the integral transform.
We are now ready to show how the integral transform can be used to instantiate popular examplesof algorithms and GNNs.
We start with the Bellman-Ford algorithm [2] (Equation 4) that wastraditionally used to demonstrate the concept of algorithmic alignment.
Bellman-Ford
Let R =
( N ∪ {∞} , + , min) be the “min-plus” semiring of extended natural numbers, with ⊗ = +and ⊕ = min .
This is the coefficient semiring over which the Bellman-Ford algorithm takes place.
Let ( V, E ) be a weighted graph with source and target maps s, t : E → V and edge weightsw : E → R .
For purely technical reasons, we also need to explicitly materialise a bias functionb : V → R , which is, in practice, a constant-zero function ( b ( v ) =
for all v ∈ V )
but will provenecessary for defining the argument pushforward.
We interpret Bellman-Ford as the following polynomial span:
[Equation]
[Equation]
Here “ + ” is the disjoint union of sets, defined as A + B = { ( a, 1) | a ∈ A } ∪ { ( b, 2) | b ∈ B } .
Notethat [ S + T, R ] ∼ =
[ S, R ] ×
[ T, R ] , i.e. specifying data on a disjoint union is equivalent to specifyingdata 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 ofnode distances ( d u ; a function in [ V, R ] ), edge weights ( w ; a function in [ E, R ] ), and thepreviously discussed bias b , a function in [ V, R ] .
Hence our overall inputs are members of[ V, R ] ×
[ E, R ] × [ V, R ] ∼ =
[ V + ( V + E ) , R ] , justifying our choice of input space.
• Arguments:
X =
( V + E ) + ( V + E ) .
Here we collect the ingredients necessary tocompute Bellman-Ford’s subproblem solutions coming from neighbouring nodes.
To do this,
we need to combine data in the nodes with data living on edges—those are the argumentsto the function.
And since they meet in the edges, we “lift” our node distances [ V, R ] to
edges they are sending from, giving us an 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 u and w uv ), and each node(containing just d u , for the purposes of access to the previous optimal solution).
Hence ourmessages are members of [ V, R ] × [ E, R ] , justifying our choice of message space.•
Output: Z = V .
Lastly, the output of one step of Bellman-Ford are updated values d ′ u ,which we can interpret as just (output) data living on V .
We now describe how to propagate data along each arrow of the diagram in turn, beginning withinputs ( f, b, w ) of node features f : V → R , a bias b : V → R , and edge weights w : E → R :
• Pullback, i ∗ : First, we can note
the input
function i : ( V + E ) + ( V + E ) → V + ( V + E )
decomposes as the sum of two arrows.
i
: V + E → V is the identity function on V and
the source function on E , and i 2 :
V + E
→ V + E is just the identity.
So we calculate the pullback i ∗ ( f, b, w ) =
( f, f ◦ s, b, w ) , giving us the arguments to compute messages.
• Argument pushforward, p ⊗ : Next, the process function p simply identifies the two copies of V + E , and sums their values.
So the argument pushforward is p ⊗ ( f, f ◦ s, b, w ) =
( f, f ◦ s ) ⊗ ( b, w ) =
( f + b, ( f ◦ s ) + w ) .
This also allows us to interpret the bias function,b , as a “self-edge” in the graph with weight 0.
• Message pushforward, o ⊕ : The output function o : V + E → V is the identity function onV and the target function on E .
So the message pushforward gives us ( o ⊕ ( f + b, ( f ◦ s ) +w ))
( u ) =
( f ( u )+
b ( u )) ⊕ (cid:76) t ( e )=
u ( f ◦ s )( e ) =
min( f ( u )+ b ( u ) , min v → u f ( v )+ w v → u ) .
Letting b ( u ) = 0 , we can see that this is exactly Equation 4.
So we have produced the formula for the
Bellman-Ford algorithm directly from the polynomial span in Diagram 6.
Note that p ⊕ is aligned with using max aggregation in neural networks—directly explaining severalprevious proposals, such as [28].
But additionally, p ⊗ , as defined, is aligned with concatenating allmessage arguments together and passing them through a linear function, which is how such a step isimplemented in GNNs’ message functions.
We now direct our polynomial span analysis at GNNs.
GNNs
We study the popular message passing neural network (MPNN) model [17], which can be interpretedusing the following polynomial span diagram:
[Equation]
Here the set 1 refers to a singleton set—sometimes also called () , or unit —which is used as a carrierfor graph-level features.
This implies the graph features will be specified as [1 , R ]
∼ = R , as expected.
Given all these features, how would we compute messages?
The natural way is to combine thefeatures of 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 argumentspace as E + ( E + E ) + E : 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 twomiddle copies of E , and the identity on the last copy of E , i.e. i ( a, b, c, d ) = { () , s ( b ) , t ( c ) , d } .
Theprocess map, p , collapses the four copies of E into just one, to hold the computed message.
Lastly,the output map, o , is 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 varioussemirings of interest) is exactly an integral transform, with an extra MLP processing step on messages: [Equation] E [Equation]
R
[Equation]
It is useful to take a moment to discuss what was just achieved: with a single abstract template (thepolynomial span), we have successfully explained both a dynamic programming algorithm, and a
GNN update rule—merely by choosing the correct support sets and latent space.
Improving GNNs with edge updates, with experimental evaluation
From now on, we will set E =
V 2 , as all our baseline GNNs will use fully connected graphs, and itwill accentuate the polynomial nature of our construction.
We now show how our polynomial span view can be used to directly propose better-aligned GNNarchitectures for certain algorithmic tasks.
Since the MPNN diagram above outputs only nodefeatures, to improve predictive performance on edge-centric algorithms, it is a natural augmentationto also update edge features, by adding edges to the output carrier (as done by, e.g., [1]):
[Equation]
But notice that there is a problem with the output arrow.
Since we are using each message twice, o isno longer a function—it’d have to send each edge message to two different objects!
To resolve this,we need to appropriately augment the messages and the arguments.
This is equivalent to specifying anew polynomial span with output V 2 , which we can then recombine with Diagram 7:
[Equation]
Most edge-centric algorithms of interest (such as the Floyd-Warshall algorithm for all-pairs shortestpaths [13]), compute edge-level outputs by reducing over a choice of “intermediate” node.
Hence, itwould be beneficial to produce messages with shape V 3 , which would then reduce to features over
V 2 .
There are three possible ways to broadcast both node and edge features into V 3 , so we proposethe following polynomial span, which materialises each of those arguments:
[Equation]
Finally, inserting this into Diagram 7 gives us a corrected polynomial span with output V + V 2 :
[Equation]
[Equation]
Here we have collapsed the copies of V 2 and V 3 in the argument position for compactness.
While Diagram 7 doesn’t make sense as a polynomial diagram of sets, we can clearly still implementit as an architecture [1], since nothing stops us from sending the same tensor to two places.
Wewant to investigate whether our proposed modification of Diagram 10, which materialises order-3messages, leads to improved algorithmic alignment on edge-centric algorithms.
To support thisevaluation, we use six edge-centric tasks from the recently proposed CLRS Algorithmic Reasoning
Benchmark [29], which evaluates how well various (G)NNs align to classical algorithms, both inand out-of-distribution.
We implemented each of these options using a two-layer MLP with embedding dimension 24 , andhidden layers of size 8 and 16 .
Our test results (out-of-distribution) are summarised in Table 1.
For convenience, we also illustrate the in-distribution performance of our models via plots given in
Appendix D.
We found that the V 3 architecture was equivalent to, or outperformed, the non-polynomial ( V 2 ) onein all cases.
In particular, the size 8 polynomial architecture has fewer parameters than the size 16non-polynomial one, but shows substantial improvements in several algorithms that benefit fromaligned edge features.
Therefore, this result directly validates our theory’s predictions, in the contextof presenting a better-aligned GNN for edge-centric algorithmic targets.
[Table]
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 deriveda 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 GNNand DP computations.
With this diagram materialised, we were able to immediately unify largequantities of prior work as simply manipulating one arrow or element in the integral transform.
Wealso 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 researchinto better-aligned neural algorithmic reasoners, especially focusing on generalising or diving intoseveral aspects of this diagram.
Lastly, it is not at all unlikely that analyses similar to ours have already been used to describe otherfields of science—beyond algorithmic reasoners.
The principal ideas of span and integral transformare central to defining Fourier series [32], and appear in the analysis of Yang-Mills equations inparticle physics [12].
Properly understanding the common ground behind all of these definitions may,in the very least, lead to interesting connections, and a shared understanding between the variousfields they span.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] We propose a 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 workis of a theoretical nature.(d)
Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of 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.
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 as a URL)?
[No] We aim torelease 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 were 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 experiments multiple times)?
[Yes](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
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 CLRSbenchmark.(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’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiableinformation or offensive content?
[N/A]
Our data is abstract and algorithmicallygenerated.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to 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]
[-Additionally, in Figure 3, we follow this construction to explicitly provide the pseudocodes for theproposed V 2 and V 
-]
A The bag and list monads
Before we conclude, we turn back to the theory behind our polynomial spans, to more preciselydetermine the restrictions on our abstract latent space R .
We found this investigation useful to includein the main paper, as it yields a strong connection to one of the most actively used concepts intheoretical computer science and functional programming.
Recall that the realisation of our pushforward operations required the existence of two aggregators: 440 (cid:78) (to fold lists) and (cid:76) (to reduce bags).
Previously, we mentioned only in passing how they can berecovered—now, we proceed to define (cid:76) axiomatically.
Given a set S , we define bag ( S ) :
= { p : S → N | # { p ( r ) ̸ = 0 } < ∞} , the natural-valued functionsof finite 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 finitely many of the n s are nonzero.
Given a function f : S → T between sets, we can define a function bag ( f ) : bag ( S )
→ bag ( T ) ,as follows: bag ( 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 = 448 (cid:80) f
( s )=
t n s .
For each S , we can also define two special functions.
The first is unit : S → bag ( S ) , sendingeach element to its indicator function (i.e. an element x ∈ S to the multiset {{ x }} ).
The second isjoin : bag ( bag ( S ))
→ bag ( S ) , which interprets a nested sum as a single sum.
These facts tell us that bag 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 forabstracting control flow, as well as gracefully handling functions with side effects [30].
It is well-known that the algebras for the monad bag are the commutative monoids , sets equippedwith a commutative and associative binary operation and a unit element.
Concretely, a commutative monoid structure on a set R is equivalent to defining an aggregatorfunction (cid:76) : bag ( R ) → R compatible with the unit and monad composition.
Here, compatibilityimplies it should correctly handle sums of singletons and sums of sums, in the sense that the followingtwo diagrams commute ; that is, they yield the same result regardless of which path is taken:
[Equation]
id
[Equation]
The first diagram explains that the outcome of aggregating a singleton multiset (i.e. the one producedby applying unit ) with (cid:76) is equivalent to the original value placed in the singleton.
The seconddiagram indicates that the (cid:76) operator yields the same results over a nested multiset, regardless ofwhether we choose to directly apply it twice (once on each level of nesting), or first perform the joinfunction to collapse the nested multiset, then aggregate the collapsed multiset with (cid:76) .
So the structure of a commutative monoid on R is exactly what we need to complete our definition ofthe message pushforward o ⊕ .
The story for the argument pushforward, p ⊗ , is remarkably similar.
Define list ( S ) :
= { ( s 1 , . . .
, s n ) | n ∈ N , s i ∈ S } , the set of all ordered lists of elements of S ,including the empty list.
Equivalently, list ( S ) =
(cid:96) n ≥ 0 S n .
We can also extend list to a functor:for a function f : S → T , list ( f ) : list ( S ) → list ( T ) is just the well-known map operation:list ( f )( s 1 , . . .
, s n ) :
= ( f ( s 1 ) , . . .
, f ( s n )) .list is also a monad, with unit : S → list ( S ) sending each x ∈ S to the singleton list ( x ) , andjoin : list ( list ( S ))
→ list ( S ) sending a list of lists to their concatenation.
The algebras for the list monads are monoids —not just commutative ones.
So R needs a secondmonoid structure, possibly noncommutative, to support our definition of the argument pushforward.
We detail how this can elegantly be done in our specific case in Appendix B.
B The monad for semirings
We have asked that R be an algebra for two monads: list and bag .
But this is an unnaturalcondition without some compatibility between the two.
It would more useful to find a single monadencapsulating both.
In general, the composition of two monads is not a monad.
For example, the composite functorlist ◦ bag does not support a monad structure.
However, the other composite bag ◦ list is actually a monad in a natural way, due to the existenceof a distributive law , which is a natural transformation λ : list ◦ bag → bag ◦ list satisfying someaxioms, see e.g. [8].
It is easy to describe λ .
Given any list of bags ( (cid:80) i 1 a i 1 , . . .
, (cid:80) i n a i n ) , we haveλ ( (cid:80) i 1 a i 1 , . . .
, (cid:80) i n a i n ) =
(cid:80) i 1 ,...,i n ( a i 1 , . .
. , a i n ) .
In other words, λ takes a list of bags andreturns the bag of all ordered selections from the list.
This is exactly how multiplication of sums works in a semiring.
For example, if I think of a polynomialas a bag of monomials, and I want to compute a product of polynomials, I interpret this product as alist of polynomials, i.e. a list of bags.
Then I expand it into a bag of lists (a sum of products), andfinally perform the products to produce the resulting bag of monomials, i.e. polynomial.
So it shouldn’t be a surprise that the algebras for the composite monad bag ◦ list are exactlysemirings, i.e. sets R equipped with a commutative monoid structure (cid:76) , another monoid structure 496 (cid:78) , and a “distributive law” (cid:78) (cid:76) → (cid:76) (cid:78) , usually written as, e.g. x ( a + b )
y = xay + xby , andextended to arbitrary sums and products by induction.
Indeed, if R is an algebra for the monad bag ◦ list , we have some “double aggregator” ev :bag ( list ( R ))
→ R .
We can recover ⊗ : list ( R )
→ R by packing our list into a singleton bag,and we can recover ⊕ : bag ( R ) → R by packing our bag into a singleton list then applying λ .
C Polynomial functors
Polynomial spans are the starting point for our integral transform, but they are also the starting pointfor polynomial functors , which arise in dendent type theory.
Let C be a locally cartesian closedcategory, and let C /A denote, for any object A of C , the category of morphisms with target A .
Apolynomial functor starts with a polynomial span:
[Equation]
And it produces a composition of three functors:
[Equation]
Here Σ o and Π p are operations called the dependent sum and dependent product respectively.
Note that there is a direct correspondence between the three arrows in each of the diagrams 5 and 11.
So it is very tempting to ask whether our integral transform is expressible as a polynomial functor.
Can our results be rephrased in those terms?
We conjecture that the answer is yes.
However, there are technical hurdles to overcome in comparingthe two definitions.
When computing with a polynomial functor, we start with a variable set E and afunction E
→ W , and produce a new set F ( E ) and a function F ( E ) →
Z .
Our integral transformstarts with a function to a fixed set, and produces a function to the same set.
To see how these two pictures are related, it is helpful to consider the example of finite sets.
Afunction π : E → W of finite sets can be expressed, up to equivalence, as a function f : W → N given by f ( w ) =
# π − 1 ( w ) .
In other words, we can describe E by describing the size of the fiberover each element of W .
Indeed, the existence of such a “decategorification” for transforms overspans was an early inspiration for our present work.
To replace N with an arbitrary semiring R , we would need to find a way to interpret a functionf : W → R as a classifying morphism for some kind of bundle E → W in a suitable category ofgeometric objects over R .
For R = N ∞ , one possibility is to define a category of R -schemes, whichshould be certain types of topological spaces equipped with sheaves of R -modules–we don’t know ofa place this theory is developed, but it is certainly related to tropical schemes, defined in [16].
For
R = R , we can consider the more familiar category of manifolds, or more generally the category oflocally compact Hausdorff spaces.
But do polynomial functors work in categories like this?
While polynomial functors were developedin type theory over locally cartesian closed categories–too strong of a condition for interestingtopology to occur–[31] has shown that polynomial functors can be defined in any category withpullbacks, as long as the “processor” morphism p :
X → Y satisfies an abstract condition calledexponentiability .
i and o can still be arbitrary morphisms.
For some intuition, we quote two results on exponentiability.
[6] shows that the exponentiablemorphisms in the category of compact Hausdorff spaces are the local homeomorphisms.
And [20]shows that a morphism R → S of commutative rings gives rise to an exponentiable morphismof affine schemes exactly when S is dualizable as an R -module.
So exponentiability seems to bestrongly linked to covering spaces in classical topology, as well as descent theory in modern algebraicgeometry.
Expanding on these ideas is far out of scope for the present work, but we hope it gives a glimpse intothe possibilities on the horizon.
D Plots of in-distribution performance on CLRS
For plots that illustrate in-distribution performance of our proposed V 3 model, against the nonpolynomial ( V 2 ) model, please refer to Figure 2.
Our findings largely mirror the ones from out-ofdistribution—with V 3 either matching the performance of the baseline or significantly outperformingit (e.g. on Insertion Sort and Floyd-Warshall).
We do note that sometimes, matched performance bythe non-polynomial V 2 baseline in-distribution can be misleading , as it significantly loses out to V 3out of distribution (cf.
Table 1).