Cooperative Graph Neural Networks
Graph neural networks are popular architectures for graph machine learning,
based on iterative computation of node representations of an input graph
through a series of invariant transformations. A large class of graph neural
networks follow a standard message-passing paradigm: at every layer, each node
state is updated based on an aggregate of messages from its neighborhood. In
this work, we propose a novel framework for training graph neural networks,
where every node is viewed as a player that can choose to either 'listen',
'broadcast', 'listen and broadcast', or to 'isolate'. The standard message
propagation scheme can then be viewed as a special case of this framework where
every node 'listens and broadcasts' to all neighbors. Our approach offers a
more flexible and dynamic message-passing paradigm, where each node can
determine its own strategy based on their state, effectively exploring the
graph topology while learning. We provide a theoretical analysis of the new
message-passing scheme which is further supported by an extensive empirical
analysis on a synthetic dataset and on real-world datasets.
Ben Finkelshtein, Xingyue Huang, Michael Bronstein, İsmail İlkan Ceylan
# Cooperative Graph Neural Networks

###### Abstract

Graph neural networks are popular architectures for graph machine learning, based on iterative computation of node representations of an input graph through a series of invariant transformations. A large class of graph neural networks follow a standard message-passing paradigm: at every layer, each node state is updated based on an aggregate of messages from its neighborhood. In this work, we propose a novel framework for training graph neural networks, where every node is viewed as a _player_ that can choose to either 'listen', 'broadcast', 'listen and broadcast', or to 'isolate'. The standard message propagation scheme can then be viewed as a special case of this framework where every node 'listens and broadcasts' to all neighbors. Our approach offers a more flexible and dynamic message-passing paradigm, where each node can determine its own strategy based on their state, effectively exploring the graph topology while learning. We provide a theoretical analysis of the new message-passing scheme which is further supported by an extensive empirical analysis on a synthetic dataset and on real-world datasets.

## 1 Introduction

Graph neural networks (GNNs) (Searselli et al., 2009; Gori et al., 2005) are a class of deep learning architectures for learning on graph-structured data. Their success in various graph machine learning tasks (Shlomi et al., 2021; Duvenaud et al., 2015; Zitnik et al., 2018) has led to a surge of different architectures (Kipf and Welling, 2017; Xu et al., 2019; Velickovic et al., 2018; Hamilton et al., 2017; Li et al., 2016). GNNs are based on an iterative computation of node representations of an input graph through a series of invariant transformations. Gilmer et al. (2017) showed that the vast majority of GNNs can be implemented through _message-passing_, where the fundamental idea is to update each node's representation based on an aggregate of messages flowing from the node's neighbors.

The message-passing paradigm has been very influential in graph ML, but it also comes with well-known limitations related to the information flow on a graph, pertaining to _long-range_ dependencies Dwivedi et al. (2022). In order to receive information from \(k\)-hop neighbors, a message-passing neural network needs at least \(k\) layers. In many types of graphs, this typically implies an exponential growth of a node's receptive field. The growing amount of information needs then to be compressed into fixed-sized node embeddings, possibly leading to information loss, referred to as _over-squashing_(Alon and Yahav, 2021).

**Motivation.** Our goal is to generalize the message-passing scheme by allowing each node to decide how to propagate information _from_ or _to_ its neighbors, thus enabling a more flexible flow of information. As a motivating example, consider the graph depicted in Figure 1 and suppose that, at every layer, the black node \(w\) needs information _only_ from the neighbors which have a yellow neighbor (node \(v\)), and hence only information from a subgraph (highlighted with gray color). Such a scenario falls outside the ability of standard message-passing schemes because there is no mechanism to condition the propagation of information on the two-hop information.

The information flow becomes more complex if the nodes have different choices across layers, since we cannot anymore view the process as merely focusing on a subgraph. This kind of message-passing is then _dynamic_ across layers. Consider the example depicted in Figure 2: the top row shows the

Figure 1: Node \(w\) listens only to \(v\). Applying the same choice iteratively, \(w\) receives information from a subgraph marked in gray.
information flow relative to the node \(u\) across three layers, and the bottom row shows the information flow relative to the node \(v\) across three layers. Node \(u\) listens to every neighbor in the first layer, only to \(v\) in the second layer, and to nodes \(s\) and \(r\) in the last layer. On the other hand, node \(v\) listens to node \(w\) for the first two layers, and to node \(u\) in the last layer.

**Main idea.** In the above examples, for each node, we need to learn whether or not to listen to a particular node in the neighborhood. To achieve this, we regard each node as a _player_ that can take the following actions in each layer:

