Agent-based Graph Neural Networks
Abstract
We present a novel graph neural network we call AgentNet, which is designed specifically for graph-level tasks.
AgentNet is inspired by sublinear algorithms, featuring a computational complexity that is independent of the graph size.
The architecture of AgentNet differs fundamentally from the architectures of traditional graph neural networks.
In AgentNet, some trained neural agents intelligently walk the graph, and then collectively decide on the output.
We provide an extensive theoretical analysis of AgentNet: We show that the agents can learn to systematically explore their neighborhood, and that AgentNet can distinguish some structures that are even indistinguishable by 3-WL.
Moreover, AgentNet is able to separate any two graphs which are sufficiently different in terms of subgraphs.
We confirm these theoretical results with synthetic experiments on hard-to-distinguish graphs and real-world graph classification tasks.
In both cases, we compare favorably not only to standard GNNs but also to computationally more expensive GNN extensions.
Introduction
Graphs and networks are prominent tools to model various kinds of data in almost every branch of science.
Due to this, graph classification problems also have a crucial role in a wide range of applications from biology to social science.
In many of these applications, the success of algorithms is often attributed to recognizing the presence or absence of specific substructures, e g. atomic groups in case of molecule and protein functions or cliques in case of social networks [9; 70; 19; 21; 60; 5].
This suggests that some parts of the graph are “more important” than others, and hence it is an essential aspect of any successful classification algorithm to find and focus on these parts.
In recent years, Graph Neural Networks (GNNs) have been established as one of the most prominent tools for graph classification tasks.
Traditionally, all successful GNNs are based on some variant of the message passing framework [3; 63].
In these GNNs, all nodes in the graph exchange messages with their neighbors for a fixed number of rounds, and then the outputs of all nodes are combined, usually by summing them [24; 48], to make the final graph-level decision.
It is natural to wonder if all this computation is actually necessary.
Furthermore, since traditional
GNNs are also known to have strong limitations in terms of expressiveness, recent works have developed a range of more expressive GNN variants; these usually come with an even higher computational complexity, while often still not being able to recognize some simple substructures.
This complexity makes the use of these expressive GNNs problematic even for graphs with hundreds of nodes, and potentially impossible when we need to process graphs with thousands or even more nodes.
However, graphs of this size are common in many applications, e.g. if we consider proteins [59; 65], large molecules [72] or social graphs [7; 5].
In light of all this, we propose to move away from traditional message-passing and approach graphlevel tasks differently.
We introduce AgentNet – a novel GNN architecture specifically focused on these tasks.
AgentNet is based on a collection of trained neural agents, that intelligently walk the graph, and then collectively classify it (see Figure 1). These agents are able to retrieve information
I see a triangle!
The Graph has triangles!
[Figure]
Update
[Figure] [Figure]
I've been there [Figure] ago!
I keep seeing these nodes!
I didn't see a triangle.
Me neither.
occupy and the state of the occupied Finally, the agent then chooses a neighboring node to transition to, based on its own state and the state of the neighboring nodes. As we will show later, even with a very naive policy, an agent can already recognize cliques and cycles, which is impossible with traditional GNNs.
One of the main advantages of AgentNet is that its computational complexity only depends on the node degree, the number of agents, and the number of steps.
This means that if a specific graph problem does not require the entire graph to be observed, then our model can often solve it using less than n operations, where n is the number of nodes.
The study of such sublinear algorithms is a popular topic in graph mining [30; 23]; it is known that many relevant tasks can be solved in a sublinear manner.
For example, our approach can recognize if one graph has more triangles than another, or estimate the frequency of certain substructures in the graph – in sublinear time!
AgentNet also has a strong advantage in settings where e.g. the relevant nodes for our task can be easily recognized based on their node features.
In these cases, an agent can learn to walk only along these nodes of the graph, hence only collecting information that is relevant to the task at hand.
The amount of collected information increases linearly with the number of steps.
In contrast to this, a standard message-passing GNN always (indirectly) processes the entire multi-hop neighborhood around each node, and hence it is often difficult to identify the useful part of the information from this neighborhood due to oversmoothing or oversquashing effects [41; 2] caused by an exponential increase in aggregated information with the number of steps.
Related work 2.1 GNN limitations
Expressiveness.
Xu et al.
[68] and Morris et al.
[49] established the equivalence of message passing
GNNs and the first order Weisfeiler-Lehman (1-WL) test.
This spurred research into more expressive
GNN architectures.
Sato et al.
[58] and Abboud et al.
[1] proposed to use random node features for unique node identification.
As pointed out by Loukas [44], message passing GNNs with truly unique identifiers are universal.
Unfortunately, such methods generalize poorly to new graphs [54].
Vignac et al.
[64] propose propagating matrices of order equal to the graph size as messages instead of vectors to achieve a permutation equivariant unique identification scheme.
Other possible expressivenessimproving node feature augmentations include distance encoding [40], spectral features [4; 22] or orbit counts [10].
However, such methods require domain knowledge to choose what structural information to encode.
Pre-computing the required information can also be quite expensive [10].
An alternative to this is directly working with higher-order graph representations [49; 47], which can directly bring k -WL expressiveness at the cost of operating on k -th order graph representations.
To improve this, methods that consider only a part of the higher-order interactions have been proposed [51; 52].
Alternatively, extensions to message passing have been proposed, which involve running a GNN on many, slightly different copies of the same graph to improve the expressiveness [54; 8; 73; 17].
Papp and Wattenhofer [53] provide a theoretical expressiveness comparison of many of these GNN extensions.
Geerts and Reutter [28] propose the use of tensor languages for theoretical analysis of GNN expressive power.
Interestingly enough, all of these expressive models focus mainly on graph-level tasks.
This suggests these are the main tasks to benefit from increased expressiveness.
Underreaching.
GNNs usually use a fixed number of layers, which only allows them to aggregate information from a few hops away.
This can cause a problem called underreaching [2], when the important information is not reached.
Historically, training GNNs with more than a few layers proved to be quite difficult due to oversmoothing [41] and oversquashing [2].
To address this, many deeper
GNN architectures have been proposed that use various normalization techniques, skip connections, adaptive information propagation, and weight tying [38; 42; 39].
While they allow for hundreds of layers, this normally comes at the expense of long training and evaluation times.
Scalability.
As traditional GNN architectures perform computations on the neighborhood of each node, their computational complexity depends on the number of nodes in the graph and the maximum degree.
To enable GNN processing of large graphs Hamilton et al.
[32] and Chen et al.
[12] propose randomly sub-sampling neighbors every batch or every layer.
To better preserve local neighborhoods
Chiang et al.
[15] and Zeng et al.
[71] use graph clustering, to construct small subgraphs for each batch.
Alternatively, Ding et al.
[20] quantize node embeddings and use a small number of quantized vectors to approximate messages from out-of-batch nodes.
In contrast to our approach, these methods focus on node-level tasks and might not be straightforward to extend to graph-level tasks, because they rely on considering only a predetermined
2.2 Sublinear algorithms
Sublinear algorithms aim to decide a graph property in much less than n time, n being the number of nodes [30].
It is possible to check if one graph has noticeably more substructures such as cliques, cycles, or even minors than another graph in constant time [27; 26; 6].
It is also possible to estimate counts of such substructures in sublinear time [23; 13; 7].
All of this can be achieved either by performing random walks [13; 7] or by investigating local neighborhoods of random nodes [27; 26; 6].
We will show that AgentNet can do both, local investigations and random walks.
2.3 Random walks in GNNs
Random walks have been previously used in graph machine learning to construct node and graph representations [56; 31; 61].
However, traditionally basic random walk strategies are used, such as choosing between dept-first or breath-first random neighborhood traversals [31].
As an alternative, Lee et al.
[37] have proposed to instead learn a graph walking procedure using reinforcement learning.
In contrast to our work, their approach is not fully differentiable; furthermore, the node features cannot be updated by the agents during the walk nor the collective neighborhood information is taken into account at each step, which both are crucial for expressiveness.
AgentNet model
The core idea of our approach is that a graph-level prediction can be made by intelligently walking on the graph many times in parallel.
To be able to learn such intelligent graph walking strategies through gradient descent we propose the AgentNet model (Figure 1).
On a high level, each agent a is assumed to be uniquely identifiable (have an ID) and is initially placed uniformly at random in the graph.
Upon an agent visiting a node v at time step t (including the initial placement), the model performs four steps: node update, neighborhood aggregation, agent update, and agent transition.
First, the embedding of each node v i that has agents a j ∈ A ( v i ) on it is updated with information from all of those agents using the node update function f v .
To ensure that the update is expressive enough, agent embeddings are processed before aggregation using a function ϕ a → v :
[Equation]
Next, in the neighborhood aggregation step, current neighborhood information is incorporated into the node state using the embeddings of neighboring nodes v j ∈ N ( v i ) .
Similarly, we apply a function ϕ N ( v ) → v on the neighbor embeddings beforehand, to ensure the update is expressive:
[Equation]
These two steps are separated so that during neighborhood aggregation a node would receive information about the current state of the neighborhood and not its state in the previous step.
The resulting node embedding is then used to update agent embeddings so that the agents know the current state of the neighborhood (e.g. how many neighbors of v i have been visited) and which other agents are on the same node.
With V ( a i ) being a mapping from agent a i to its current node:
[Equation]
Now the agent has collected all possible information at the current node it is ready to make a transition to another node.
For simplicity, we denote the potential next positions of agent a i by
N t ( a i ) :
=
N ( v tV ( a i ) ) ∪ v tV ( a i ) .
First, probability logits z a i → v j are estimated for each v tj ∈ N t ( a i ) .
Then this distribution is sampled to select the next agent position V ( a i ) :
[Equation]
a i ) ← GumbelSoftmax (cid:0)(cid:8) z a i → v j for v tj ∈ N t ( a i ) (cid:9)(cid:1) .
As we need a categorical sample for the next node position, we use the straight-through Gumbel softmax estimator [34; 46].
In the practical implementation, we use dot-product attention [62] to determine the logits z a i → v j used for sampling the transition.
To make the final graph-level prediction, pooling is applied on the agent embeddings, followed by a readout function.
Edge features can also be trivially incorporated into this model.
For details on this and the practical implementation of the model see Appendix C.
There we also discuss other simple extensions to the model, such as providing node count to the agents or global communication between them.
Let’s consider a Simplified AgentNet version of this model, where the agent can only decide if it prefers transitions to explored nodes, unexplored nodes, going back to the previous node, or staying on the current one.
Such an agent can already recognize cycles.
If the agent always prefers unexplored nodes and marks every node it visits, which it can do uniquely as agents are uniquely identifiable, it will know once it has completed a cycle (Figure 1 (c)).
Similarly, if for c steps an agent always visits a new node, and it always sees all the previously visited nodes, it has observed a c -clique (Figure(d)).
Naturally, this clique detection can be improved if the agent systematically visits all neighbors of a node v , going back to v after each step, until a clique is found.
This implies that even this simple model could distinguish the Rook’s 4 × 4 and Shrikhande graphs, which are indistinguishable by 3-WL, but can be distinguished by detecting 4-cliques.
If we have more than one agent they could successfully run such algorithms in parallel to increase interesting substructure recognition chances.
Theoretical analysis
We first analyze the theoretical expressiveness of the AgentNet model described above.
We begin with the case of a single agent, and then we move on to having multiple agents simultaneously.
The proofs and more detailed discussion of our theorems are available in Appendices A and B.
4.1 Expressiveness with a single agent
We first study the expressiveness of AgentNet with a single agent only.
That is, assume that an agent is placed on the graph at a specific node v .
In a theoretical sense, how much is an agent in our model able to explore and understand the region of the graph surrounding v ?
Let us define the r -hop neighborhood of v (denoted N r ( v ) ) as the subgraph induced by the nodes that are at distance at most r from v .
One of the most fundamental observations is that AgentNets are powerful enough to efficiently explore and navigate the r -hop neighborhood around the starting node v .
In particular, AgentNets can compute (and maintain) their current distance from v , recognize the nodes that they have already visited, and also consciously select the next neighbor to visit in each step; this already allows them to efficiently traverse all the nodes in N r ( v ) .
Lemma 1.
An AgentNet can learn to execute a iteratively deepening depth-first search of its r -hop neighborhood, visiting each node of N r ( v ) in ℓ ≤ O ( r · | N r ( v ) | ) steps altogether.
We note that in many practical cases, e.g. if the nodes of interest for a specific application around v can already be identified from their node features, then AgentNets can traverse the neighborhood even more efficiently with a depth-first search.
We discuss this special case (together with the rest of the proofs) in more detail in Appendix A.
While using these methods to traverse the neighborhood around v , an AgentNet can also identify all the edges between any two visited nodes; intuitively speaking, this allows for a complete understanding of the neighborhood in question.
More specifically, if the transition functions of the agent are implemented with a sufficiently expressive method (e.g. with multi-layer perceptrons), then one can show that it is possible to develop a maximally expressive AgentNet implementation, i.e. where each transition function is injective, similarly to GIN in a standard GNN setting [68].
Together with
Lemma 1, this provides a strong statement: it allows for an AgentNet that can distinguish any pair of different r -hop neighborhoods around v (for any finite r and ∆ , where ∆ denotes the maximal degree in the graph).
Theorem 2.
There exists an injective implementation of an AgentNet (with ℓ ≤ O ( r · | N r ( v ) | ) steps) which computes a different final embedding for every non-isomorphic r -hop neighborhood that can occur around a node v .
Note that this already implies that with a sufficient number of steps, running an AgentNets from node v can be strictly more expressive than a standard GNN from v with r layers, or as a matter of fact, any GNN extension that is unable to distinguish every pair of non-isomorphic r -hop neighborhoods around v .
In particular, consider the Rook’s 4 × 4 and Shrikhande graphs, which are fundamental examples of graphs that are challenging to distinguish even for some sophisticated GNN extensions; an AgentNet can learn to distinguish these two graphs in as little as ℓ = 11 steps.
Since the comparison of the theoretical expressiveness of different GNN variants is a heavily studied topic, we also add these observations as an explicit theorem.
Corollary 3.
The AgentNet approach can also distinguish graphs that are not separable by standard
GNNs or even more powerful GNN extensions such as PPGN [47], GSN [10] or DropGNN [54].
Intuitively, Theorem 2 also means that AgentNets are expressive enough to learn any property that is a deterministic function of the r -hop neighborhood around v .
In particular, for any subgraph H that is contained within distance r of one of its nodes v 0 , there is an AgentNet that can compute (in
O ( r · | N r ( v ) | ) steps) the number of occurrences of H around a node v of G , i.e. the number of times
H appears as an induced subgraph of G such that node v takes the role of v 0 .
Lemma 4.
Let H be any subgraph of radius at most r around a specific node v 0 .
Then there exists an AgentNet that can compute the number of occurrences of H around a specific node v of G .
For several applications, cliques and cycles are often mentioned as some of the most relevant substructures in practice.
As such, we also include more specialized lemmas that consider these two structures in particular.
In these lemmas, we say that an event happens with high probability ( w.h.p ) if an agent can ensure that it happens with probability p for an arbitrarily high constant p < 1 .
Lemma 5.
There exists an AgentNet that can count cliques (of any size) in ℓ = 2 · ∆ − 1 steps, but there is no AgentNet that can count them w.h.p.
in less steps.
Lemma 6.
There exists an AgentNet that can count c -cycles in ℓ =
Θ( r · | N r ( v ) | ) steps, but there is no AgentNet that can count them w.h.p.
in less than 2 · | N r ( v ) | − ⌊ c 2 ⌋ steps.
In general, these theorems will allow us to show that if a specific subgraph appears much more frequently in some graph G 1 than in another graph G 2 , then we can already distinguish the two graphs with only a few agents (we formalize this in Theorem 9 for the multi-agent setting).
Note that all of our theorems so far rely on the agents recognizing the structures around their starting point v .
This is a natural approach, since in the general case, without prior knowledge of the structure of the graph, the agents cannot make more sophisticated (global) decisions regarding the directions in which they should move in order to e.g. find a specific structure.
However, another natural idea is to consider the case when the agents do not learn an intelligent transition strategy at all, but instead keep selecting from their neighbors uniformly at random .
We do not study this case in detail, since the topic of such random walks in graphs has already been studied exhaustively from a theoretical perspective.
In contrast to this, the main goal of our paper is to study agents that learn to make more sophisticated navigation decisions.
In particular, a clever initialization of the Simplified AgentNet transition function already ensures in the beginning that the movement of the agent is more of a conscious exploration strategy than a uniform random walk.
Nonetheless, we point out that many of the known results on random walks automatically generalize to the AgentNets setting.
For example, Dasgupta et al.
and Chierichetti et al.
[16; 18] defined a random walk access model where an algorithm begins from a seed vertex v and has very limited access to the graph: it can (i) move to a uniform random neighbor, (ii) query the degree of the current node, and (iii) query whether two already discovered nodes are adjacent.
Several algorithms have been developed and studied in this model, e.g. for counting triangles in the graph efficiently [7].
Since an AgentNet can execute all of these fundamental steps, it is also able to simulate any of the algorithms that were developed for this model.
Theorem 7.
An AgentNet can simulate any algorithm developed in the random walk access model.
4.2 Multiple agents
We now analyze the expressiveness of AgentNet with multiple agents.
Note that the main motivation for using multiple agents concurrently is that it allows for a significant amount of parallelization, hence making the approach much more efficient in practice.
Nonetheless, having multiple agents also comes with some advantages (and drawbacks) in terms of expressive power.
We first note that adding more agents can never reduce the theoretical expressiveness of the AgentNet framework; intuitively speaking, with unique IDs, the agents are expressive enough to disentangle their own information (i.e. the markings they leave at specific nodes) from that of other agents.
Lemma 8.
Given an upper bound b on agent IDs, there is an AgentNet implementation that always computes the same final embedding starting from v , regardless of the actions of the remaining agents.
This shows that even if we have multiple agents, they always have the option to operate independently from each other if desired.
Together with Theorem 2, this allows us to show that if two graphs differ significantly in the frequency of some subgraph, then they can already be distinguished by constantly many agents.
More specifically, let G 1 and G 2 be two graphs on n nodes, let H be a subgraph that can be traversed in ℓ steps as discussed above, and let γ H ( G i ) denote the number of nodes in G i that are incident to at least one induced subgraph H .
Theorem 9.
Let G 1 and G 2 be graphs such that γ H ( G 1 ) − γ H ( G 2 ) ≥ δ · n for some constant δ > 0 .
Then already with k ∈ O (1) agents an AgentNet can distinguish the two graphs w.h.p.
In general, Lemma 8 shows that if the number of steps ℓ is fixed, then we strictly increase the expressiveness by adding more agents.
However, another (in some sense more reasonable) approach is to compare two AgentNet settings with the same number of total steps: that is, if we consider an AgentNet with k distinct agents, each running for ℓ steps, then is this more expressive than an
AgentNet with a single agent that runs for k · ℓ steps?
We point out that in contrast to Lemmas 1-6 (which hold in every graph), our results for this question always describe a specific graph construction, i.e. they state that we can embed a subgraph H in a specific kind of graph G such that it can be recognized by a given AgentNet configuration, but not by another one.
A more detailed discussion of the theorems is B.
On the one hand, it is easy to find an example where having a single agent with k · ℓ steps is more beneficial: if we want to recognize a very large structure, e.g. a path on more than ℓ nodes, then this might be straightforward in a single-agent case, but not possible in the case of multiple agents.
Lemma 10.
There is a subgraph H (of radius larger than ℓ ) that can be recognized by a single agent with k · ℓ steps, but not by k distinct agents running for ℓ steps each.
On the other hand, it is also easy to find an example where the multi-agent case is more powerful at distinguishing two graphs, e.g. if this requires us to combine information from two distant parts of the graph that are more than k · ℓ steps away from each other.
Lemma 11.
There is a pair of non-isomorphic graphs (of radius larger than k · ℓ ) that can be distinguished by k distinct agents with ℓ steps each, but not by a single agent with k · ℓ steps.
However, neither of these lemmas cover the simplest (and most interesting) case when we want to find and recognize a single substructure H of small size, i.e. that has a radius of at most ℓ .
In this case, one might first think that the single-agent setting is more powerful, because it can essentially simulate the multi-agent case by consecutively following the path of each separate agent.
For example, if two agents a 1 , a 2 meet at a given node v while identifying a substructure in the multi-agent setting, then a single agent could first traverse the path of a 1 , then move back to v and then traverse the path of a 2 from here.
However, it turns out that this is not the case in general: one can design a construction where a subgraph H can be efficiently explored by multiple agents, but not by a single agent.
Intuitively, the main idea of this construction is to develop a one-way tree structure that can only be efficiently traversed in one direction; that is, in the correct direction, the next step of a specific path is always clearly distinguishable from the node features, but when traversing it the opposite way, there are many identical-looking neighbors in each step.
Theorem 12.
There exists a subgraph H that can be recognized w.h.p.
by 2 agents in ℓ steps, but it cannot be recognized by 1 agent in c · ℓ steps (for any constant c ).
Finally, we note that as an edge case, AgentNets can trivially simulate traditional message-passing
GNNs with n agents if a separate agent is placed on each node, and performs the node update and neighborhood aggregation steps, while never transitioning to another node.
However, if the problem at hand requires such an approach, it is likely more reasonable to use a traditional GNN instead.
Experiments
In this section, we aim to demonstrate that AgentNet is able to recognize various subgraphs, correctly classifies large graphs if class-defining subgraphs are common, performs comparably on real-world tasks to other expressive models which are much more computationally expensive, and can also successfully solve said tasks better than GIN while performing much fewer than n steps on the graph.
In all of the experiments, we parameterize all of the AgentNet update functions with 2-layer MLPs with a skip connection from the input to the output.
For details on model implementation see
Appendix C. Unless stated otherwise, to be comparable in complexity to standard GNNs (e.g. GIN) we use the number of agents equal to the mean number of nodes n in graphs of a given dataset.
For details on the exact setup of all the experiments and the hyperparameters considered see Appendix D.
5.1 Synthetic datasets
First, to ensure that AgentNet can indeed recognize substructures in challenging scenarios we test it on synthetic GNN expressiveness benchmarks (Table 1).
There we consider three AgentNet versions.
The random walk AgentNet, where the exploration is purely random.
The Simplified AgentNet, where each agent only predicts the logits for staying on the current node, going to the previous node, or going to any unexplored or any explored node, and finally the full AgentNet, which uses dot-product attention for transitions.
Notice that the random walk AgentNet is not able to solve the 3-WL task, as a random walk is a comparatively limited and sample-intensive exploration strategy.
The other two
AgentNet models successfully solve all tasks, which means they can detect cycles and cliques.
Since
[Table]
[Table]
[Table]
enough, while 1-2-3 GNN should have a similar expressive power to the 3-WL test, it solves this task.
This is likely due to the fact that while all possible triplets are needed to simulate a 3-WL test, in practice only a subset of them is considered to reduce computational complexity [49].
the simplified exploration strategy does not account for neighboring node features, in the rest of the experiments we will only consider the full AgentNet model.
Secondly, to test if AgentNet is indeed able to separate two graphs when the defining substructure is prevalent in the graph we perform the ablation in Figure 2a.
We see, that indeed when there are many defining subgraphs AgentNet can successfully differentiate between the graphs.
[Figure]
If the subgraph of interest is common in the graph (blue line) AgentNet successfully learns to differentiate the graphs, even when k · ℓ < n (gray line). Secondly, to test how this transfers to real-world tasks we test AgentNet on a large protein DD dataset using a different number of agents and steps (d).
Grey-bordered cells mark configurations that perform fewer node visits than GIN ( k · ℓ < 4 · n ) and black-bordered cells mark configurations that perform less than n node visits ( k · ℓ < n ). GIN achieved 77% accuracy on this task. AgentNet outperforms it, even when using just n 8 node visits and matches GIN while using only n 16 node visits.
5.2 Real-world graph classification datasets
To verify that this novel architecture works well on real-world graph classification datasets, following
Papp et al.
[54] we use three bioinformatics datasets (MUTAG, PTC, PROTEINS) and two social network datasets (IMDB-BINARY and IMDB-MULTI)
As graphs in these datasets only have 10s of nodes per graph on average, we also include one large-graph protein dataset (DD) [21] and
[Table]
and ESAN models use r ≈ O ( n ) versions of each graph, which makes them of quadratic complexity in practice.
We set k ≈ n for AgentNet models to have a comparable setting to GIN.
We also compare to a 1-AgentNet, which uses only one agent and cannot visit the whole graph.
* We report the best result achieved by any of the different model versions.
PPGN has 3 different versions and ESAN has 8 different versions.
† Originally 1-2-3 GNN used a slightly different evaluation setup.
We re-trained it to follow the same experimental setup as the other baselines [68] (see Appendix D). All other results come from the respective papers. one large-graph social dataset (REDDIT-BINARY) [69]. In these datasets, graphs have hundreds of nodes on average and the largest graphs have thousands of nodes (Appendix E).
We compare our model to GIN, which has maximal 1-WL expressive power achievable by a standard message-passing
GNN [68], and more expressive GNN architectures, which do not require pre-computed features [54; 8; 49; 47].
As you can see in Table 2 our novel approach usually outperforms at least half of the more expressive baselines.
Even the model that uses just a single agent matches the performance of
GIN, even though in this case the agent cannot visit the whole graph, as for these experiments we only consider ℓ ∈ { 8 , 16 } (Appendix D) and using only one agent reduces expressiveness.
Naturally, 1-AgnetNet performance deteriorates on the large graphs, while AgentNet does very well.
The higher-order methods cannot even train on these large graphs on a 24GB GPU.
To train the ESAN and DropGIN models on them we have to only use 5% of the usual n different versions of the same graph [8].
For DropGIN, this results in a loss of accuracy.
While ESAN performs well in this scenario, it requires lengthy pre-processing and over 120GB of main memory to load the pre-processed dataset (original datasets are < 150 MB), which can become prohibitive if we need to test even larger graphs.
[Table]
on the OGB-MolHIV dataset with edge features.
To test in how AgentNet performance depends on the number of agents and steps in large real-world graphs we perform an ablation on the DD dataset (Figure 2d).
We see that many configurations perform well, even when just a fraction of the graph is explored.
We also observe that unlike with traditional GNNs [41; 38], using more steps (layers) does not cause a noticeable decrease in performance.
Especially when fewer than n agents are used ( n ≈ 284 ).
This can indicate higher robustness to oversquashing and oversmoothing.
The previous datasets do not use edge features.
As as mentioned in Section 3, this is a straightforward extension (Appendix C).
We test this extension on the OGB-MolHIV molecule classification dataset, which uses edge features [33].
In Table 3 we can see that AgentNet performs well in this scenario and outperforms even the best ESAN model.
Conclusion
In this work, we presented a novel AgentNet architecture for graph-level tasks.
We provide an extensive theoretical analysis, which shows that this architecture is able to distinguish various substructures that are impossible to distinguish with traditional GNNs.
We show that AgentNet also brings improvements to real-world datasets.
Furthermore, if features necessary to determine the graph class are frequent in the graph, AgentNet allows for classification in sublinear or even a constant number of rounds.
To our knowledge, this is the first fully differentiable GNN computational model capable of graph classification which inherently has this feature, without requiring explicit graph sparsification – it learns which substructures are worth exploring on its own.
A Proofs for Section 4. A.1 Injective implementation
The fundamental idea behind developing a maximally expressive AgentNet implementation is to ensure that the functions learned by the agent are injective.
A similar proof technique for developing injective GNN implementations has already been applied to analyze the expressive power of standard
GNNs [68] and also some GNN extensions [54]; we refer the reader to these papers for more technical details on this proof approach.
As in the rest of these proofs, we assume that the range of the node features is a finite set, and we also have a finite upper bound on the maximal degree ∆ of the graph.
With an induction, this already shows that after any finite number of time steps t , the range of possible embeddings (of agents or nodes) is still finite.
The main tool we will use is the following: assume we have a finite number of embeddings h 0 , ..., h j − 1 ∈ R , and assume for simplicity that they are from the interval (0 , 1) .
Then there exists a function f : (0 , 1) j
→ (0 , 1) that is injective on its entire domain.
; that is, if we have ˆ h 1 , ..., ˆ h j such that h i ̸
= ˆ h
i for at least one i ∈ { 0 , ..., j − 1 } , then f ( h 0 , ..., h j − 1 ) ̸
= f
(ˆ h 0 , ..., ˆ h j − 1 ) .
One possible way to construct such an f is as follows: we consider the binary representation of each h i , and encode it in the bits of f ( h 0 , ..., h j − 1 ) at positions that are equal to i modulo j .
That is, if i =
z 1 · j + z 2 for some 0 ≤ z 2 < j , then we define the i -th bit of f ( h 0 , ..., h j − 1 ) (after the decimal point) to be equal to the z 1 -th bit of h z 2 (after the decimal point).
Note that this function is indeed injective, since for any f ( h 0 , ..., h j − 1 ) , we can uniquely reconstruct all the values h 0 , ..., h j − 1 .
Note that the numbers are only restricted to the interval (0 , 1) for the sake of simplicity; if h i > 1 , then we can encode the bits of h i by alternatingly moving in both directions from the decimal point.
As such, for any finite number of values, there exists an injective function f that combines these into a single real embedding.
Such a function f can be approximated arbitrarily closely with a multi-layer perceptron according to the universal approximation theorem, and hence an AgentNet can also have the theoretical expressiveness to learn any such function.
From here, the concrete design of the injective AgentNet implementation only requires the repeated application of this idea.
As a first step, we need to create an injective node update function f v for the single agent case; for this, we encode the values v t − 1 i and a t − 1 (and for convenience, also the value of t ) with the method above.
That is, whenever we have either v t − 1 i
[Equation]
then our update function will ensure f v ( v t − {+
+}
i
[Equation]
As described above, such a function f v exists, and hence can be learned by a sufficiently powerful AgentNet.
As a next step, the node essentially executes a message-passing phase around its neighborhood; this can again be done injectively.
In particular, the work of [68] describes how to design an injective multiset function, i.e. an aggregation of neighbors that returns a different embedding for every possible multiset of neighboring embeddings.
We can once again combine this with the embedding v ti of the node as discussed above, creating an f n that is injective in v ti and the multiset v tj ∈ N ( v i ) .
In the single-agent case, we can simply select f a ( a t − 1 i ) =
v t i , since this already contains all the information available to the agent at this point.
Intuitively, the injective implementation means that any decision that can be made by a deterministic algorithm from the same information can also be executed in our setting.
In particular, this also holds for the node transition functions: in each step, we can select the transition function f p to model almost any categorical distribution π over the closed neighborhood { v 0 } ∪ N ( v 0 ) of the current node v 0 , based on the current node state a tv 0 .
That is, given the desired probabilities π ( v j ) , inverting the softmax function determines the set of required logits, and we need to define f p accordingly.
Recall that in the injective implementation, a tv 0 already determines both v t 0 and the multiset of v tj in N ( v 0 ) ; hence when defining f p ( a ti , v tj ) , we can indeed already determine the set of desired logits from a ti alone, and assign the appropriate value for each v tj .
For a categorical distribution π to be obtainable as a transition function this way, it needs to satisfy two simple properties: (i) π ( u ) ∈ (0 , 1) for any u ∈ { v 0 } ∪ N ( v 0 ) , and (ii) if u 1 , u 2 has the same current embedding, then π ( u 1 ) =
π ( u 2 ) .
Furthermore, property (i) is also not strict in the sense that if a probability of 0 is desired, then we can also select a sufficiently small transition value such that the probability is essentially 0 after the softmax; that is, if ℓ and ∆ are both finite, then we can always set π ( u ) small enough to ensure that the probability of executing any of these “ 0 -probability transitions” at any point during the entire ℓ steps is below an arbitrary constant ϵ .
As such, the only real restriction on the set of obtainable categorical distributions is that two nodes with the same embedding must receive the same transition probabilities.
Since an injective agent ensures that every visited node has a different embedding, this only means the following restriction in practice: if v 0 has multiple unvisited neighbors with the same initial node features, then the transition probabilities to these nodes will be identical.
If any categorical distribution over the closed neighborhood satisfies this property, then there exists an f p function corresponding to this distribution, and an agent can learn to approximate this.
For completeness, we also introduce a technical modification to our injective agent implementation to ensure that the range of the function f v can never accidentally coincide with any of the (finite) possible initial values of the node features.
To do this, we encode one more finite value in the representation of f v (in every j -th bit, as before).
Since there are only finitely many possible node features, we can e.g. use the i -th of these extra bits to ensure that that f v is not equal to the i -th possible initial feature (this approach is also known as the “diagonalization method”).
This ensures that any node that is not yet visited will have a different embedding than all the already visited nodes.
Finally, let us make some simple observations about this injective implementation that can serve as building blocks for our more complex algorithms later.
First of all, we note that each node encodes the entire history of the agent from all of the time steps when the node was visited; in particular, a node is aware of each time unit when it was visited by the agent.
Furthermore, in each step, an agent can uniquely determine the neighbor of the current node that it has arrived from.
Also, for each neighbor of the current node, an agent can determine whether it has already been visited or not.
A.2 Graph traversal methods
With the injective implementation discussed above, agents can already learn to intelligently traverse the r -hop neighborhood of their starting node v in the graph.
We first discuss the iterative depth-first search approach that provides Lemma 1, and then also comment on the simpler case when a depth-first search is sufficient.
Proof of Lemma 1.
Since the agent is not aware of the size and shape of its r -hop neighborhood in advance, it needs to execute an iteratively deepening depth-first search (IDDFS) in the neighborhood to ensure that it discovers all nodes from N r ( v ) , but only these nodes.
The IDDFS algorithm iterates a depth limit d from 1 to r , and in each iteration, it executes a depth-first search from v up to depth d only.
This ensures that all the nodes at distance d from v are already identified in iteration d , but no nodes that are farther away from v are visited.
Note that we cannot achieve the same with a regular depth-first search with depth limit r : it might happen that some parts of the search tree from node u are discarded due to this depth limit, and we only find a shorter path to u later which shows that the distance from v to these discarded nodes is, in fact, smaller than r .
It is easy to see that an injective AgentNet (as described above) can indeed execute such an IDDFS algorithm; we just need to select the transition function f p appropriately.
Note that the current embedding a j of the agent already uniquely determines the current iteration d that the agent is executing; alternatively, we can also save this current value d on separate bits when constructing the injective functions of the agent (recall that we can encode any finite number of values with the same approach).
This also implies that the current embedding of each node u determines whether u was already visited in iteration d or not.
Besides this, the embedding of each (already visited) node u determines the distance of u from v : in the iteratively deepening setting, this is simply the index of the first iteration where this node was visited.
Finally, when the agent is on a node u , it can determine the predecessor of u in the depth-first search tree of the current iteration: it can consider the first time step t when u was visited in iteration d , and the predecessor of u is the node that was visited in time step ( t − 1) .
Based on these observations, we can describe the transition function f p with the following few simple rules.
When the agent is at distance at most ( d − 1) from v , and there are still neighbors that have not been visited in the current iteration, then we assign a large constant value of c to these nodes in f p , and a value of (essentially) 0 to all other nodes.
When the agent is at distance d from v , or all of its neighbors have already been visited in iteration d , then the agent moves backward: we set a value of c for the predecessor of the current node, and a value of 0 to all other nodes.
Finally, if the agent is at v and all of its neighbors have been visited in iteration d already, then it begins iteration ( d + 1) and sets the value of all neighbors to c (it selects the next neighbor uniformly at random).
Any iteration of the IDDFS traverses a depth-first search tree with at most N r ( v ) nodes; each edge of the tree is traversed at most twice, resulting in at most 2 ·| N r ( v ) | steps.
Since the number of iterations is r , our agent indeed needs ℓ ≤ O ( r · | N r ( v ) | ) .
We note that in most practical cases, we can even drop the factor r from this expression: e.g. when the graph region around v is a ∆ -regular tree, then we will have most nodes in N r ( v ) at distance exactly r from v , and thus all iterations except for the last one will become asymptotically irrelevant.
This iteratively deepening search method is only required because we always need to be aware of the distance of the current node from v in order to decide whether we need to explore the graph further in a specific direction.
However, in some special cases, we can apply a much more efficient depth-first search (DFS) traversal of the neighborhood; this removes the factor r from the required number of steps.
One simple example of such a setting is when the connected component of the graph containing v is very small.
In particular, assume that the entire component only consists of n 0 nodes; while this might be unusual in actual applications, it is relatively frequent in the synthetic datasets that are often used to measure the expressive power of different GNN extensions.
For another example, consider applications where the nodes that are important for our purpose have some specific features that make them easy to distinguish from other nodes; in this case, an injective AgentNet can learn to only consider these nodes while processing the graph (i.e. set the transition probability to all other nodes to 0 ).
As such, we can fictitiously remove every from the graph that does not have the appropriate features, and if the connected component of v in the remaining graph (the “interesting part” of the region around v ) is relatively small with only n 0 nodes, then we can again restrict our exploration strategy to this subgraph.
In the cases mentioned above, the agent can traverse the entire subgraph of size n 0 around v with a depth-first search approach: it becomes unnecessary to maintain the distances from v anymore since the agent traverses the entire connected component anyway.
This is a much more efficient method in terms of the number of steps, not requiring an extra factor r : the entire connected component around v can be traversed in 2 · n 0 − 3 steps (or in terms of the radius r of the connected component, in 2 · | N r ( v ) | − 3 steps).
This is because the DFS tree has n 0 − 1 edges, each of these is traversed at most twice, and the edges leading to the last discovered node are traversed only once (but for this we can only subtract a single edge in the worst case).
We note that this bound is indeed tight, i.e. even in this DFS-based setting, we cannot explore every neighborhood in ℓ = 2 · n 0 − 4 steps.
This also allows us to show that there are subgraphs H of radius r such that no AgentNet can decide w.h.p.
in 2 · n 0 − 4 steps whether H occurs as an induced substructure around v .
A concrete example is shown for this in the proof of the negative result in Lemma 5: the neighborhood of interest (that can contain a triangle) in this case consists [-ofn 
-]of n 0 = 1 + ∆ = 4 nodes, but it is not possible w.h.p.
in 2 · n 0 − 4 = 4 steps to decide whether v has an incident triangle.
A.3 Recognizing structures
Given an agent that systematically traverses its neighborhood, we can now consider the claims of recognizing specific substructures.
We begin with the proof of Theorem 2.
We point out that in the context of graphs with node features, we when we say that two graphs G 1 ( V 1 , E 1 ) and G 2 ( V 2 , E 2 ) are isomorphic, then besides the regular graph-theoretic definition of having σ : V 1 → V 2 such that ( v 1 , v 2 ) ∈
E 1 ⇐⇒ ( σ ( v 1 ) , σ ( v 2 ))
∈ E 2 , we also require that v 1 and σ ( v 1 ) have identical features for all v 1 ∈ V 1 .
Proof of Theorem 2.
According to Lemma 1, there exists an AgentNet that explores every node in the r -hop neighborhood of v ; furthermore, during this traversal, the agent also becomes aware of the features of every node in N r ( v ) , and all edges between pairs of nodes in N r ( v ) when the second endpoint of the edge is visited.
This allows an agent to uniquely identify the entire r -hop neighborhood around v .
More specifically, if an injective agent produces the same output embedding for two neighborhoods, then this implies that for the two traversals T 1 , T 2 , the following properties must all hold: (i) the nodes visited in the t -th and t ′ -th steps of T 1 are the same node if and only if the nodes visited in the t -th and t ′ -th steps of T 2 are the same node, (ii) the t -th visited nodes of Tand T 2 have the same node features, and (iii) the t -th and t ′ -th visited nodes of T 1 are adjacent if and only if the t -th and t ′ -th visited nodes of T 2 are adjacent.
These properties provide a clear bijection between the nodes of the two neighborhoods, also preserving node features and edges; this implies that the two neighborhoods are isomorphic.
This shows that an injective AgentNet always assigns a different embedding to non-isomorphic neighborhoods.
However, we also need to ensure that isomorphic neighborhoods, on the other hand, obtain the same embedding.
Indeed, with the injective AgentNet described so far, it could easily happen that two nodes have isomorphic neighborhoods, but an agent traverses these neighborhoods in a different order, and hence computes a different embedding in the two cases.
In order to resolve this, we only need to observe that there is a function assigning every possible
IDDFS traversal to the isomorphism class of N r ( v ) , and a sufficiently powerful agent can learn to apply this function on the final embedding.
More specifically, let G denote the set of all different (non-isomorphic) graphs of radius at most r (and degree at most ∆ ) around v .
We have already seen that if two r -hop neighborhoods are non-isomorphic, then our injective agent always produces a different final embedding for them.
This implies that the final embedding uniquely determines the graph induced by N r ( v ) , i.e. there exists a well-defined function ψ 1 : R → G which assigns the graph induced by N r ( v ) to every possible final embedding generated by the agent.
Finally, consider a function ψ 2 : G → R which assigns the numbers { 1 , . .
. , |G|} to the graphs in G in arbitrary order.
Then ψ := ψ 2 ◦ ψ 1 is simply a function ψ : R → R , and according to the universal approximation theorem, it can be learnt e.g. by an MLP implementation.
Applying this function ψ on the final embedding (in the last step of the traversal) ensures that the final output of the agent describes the isomorphism class of N r ( v ) , and hence two starting nodes receive the same final embedding if and only if their r -hop neighborhoods are isomorphic.
Proof of Corollary 3.
Corollary 3 follows easily from Theorem 2 and the fact that there are known limitations to the expressiveness of each of the listed GNN extensions.
In the case of a standard
GNN, two small cycles of different lengths are already a simple example of indistinguishability [68].
For the 3 -WL algorithm (and hence equivalently, PPGN), the Rook 4 × 4 and Shrikhande graphs are a well-known example on 16 nodes that are not distinguishable.
For GSN and DropGNN, the analysis of [53] describes an example construction that cannot be distinguished without preprocessing structures of size Θ(∆) or removing Θ(∆) nodes.
Hence for any fixed parametrization of GSN or
DropGNN (preprocessing substructures of fixed size, or removing a fixed number of nodes), there exists a pair of neighborhoods that cannot be distinguished by GSN or DropGNN.
These example graphs have ∆ + 1 nodes, so they can still be distinguished by AgentNet in ℓ = 2 · ∆ − 1 steps.
Lemma 4 follows easily from Theorem 2.
Proof of Lemma 4.
The r -hop neighborhood around v already determines the number of occurrences of any structure H of radius at most r from v .
That is, if we denote the set of all possible nonisomorphic r -hop neighborhoods around v by N r , then there is a deterministic function f H : N r → N that describes the number of occurrences of H for each neighborhood in N r .
Theorem 2 shows that we can learn an injective function
f A : N r
→
R that describes the final embedding of an agent; hence we can define the function g =
f
H ◦ f − 1 A that assigns the appropriate number of occurrences of H to any final embedding.
This function g can also be learned according to the universal approximation theorem, and hence an AgentNet can learn to execute this function in its last step.
We discuss these substructure-related claims for cliques and cycles explicitly, which are both important substructures for specific applications.
Note that the positive parts of the lemmas are deterministic, i.e. given a specific number of steps, there is an agent implementation that always (with probability 1 ) returns the desired number of subgraphs as its final embedding; meanwhile, the negative parts claim that if the number of steps is small, then no agent can return the correct number w.h.p.
Proof of Lemma 5.
As a special case of the IDDFS discussed before (with r = 1 ), an injective
AgentNet can learn to mark the starting node v , then transition to an unvisited neighbor of v in each odd step, and then move back to v in even steps.
This allows it to visit all neighbors of v in 2 · ∆ − v vvvv ′v vvvv ′v ′v ′ steps, and identify the entire 1 -hop neighborhood of v as discussed in Lemma 1.
Since every clique containing v is completely included in this induced subgraph, this AgentNet can compute the number of cliques incident to v for any clique size.
On the other hand, if ℓ ≤ 2 · ∆ − 2 , then the agent does not have enough steps to visit every neighbor of v ; as such, it might not detect a clique if the last unvisited neighbor of v is contained in it.
For a concrete example, consider two graphs where v has degree 3 (and also ∆ = 3 ); in G 1 , we have edges { ( v, v 1 ) , ( v, v 2 ) , ( v, v 3 ) , ( v 1 , v 2 ) , ( v 3 , v ′ 3 ) } , i.e. a triangle and a path of length 2 incident to v , while in G 2 , we have edges { ( v, v 1 ) , ( v, v 2 ) , ( v, v 3 ) , ( v 1 , v ′ 1 ) , ( v 2 , v ′ 2 ) , ( v 3 , v ′ 3 ) } , i.e. three paths of length 2 incident to v ; see Figure 3 for an illustration.
Assume that all nodes begin with identical node features.
In the first step, any agent can only move to a uniform random neighbor u of v (staying at the current node is not a reasonable action in this setting).
In either of the graphs, the agent observes the same situation after the first step, so it must move to the other neighbor of u with a fixed probability p 1 , and back to v with probability (1 − p 1 ) .
Note that if p 1 ≥ 12 , and if the current node u happens to be v 3 (this has probability 13 ), then the agent cannot distinguish the two graphs in the remaining steps (with probability at least 16 ).
If the agent returns to v in the second step, then in the remaining two steps, it gains the highest possible amount of information by visiting another neighbor u ′ of v and then moving to the other neighbor of u ′ .
However, once again, if u = v 3 (this has probability 13 ), then the agent cannot distinguish the two graphs from the path it has traversed.
Thus (1 − p 1 ) ≥ 12 implies a failure probability of at leastagain.
This shows for any choice of p 1 , the agent fails to distinguish the two graphs with an arbitrarily high constant probability (i.e. it higher than 56 ).
Proof of Lemma 6.
A cycle of c nodes has radius ⌊ c 2 ⌋ , so according to Lemma 4, an AgentNet can indeed count the number of incident cycles of length c in O ( c · | N ⌊ c 2 ⌋ ( v ) | ) steps.
On the other hand, if the ⌊ c 2 ⌋ -hop neighborhood of v consists only of paths of length ⌊ c 2 ⌋ starting from v , and we need to visit each of these paths to see if the endpoints of two paths are adjacent (thus forming a c -cycle), then ℓ = 2 · | N ⌊ c 2 ⌋ ( v ) | − ⌊ c 2 ⌋ − 2 steps can indeed be required: intuitively, we need to visit each path starting from v up to a distance ⌊ c 2 ⌋ , and then return to v on each occasion except for the last one.
For a concrete example, we can consider the graphs G 1 , G 2 from the proof of
Lemma 5 again: a triangle is a cycle for c = 3 , we have 2 · | N ⌊ c 2 ⌋ ( v ) | − ⌊ c 2 ⌋ − 2 = 2 · 4 − 1 − 2 =in this case, and we have already seen that that no AgentNet can distinguish these graphs w.h.p.
in less than 5 steps.
We note that for a tight lower bound in this case, we would also need to incorporate a further factor r =
⌊ c 2 ⌋ (to account for the iterative phases of the IDDFS); however, this would require a much more complex construction, since here intuitively we would also need to ensure that leaving out any of the iterative phases would result in an incorrect traversal of the neighborhood.
A.4 Random walk access model
Finally, we have already outlined the proof of Theorem 7 in Section 4.1.
Proof of Theorem 7.
The injective version of AgentNet can easily carry out all three fundamental steps of the random walk access model.
Moving to a uniform random neighbor in each step can be implemented by a simple transition function f p that assigns the same value to each neighbor.
Besides this, the agent ensures (by encoding the time step t ) that it leaves a unique node embedding at each of the visited nodes, so an injective AgentNet can recognize any of these nodes (or in fact, any possible subset of these nodes) in all subsequent steps.
In this case, an injective AgentNet can ensure in general that in the i -th step (for any finite i ), it computes an embedding that is injective in the (i) degree of the current node and (ii) set of previously visited nodes that are adjacent to the current node.
That is, in the injective implementation, after any number of steps, the agent is aware of the degree sequence observed so far and the adjacency relations between all the already discovered nodes.
Such an agent is already in possession of all the information that can be queried in the random walk access model.
As such, the output of any algorithm in this model can be expressed as a deterministic function of the final embedding of the agent, and hence due to the universal approximation theorem, an agent can approximate this output if implemented by a sufficiently expressive method (e.g. multi-layer perceptron).
B Proofs for Section 4. B.1 Simpler claims in the multi-agent setting
We now discuss the theorems on expressiveness with multiple agents.
We begin with Lemma 8, which only requires a further extension of the injective implementation provided earlier.
Proof of Lemma 8.
The lemma assumes that each agent has a unique ID from 1 to some known upper bound b (alternatively, we can also select the agent IDs from a predetermined finite set).
In this case, we can modify the injective construction described in Appendix A to ensure that the i -th agent only uses the bits at positions i modulo b for its own encoding.
This means that when summing the embeddings of different agents in f v , we again have an injective function, even if all the agents visit the same node at the same time.
As such, any single agent can determine its final embedding based on only its own part of the node embeddings, and disregard the embeddings of other agents.
After this proof, Theorem 9 only requires a few more steps.
Note that the rest of the claims in this section discuss the ability of AgentNets (with a specific k and ℓ ) to distinguish two graphs G 1 and
G 2 .
To formalize this concept, we say that an AgentNet can distinguish G 1 and G 2 if it returns a different output (in the final readout function) for the two graphs w.h.p.
, i.e. it returns a given value α 1 w.h.p.
in case of G 1 , and another value α 2 w.h.p.in case of G 2 .
We say that it cannot distinguish two graphs if this is not possible, i.e. if there is a constant upper bound on the success probability.
Proof of Theorem 9.
From Lemmas 4 and 8, it follows that there exists an AgentNet implementation such that if an agent is placed on a node v , then it computes a final embedding of 1 if v is incident to a copy of H in G , and 0 otherwise (note that we can include it in the last transition function to also convert the injective outputs into these more convenient 0 - 1 values).
Consider an aggregation of the agents that simply sums up their final embeddings; this results in a final embedding that equals to the number of starting nodes v that were incident to a copy of H .
Let γ i =
γ H
G i ) n for i ∈ { 1 , 2 } , and γ ′ =
12 · ( γ 1 − γ 2 ) .
Recall that γ 1 − γ ′ ≥ δ 2 and γ ′ − γ 2 ≥ δ 2 .
The starting point of the agents is chosen uniformly at random and independently from each other; hence if we run an AgentNet with k agents in G i , then each of these will output 1 with probability γ i and 0 otherwise, and thus our final embedding (after aggregating the agents) will follow a binomial distribution X with parameters k and γ i .
We can then use a Chernoff bound to upper bound the probability that the sum falls below (above) γ ′ · k in G 1 ( G 2 , respectively).
If the value is above γ ′ · k in G 1 and below γ ′ · k in G 2 w.h.p.
, then a simple classification function that compares the sum of embeddings to this value γ ′ · k can already distinguish the two graphs.
Let p denote the constant probability we require for our definition of w.h.p.
Let ϵ :
=
δ 2 · γ i (for a fixed i ∈ { 1 , 2 } ); then for the expected value E X of the binomial distribution, we have ϵ · E X =
δ 2 · γ i · γ i · k = δ 2 · k .
According to the Chernoff bound, we have
[Equation]
While this looks like a complicated expression, δ 2 4 · γ i 2 · γ i is simply a constant in our case.
This implies that for any choice of p , there exists a high enough constant value k such that the expression on the right-hand side is smaller than (1 − p ) .
By our choice of ϵ , having | X − E X | < ϵ · E X implies
X > γ ′ in G 1 and X < γ ′ in G 2 .
This means that the probability of having X ≤ γ ′ in G 1 is at most (1 − p ) , and similarly, the probability of having X ≥ γ ′ in G 2 is at most (1 − p ) .
We next present the proofs of Lemmas 10 and 11 that compare the single-agent and multi-agent settings.
Proof of Lemma 10.
A simple example for such a subgraph H is a path on ℓ nodes; let us assume for simplicity that the nodes have a single feature, and the value of this feature is 0 for all nodes of the path.
Consider a graph G 1 that consists of such a path H , and n − ℓ further nodes (with feature 1 ) that are all connected to the first node of the path.
If n is large enough (and k is smaller, e.g. a constant), then w.h.p.
all of the agents will start on a node with feature 1 .
As such, in ℓ steps, neither of them is able to reach the other end of the path.
This means that if we consider another graph G 2 where the only difference is that the last node of the path also has feature 0 , then any AgentNet with k agents fails to distinguish the graphs w.h.p.
In contrast to this, a single agent with at least ℓ + 1 steps can easily learn to walk to the other end of the path (by simply moving to the neighbor with feature 1 in the first step, and then always to the unmarked neighbor of feature 1 ), and hence distinguish the two graphs.
Note that if we want a version of this construction that has ∆ =
O (1) , we can simply replace the nodes added to the beginning of the path by a complete (∆ − 1) -ary tree of depth h : the root node is the first node of the path, and every node (apart from the leaves) has (∆ − 1) children.
By assigning a different feature to each level of the tree, we can ensure that wherever an agent begins in the tree, it can learn to directly walk to the root in at most h steps (we discuss these one-way trees in detail later).
Note that we have h ≤ O (log n ) in this graph.
Hence if we select ℓ = h , then we can again ensure w.h.p.
that all nodes start in the tree, and hence cannot reach the end of the path after ℓ steps.
In contrast to this, with k · ℓ steps (assuming k ≥ 2 ), a single agent can always reach the end of the path and hence distinguish the two graphs.
Proof of Lemma 11.
Let both G 1 and G 2 consist of two identical connected components of equal size; in the first component, both graphs have a node feature of 0 on all nodes, whereas, in the second component, the node features are 0 and 1 in G 1 and G 2 , respectively.
If having multiple components are undesired, we can also connect the two components with a path of length k · ℓ + 1 (with node features of, say, 2 ).
In the case of a single agent, the agent appears in the first part of the graph with probability (almost) 12 and is hence unable to distinguish the two graphs w.h.p.
, regardless of its actions.
However, in the case of a sufficiently high number of agents k (i.e. if 1 − 2 k > p is satisfied, and assuming that the connecting path is an asymptotically irrelevant part of the graph), at least one of the agents begins in the second component; in this case, they can easily recognize a node with feature 1 , and hence distinguish the two graphs.
B.2 Proof of Theorem
It remains to prove the more involved Theorem 12 regarding the comparison of the two settings.
For completeness, we first state the theorem in a more precise form:
Theorem 12 ( detailed ) .
There exists a pair of non-isomorphic graphs G 1 , G 2 and a connected subgraph H such that
• H appears as a subgraph in G 1 , but not in G 2 , • there exists an AgentNet with 2 agents and ℓ steps that can distinguish the two graphs w.h.p.
by walking through the nodes of H in G 1 , • there exists no AgentNet with 1 agent and c · ℓ steps (for any constant c ) that can distinguish the two graphs w.h.p.
That is, our proof here is a more surprising construction than that of Lemma 11: the whole graph has a diameter 2 · ℓ only, and there is a specific structure H that distinguishes the two graphs, but a single agent is unable to navigate this structure.
We first present the key idea to develop such constructions, which we call one-way trees.
Definition 13.
A one-way tree is a complete (∆ − 1) -ary tree of depth h , i.e. a rooted tree where every node (apart from the leaves) has exactly (∆ − 1) children.
The single node feature of every node in the i -th level is the integer
i (for i ∈ { 1 , ..., h } ).
Note that the node degrees are still bounded by ∆ this way.
The key idea of the tree is that an
AgentNet can easily lean to walk towards the root in this tree, but finding a specific leaf is not possible without traversing a significant portion of the tree.
Lemma 14.
Consider a one-way tree, select one of its leaf nodes u , and let us add an extra neighbor (with feature ( h + 1) ) to this leaf u .
• There exists an AgentNet implementation which, if placed on u initially, can reach the root of the tree in ℓ = h steps.
• There exists no AgentNet implementation which, if placed on the root node initially, can reach u in ℓ = 12 · (∆ − 1) h steps with probability larger than 12 .
Proof.
If an agent is placed at u initially, then in each step, it can easily distinguish the parent node of its current node in the tree, since this is the only node that has a smaller feature.
That is if the agent learns to transition to the node with feature h in the first step, with feature ( h − 1) in the second step, and so on, then it can reach the root of the tree in h steps without ambiguity (i.e. with probability 1 ).
On the other hand, when starting from the root, the agent has no way to distinguish the different children of the current node from each other.
The agent can only proceed by visiting each of the (∆ − 1) h leaves of the tree sequentially to find the leaf which has a neighbor with feature ( h + 1) .
In 12 · (∆ − 1) h steps, the agent can visit strictly less than half of the leaves, so it will only find u with probability less than 12 .
With this tool, we can already move on to the proof of Theorem 12.
Proof of Theorem 12.
Let b be a constant parameter (to be discussed later), and consider a node v with feature 0 .
Let us consider b distinct one-way trees of depth
h 1 (we will call these primary trees ), and connect the root of all these trees to v .
So far we do this in both of our graphs.
Then in the case of G 2 , let us further select one of the leaf nodes v 0 in one of the primary trees, and attach the root node of another one-way tree of depth h 2 (called the secondary tree ) to v 0 .
In contrast to this, in G 1 , we select an arbitrary leaf node in every primary tree, and we attach another one-way tree of depth h(a secondary tree) to each of these leaf nodes.
For our substructure H , let us select a path on 2 · h 1 + 3 nodes, such that the node features on the path (in order) are (1 , h 1 , h 1 − 1 , ..., 2 , 1 , 0 , 1 , 2 , ..., h 1 − 1 , h 1 , 1) .
Note that this structure appears in G 1 : we start from a root node of a secondary tree, then move to the attached leaf of the primary tree, then to the root of the primary tree, then to v , then to the root of another primary tree, then all the way down to another leaf node that has a secondary tree attached, and finally to the root of this secondary tree.
On the other hand, H does not appear as a substructure in G 2 , since v 0 is the only node in the graph which has feature h 1 and has a neighbor with feature 1 .
We select h 2 significantly larger than h 1 ; with this, we can ensure that with arbitrarily high probability, the agents are always placed in one of the secondary trees (even in G 2 , since b is only a constant).
Furthermore, note that if we have k = 2 agents, then with probability 1 − 1 b they are placed into different secondary trees initially in G 1 .
That is, by selecting the constant b large enough, we can ensure that the two agents are placed into two distinct secondary trees with a probability of at least p (i.e. w.h.p.).
Finally, we select ℓ =
h 2 + h 1 .
In the case of 2 agents, the agents can learn to always move to the root of the trees as discussed in the proof of Lemma 14: first to the root of the secondary tree, and then to the root of the primary tree.
In ℓ steps, both agents can reach v .
In the case of G 1 , if the two agents begin in different secondary trees, then they only meet at node v ; in G 2 , they will always already meet at the root of the single secondary tree in the graph.
Since the agents can learn to check if they share their location with the other agent in a given step, they can easily distinguish these two cases.
Hence they can separate the two graphs w.h.p.
: in G 1 they are correct with probability at least 1 − 1 b ≥ p , and in G 2 they are correct with probability 1 .
However, a single agent is unable to do the same.
It can walk to node v in ℓ steps in either of the graphs, but from there it cannot decide if there exists another secondary tree attached to the graph in one of the other primary trees.
More specifically, according to the second part of Lemma 14, the agent cannot find another secondary tree with probability larger than 12 in 12 · (∆ − 1) h 1 steps, so if 12 · (∆ − 1) h 1 > c · ℓ (for any given c and ∆ ≥ 3 , we can ensure this when choosing the value of h 1 ), then the agent only finds a copy of H with probability 12 at most, even in G 1 where such a subgraph exists.
Hence the agent encounters this situation (not finding another secondary tree) with a probability of at least 12 in both graphs.
This implies that it cannot distinguish the two graphs w.h.p. : whichever classification ( G 1 or G 2 ) the agent assigns to this situation, it is wrong in one of the two graphs with a probability of at least 12 .
Note that if we want an example for Theorem 12 where the two graphs have the same number of nodes (e.g. to extend it to the case when the agents are aware of n ), we need to adjust it slightly.
We can define
G ′ 1 to consist of b independent copies of our original G 1 .
Now both graphs contain b secondary trees.
Since h 2 is significantly larger than h 1 , their number of nodes is essentially equal: that is, we can add a small independent component G +2 to our original G 2 in order to ensure that G ′ 1 and G ′ 2 :
= G 2 ∪ G + have the same size (i.e. G +2 is just an arbitrary graph on ( b − 1) · (1 + b · (cid:80) h 1 − 1 i [-=
-]=0 (∆ − 1) i ) nodes).
By choosing h 2 much larger than h 1 , we can ensure that G +2 is an asymptotically irrelevant part of
G ′ 2 , i.e. it holds with an arbitrarily high constant probability that neither of the agents is placed within
G +2 .
Now given these graphs G ′ 1 and G ′ 2 , a single agent is still unable to distinguish the two cases, since they observe the same situation in both graphs with a probability of at least 12 .
However, when we have two agents, in G ′ 2 they meet w.h.p.
only at v , whereas in G ′ 1 , they always meet either already at the root of the secondary tree, or not at all; hence we can distinguish the two graphs.
C Model implementation
In this section, we will discuss all aspects of the practical implementation of our model and its possible extensions.
The model is implemented using PyTorch [55] and PyTorch Geometric [25].
C.1 General notes and initialization
We use 2-layer MLPs to parameterize all of the model update functions (node update f v , neighborhood aggregation f n , agent update f a ) and the different inputs are concatenated.
We use a Leaky ReLU [45] activation function with a negative slope of 0 . 01 .
As you will see below, we use skip connections for both node and agent embeddings.
This necessitates the use of Layer Normalization (LN) for the MLP inputs.
This makes each MLP a Pre-LN residual block [67].
Both agents and nodes have h -dimensional embedding vectors.
To ensure starting node embeddings are of the correct dimension and that their sum is injective we first pre-process all node features with an MLP.
Since for the majority of the graphs we consider we expect the agents to observe a large fraction of the graph, we process all nodes once in parallel.
To stick to the strictly sub-linear setting the model can be modified to only process node embeddings the first time they are observed.
As we need our agents to be uniquely identifiable, we use k randomly initialized and learnable h -dimensional embedding vectors, which are used as the starting agent state.
We initially place agents uniformly at random on the graph.
C.2 Model steps and possible extensions
Let’s now see, how the four model steps performed upon visiting a node we described in Sectionare implemented in practice.
Note that the same functions (MLPs) are used for every agent.
However, because agents have different embeddings, the functions can produce different outcomes for different agents.
Here, we will also show how the model can be extended with edge features and global communication between the agents.
In practice, we always include global agent communication.
If the given dataset does not make use of edge features, we use a simplified model where the pre-aggregation update functions ϕ N ( v ) → v and ϕ a → v are set to identity.
As MLPs are universal function approximators, previous MLP can already directly approximate ϕ N ( v ) → v and ϕ a → v , thus we do not lose expressiveness by excluding them [68].
One could also imagine, that providing the agents with the total number of nodes in the graph could help.
If the task, for example, is to estimate the density of triangles in the graph, by observing just a fraction of it.
Such conditioning can be trivially achieved by including n as input to every update function.
However, in our experiments, we did not consider such augmentation.
Node update.
As discussed earlier we use a skip connection.
If we have edge features, then ϕ a → v includes both the agent state and the edge e a j it took to arrive at the current node ϕ a → v (cid:0) a t − 1 j [Equation] j (cid:1) =
LeakyReLU (cid:0) LinearLayer (cid:0) LN (cid:0) a t − 1 j , e a j (cid:1)(cid:1)(cid:1) .
In this case, we use a negative slope of 0 . 2 .
If there are no edge features ϕ a → v (cid:0) a t − 1 j
[Equation]
[Equation]
If we want to have global information between the agents, we can include mean of their embeddingsk (cid:80) a j ∈ A a j as another input to f v .
As all of the subsequent operations make use of the current node state v ti , this global information can impact every other update function.
Note that since k is fixed, this mean is also injective.
To ensure all of these sum pooling operations do not cause value explosion and training problems but retain their expressiveness we implement them as mean scaled by the log of summand count.
This is also true for the neighbor aggregation.
Neighborhood aggregation.
In the same fashion, a skip connection is used, while ϕ N ( v ) → v is either identity or a linear layer followed by a non-linearity, depending on whether the graph has edge features or not.
e v j → v i is used to denote an edge from node v j to node v i .
This results in a GIN-like convolution [68] with a skip connection:
[Equation]
Agent update.
Similarly, the agent update is straightforward and can take into account the edge e a i the agent used to reach the current node if edge features are used:
[Equation]
Agent transition.
First, let us consider the Simplified AgentNet.
In this case, for each agent a i we track which nodes v j have been explored by it x ( a i , v j ) ∈ [0 , 1] .
Every time step we decay the values x = 0 .
9 · x and set the values for the current nodes V ( a i ) of all of the agents a i ∈
A to x ( a i , V ( a i ))
= 1 .
We can now use these exploration values x to construct the simplified transition function f ps .
Using agent embedding an MLP g ( a ti ) produces four logits [ g p , g c , g e , g u ] = g
( a ti ) , respectively for the previous node, the current node, explored nodes and unexplored nodes.
Then, for each neighboring node, we determine its final logits as a weighted sum of these four values.
To check if a given node v j was the agent’s a i previous node V t − 1 ( a i ) we use an indicator variablev j = V t − 1 ( a i ) .
Another indicator variable checks if the node is the agent’s current node 1 v j = V ( a i ) .
The explored and unexplored node logits g e and g u are interpolated using the node’s exploration value x ( a i , v j ) :
[Equation]
[Equation]
To ensure efficient training in the beginning we want as many agents as possible to observe the defining subgraphs.
This requires exploration.
Unfortunately, random walks are sample inefficient and require many steps to even walk a whole simple connected component e.g. a cycle.
It is known, that just preventing the random walk from backtracking already greatly increases the sample efficiency [36].
In the same vein, we initialize the learnable bias of the g ( a ti ) output layer such that [ g p , g c , g e , g u ] ≈ [0 , − 1 , 0 , 5] and the model is initially biased to focus on the yet unexplored nodes.
Note that in principle one can also restrict g ( a ti ) to be just a set of learnable global bias parameters that do not depend on the agent to produce an even simpler model.
For the full AgentNet, we use dot-product attention f p to determine the next node, where the query vector Q ( a ti ) is a linear projection of the agent embedding and the key vector K ( v ti , v tj , e v j → v i ) is a linear projection of the source node embedding, target node embedding, and any edge features.
[Equation]
h
To still ensure that the exploration is efficient at the beginning of the training, we combine dot-product attention with values produced by the Simplified AgentNet transition function f ps .
In this case, the transition logits g p , g c , g e , and g u are just global learnable bias parameters that are independent of the agent and are initialized the same way as before:
[Equation]
The resulting logits z a i → v j are used to sample a node from the neighborhood using straight-through
Gumbel Softmax:
[Equation]
To ensure the gradients flow through the Gumbel Softmax sampling, we interpret its output as a sparse one-hot vector (where only the 1s are present).
We use the resulting agent → node adjacency matrix for agent pooling in the node update step and for selecting the appropriate node in the agent update step, thus multiplying the 1s carrying the gradients with the correct embedding vectors.
C.3 Readout
As we have stated in Section 3, agent embeddings are pooled together to make the final graph-level decision.
In practice, following Xu et al.
[68] we pool agent embeddings after each agent update step and sum the individual step predictions to make the final prediction:
[Equation]
Here ϕ o is a 2-layer MLP used to project agent embeddings before the pooling.
We use both mean and max pooling because in theory there could be two kinds of problems, respectively: problems where agents collectively need to decide how commonplace certain features are and problems where it is sufficient that just one agent finds a class-defining feature.
As the final readout f o we use a simple linear layer.
C.4 Time-step conditioning
Technically, the agents are able to track time themselves, as the model is injective and they act every time step.
However, to make the task easier for the model, we condition all of the update functions (node update f v , neighborhood aggregation f n , agent update f a ) and the readout function ϕ o on the current time step.
We achieve this by adding the Transformer sinusoidal position (time step) embedding [62] to the inputs of each function (MLP).
C.5 Possible model simplifications
We implemented the model to be as flexible as possible and to match the theoretical analysis.
However, the resulting model performs quite a few operations every step.
We could simplify it at a loss of some expressiveness.
For example, node update and neighborhood aggregation steps can be merged, but this will cause the neighborhood aggregation to see the neighbors as they were in the previous time step.
Similarly, the agent update step could be merged into this unified update step, if we are content with producing the same delta update for every agent that is on the same node.
Even the agent transition probabilities could be incorporated in this unified step if for example we would model the node aggregation after the GATv2 [11] convolution and would use the final attention head to produce transition probabilities for every edge.
The same transition probabilities would then be used by every agent.
Naturally, this would reduce the model’s expressiveness, but it would result in a simpler and fully parallelized model.
However, in this work, we aimed to provide an expressive, theoretically motivated model and show that it works well in practice.
We leave the investigation of various model simplifications for future work.
D Experimental setup
For AgentNet in all of the tasks, we use AdamW optimizer [43] with a weight decay of 0 .
We set the initial learning rate to 10 − 4 and decay it over the whole training to 10 − 11 using a cosine schedule.
We also clip the global gradient norm to 1 .
In all of the cases, Gumbel-Softmax temperature is set to 23 as this has been suggested as a robust choice when the distribution has only a few categories [46].
D.1 Synthetic datasets
Expressiveness benchmarks.
To ensure all of the baseline models as well as the different AgentNet versions have a fair shot at solving the given tasks we perform a generous grid search for all of them.
We consider batch size ∈ { 50 , 300 } and hidden units ∈ { 64 , 128 } and learning rate ∈ { 0 .
For AgentNet we set number of agents k = n (16 for 4-Cycles and 3-WL, 41 for Circular Skip Links) and consider ℓ ∈ { 16 , 64 } .
For other GNN architectures, we also include ℓ ∈ { 4 , 8 } as they tend to perform worse with high depth.
On top of this for AgentNet we consider the number of agents k ∈ { 2 , n } as these problems should be solvable even with few agents.
In fact, we did observe that having k = n agents can make the model convergence slower, and for 3-WL dataset results in mean accuracy of 98 ± 2% .
We train each configuration using 10 random seeds and report the mean and the standard deviation of the best configuration after 10 thousand training steps.
For the baseline models, we use the same training setup as for AgentNet.
The baselines themselves were chosen as the most expressive models among their class of expressive GNNs.
PPGN [47] represents the higher-order GNNs and matches 3-WL in expressive power, GIN with random features [58; 1] represents non-equivariant node identification, SMP [64] represents the equivariant node identification scheme, while DropGNN [54] represents the models that use many different versions of a graph to make a prediction [54; 8; 17].
As described in Appendix C.2 the AgentNet and the Simplified AgentNet use initial attention weights biased for exploration.
We also tested these models on these synthetic benchmarks without this bias and they still successfully solved the task.
However, we kept this bias in the other experiments due to the theoretically better sample complexity.
Subgraph density ablation.
As discussed in Figure 2, we set k = 16 , ℓ = 16 , use 128 hidden units and a batch size of 200 and train using 10 random seeds for 10 thousands steps.
We report the mean and the standard deviation of accuracy at the end of the training.
While such a large batch size and a large number of hidden units are not necessary for this task, we aimed to make sure these hyperparameters are large enough to rule out any training or capacity issues in the results of this ablation study.
D.2 Real-world graph classification datasets
TU datasets.
We follow the evaluation setup by Xu et al.
[68], as was done by all of the baseline models.
We perform a 10 -fold cross-validation and report the mean and the standard deviation.
In line with baseline models, we train for 350 epochs, and for each of the datasets we perform a grid search over the batch size ∈ { 32 , 128 } , hidden units ∈ { 32 , 64 , 128 } and number of steps ℓ ∈ { 8 , 16 } .
We always set the number of agents to the mean number of nodes in the graphs of the corresponding dataset (see Table 4).
As REDDIT-BINARY has some very high-degree nodes with thousands of neighbors (Table 4) this can cause memory issues when many agents end up on the same high-degree node at the same time and compute their transition probabilities.
To avoid this, for REDDIT-BINARY we use k = 350 agents instead of 430 and consider ℓ ∈ { 4 , 8 } .
For the DD dataset, we use the usual setup, as it does not have such high-degree nodes.
Neither of the baselines was trained on DD or
REDDIT-BINARY datasets (except ESAN on REDDIT-BINARY)
, thus we trained them using the original code and hyperparameter tuning.
We also re-trained 1-2-3 GNN [49] to follow the experimental setup by Xu et al.
We use the original architecture 1 including the layer counts and a similar hyperparameter search as used for our model and by Xu et al.
The search is performed over the batch size ∈ { 32 , 64 } , the hidden unit count ∈ { 16 , 32 , 64 , 128 } and dropout ratio ∈ { 0 . 0 , 0 . 5 } .
The model is trained with Adam optimizer [35], a learning rate of 0 . 001 , and no weight decay as done originally [49]. We used a learning rate decay of 50% every 50 epochs as done by Xu et al. We found this schedule to perform slightly better than the original decay on plateau used by Morris et al. OGB.
We follow the standard evaluation setup proposed by Hu et al.
Similarly to the previous graph classification tasks, we set k to mean number of nodes ( n ≈ 26 ), considered number of steps ℓ ∈ { 8 , 16 } , batch size ∈ { 32 , 64 } and hidden units ∈ { 64 , 128 } .
We train the model with 10 random seeds for 100 epochs, select the model with the best validation ROC-AUC and report the mean and the standard deviation.
The best setup proved to be using a batch size of 64 , 128 hidden units, and ℓ = 16 steps.
In general, over most of the tasks we tested, we observed that a larger batch size tends to improve the AgentNet training and that having around 128 hidden units is a good choice, at least with our considered range of values for the number of agents k .
E Graph statistics
In Table 4 we provide the graph statistics for all of the real-world datasets used in Section 5.2.
As you can see, most of the commonly used TU datasets [50] and OGB-MolHIV [33] have small graphs.
The two large datasets we include (DD and REDDIT-BINARY) have much larger graphs, especially when considering the largest examples.
REDDIT-BINARY also has some very high-degree nodes.
To an extent, this is also true for the smaller social-graph datasets (IMDB-BINARY and IMDB-MULTI).
[Table]
F Performance on poorly aligned tasks
We want to check how our model performs in graph-level tasks it is not well aligned with.
To this end, we test our AgentNet on the molecule property regression task on the QM9 dataset [57].
Conceptually, for this molecule property regression task, we probably want to perform message passing on each node every time, as we likely need to learn the exact geometric structure of the molecule (position of every node, relative to other nodes, taking charges and bonds into account).
This would make
AgentNet (with randomly placed agents) not well suited for this task.
However, in Table 5 we can see that AgentNet still outperforms the non-expressive baselines (MPNN and 1-GNN) of similar computational complexity while performing comparably to the expressive baselines (1-2-3 GNN,
PPGN, DropMPNN, and Drop-1-GNN) which have much higher computational complexity.
This means that even on unfavorable tasks it can perform sufficiently well. [Table]
For this task, we changed the neighborhood aggregation function f n to align better with the task, by using the continuous kernel-based convolutional operator proposed by Gilmer et al.
[29] for the message pre-processing function ϕ N ( v ) → v .
This convolution is used by all of the baseline models we consider.
Similarly, to stay close to baseline models we parameterize both the final readout f
o
and the ϕ a → v function used for agent aggregation in the node update by 2-layer MLPs.
We also use the global information exchange between the agents, both in the node update step and in the agent update step.
Otherwise, the training procedure is the same as for the other tasks.