Learning How to Propagate Messages in Graph Neural Networks
This paper studies the problem of learning message propagation strategies for
graph neural networks (GNNs). One of the challenges for graph neural networks
is that of defining the propagation strategy. For instance, the choices of
propagation steps are often specialized to a single graph and are not
personalized to different nodes. To compensate for this, in this paper, we
present learning to propagate, a general learning framework that not only
learns the GNN parameters for prediction but more importantly, can explicitly
learn the interpretable and personalized propagate strategies for different
nodes and various types of graphs. We introduce the optimal propagation steps
as latent variables to help find the maximum-likelihood estimation of the GNN
parameters in a variational Expectation-Maximization (VEM) framework. Extensive
experiments on various types of graph benchmarks demonstrate that our proposed
framework can significantly achieve better performance compared with the
state-of-the-art methods, and can effectively learn personalized and
interpretable propagate strategies of messages in GNNs.
Teng Xiao, Zhengyu Chen, Donglin Wang, Suhang Wang
# Learning How to Propagate Messages in Graph Neural Networks

###### Abstract.

This paper studies the problem of learning message propagation strategies for graph neural networks (GNNs). One of the challenges for graph neural networks is that of defining the propagation strategy. For instance, the choices of propagation steps are often specialized to a single graph and are not personalized to different nodes. To compensate for this, in this paper, we present learning to propagate, a general learning framework that not only learns the GNN parameters for prediction but more importantly, can explicitly learn the interpretable and personalized propagate strategies for different nodes and various types of graphs. We introduce the optimal propagation steps as latent variables to help find the maximum-likelihood estimation of the GNN parameters in a variational Expectation-Maximization (VEM) framework. Extensive experiments on various types of graph benchmarks demonstrate that our proposed framework can significantly achieve better performance compared with the state-of-the-art methods, and can effectively learn personalized and interpretable propagate strategies of messages in GNNs. Code is available at [https://github.com/tengxiao1/L2P](https://github.com/tengxiao1/L2P).

Graph Neural Networks; Graph Representation Learning +
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
Footnote †: journal: Computing methodologies Machine learning

+
those heterophily graphs (Srivastava et al., 2017) (where connected nodes may have different class labels and dissimilar features), message propagation steps may hurt the node similarity, and stacking deep layers cannot achieve better performance compared with homophily graphs (Kumar et al., 2017; Li et al., 2018; Li et al., 2019). Since there is no strategy to learn how to propagate the message, existing GNNs need a hand-crafted layer number depending on different types of nodes and graphs. This requires expert domain knowledge and careful parameter tuning and will be sub-optimal. However, whether it is possible to learn personalized strategies while optimizing GNNs remains an open problem.

Motivated by the discussion above, in this paper, we investigate whether one can automatically learn personalized propagation strategies to help GNNs learn interpretable prediction and improve generalization. In essence, we are faced with several challenges: (i) The graph data is very complex, thus building hand-crafted and heuristic rules for propagation steps for each node tends to be infeasible when we know little knowledge underlying graph or the node distributions are too complicated. (ii) In practice, there is no way to directly access the optimal strategy of propagation. The lack of supervision about how to propagate obstructs models from modeling the distribution of propagation for each node. (iii) GNNs are also prone to be over-fitting (Srivastava et al., 2017), where the GNNs fit the training data very well but generalizes poorly to the testing data. It will suffer from the over-fitting issue more severely when we utilize an addition parameterized model for each node to learn how to propagate given limited labeled data in the real world.

To address the challenges mentioned above, we propose a simple yet effective framework called learning to propagate (L2P) to simultaneously learn the optimal propagation strategy and GNN parameters to achieve personalized and adaptive propagation. Our framework requires no heuristics and is generalizable to various types of nodes and graphs. Since there is no supervision, we adopt the principle of the probabilistic generative model and introduce the optimal propagation steps as latent variables to help find the maximum-likelihood estimation of GNN parameters in a variational Expectation-Maximization (VEM) framework. To further alleviate the over-fitting, we introduce an efficient bi-level optimization algorithm. The bi-level optimization closely matches the definition of generalization since validation data can provide accurate estimation of the generalization. The main contributions of this work are:

\(\bullet\) We study a new problem of learning propagation strategies for GNNs. To address this problem, we propose a general L2P framework which can learn personalized and interpretable propagation strategies, and achieve better performance simultaneously.

\(\bullet\) We propose an effective stochastic algorithm based on variational inference and bi-level optimization for the L2P framework, which enables simultaneously learning the optimal propagation strategies and GNN parameters, and avoiding the over-fitting issue.

\(\bullet\) We conduct experiments on homophily and heterophily graphs and the results demonstrate the effectiveness of our framework.

## 2. Related Work

### Graph Neural Networks

GNNs have achieved great success in modeling graph-structured data. Generally, GNNs can be categorized into two categories, i.e., spectral-based and spatial-based. Spectral-based GNNs define graph convolution based on spectral graph theory (Golovolov et al., 2010; Kipf and Welling, 2014; Kipf and Welling, 2014). GCN (Gori et al., 2015) further simplifies graph convolutions by stacking layers of first-order Chebyshev polynomial filters together with some approximations. Spatial-based methods directly define updating rules in the spatial space. For instance, GAT (Golovolov et al., 2010) introduces the self-attention strategy into aggregation to assign different importance scores of neighborhoods. We refer interested readers to the recent survey (Kipf and Welling, 2014) for more variants of GNN architectures. Despite the success of variants GNNs, the majority of existing GNNs aggregate neighbors' information for representation learning, which are shown to suffer from the over-smoothing (Kipf and Welling, 2014; Li et al., 2019) issue when many propagation layers are stacked, the representations of all nodes become the same.

To tackle the over-smoothing issue, some works (Gori et al., 2015; Li et al., 2019) try to add residual or dense connections (Kipf and Welling, 2014) in propagation steps for preserving the locality of the node representations. Other works (Kipf and Welling, 2014; Srivastava et al., 2017) augment the graph by randomly removing a certain number of edges or nodes to prevent the over-smoothing issue. Recently, GCNII (Kipf and Welling, 2014) introduces initial residual and identity mapping techniques for GCN and achieves promising performance. Since the feature propagation and transformation steps are commonly coupled with each other in standard GNNs, several works (Li et al., 2019; Li et al., 2019) separate this into two steps to reduce the risk of over-smoothing. We differ from these methods as (1) instead of focus on alleviating over-smoothing, we argue that different nodes and graphs may need a different number of propagation layers, and propose a framework of learning propagation strategies that generalizable to various types of graphs and backbones, and (2) we propose the bilevel optimization to utilize validation error to guide learning propagation strategy for improving the generalization ability of graph neural networks.

### The Bi-level Optimization

Bi-level optimization (Kipf and Welling, 2014), which performs upper-level learning subject to the optimality of lower-level learning, has been applied to different tasks such as few-shot learning (Kipf and Welling, 2014; Kipf and Welling, 2014; Kipf and Welling, 2014), searching architectures (Li et al., 2019), and reinforcement learning (Li et al., 2019). For the graph domain, Franceschi et al. propose a bi-level optimization objective to learn the structures of graphs. Some works (Kipf and Welling, 2014; Kipf and Welling, 2014) optimize a bi-level objective via reinforcement learning to search the architectures of GNNs. Moreover, Meta-attack (Zhang et al., 2017) adopts the principle of meta-learning to conduct the poisoning attack on the graphs by optimizing a bi-level objective. Recently, Hwang et al. propose SELAR (Huang et al., 2019) which learns the weighting function for self-supervised tasks to help the primary task on the graph with a bi-level objective. To conduct the few-shot learning in graphs, the work (Huang et al., 2019), inspired by MAML (Mamlak et al., 2017), try to obtain a parameter initialization that can adapt to unseen tasks quickly, using gradients information from the bi-level optimization. By contrast, in this paper, our main concern is the generalization, and we propose a bilevel programming with variational inference to develop a framework for learning propagation strategies, while avoiding the over-fitting issues.

## 3. Preliminaries

### Notations and Problem Definition

Let \(G=(\mathcal{V},\mathcal{E})\) denote a graph, where \(\mathcal{V}\) is a set of \(|\mathcal{V}|=N\) nodes and \(\mathcal{E}\subseteq\mathcal{V}\times\mathcal{V}\) is a set of \(|\mathcal{E}|\) edges between nodes. \(\mathbf{A}\in\{0,1\}^{N\times N}\) is the adjacency matrix of \(G\). The \((i,j)\)-th element
1 if there exists an edge between node \(v_{i}\) and \(v_{j}\), otherwise \(\mathbf{A}_{ij}=0\). Furthermore, we use \(\mathbf{X}=\left[\mathbf{x}_{1},\mathbf{x}_{2},\ldots,\mathbf{x}_{N}\right]\in \mathbb{R}^{NN\times d}\) to denote the features of nodes, where \(\mathbf{x}_{n}\) is the \(d\)-dimensional feature vector of node \(v_{n}\). Following the common semi-supervised node classification setting (Wang et al., 2017; Wang et al., 2018), only a small portion of nodes \(\mathcal{V}_{o}=\left\{v_{1},v_{2},\ldots,v_{o}\right\}\) are associated with observed labels \(\mathcal{Y}^{o}=\left\{y_{1},y_{2},\ldots,y_{o}\right\}\), where \(y_{n}\) denotes the label of \(v_{n}\). \(\mathcal{V}_{u}=\mathcal{V}\backslash\mathcal{V}_{o}\) is the set of unlabeled nodes. Given the adjacency matrix \(\mathbf{A}\), features \(\mathbf{X}\) and the observed labels \(\mathcal{Y}^{o}\), the task of node classification is to learn a function \(f_{\theta}\) which can accurately predict the labels \(\mathcal{Y}^{u}\) of unlabeled nodes \(\mathcal{V}_{u}\).

### Message Propagation

Generally, GNNs adopt the message propagation process, which iteratively aggregates the neighborhood information. Formally, the propagation process of the \(k\)-th layer in GNN is two steps:

\[\mathbf{m}_{k,n}=\textsc{AGGREGATE}\ \left(\left(\mathbf{h}_{k-1,u}:u \in\mathcal{N}(n)\right)\right) \tag{2}\] \[\mathbf{h}_{k,n}=\textsc{UPDATE}\left(\mathbf{h}_{k-1,u},\mathbf{ m}_{k,u}\mathbf{h}_{0,n}\right) \tag{1}\]

where \(\mathcal{N}_{n}\) is the set of neighbors of node \(v_{n}\), AGGREGATE is a permutation invariant function. After \(K\) message-passing layers, the final node embeddings \(\mathbf{H}_{K}\) are used to perform a given task. In general, most state-of-the-art GNN backbones (Wang et al., 2017; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018) follow this message propagation form with different AGGREGATE functions, UPDATE function, or initial feature \(\mathbf{h}_{0,n}\). For instance, APPNP (Wang et al., 2018) and GCNII (Wang et al., 2018) add the initial feature \(\mathbf{h}_{0,n}=\textsc{MLP}(\mathbf{x}_{n};\theta)\) to each layer in the UPDATE function. In general, GNN consists of several message propagation layers. We abstract the message propagation with \(K\) layers as one parameterized function \(GNN(\mathbf{X},\mathbf{A},K)\).

## 4. Learning to Propagate

In this section, we introduce the Learning to Propagate (L2P) framework, which can perform personalized message propagation for better node representation learning and a more interpretable prediction process. The key idea is to introduce a discrete latent variable \(t_{n}\) for each node \(v_{n}\), which denotes the personalized optimal propagation step of \(v_{n}\). How to learn \(t_{n}\) is challenging given no explicit supervision on the optimal propagation step of \(v_{n}\). To address the challenge, we propose a generative model for modeling the joint distribution of node labels and propagation steps conditioned on node attributes and graphs, i.e., \(p(y_{n},t_{n}|\mathbf{X},\mathbf{A})\), and formulate the Learning to Propagate framework as a variational objective, where the goal is to find the parameters of GNNs and the optimal propagation distribution, by iteratively approximating and maximizing the log-likelihood function. To alleviate the over-fitting issue, we further frame the variational process as a bi-level optimization problem, and optimize the variational parameters of learning the propagation strategies in an outer loop to maximize generalization performance of GNNs trained based on the learned strategies.

### The Generative Process

Generally, we can consider the designing process of graph neural networks as follows: we first choose the number of propagation layers \(K\) for all nodes and the type of the aggregation function parameterized by \(\theta\). Then for each training label \(y_{n}\) of node \(n\), we typically conduct the Maximum Likelihood Estimation (MLE) of the marginal log-likelihood over the observed labels as:

\[\max_{\theta}\mathcal{L}\left(\theta;\mathbf{A},\mathbf{X},\mathcal{Y}^{o} \right)=\sum\nolimits_{y_{n}\in\mathcal{Y}^{o}}\log p_{\theta}\left(y_{n}|GNN \left(\mathbf{X},\mathbf{A},K\right)\right), \tag{3}\]

where \(p_{\theta}\left(y_{n}|GNN(\mathbf{X},\mathbf{A},K)\right)=p\left(y_{n}| \mathbf{H}_{K}\right)\) is the predicted probability of node \(v_{n}\) having label \(y_{n}\) using \(\mathbf{H}_{K,n}\). \(\mathbf{H}_{K,n}\) is the node representation of \(v_{n}\) after stacking \(K\) propagation steps (see SS 3.2). Generally, a softmax is applied on \(\mathbf{H}_{K,n}\) for predicting label \(y_{n}\).

Although the message propagation strategy above has achieved promising performance, it has two drawbacks: (i) The above strategy treats each node equally, i.e., each node stacks \(K\)-layers; while in practice, different nodes may need different propagation steps/layers. Simply using a one-for-all strategy could potentially lead to sub-optimal decision boundaries and is less interpretable, and (ii) Different datasets/graphs may also have different optimal propagation steps. Existing GNNs require a hand-crafted number of propagation steps, which requires expert domain knowledge, careful parameter tuning, and is time-consuming. Thus, it would be desirable if we could learn the personalized and adaptive propagation strategy which is applicable to various types of graphs and GNN backbones.

Based on the motivation above, we propose to learn a personalized propagate distribution from the given labeled nodes and utilize the learned propagate distribution at test time, such that each test node would automatically find the optimal propagate step to explicitly improve the performance. A natural idea of learning optimal propagate distribution is supervised learning. However, there is no direct supervision of the optimal propagate strategy for each node. To solve this challenge, we treat the optimal propagation layer of each node as a discrete latent variable and adopt the principle of the probabilistic generative model, which has shown to be effective in estimating the underlying data distribution (Wang et al., 2017; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018).

Specifically, for each node \(v_{n}\), we introduce a latent discrete variable \(t_{n}\in\left\{0,1,2,\cdots,K\right\}\) to denote its optimal propagation step, where \(K\) is the predefined maximum step. Note that \(t_{n}\) can be 0, which corresponds to use non-aggregated features for prediction. We allow \(t_{n}\) to be 0, because for some nodes in heterophily graphs, the neighborhood information is noisy, aggregating the neighborhood information may result in worse performance (Wang et al., 2018; Wang et al., 2018). \(t_{n}\) is node-wise because the optimal propagation step for each node may vary largely from one node to another. With the latent variable \(\left\{t_{n}\right\}_{n=1}^{|\mathcal{V}|}\), we propose the following generative model with modeling the joint distribution of each observed label \(y_{n}\) and latent \(t_{n}\):

\[p_{\theta}(y_{n},t_{n}|\mathbf{X},\mathbf{A})=p_{\theta}(y_{n}|GNN(\mathbf{X}, \mathbf{A},t_{n}))p(t_{n}), \tag{4}\]

where \(p(t_{n})\) is the prior of propagation variable and \(\theta\) is the parameter shared by all nodes. \(p_{\theta}(y_{n}|GNN(\mathbf{X},\mathbf{A},t_{n}))\) represents the label prediction probability using \(v_{n}\)'s representation from the \(t_{n}\)-th layer, i.e., \(\mathbf{H}_{t_{n},n}\). Since we do not have any prior of how to propagate, \(p(t_{n})=\frac{1}{K+1}\) is defined as uniform distribution on all layers of all nodes in this paper. We can also use an alternative prior with lower probability on the deeper layers, if we want to encourage shallower GNNs. Given the generative model in Eq. (4) and from the Bayesian perspective, what we are interested in are two folds:

(1) Learning the parameter \(\theta\) of the GNN by maximizing the follow likelihood which helps make label prediction in the testing phase:

\[\log p_{\theta}\left(y_{n}|\mathbf{X},\mathbf{A}\right)=\log\sum\nolimits_{t_{n }\to\theta}^{K}p_{\theta}(y_{n}|GNN(\mathbf{X},\mathbf{A},t_{n}))p(t_{n}). \tag{5}\]
(2) Inferring the following posterior \(p(t_{n}|\mathbf{X},\mathbf{A},y_{n})\) of the latent variable \(t\), which is related to the optimal propagation distribution.

\[p(t_{n}=k|\mathbf{X},\mathbf{A},y_{n})=\frac{p_{\theta}(y_{n}|GNN(\mathbf{X}, \mathbf{A},k))}{\sum_{k^{\prime}=0}^{K}p_{\theta}(y_{n}|GNN(\mathbf{X},\mathbf{ A},k^{\prime}))}. \tag{6}\]

Intuitively, this posterior can be understood as we choose the propagation step \(t_{n}\) of node \(v_{n}\) according to the largest likelihood (i.e., the smallest loss) in the defined propagation steps.

However, there are several challenges to solve these two problems. For learning, we cannot directly learn the parameter \(\theta\), since it involves marginalizing the latent variable, which is generally time-consuming and intractable (Kang and Hinton, 2015). In terms of the inference, since we do not have labels for test nodes, i.e, \(y_{n}\) for \(v_{n}\in\mathcal{V}_{n}\), the non-parametric true posterior in Eq. (6), which involves evaluating the likelihood \(p_{\theta}(y_{n}|GNN(\mathbf{X},\mathbf{A},k))\) of test nodes, is not applicable. To solve the challenges in the learning and inference, we adopt the variational inference principle (Kang and Hinton, 2015; Hinton et al., 2015), and instead consider the following lower bound of the marginal log-likelihood in Eq. (5) which gives rise to our following formal variational objective:

\[\mathcal{L}(\theta,q)=\mathbb{E}_{q(t_{n})}\left[\log p_{\theta} \left(y_{n}|\mathbf{X},\mathbf{A},t_{n}\right)\right]-\text{KL}(q(t_{n})||p(t _{n})), \tag{7}\]

where the derivations are given in Appendix A.1 and \(q(t_{n})\) is the introduced variational distribution. Maximizing the ELBO \(\mathcal{L}(\theta,q)\) is equivalent to (i) maximize Eq. (5) and to (ii) make the variational distributions \(q(t_{n})\) of each node be close to its intractable true posteriors \(p(t_{n}|\mathbf{X},\mathbf{A},y_{n})\). Note that the ELBO holds for any type of variational distribution \(q(t_{n})\). We defer discussion of the learning and inference process until the next section. Here, we first introduce two ways to show how to exactly parameterize the variational distribution \(q(t_{n})\), resulting in two instances of our L2P framework.

### Learning to Select

In the variational inference principle, we can introduce a variational distribution \(q_{\phi}(t_{n}|\mathbf{v}_{n})\) parameterized by \(\mathbf{v}_{n}\in\mathbb{R}^{K}\). However, we cannot fit each \(q_{\phi}(t_{n}|\mathbf{v}_{n})\) individually by solving \(N\cdot K\) parameters, which increases the over-fitting risk given the limited labels in the graphs. Thus, we consider the amortization inference (Kang and Hinton, 2015) which avoids the optimization of the parameter \(\mathbf{v}_{n}\) for each local variational distribution \(q_{\phi}(t_{n}|\mathbf{v}_{n})\). Instead, it fits a shared neural network to calculate each local parameter \(\mathbf{v}_{n}\). Since the latent variable \(t_{n}\) is a discrete multinomial variable, the simplest and most naive way to represent categorical variable is the softmax function. Thus, we pass the features of nodes through a softmax function to parameterize the categorical propagation distribution as:

\[q_{\phi}(t_{n}=k|\mathbf{X},\mathbf{A})=\frac{\exp(\mathbf{w}_{k}^{ \top}\mathbf{H}_{k,n})}{\sum_{k^{\prime}=0}^{K}\exp(\mathbf{w}_{k^{\prime}}^{ \top}\mathbf{H}_{k^{\prime},n})}, \tag{8}\]

where \(\mathbf{w}_{k}\) represents the trainable linear transformation for the \(k\)-th layer. \(\mathbf{H}_{k,n}\) is the representation of node \(n\) at the \(k\)-th layer and \(\phi\) represents the set of parameters. The main insight behind this amortization is to reuse the propagation representation of each layer, leveraging the accumulated knowledge of representation to quickly infer propagation distribution. With amortization, we reduce the number of parameters to \((K+1)\cdot D\), where \(K\) is the predefined maximum propagation step and \(D\) is the dimension of the representation of nodes. Since this formulation directly models the selection probability overall propagation steps of each node, we refer to this method as _Learning to Select_ (L2S). Figure 2(b) gives an illustration of L2S. We adopt the node representation of \(v_{n}\) in each layer to calculate \(q_{\phi}(t_{n}=k|\mathbf{X},\mathbf{A})\), which makes it able to personally and adaptively decide which propagation layer is best for \(v_{n}\). It also allows each graph to learn its own form of propagation with its own decay form from the validation signal (see SS 5.1 for details).

### Learning to Quit

Instead of directly modeling the selection probability over every propagation step, we can model the probability of exiting the propagation process and transform the modeling of multinomial probability parameters into the modeling of the logits of binomial probability parameters. More specifically, we consider modeling the quit probability at each propagation layer for each node \(n\) as follows:

\[\alpha_{k,n}=\frac{1}{1+\exp(-\mathbf{w}_{k}^{\top}\mathbf{H}_{k,n})}, \tag{9}\]

where \(\alpha_{k,n}\) denotes the probability that node \(v_{n}\) quits propagating at the \(k\)-th layer. The challenge is how to transfer the logits \(\alpha_{k,n}\) to the multinomial variational distribution \(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})\). In this paper, we consider the stick breaking (a non-parametric Bayesian) process (Kang and Hinton, 2015) to address this challenge. Specifically, the probability of the first step (no propagation), i.e., \(q(t_{n}=0)\) is modeled as a break of proportion \(\alpha_{0,n}\) (quit at the first-step), while the length of the remainder of the propagation is left for the next break. Each probability of propagation step can be deterministically computed by the quit probability \(q(t_{n}=k)=\alpha_{k,n}\prod_{k^{\prime}=0}^{k-1}\left(1-\alpha_{k^{\prime},n}\right)\) until \(K-1\), and the probability of last propagation step is \(q(t_{n}=K)=\prod_{k^{\prime}=0}^{K-1}(1-\alpha_{k^{\prime},n})\). Assume the maximum propagation step \(K=2\), then the propagation probability is generated by 2 breaks where \(p(t_{n}=0)=\alpha_{0,n}\), \(p(t_{n}=1)=\alpha_{1,n}\left(1-\alpha_{0,n}\right)\) and the last propagation step \(p(t_{n}=2)=(1-\alpha_{1,n})\left(1-\alpha_{0,n}\right)\) (not quit until the end). Hence, for different values of \(K\), this non-parametric breaking process always satisfies \(\sum_{k=0}^{K}q(t_{n}=k)=1\). We call this method _Learning to Quit_ (L2Q). Compared with L2S, L2Q models the quit probability of each node at each propagation step via the stick-breaking process which naturally induces the sparsity property of the modeling propagation step for each node. The deeper layers are less likely to be sampled. Figure 2(c) shows the architecture of L2Q.