* Standard (S): Broadcast to neighbors that listen _and_ listen to neighbors that broadcast.
* Listen (L): Listen to neighbors that broadcast.
* Broadcast (B): Broadcast to neighbors that listen.
* Isolate (I): Neither listen nor broadcast, effectively isolating the node.

The case where all nodes perform the action Standard corresponds to the standard message-passing scheme used in GNNs. Conversely, having all the nodes Isolate corresponds to removing all the edges from the graph implying node-wise predictions. The interplay between these actions and the ability to change them _locally_ and _dynamically_ makes the overall approach richer and allows to decouple the input graph from the computational one and incorporate directionality into message-passing: a node can only listen to those neighbors that are currently broadcasting, and vice versa. We can emulate the example from Figure 2 by making \(u\) choose the actions \(\langle\mathrm{L},\mathrm{L},\mathrm{S}\rangle\), \(v\) and \(w\) the actions \(\langle\mathrm{S},\mathrm{S},\mathrm{L}\rangle\), and \(s\) and \(r\) the actions \(\langle\mathrm{S},\mathrm{I},\mathrm{S}\rangle\).

**Contributions.** In this paper, we develop a new class of architectures, dubbed _cooperative graph neural networks_ (Co-GNNs), where every node in the graph is viewed as a player that can perform one of the aforementioned actions. Co-GNNs comprise two jointly trained cooperating message-passing neural networks: an _environment network_\(\eta\) (for solving the given task), and an _action network_\(\pi\) (for choosing the best actions). Our contributions can be summarized as follows:

* We propose a novel, flexible message-passing mechanism for graph neural networks, which leads to Co-GNN architectures that effectively explore the graph topology while learning (Section 4).
* We provide a detailed discussion on the properties of Co-GNNs (Section 5.1) and show that they are more expressive than 1-dimensional Weisfeiler-Leman algorithm (Section 5.2), and more importantly, better suited for long-range tasks due to their adaptive nature (Section 5.3).
* Empirically, we focus on Co-GNNs with basic action and environment networks to carefully assess the virtue of the new message-passing paradigm. We first experimentally validate the strength of our approach on a synthetic task (Section 6.1). Afterwards, we conduct experiments on real-world datasets, and observe that Co-GNNs always improve compared to their baseline models, and further yield multiple state-of-the-art results (Sections 6.2 and 6.3). We complement these with further experiments reported in the Appendix. Importantly, we illustrate the dynamic and adaptive nature of Co-GNNs (Appendix B) and also provide experiments to evaluate Co-GNNs on long-range tasks (Appendix D).

Proofs of the technical results and additional experimental details can be found in the Appendix.

Figure 2: Example information flow for two nodes \(u,v\). **Top**: information flow relative to \(u\) across three layers. Node \(u\) listens to every neighbor in the first layer, but only to \(v\) in the second layer, and only to \(s\) and \(r\) in the last layer. **Bottom**: information flow relative to \(v\) across three layers. The node \(v\) listens only to \(w\) in the first two layers, and only to \(u\) in the last layer.
Background

**Graph neural networks.** We consider simple, undirected attributed graphs \(G=(V,E,\mathbf{X})\), where \(\mathbf{X}\in\mathbb{R}^{|V|\times d}\) is a matrix of (input) node features, and \(\mathbf{x}_{v}\in\mathbb{R}^{d}\) denotes the feature of a node \(v\in V\). We focus on _message-passing neural networks (MPNNs)_(Gilmer et al., 2017) that encapsulate the vast majority of GNNs. An MPNN updates the initial node representations \(\mathbf{h}_{v}^{(0)}=\mathbf{x}_{v}\) of each node \(v\) for \(0\leq\ell\leq L-1\) iterations based on its own state and the state of its neighbors \(\mathcal{N}_{v}\) as:

\[\mathbf{h}_{v}^{(\ell+1)}=\phi^{(\ell)}\left(\mathbf{h}_{v}^{(\ell)},\psi^{(\ell)} \left(\mathbf{h}_{v}^{(\ell)},\{\mathbf{h}_{u}^{(\ell)}\mid u\in\mathcal{N}_{v}\} \right)\right),\]

where \(\{\!\!\{\cdot\}\!\}\) denotes a multiset and \(\phi^{(\ell)}\) and \(\psi^{(\ell)}\) are differentiable _update_ and _aggregation_ functions, respectively. We denote by \(d^{(\ell)}\) the dimension of the node embeddings at iteration (layer) \(\ell\). The final representations \(\mathbf{h}_{v}^{(L)}\) of each node \(v\) can be used for predicting node-level properties or they can be pooled to form a graph embedding vector \(\mathbf{z}_{G}^{(L)}\), which can be used for predicting graph-level properties. The pooling often takes the form of simple averaging, summation, or element-wise maximum. In this paper, we largely focus on the basic MPNNs of the following form:

