A SIMPLE AND GENERAL GRAPH NEURAL NETWORK WITH STOCHASTIC MESSAGE PASSING

Abstract

Graph neural networks (GNNs) are emerging machine learning models on graphs. One key property behind the expressiveness of existing GNNs is that the learned 2 node representations are permutation-equivariant. Though being a desirable prop-3 erty for certain tasks, however, permutation-equivariance prevents GNNs from 4



being proximity-aware, i.e., preserving the walk-based proximities between pairs 5 of nodes, which is another critical property for graph analytical tasks. On 6 the other hand, some variants of GNNs are proposed to preserve node prox- proximities between nodes such as the shortest distance or high-order proximities (see Theorem 1).

29

Pairwise proximities between nodes are crucial for graph analytical tasks such as link predic-30 tion (Hu et al., 2020; You et al., 2019) . To enable a proximity-aware GNN, Position-aware GNN 31 (P-GNN) (You et al., 2019)foot_0 proposes a sophisticated GNN architecture and shows better perfor-32 mance for proximity-aware tasks. But P-GNN needs to explicitly calculate the shortest distance be-33 tween nodes and its computational complexity is unaffordable for large graphs. Moreover, P-GNN 34 completely ignores the permutation-equivariance property. Therefore, it cannot produce satisfactory 35 results when permutation-equivariance is helpful.

36

In real-world scenarios, both proximity-awareness and permutation-equivariance are indispensable 37 properties for GNNs. Firstly, different tasks may require different properties. For example, recom-38 mendation applications usually require the model to be proximity-aware (Konstas et al., 2009) while 39 permutation-equivariance is a basic assumption in centrality measurements (Borgatti, 2005) . Even



In (You et al., 2019), the authors consider the special case of shortest distance between nodes and name such property as "position-aware". In this paper, we consider a more general case of any walk-based proximity.



imities, but they fail to maintain permutation-equivariance. How to empower 8 GNNs to be proximity-aware while maintaining permutation-equivariance re-9 mains an open problem. In this paper, we propose Stochastic Message Passing 10 (SMP), a general and simple GNN to maintain both proximity-awareness and 11 permutation-equivariance properties. Specifically, we augment the existing GNNs 12 with stochastic node representations learned to preserve node proximities. Though 13 seemingly simple, we prove that such a mechanism can enable GNNs to preserve 14 node proximities in theory while maintaining permutation-equivariance with cer-15 tain parametrization. Extensive experimental results demonstrate the effectiveness 16 and efficiency of SMP for tasks including node classification and link prediction.

17

networks (GNNs), as generalizations of neural networks in analyzing graphs, have 19 attracted considerable research attention. GNNs have been widely applied to various applications 20 such as social recommendation (Ma et al., 2019), physical simulation (Kipf et al., 2018), and protein 21 interaction prediction (Zitnik & Leskovec, 2017). 22 One key property of most existing GNNs is permutation-equivariance, i.e., if we randomly permu-23 tate the IDs of nodes while maintaining the graph structure, the representations of nodes in GNNs 24 are permutated accordingly. Mathematically, permutation-equivariance reflects one basic symmet-25 ric group of graph structures. Although it is a desirable property for tasks such as node or graph 26 classification (Keriven & Peyré, 2019; Maron et al., 2019b), permutation-equivariance also prevents 27 GNNs from being proximity-aware, i.e., permutation-equivariant GNNs cannot preserve walk-based 28

annex

link prediction as an example, we observe that permutation-equivariant GNNs such as GCN (Kipf & 42 Welling, 2017) or GAT (Velickovic et al., 2018) show better results than P-GNN in coauthor graphs, 43 but the opposite in biological graphs (please see Section 5.2 for details). Unfortunately, in the current 44 GNN frameworks, these two properties are contradicting, as we show in Theorem 1. Whether there GNN (You et al., 2019) , which proposes to capture the positions of nodes using the relative 85 distance between the target node and some randomly chosen anchor nodes. However, P-GNN cannot 86 satisfy permutation-equivariance and is computationally expensive.

87

Very recently, motivated by enhancing the expressive power of GNNs in graph isomorphism tests 88 and distributed computing literature (Angluin, 1980; Linial, 1992; Naor & Stockmeyer, 1995) , 89 some studies suggest assigning unique node identifiers for GNNs (Loukas, 2020) such as one-hot 90 IDs (Murphy et al., 2019) or random numbers (Dasoulas et al., 2019; Sato et al., 2020; Corso et al., 