### Learning and Inference

Maximizing the ELBO in Eq. (7) is challenging. The lack of labels for test data further exacerbates the difficulty. Thus, in this paper, we propose two algorithms: the alternate expectation maximization

Figure 2. Illustrations of our L2P framework. (a) The vanilla GNN architecture. (b) L2S predicts the selection probability over all propagation steps for each node. (c) L2Q forces each node to personally quit its propagation process.
and iterative variational inference algorithms to maximize it.

**Alternate expectation maximization.** Minimization of the negative ELBO in Eq. (7) can be solved by the expectation maximization (EM) algorithm, which iteratively infers \(q(t_{n})\) at E-step and learns \(\theta\) at M-step. More specifically, at each iteration \(i\), given the current status parameters \(\theta^{(i)}\), the E-step that maximizes \(\mathcal{L}(\theta^{(i)},q)\) w.r.t q has a closed-form solution the same as Eq. (6):

\[q^{(i+1)}(t_{n})=q(t_{n}|\mathbf{X},\mathbf{A},y_{n})=\frac{p_{\theta^{(i)}}(y _{n}|\mathbf{X},\mathbf{A},t_{n})}{\sum_{t_{n}=0}^{K}p_{\theta^{(i)}}(y_{n}| \mathbf{X},\mathbf{A},t_{n})}. \tag{10}\]