\[\mathbf{h}_{v}^{(\ell+1)}=\sigma\left(\mathbf{W}_{s}^{(\ell)}\mathbf{h}_{v}^{(\ell)}+\mathbf{ W}_{n}^{(\ell)}\psi\left(\left\{\mathbf{h}_{u}^{(\ell)}\mid u\in\mathcal{N}_{v}\mathbf{ \mathbbm{1}}\right\}\right)\right),\]

where \(\mathbf{W}_{s}^{(\ell)}\) and \(\mathbf{W}_{n}^{(\ell)}\) are \(d^{(\ell)}\times d^{(\ell+1)}\) learnable parameter matrices acting on the node's self-representation and on the aggregated representation of its neighbors, respectively, \(\sigma\) is a non-linearity, and \(\psi\) is either _mean_ or _sum_ aggregation function. We refer to the architecture with mean aggregation as MeanGNNs and to the architecture with sum aggregation as SumGNNs (Hamilton, 2020). We also consider prominent models such as GCN (Kipf and Welling, 2017) and GIN (Xu et al., 2019).

**Straight-through Gumbel-softmax estimator.** In our approach, we rely on an action network for predicting categorical actions for the nodes in the graph, which is not differentiable and poses a challenge for gradient-based optimization. One prominent approach to address this is given by the Gumbel-softmax estimator (Jang et al., 2017; Maddison et al., 2017) which effectively provides a differentiable, continuous approximation of discrete action sampling. Consider a finite set \(\Omega\) of actions. We are interested in learning a categorical distribution over \(\Omega\), which can be represented in terms of a probability vector \(\mathbf{p}\in\mathbb{R}^{|\Omega|}\), whose elements store the probabilities of different actions. Let us denote by \(\mathbf{p}(a)\) the probability of an action \(a\in\Omega\). Gumbel-softmax is a special reparametrization trick that estimates the categorical distribution \(\mathbf{p}\in\mathbb{R}^{|\Omega|}\) with the help of a Gumbel-distributed vector \(\mathbf{g}\in\mathbb{R}^{|\Omega|}\), which stores an i.i.d. sample \(\mathbf{g}(a)\sim\textsc{Gumbel}(0,1)\) for each action \(a\). Given a categorical distribution \(\mathbf{p}\) and a temperature parameter \(\tau\), Gumbel-softmax scores can be computed as follows:

\[\mathrm{Gumbel-softmax}\left(\mathbf{p};\tau\right)=\frac{\exp\left((\log(\mathbf{p})+ \mathbf{g})/\tau\right)}{\sum_{a\in\Omega}\exp\left((\log(\mathbf{p}(a))+\mathbf{g}(a))/ \tau\right)}\]

As the softmax temperature \(\tau\) decreases, the resulting vector tends to a _one-hot_ vector. Straight-through Gumbel-softmax estimator utilizes the Gumbel-softmax estimator during the backward pass only (for a differentiable update), while during the forward pass, it employs an ordinary sampling.

## 3 Related work

Most of GNNs used in practice are instances of MPNNs (Gilmer et al., 2017) based on the message-passing approach, which has roots in classical GNN architectures (Scarselli et al., 2009; Gori et al., 2005) and their modern variations (Kipf and Welling, 2017; Xu et al., 2019; Velickovic et al., 2018; Hamilton et al., 2017; Li et al., 2016).

Despite their success, MPNNs have some known limitations. First of all, their expressive power is upper bounded by the 1-dimensional Weisfeiler-Leman graph isomorphism test (1-WL) (Xu et al., 2019; Morris et al., 2019) in that MPNNs cannot distinguish any pair of graphs which cannot be distinguished by 1-WL. This drawback has motivated the study of more expressive architectures, based on higher-order graph neural networks (Morris et al., 2019; Maron et al., 2019; Keriven and Peyre, 2019), subgraph sampling approaches (Bevilacqua et al., 2022; Thiede et al., 2021), lifting graphs to higher-dimensional topological spaces (Bodnar et al., 2021), enriching the node features with unique
identifiers (Bouritsas et al., 2022; Loukas, 2020; You et al., 2021) or random features (Abboud et al., 2021; Sato et al., 2021). Second, MPNNs generally perform poorly on long-range tasks due to their information propagation bottlenecks (Li et al., 2018; Alon and Yahav, 2021). This motivated approaches based on rewiring the input graph (Klicpera et al., 2019; Topping et al., 2022; Karhadkar et al., 2023) by connecting relevant nodes and shortening propagation distances to minimize bottlenecks, or designing new message-passing architectures that act on distant nodes directly, e.g., using shortest-path distances (Abboud et al., 2022; Ying et al., 2021). Lately, there has been a surging interest in the advancement of Transformer-based approaches for graphs (Ma et al., 2023; Ying et al., 2021; Yun et al., 2019; Kreuzer et al., 2021; Dwivedi and Bresson, 2021), which can encompass complete node connectivity beyond the local information classical MPNNs capture, which in return, allows for more effective modeling of long-range interactions. Finally, classical message passing updates the nodes in a synchronous manner, which does not allow the nodes to react to messages from their neighbors individually. This has been recently argued as yet another limitation of classical message passing from the perspective of algorithmic alignment (Faber and Wattenhofer, 2022).