However, we can not utilize this non-parametric posterior since the label \(y_{n}\) is not available for the test nodes. We need to let the training and testing pipeline be consistent. Thus, we consider projecting the non-parametric posterior to a parametric posterior \(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})\) (i.e., L2S or L2Q). We adopt an approximation, which has also been used in the classical wake-sleep algorithm (Srivastava et al., 2015) by minimizing the forward KL divergence \(KL(q^{(i+1)}(t_{n})||q_{\phi}(t_{n}|\mathbf{X},\mathbf{A}))\). Then we can get the following pseudo maximizing likelihood objective:

\[\phi^{(i+1)}=\arg\max_{\phi}\mathbb{E}_{q^{(i+1)}(t_{n})}\left[\log q_{\phi}(t _{n}|\mathbf{X},\mathbf{A})\right]. \tag{11}\]

Given the parametric posterior \(q_{\phi^{(i+1)}}(t_{n}|\mathbf{X},\mathbf{A})\), the M-step optimizes \(\mathcal{L}(\theta,\,q_{\phi^{(i+1)}}(t_{n}|\mathbf{X},\mathbf{A}))\) w.r.t \(\theta\). Since there is no analytical solution for deep neural networks, we update the model parameters \(\theta\) with respect to the ELBO by one step of gradient descent.

**Iterative variational inference.** Although the alternate expectation maximization algorithm is effective to infer the optimal propagation variable, the alternate EM steps are time-consuming and we need calculating the loss at every layer for each training node, i.e., the \(O(N*(K+1))\) complexity. Thus, we propose an end-to-end iterative algorithm to minimize negative ELBO. Specifically, we introduce the parameterized posterior \(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})\) (i.e., L2S or L2Q) into Eq. (7) and directly optimize ELBO using reparameterization trick (Kang et al., 2015). We infer the optimal propagation distribution \(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})\) and learn GNN weights \(\theta\) jointly through standard back-propagation from the ELBO in Eq. (7). However, the optimal propagation steps \(t\) is discrete and non-differentiable which makes direct optimization difficult. Therefore, we adopt Gumbel-Softmax Sampling (Gumbel and Softmax, 1995; Gumbel and Softmax, 1995), which is a simple yet effective way to substitutes the original non-differentiable sample from a discrete distribution with a differentiable sample from a corresponding Gumbel-Softmax distribution. Specifically, we minimize the following negative ELBO in Eq. (7) with the reparameterization trick (Kang et al., 2015):

\[\mathcal{L}(\theta,\phi)=-\log p_{\theta}(\mathbf{y}|GNN(\mathbf{X},\mathbf{A},\hat{t}))+\mathrm{KL}(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})||p(t_{n})), \tag{12}\]