Through a new message-passing scheme, our work presents new perspectives on these limitations by dynamically changing the information flow depending on the task, and resulting in more flexible architectures than the classical MPNNs, while also allowing asynchronous message passing across nodes. Our work is related to the earlier work of Lai et al. (2020), where the goal is to update each node using a potentially different number of layers, which is achieved by learning the optimal aggregation depth for each node through a reinforcement learning approach. Co-GNNs are orthogonal to this study both in terms of the objectives and the approach (as detailed in Section 5).

## 4 Cooperative graph neural networks

Co-GNNs view each node in a graph as _players_ of a multiplayer environment, where the state of each player is given in terms of the representation (or _state_) of its corresponding node. In this environment, every node is updated following a two-stage process. In the first stage, each chooses an action from the set of actions given their current state and the states of their neighboring nodes. In the second stage, every node state gets updated based on their current state and the states of a _subset_ of the neighboring nodes, as determined by the actions in the first stage. As a result, every node can determine how to propagate information from or to its neighbors.

Formally, a Co-GNN \((\pi,\eta)\) architecture is given in terms of two cooperating GNNs: (i) an action network \(\pi\) for choosing the best actions, and (ii) an environment network \(\eta\) for updating the node representations. A Co-GNN layer updates the representations \(\mathbf{h}_{v}^{(\ell)}\) of each node \(v\) as follows. First, an action network \(\pi\) predicts, for each node \(v\), a probability distribution \(\mathbf{p}_{v}^{(\ell)}\in\mathbb{R}^{4}\) over the actions \(\{\mathrm{S},\mathrm{L},\mathrm{B},\mathrm{I}\}\) that \(v\) can take, given its state and the state of its neighbors \(\mathcal{N}_{v}\), as follows:

\[\mathbf{p}_{v}^{(\ell)}=\pi\left(\mathbf{h}_{v}^{(\ell)},\{\!\!\{\mathbf{h}_{u}^{(\ell)} \mid u\in\mathcal{N}_{v}\}\!\!\}\right). \tag{1}\]

Then, for each node \(v\), an action is sampled \(a_{v}^{(\ell)}\sim\mathbf{p}_{v}^{(\ell)}\) using the Straight-through Gumbel-softmax estimator, and an environment network \(\eta\) is utilized to update the state of each node in accordance with the sampled actions, as follows:

\[\mathbf{h}_{v}^{(\ell+1)}=\begin{cases}\eta^{(\ell)}\big{(}\mathbf{h}_{v}^{(\ell)}, \{\!\!\{\mathbf{h}_{u}^{(\ell)},\!\!\{\mathbf{h}_{u}^{(\ell)}\mid u\in\mathcal{N}_{v}, a_{v}^{(\ell)}=\mathrm{S}\setminus\mathrm{B}\}\!\!\}\big{)},&a_{v}^{(\ell)}= \mathrm{I}\lor\mathrm{B}\\ \eta^{(\ell)}\big{(}\mathbf{h}_{v}^{(\ell)},\{\!\!\{\mathbf{h}_{u}^{(\ell)}\mid u\in \mathcal{N}_{v},a_{u}^{(\ell)}=\mathrm{S}\setminus\mathrm{B}\}\!\!\}\big{)},&a_{ v}^{(\ell)}=\mathrm{L}\lor\mathrm{S}.\end{cases} \tag{2}\]

This corresponds to a single layer update, and, as usual, by stacking \(L\geq 1\) layers, we obtain the representations \(\mathbf{h}_{v}^{(L)}\) for each node \(v\). In its full generality, a Co-GNN \((\pi,\eta)\) architecture can use any GNN architecture in place of the action network \(\pi\) and the environment network \(\eta\). For the sake of our study, we focus on simple models such as SumGNNs, MeanGNNs, GCN, and GIN, which are respectively denoted as \(\sum\), \(\mu\), \(*\) and \(\epsilon\). For example, we write Co-GNN\((\Sigma,\mu)\) to denote a Co-GNN architecture which uses SumGNN as its action network and MeanGNN as its environment network.

Fundamentally, Co-GNNs update the node states in a fine-grained manner as formalized in Equation (2): if a node \(v\) chooses to Isolate or to Broadcast then it gets updated only based on its previous state, which corresponds to a node-wise update function. On the other hand, if a node \(v\) chooses the action Listen or Standard then it gets updated based on its previous state as well as the state of its neighbors which perform the actions Broadcast or Standard at this layer.
## 5 Model properties

In this section, we provide a detailed analysis of Co-GNNs, focusing on its conceptual novelty, its expressive power, and its suitability to long-range tasks.

### Conceptual properties of the learnable message-passing paradigm

**Task-specific**: Standard message-passing updates nodes based on their local neighborhood, which is completely task-agnostic. By allowing each node to listen to the information only from'relevant' neighbors, Co-GNNs can determine a computation graph which is best suited for the target task. For example, if the task requires information only from the neighbors with a certain degree then the action network can learn to listen only to these nodes, as we experimentally validate in Section 6.1.

**Directed**: The outcome of the actions that the nodes can take amounts to a special form of _'directed rewiring'_ of the input graph: an edge can be _dropped_ (e.g., if two neighbors listen without broadcasting); an edge can remain _undirected_ (e.g., if both neighbors apply the standard action); or, an edge can _become directed_ implying directional information flow (e.g., if one neighbor listens while its neighbor broadcasts). Taking this perspective, the proposed message-passing can be seen as operating on a potentially different directed graph induced by the choice of actions at every layer. Formally, given a graph \(G=(V,E)\), let us denote by \(G^{(\ell)}=(V,E^{(\ell)})\) the directed computational graphs induced by the actions chosen at layer \(\ell\), where \(E^{(\ell)}\) is the set of directed edges at layer \(\ell\). We can rewrite the update given in Equation (2) concisely as follows:

\[\mathbf{h}_{v}^{(\ell+1)}=\eta^{(\ell)}\left(\mathbf{h}_{v}^{(\ell)},\{\!\!\{\mathbf{h}_ {u}^{(\ell)}\mid(u,v)\in E^{(\ell)}\}\!\!\}\right).\]

Consider the input graph \(H\) from Figure 3: \(u\) gets messages from \(v\) only in the first two layers, and \(v\) gets messages from \(u\) only in the last layer, illustrating a directional message-passing between these nodes. This abstraction allows for a direct implementation of Co-GNNs by simply considering the induced graph adjacency matrix at every layer.

**Dynamic**: In this setup, each node learns to interact with the'relevant' neighbors and does so only as long as they remain relevant: Co-GNNs do not operate on a pre-fixed computational graph, but rather on a learned computational graph, which is dynamic across layers. In our running example, observe that the computational graph is a different one at every layer (depicted in Figure 4). This brings advantages for the information flow as we outline in Section 5.3.

**Feature and structure based**: Standard message-passing is completely determined by the structure of the graph: two nodes with the same neighborhood get the same aggregated message. This is not necessarily the case in our setup, since the action network can learn different actions for two nodes with different node features, e.g., by choosing different actions for a _red_ node and a _blue_ node. This enables different messages for different nodes even if their neighborhoods are identical.

Figure 4: The computational graph of \(H\).

Figure 3: The input graph \(H\) and its computation graphs \(H^{(0)}\), \(H^{(1)}\), \(H^{(2)}\) at the respective layers. The computation graphs are a result of applying the following actions: \(\langle\operatorname{L},\operatorname{L},\operatorname{S}\rangle\) for the node \(u\); \(\langle\operatorname{S},\operatorname{S},\operatorname{L}\rangle\) for the nodes \(v\) and \(w\); \(\langle\operatorname{S},\operatorname{I},\operatorname{S}\rangle\) for the nodes \(s\) and \(r\); \(\langle\operatorname{S},\operatorname{S},\operatorname{S}\rangle\) for all other nodes.
Asynchronous**: Standard message-passing updates all nodes synchronously at every iteration, which is not always optimal as argued by Faber and Wattenhofer (2022), especially when the task requires to treat the nodes non-uniformly. By design, Co-GNNs enable asynchronous updates across nodes.

EfficientWhile being more sophisticated, the proposed message-passing algorithm is efficient in terms of runtime, as we detail in Appendix C. Co-GNNs are also parameter-efficient: they use the same policy network and as a result a comparable number of parameters to their baseline models.

### Expressive power of cooperative graph neural networks

The environment and action networks of Co-GNN architectures are parameterized by standard MPNNs. This raises an obvious question regarding the expressive power of Co-GNN architectures: _are_ Co-GNNs _also bounded by 1-WL?_ Consider, for instance, the non-isomorphic graphs \(G_{1}\) and \(G_{2}\) depicted in Figure 5. Standard MPNNs cannot distinguish these graphs, while Co-GNNs can:

**Proposition 5.1**.: _Let \(G_{1}=(V_{1},E_{1},\mathbf{X}_{1})\) and \(G_{2}=(V_{2},E_{2},\mathbf{X}_{2})\) be two non-isomorphic graphs. Then, for any threshold \(0<\delta<1\), there exists a parametrization of a Co-GNN architecture using sufficiently many layers \(L\), satisfying \(\mathbb{P}(\mathbf{z}_{G_{1}}^{(L)}\neq\mathbf{z}_{G_{2}}^{(L)})\geq 1-\delta\)._

The explanation for this result is the following: Co-GNN architectures learn, at every layer, and for each node \(u\), a probability distribution over the actions. These learned distributions are identical for two isomorphic nodes. However, the process relies on _sampling_ actions from these distributions, and clearly, the samples from identical distributions can differ. This makes Co-GNN models _invariant in expectation_, and the variance introduced by the sampling process helps to discriminate nodes that are 1-WL indistinguishable. Thus, for two nodes indistinguishable by 1-WL, there is a non-trivial probability of sampling a different action for the respective nodes, which in turn makes their direct neighborhood differ. This yields unique node identifiers (see, e.g., Loukas, 2020) with high probability and allows us to distinguish any pair of graphs assuming an injective graph pooling function (Xu et al., 2019). Our result is analogous to GNNs with random node features (Abboud et al., 2021; Sato et al., 2021), which are more expressive than their classical counterparts. Another relation is to subgraph GNNs (Bevilacqua et al., 2022; Papp et al., 2021), which bypass 1-WL expressiveness limitation by considering subgraphs extracted from \(G\). In our framework, the sampling of different actions has an analogous effect. Note that Co-GNNs are _not_ designed for more expressive power, and our result relies merely on variations in the sampling process, which should be noted as a limitation. We validate the stated expressiveness gain on a synthetic experiment in Appendix D.

### Dynamic message-passing for long-range tasks

Long-range tasks necessitate to propagate information between distant nodes: we argue that a dynamic message-passing paradigm is highly effective for such tasks since it becomes possible to propagate only relevant task-specific information. Suppose, for instance, that we are interested in transmitting information from a source node to a distant target node: our message-passing paradigm can efficiently filter irrelevant information by learning to focus on the shortest path connecting these two nodes, hence maximizing the information flow to the target node. We can generalize this observation towards receiving information from multiple distant nodes and prove the following:

**Theorem 5.2**.: _Let \(G=(V,E,\mathbf{X})\) be a connected graph with node features. For some \(k>0\), for any target node \(v\in V\), for any \(k\) source nodes \(u_{1},\ldots,u_{k}\in V\), and for any compact, differentiable function \(f:\mathbb{R}^{d^{(0)}}\times\ldots\times\mathbb{R}^{d^{(0)}}\rightarrow\mathbb{ R}^{d}\), there exists an \(L\)-layer Co-GNN computing final node representations such that for any \(\epsilon,\delta>0\) it holds that \(\mathbb{P}(|\mathbf{h}_{v}^{(L)}-f(\mathbf{x}_{u_{1}},\ldots\mathbf{x}_{u_{k}})|<\epsilon )\geq 1-\delta\)._

This means that if a property of a node \(v\) is a function of \(k\) distant nodes then Co-GNNs can approximate this function. This follows from two findings: (i) the features of \(k\) nodes can be transmitted to the source node without loss of information in Co-GNNs and (ii) the final layer of a Co-GNN architecture, e.g., an MLP, can approximate any differentiable function over \(k\) node features (Ornik, 1991; Cybenko, 1989). We validate these findings empirically on long-range interactions datasets (Dwivedi et al., 2022) in Appendix D.

Figure 5: \(G_{1}\) and \(G_{2}\) are indistinguishable by a wide class of GNNs. Co-GNNs can distinguish these graphs.
### Over-squashing and over-smoothing

**Over-squashing** refers to the failure of message passing to propagate information on the graph. Topping et al. (2022); Di Giovanni et al. (2023) formalized over-squashing as the insensitivity of an \(r\)-layer MPNN output at node \(u\) to the input features of a distant node \(v\), expressed through a bound on the Jacobian \(\|\partial\mathbf{h}_{v}^{(r)}/\partial\mathbf{x}_{u}\|\leq C^{r}(\hat{\mathbf{A}}^{r})_{vu}\), where \(C\) encapsulated architecture-related constants (e.g., width, smoothness of the activation function, etc.) and the normalized adjacency matrix \(\hat{\mathbf{A}}\) captures the effect of the graph. Graph rewiring techniques amount to modifying \(\hat{\mathbf{A}}\) so as to increase the upper bound and thereby reduce the effect of over-squashing.

Since the actions of every node in Co-GNNs result in an effective graph rewiring (different at every layer), the action network can choose actions that transmit the features of node \(u\in V\) to node \(v\in V\) as shown in Theorem 5.2, resulting in the maximization of the bound on the Jacobian.

**Over-smoothing** refers to the tendency of node embeddings to become increasingly similar across the graph with the increase in the number of message passing layers (Li et al., 2018). Recently, Rusch et al. (2022) showed that over-smoothing can be mitigated through the gradient gating mechanism, which adaptively disables the update of a node from neighbors with similar features. Our architecture, through the choice of Broadcast or Isolate actions, allows to mimic this mechanism.

## 6 Experimental results

We evaluate Co-GNNs on (i) a synthetic experiment comparing with classical MPNNs, (ii) real-world node classification datasets (Platonov et al., 2023), and (iii) real-world graph classification datasets (Morris et al., 2020). We provide a detailed analysis regarding the actions learned by Co-GNNs on different datasets in Appendix B, illustrating the task-specific nature of Co-GNNs. Finally, we report a synthetic expressiveness experiment and an experiment on long-range interactions datasets (Dwivedi et al., 2022) in Appendix D.

**Implementation.** We trained and tested our model on NVidia GTX V100 GPUs with Python 3.9, Cuda 11.8, PyTorch 2.0.0, and PyTorch Geometric (Fey and Lenssen, 2019).

### Synthetic experiment on RootNeighbors

**Task.** In this experiment, we compare Co-GNNs to a class of MPNNs on a new dataset: RootNeighbors. Specifically, we consider the following regression task: _given a rooted tree, predict the average of the features of root-neighbors of degree \(6\)_. This is an intricate task since it requires first identifying the neighbors of the root node with degree \(6\) and then returning the average feature of these nodes as the target value. RootNeighbors consists of trees of depth \(2\) with random node features of dimension \(d=5\). The data generation ensures, for each root in each tree has at least one neighbor which is of degree \(6\) (detailed in Appendix E.2). One example tree is shown in Figure 5(a): the root node \(r\) has only two neighbors with degree \(6\) (\(u\) and \(v\)) and the target prediction value is \((\mathbf{x}_{u}+\mathbf{x}_{v})/2\).

**Setup.** We trained GCN (Kipf and Welling, 2017), GraphSAGE (Hamilton et al., 2017), GAT (Velickovic et al., 2018), SumGNN, MeanGNN, as baselines and compared them to Co-GNN\((\Sigma,\Sigma)\), Co-GNN\((\mu,\mu)\) and Co-GNN\((\Sigma,\mu)\). We used the Adam optimizer and reported all hyperparameters in the appendix (Table 8). We report the Mean Average Error (MAE).

**Results for MPNNs.** The results are presented in Table 1, which includes the random baseline: the average MAE obtained via a random prediction. All MPNNs perform poorly on this task. GCN, GAT, and MeanGNN fail to identify node degrees, making it impossible to detect nodes with a

Figure 6: RootNeighbors example.
specific degree, which is crucial for the task. GCN and GAT are only marginally better than the random baseline, whereas MeanGNN performs substantially better than the random baseline. The latter can be explained by the fact that MeanGNN employs a different transformation on the source node rather than treating it as a neighbor (unlike the self-loop in GCN/GAT) and this yields better average MAE. SAGE and SumGNN use sum aggregation and can identify the node degrees, but they struggle in averaging the node features, which yield comparable MAE results to that of MeanGNN.

**Results for Co-GNNs.** The ideal mode of operation for Co-GNNs would be as follows:

1. The action network chooses either the action Listen or Standard for the root node, and the action Broadcast or Standard for the root-neighbors which have a degree \(6\),
2. The action network chooses either the action Listen or the action Isolate for all the remaining root-neighbors, and
3. The environment network updates the root node by averaging the features of its neighbors which are currently broadcasting.

This is depicted in Figure 6 and can be achieved with a single-layer Co-GNN assuming (1)-(3) can be accomplished. The best result is achieved by Co-GNN\((\Sigma,\mu)\), because SumGNN (as the action network) can accomplish (1) and (2), and MeanGNN (as the environment network) can accomplish (3). This model leverages the strengths of the SumGNN model and the MeanGNN model to cater to the different roles of the action and environment networks, making it the most natural Co-GNN model for the regression task. The next best model is Co-GNN\((\Sigma,\Sigma)\), which also uses SumGNN as the action network, accomplishing (1) and (2). However, it uses another SumGNN as the environment network which cannot easily mimic the averaging of the neighbor's features. Finally, Co-GNN\((\mu,\mu)\) performs weakly, since MeanGNN as an action network cannot achieve (1) hindering the performance of the whole task. Indeed, Co-GNN\((\mu,\mu)\) performs comparably to MeanGNN suggesting that the action network is not useful in this case.

To shed light on the performance of Co-GNN models, we computed the percentage of edges which are accurately retained or removed by the action network in a single layer Co-GNN model. We observe an accuracy of 57.20% for Co-GNN\((\mu,\mu)\), 99.55% for Co-GNN\((\Sigma,\Sigma)\), and 99.71% for Co-GNN\((\Sigma,\mu)\), which empirically confirms the expected behavior of Co-GNNs. In fact, the example tree shown in Figure 6 is taken from the RootNeighbors, and reassuringly, Co-GNN\((\Sigma,\mu)\) learns precisely the actions that induce the shown optimal subgraph.

### Node classification

One of the strengths of Co-GNNs is their capability to utilize task-specific information propagation, which raises an obvious question: could Co-GNNs outperform the baselines on heterophilious graphs, where standard message passing is known to suffer? To answer this question, we assess the performance of Co-GNNs on heterophilic node classification datasets from (Platonov et al., 2023).

**Setup.** We evaluate Co-GNN\((\Sigma,\Sigma)\) and Co-GNN\((\mu,\mu)\) on the 5 heterophilic graphs, following the 10 data splits and the methodology of Platonov et al. (2023) and report the accuracy and standard deviation for roman-empire and amazon-ratings, and mean ROC AUC and standard deviation for minesweeper, tokokers, and questions. The classical baselines GCN (Kipf and Welling, 2017), GraphSAGE (Hamilton et al., 2017), GAT (Velickovic et al., 2018), GAT-sep, GT (Shi et al., 2021) and GT-sep are from (Platonov et al., 2023). We use the Adam optimizer and report all hyperparameters in the Appendix E.4.

**Results.** All results are reported in Table 2: Co-GNNs achieve state-of-the-art across results the board, despite using relatively simple architectures as their action and environment networks. Overall, Co-GNNs demonstrate an average accuracy improvement of 2.23% compared to all baseline methods, across all datasets, surpassing the performance of more complex models such as GT. These results are reassuring as they establish Co-GNNs as a strong method in the heterophilic setting.

\begin{table}
\begin{tabular}{l c} \hline \hline Model & MAE \\ \hline Random & 0.474 \\ GCN & 0.468 \\ SAGE & 0.336 \\ GAT & 0.442 \\ SumGNN & 0.370 \\ MeanGNN & 0.329 \\ \hline Co-GNN\((\Sigma,\Sigma)\) & 0.196 \\ Co-GNN\((\mu,\mu)\) & 0.339 \\ Co-GNN\((\Sigma,\mu)\) & 0.079 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Results on RootNeighbors. Top three models are colored by First, Second, Third.
### Graph classification

In this experiment, we evaluate Co-GNNs on the TUDataset (Morris et al., 2020) graph classification benchmark.

**Setup.** We evaluate Co-GNN\((\Sigma,\Sigma)\) and Co-GNN\((\mu,\mu)\) on the 7 graph classification benchmarks, following the risk assessment protocol of Errica et al. (2020), and report the mean accuracy and standard deviation. The results for the baselines DGCNN (Wang et al., 2019), DiffPool (Ying et al., 2018), Edge-Conditioned Convolution (ECC) (Simonovsky and Komodakis, 2017), GIN (Xu et al., 2019), GraphSAGE (Hamilton et al., 2017) are from Errica et al. (2020). We also include ICGMM\({}_{f}\)(Castellana et al., 2022), and SPN\((k=5)\)(Abboud et al., 2022) as more recent baselines. OOR (Out of Resources) implies extremely long training time or GPU memory usage. We use Adam optimizer and StepLR learn rate scheduler, and report all hyperparameters in the appendix (Table 11).

**Results.** Co-GNN models achieve the highest accuracy on three datasets out of six as reported in Table 3 and remain competitive on the other datasets. Co-GNN yield these performance improvements, despite using relatively simple action and environment networks, which is intriguing as Co-GNNs unlock a large design space which includes a large class of model variations.
## 8 Reproducibility statement

To ensure the reproducibility of this paper, we include the Appendix with five main sections. Appendix A includes detailed proofs for the technical statements presented in the paper. Appendix E provides data generation protocol for RootNeighbors and further details of the real-world datasets that are used in Section 6. The experimental results in the paper are reproducible with the hyperparameter settings for all results contained in Tables 8 to 11 in Appendix E.4. The code to reproduce all experiments, along with the code to generate the datasets and tasks we propose is released at [https://anonymous.4open.science/r/CoGNN](https://anonymous.4open.science/r/CoGNN).