where \(\hat{t}\) is drawn from a categorical distribution with the discrete variational distribution \(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})\) parameterized by \(\phi\):

\[\hat{t}_{k}=\frac{\exp((\log(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})\left[a_{k} \right])+g_{k})/y_{g})}{\sum_{k^{\prime}=0}^{K}\exp((\log(q_{\phi}(t_{n}| \mathbf{X},\mathbf{A})\left[a_{k^{\prime}}\right])+g_{k^{\prime}})/y_{g})}, \tag{13}\]

where \((g_{k^{\prime}})_{k=0}^{K}\) are i.i.d. samples drawn from the Gumbel (0, 1) distribution, \(y_{g}\) is the softmax temperature, \(\hat{t}_{k}\) is the \(k\)-th value of sample \(\hat{t}\) and \(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})[a_{k}]\) indicates the \(a_{k}\)-th index of \(q_{\phi}(t_{n}|\mathbf{X},\mathbf{A})\), i.e., the logit corresponding the \((a_{k}-1)\)-th layer. Clearly, when \(\tau>0\), the Gumbel-Softmax distribution is smooth so \(\phi\) can be optimized by standard back-propagation. The KL term in Eq. (12) is respect to two categorical distributions, thus it has a closed form.

## 5. Bi-level variational inference

So far, we have proposed the L2P framework and shown how to solve it via variational inference. However, as suggested by previous work (Kang et al., 2015; Wang et al., 2015), GNNs suffer from over-fitting due to the scarce label information in the graph domain. In this section, we propose the bilevel variational inference to alleviate the over-fitting issue.

### The Bi-level Objective

For our L2P framework, the introduced inference network for joint learning optimal propagation steps in L2S and L2Q also increases the risk of over-fitting as shown in experiments (SS 6.4). To solve the over-fitting issue, we draw inspiration from gradient-based meta-learning (learning to learn) (Kang et al., 2015; Wang et al., 2015). Briefly, the objective of \(\phi\) is to maximize the ultimate measure of the performance of GNN model \(p_{\theta}(y|GNN(\mathbf{X},\mathbf{A},t))\), which is the model performance on a held-out validation set. Formally, this goal can be formulated as the following bi-level optimization problem:

\[\min_{\phi}\mathcal{L}_{\text{val}}\left(\theta^{*}(\phi),\phi\right)\text{ s.t. }\theta^{*}(\phi)=\arg\min_{\theta}\mathcal{L}_{\text{train}}(\theta,\phi), \tag{14}\]

where \(\mathcal{L}_{\text{val}}\left(\theta^{*}(\phi),\phi\right)\) and \(\mathcal{L}_{\text{train}}(\theta,\phi)\) are called upper-level and lower-level objectives on validation and training sets, respectively. For our L2P framework, the objective is the negative ELBO \(\mathcal{L}(\theta,\phi)\) in Eq. (7). This bi-level update is to optimize the propagation strategies of each node so that the GNN model performs best on the validation set. Instead of using fixed propagation steps, it learns to assign adaptive steps while regularizing the training of a GNN model to improve the generalization. Generally, the bi-level optimization problem has to solve each level to reach a local minimum. However, calculating the optimal \(\phi\) requires two nested loops of optimization, i.e., we need to compute the optimal parameter \(\theta^{*}(\phi)\) for each \(\phi\). Thus, in order to control the computational complexity, we propose an approximate alternating optimization method by updating \(\theta\) and \(\phi\) iteratively in the next section.

### Bi-level Training Algorithm

Indeed, in general, there is no closed-form expression of \(\theta\), so it is not possible to directly optimize the upper-level objective function in Eq. (14). To tackle this challenge, we propose an alternating approximation algorithm to speed up computation in this section.

**Updating lower level \(\theta\).** Instead of solving the lower level problem completely per outer iteration, we fix \(\phi\) and only take the following gradient steps over mode parameter \(\theta\) at the \(i\)-th iteration:

\[\theta^{(i)}=\theta^{(i-1)}-\eta_{\theta}\nabla_{\theta}\mathcal{L}_{\text{ train}}(\theta^{(i-1)},\phi^{(i-1)}), \tag{15}\]

where \(\eta_{\theta}\) is the learning rate for \(\theta\).

**Updating upper level \(\phi\).** After receiving the parameter \(\theta^{(i)}\) (a reasonable approximation of \(\theta^{(*)}(\phi)\)), we can calculate the upper level objective, and update \(\phi\) through:

\[\phi^{(i)}=\phi^{(i-1)}-\eta_{\phi}\nabla_{\phi}\mathcal{L}_{\text{ val}}(\theta^{(i)},\phi^{(i-1)}). \tag{16}\]

Note that \(\theta^{(i)}\) is a function of \(\phi\) due to Eq. (15), we can directly back-propagate the gradient through \(\theta^{(i)}\) to \(\phi\). the \(\nabla_{\phi}\mathcal{L}_{\text{val}}(\theta^{(i)},\phi^{(i-1)})\) can be approximated as (see Appendix A.2 for detailed derivations):

\[\nabla_{\phi}\mathcal{L}_{\text{val}}(\theta^{(i)},\phi^{(i-1)})= \nabla_{\phi}\mathcal{L}_{\text{val}}(\bar{\theta}^{(i)},\phi^{(i-1)})\] \[-\eta_{\theta}\frac{1}{\epsilon}(\nabla_{\phi}\mathcal{L}_{\text {train}}(\theta^{(i-1)}+\epsilon\epsilon,\phi^{(i-1)})-\nabla_{\phi}\mathcal{L}_{ \text{train}}(\theta^{(i-1)},\phi^{(i-1)})), \tag{17}\]
where \(v=\nabla_{\theta}\mathcal{L}_{\text{val}}(\theta^{(i)},\hat{\phi}^{(i-1)})\), and \(\bar{\theta}^{(i)}\) and \(\hat{\phi}^{(i-1)}\) means stopping the gradient. This can be easily implemented by maintaining a shadow version of \(\theta^{(i-1)}\) at last step, catching the training loss \(\mathcal{L}_{\text{train}}(\theta^{(i-1)},\hat{\phi}^{(i-1)})\) and computing the new loss \(\mathcal{L}_{\text{train}}(\theta^{(i-1)}+\epsilon v,\phi^{(i-1)})\). When \(\eta_{\theta}\) is set to 0 in Eq. (17), the second-order derivative will disappear, resulting in a first-order approximate. In experiments in SS 6.4, we study the effect of bi-level optimization, and the first- and second-order approximates.

Given the above derivations of gradients, we have the complete L2P algorithm by alternating the update rules in Eqs. (15) and (16). The time complexity mainly depends on the bi-level optimization. For the first-order approximate, the complexity is the same as vanilla GNN methods. L2P needs approximately 3 \(\times\) training time for the second-order approximate since it needs extra forward and backward passes of the weight to compute the bilevel gradient. However, as the experiments in SS 6.4 show, the first-order approximate is sufficient to achieve the best performance.

## 6. Experiment

In this section, we conduct experiments to evaluate the effectiveness of the proposed frameworks with comparison to state-of-the-art GNNs. Specifically, we aim to answer the following questions:

* How effective is the proposed L2P framework for the node classification task on both heterophily and homophily graphs?
* Could the proposed L2P alleviate over-smoothing?
* How do the proposed learning algorithms work? Could the bi-level optimization alleviate the over-fitting issue?
* Could the proposed framework adaptively learn propagation strategies for better understanding the graph structure?
* Could the proposed L2P framework effectively the personalized and interpretable propagation strategies for each node?

### Experimental Settings

**Datasets.** We conduct experiments on both homophily and heterophily graphs. For homophily graphs, we adopts three standard citation networks for semi-supervised node classification, i.e., Cora, CiteSeer, and PubMed (Yang et al., 2017). Recent studies (Zhou et al., 2018; Liu et al., 2019; Liu et al., 2019) show that the performance of GNNs can significantly drop on heterophily graphs, we also include heterophily benchmark in our experiments, including Actor, Cornell, Texas, and Wisconsin (Zhou et al., 2018; Liu et al., 2019). The descriptions and statistics of these datasets are provided in Appendix A.3.

**Baselines.** To evaluate the effectiveness of the proposed framework, we consider the following representative and state-of-the-art GNN models on the semi-supervised node classification task. GCN (Liu et al., 2019), GAT (Liu et al., 2019), JK-Net (Liu et al., 2019), APPNP (Liu et al., 2019), DAGNN (Liu et al., 2019), IncepGCN (Liu et al., 2019), and GCNII* (Zhou et al., 2019). We also compare our proposed methods with GCN(DropEdge), ResGCN(DropEdge), JKNet(DropEdge) and IncepGCN(DropEdge) by utilizing the drop-edge trick (Liu et al., 2019). The details and implementations of baselines are given in Appendix A.4.

**Setup.** For our L2P framework, we consider APPNP as our backbone unless otherwise stated, but note that our framework is broadly applicable to more complex GNN backbones (Zhou et al., 2018; Liu et al., 2019; Liu et al., 2019). We randomly initialize the model parameters. We utilize the first-order approximate for our methods due to its efficiency and study the effect of second-order approximate separately in SS 6.4. For each search of hyper-parameter configuration, we run the experiments with 20 random seeds and select the best configuration of hyper-parameters based on average accuracy on the validation set. Hyper-parameter settings and the splitting of datasets are given in Appendix A.5.

### RQ1. Performance Comparison

To answer RQ1, we conduct experiments on both homophily and heterophily graphs with comparison to state-of-the-art methods.

**Performance on homophily graphs.** Table 1 reports the mean classification accuracy with the standard deviation on the test nodes after 10 runs. From Table 1, we have the following findings: (1) Our L2S and L2Q improve the performance of the APPAP backbone consistently and significantly in all settings. This is because that our framework has the advantage of adaptively learning personalized strategies via bi-level training. This observation demonstrates our motivation and the effectiveness of our framework. (2) Our L2S and L2Q can achieve comparable performance with state-of-the-art methods such as DAGNN and GCNII* on Cora and PubMed, and outperform them on CiteSeer. This once again demonstrates the effectiveness of our L2P framework on the node classification task. (3) In terms of our methods, the L2Q performs better than L2S, indicting that the simple softmax is not the best parameterization for the variational distribution of the latent propagation variable.

**Performance on heterophily graphs.** Besides the previously

\begin{table}
\begin{tabular}{l c c c} \hline \hline
**Method** & **Cora** & **CiteSeer** & **PubMed** \\ \hline GCN & 81.3 \(\pm\) 0.8 & 71.1 \(\pm\) 0.7 & 78.8 \(\pm\) 0.6 \\ GAT & 83.0 \(\pm\) 0.7 & 72.5 \(\pm\) 0.7 & 79.0 \(\pm\) 0.3 \\ APPNP & 83.3 \(\pm\) 0.5 & 71.8 \(\pm\) 0.5 & 79.7 \(\pm\) 0.3 \\ JKNet & 80.6 \(\pm\) 0.5 & 69.6 \(\pm\) 0.2 & 77.8 \(\pm\) 0.3 \\ JKNet(Drop) & 83.0 \(\pm\) 0.3 & 72.2 \(\pm\) 0.7 & 78.9 \(\pm\) 0.4 \\ Incep(Drop) & 83.0 \(\pm\) 0.5 & 72.3 \(\pm\) 0.4 & 79.3 \(\pm\) 0.3 \\ DAGNN & 84.2 \(\pm\) 0.5 & 73.3 \(\pm\) 0.6 & 80.3 \(\pm\) 0.4 \\ GCNII* & **85.3 \(\pm\) 0.2** & 73.2 \(\pm\) 0.8 & 80.3 \(\pm\) 0.4 \\ \hline L2S & 84.9 \(\pm\) 0.3 & 74.2 \(\pm\) 0.5 & 80.2 \(\pm\) 0.5 \\ L2Q & 85.2 \(\pm\) 0.5 & **74.6 \(\pm\) 0.4** & **80.4 \(\pm\) 0.4** \\ \hline \hline \end{tabular}
\end{table}
Table 1. Summary of results on homophily graphs. Note our results can be easily improved by using a more complex backbone. For example, by using GCNII* as our backbone, L2S can achieve 85.6 \(\pm\) 0.2 on Cora and 80.9 \(\pm\) 0.3 on PubMed.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline
**Method** & **Actor** & **Cornell** & **Texas** & **Wisconsin** \\ \hline GCN & 26.86 & 52.71 & 52.16 & 45.88 \\ GAT & 28.45 & 54.32 & 58.38 & 49.41 \\ Geom-GCN-I & 29.09 & 56.76 & 57.58 & 58.24 \\ Geom-GCN-P & 31.63 & 60.81 & 67.57 & 64.12 \\ Geom-GCN-S & 30.30 & 55.68 & 59.73 & 56.67 \\ APPNP & 32.41 & 73.51 & 65.41 & 69.02 \\ JKNet & 27.41 & 57.30 & 56.49 & 48.82 \\ JKNet(Drop) & 29.21 & 61.08 & 57.30 & 50.59 \\ Incep(Drop) & 30.13 & 61.62 & 57.84 & 50.20 \\ GCNII* & 35.18 & 76.49 & 77.84 & 81.57 \\ \hline L2S & 36.58 & 80.54 & 84.12 & 84.31 \\ L2Q & **36.97** & **81.08** & **84.56** & **84.70** \\ \hline \hline \end{tabular}
\end{table}
Table 2. Node classification accuracy on heterophily graphs.
mentioned baselines, we also compare our methods with three variants of Geom-GCN (Wang et al., 2019): Geom-GCN-I, Geom-GCN-P, and Geom-GCN-S. Table 2 reports the results. (1) We can observe that L2S and L2Q outperform the APPNP backbone on four heterophily graphs, which indicates our framework can still work well on the heterophily graphs. (2) L2S and L2Q consistently improve GCNII* by a large margin and achieve new state-of-the-art results on four heterophily graphs. (3) We can find that the improvement on heterophily graphs is usually larger than that on homophily graphs (Table 1). This is because the neighborhood information is noisy, aggregating the neighborhood information may result in worse performance for GCNII*. In contrast, our L2S and L2Q adaptively learn the process of propagation which can avoid utilizing the structure information which maybe not helpful for heterophily graphs.
the best accuracy under a depth beyond 2, which again verifies the impact of L2P on formulating graph neural networks. Notably, our methods achieve the best performance as we increase the network depth to 64 and the results of our methods remain stable with stacking many layers. On the other hand, the performance of GCN with DropEdge and JKNet drops rapidly as the number of layers exceeds 32, which represents that they suffer from over-smoothing. This phenomenon suggests that with an adaptive and personalized message propagation strategy, L2P can effectively resolve the over-smoothing problem and achieve better performance.

### RQ3. The Effect of Learning Algorithms

To answer RQ3, We first compare the performance of alternate expectation maximization (EM) and iterative variational inference (VI). From Figure 4, we can find that our methods with two learning algorithms both achieve better performance compared to the best results of APPNP, which verifies the effectiveness of our learning algorithm. In general, the iterative VI achieves better performance than the EM algorithm. We then analyze the model loss of stochastic bi-level variational inference with _training_ (we optimize \(\phi\) simultaneously with \(\theta\) on training data without validation), _first-order_ and _second-order_ approximates. Figures 3 show the learning curves of training loss and validation loss on the Texas and PubMed datasets of L2Q. We can observe that the _training_ gets stuck in the overfitting issue attaining low training loss but high validation loss. The gap between training and validation losses is much smaller for first-order and second-order. This demonstrates that the bilevel optimization can significantly improve generalization capability and the first-order approximate is sufficient to prevent the overfitting.

### RQ4. Adaptive Propagation Strategies.

One of the most important properties of our framework is that the learned propagation strategy is interpretable and is different for different types of graphs and nodes. Thus, in this subsection, we investigate if the proposed framework can learn adaptive propagation strategies, which aims to answer RQ4. We visualize the average propagation distribution (via averaging propagation distributions of all nodes) for seven graphs learned by L2S and L2Q with K=16 in Figure 5. The darkness of a step represents the probability that the step is selected for propagation. From Figure 5, we can find that (1) different types of graphs exhibit different propagation distributions although the pre-defined step is 16 for all of them. For instance, the 0-th step probability in heterophily graphs is much larger than that of homophily graphs. This is because that the feature information in those heterophily graphs is much more important than the structure information. (2) The propagation distribution learned by L2Q is much sparse, and the layers on the tail are less likely to be sampled. In Figure 6, we also provide the correlation, i.e. the cosine similarity of learned propagation distributions of different graphs. We clearly observe the correlations between the same types of graphs are large while the correlation between homophily and heterophily graphs is small, which meets our expectation that similar types of graphs should generally have similar propagation strategies.

### RQ5. Personalized Propagation Strategies

To evaluate if L2P can learn good personalized and interpretable propagation for RQ5, we study the propagation strategy for individual nodes. Figures 7 and 8 show the case studies of personalized propagation on homophily and heterophily graphs. In Figures 7 and 8, we plot the 3-hop neighborhood of each test node and use different colors to indicate different labels. We find that a test node with more same class neighbors tends to propagate few steps. In contrast, a test node with fewer same class nodes will probably have more propagation steps to predict truly its label. This observation matches our intuition that different nodes need different propagation strategies, and the prediction of a node will be confused if it has too many propagation steps and thus can not benefit much from message propagation. Additionally, we can find that our framework successfully identifies the propagation steps that are important for predicting the class of nodes on both homophily and heterophily graphs and has a more interpretable prediction process.

Figure 8. Case studies of the personalized propagation on two heterophily datasets. The bigger node in each sub-graph is the test node. The propagation distributions learned by L2Q of the test nodes are visualized with heatmaps (bottom).

Figure 7. Case studies of the personalized propagation on two homophily datasets. The bigger node in each sub-graph is the test node. The propagation distributions learned by L2Q of the test nodes are visualized with heatmaps (bottom).
## 7. Conclusion

In this paper, we study the problem of learning the propagation strategy in GNNs. We propose learning to propagate (L2P), a general framework to address this problem. Specifically, we introduce the optimal propagation steps as latent variables to help find the maximum-likelihood estimation of the GNN parameters and infer the optimal propagation step for each node via the VEM. Furthermore, we propose L2S and L2Q, two instances to parameterize the variational propagation distribution and frame the variational inference process as a bi-level optimization problem to alleviate the over-fitting problem. Extensive experiments demonstrate that our L2P can achieve state-of-the-art performance on seven benchmark datasets and adaptively capture the personalized and interpretable propagation strategies of different nodes and various graphs.

###### Acknowledgements.

 The authors would like to thank the Westlake University and Bright Dream Robotics Joint Institute for the funding support. Suhang Wang is supported by the National Science Foundation under grant number IIS-1909702, IIS1955851, and Army Research Office (ARO) under grant number W911NF-21-1-0198.
