GraphNorm: A Principled Approach to Accelerating Graph Neural Network Training

Introduction

Recently, there has been a surge of interest in Graph Neural Networks (GNNs) for learning with graphs (Gori et al., 2005; Scarselli et al., 2008; Hamilton et al., 2017; Kipf & Welling, 2017; Velickovic et al., 2018; Xu et al., 2018; Ying et al., 2021). GNNs learn node and graph representations by recursively aggregating and updating the node representations from neighbor representations (Gilmer et al., 2017). Empirically, GNNs have succeeded in a variety of tasks such as computational chemistry (Stokes et al., 2020), recommendation systems (Ying et al., 2018), and visual question answering (Santoro et al., 2017). Theoretically, existing works have studied GNNs through the lens of expressive power (Keriven & Peyré, 2019; Xu et al., 2019; Sato et al., 2019; Loukas, 2020; Ying et al., 2021), generalization (Scarselli et al., 2018; Du et al., 2019b; Xu et al., 2020), and extrapolation (Xu et al., 2021). However, the optimization of GNNs is less well understood, and in practice, the training of GNNs is often unstable and the convergence is slow (Xu et al., 2019).

In this paper, we study how to improve the training of GNNs via normalization. Normalization methods shift and scale the hidden representations and are shown to help the optimization for deep neural networks (Ioffe & Szegedy, 2015; Ulyanov et al., 2016; Ba et al., 2016; Salimans & Kingma, 2016; Xiong et al., 2020; Salimans et al., 2016; Miyato et al., 2018; Wu & He, 2018; Santurkar et al., 2018). Curiously, no single normalization helps in every domain, and different architectures require specialized methods. For example, Batch normalization (BatchNorm) is a standard component in computer vision (Ioffe & Szegedy, 2015); Layer normalization (LayerNorm) is popular in natural language processing (Ba et al., 2016; Xiong et al., 2020); Instance normalization (InstanceNorm) has been found effective for style transfer tasks (Ulyanov et al., 2016) . This motivates the question: What normalization methods are effective for GNNs?

We take an initial step towards answering the question above. First, we adapt the existing methods from other domains, including BatchNorm, LayerNorm, and InstanceNorm, to GNNs and evaluate their performance with extensive experiments on graph classification tasks. We observe that our adaptation of InstanceNorm to GNNs, which for each individual graph normalizes its node hidden representations, obtains much faster convergence compared to BatchNorm and LayerNorm. We provide an explanation for the success of InstanceNorm by showing that the shift operation in InstanceNorm serves as a preconditioner of the graph aggregation operation. Empirically, such preconditioning makes the optimization curvature smoother and makes the training more efficient. We also explain why the widely used BatchNorm does not bring the same level of acceleration. The variance of the batch-level statistics on graph datasets is much larger if we apply the normalization across graphs in a batch instead of across individual graphs. The noisy statistics during training may lead to unstable optimization.

Second, we show that the adaptation of InstanceNorm to GNNs, while being helpful in general, has limitations. The shift operation in InstanceNorm, which subtracts the mean statistics from node hidden representations, may lead to an expressiveness degradation for GNNs. Specifically, for highly regular graphs, the mean statistics contain graph structural information, and thus removing them could hurt the performance. Based on our analysis, we propose GraphNorm to address the issue of InstanceNorm with a learnable shift (Step 2 in Figure 1). The learnable shift could learn to control the ideal amount of information to preserve for mean statistics. Together, GraphNorm normalizes the hidden representations across nodes in each individual graph with a learnable shift to avoid the expressiveness degradation while inheriting the acceleration effect of the shift operation.

We validate the effectiveness of GraphNorm on eight popular graph classification benchmarks. Empirical results confirm that GraphNorm consistently improves the speed of converge and stability of training for GNNs compared to those with BatchNorm, InstanceNorm, LayerNorm, and those without normalization. Furthermore, GraphNorm helps GNNs achieve better generalization performance on most benchmarks.

Closely related to our work, InstanceNorm (Ulyanov et al., 2016) is originally proposed for real-time image generation. Variants of InstanceNorm are also studied in permutation equivalent data processing (Yi et al., 2018; Sun et al., 2020). We instead adapt InstanceNorm to GNNs and find it helpful for the training of GNNs. Our proposed GraphNorm builds on and improves InstanceNorm by addressing its expressiveness degradation with a learnable shift.

Few works have studied normalization in the GNN literature. Xu et al. (2019) adapts BatchNorm to GIN as a plug-in component. A preliminary version of Dwivedi et al. (2020) normalizes the node features with respect to the graph size. Our GraphNorm is size-agnostic and significantly differs from the graph size normalization. More discussions on other normalization methods are in Appendix E.

The reason behind the effectiveness of normalization has been intensively studied. While scale and shift are the main components of normalization, most existing works focus on the scale operation and the “scale-invariant” property: With a normalization layer after a linear (or convolutional) layer, the output values remain the same as the weights are scaled. Hence, normalization decouples the optimization of direction and length of the parameters (Kohler et al., 2019), implicitly tunes the learning rate (Ioffe & Szegedy, 2015; Hoffer et al., 2018; Arora et al., 2018b; Li & Arora, 2019), and smooths the optimization landscape (Santurkar et al., 2018). Our work offers a different view by instead showing specific shift operation has the preconditioning effect and can accelerate the training of GNNs.

Preliminaries

Graph Neural Networks. GNNs use the graph structure and node features to learn the representations of nodes and graphs. Modern GNNs follow a neighborhood aggregation strategy (Sukhbaatar et al., 2016; Kipf & Welling, 2017; Hamilton et al., 2017; Velickovic et al., 2018; Monti et al., 2017; Ying et al., 2021), where the representation of a node is iteratively updated by aggregating the representation of its neighbors. To be concrete, we denote $h^{(k)}_{i}$ as the representation of $v_{i}$ at the $k$-th layer and define $h_{i}^{(0)}=X_{i}$. We use AGGREGATE to denote the aggregation function in the $k$-th layer:

where $\mathcal{N}(v_{i})$ is the set of nodes adjacent to $v_{i}$. Different GNNs can be obtained by choosing different AGGREGATE functions. Graph Convolutional Networks (GCN) (Kipf & Welling, 2017) can be defined in matrix form as:

Graph Isomorphism Network (GIN) (Xu et al., 2019) is defined in matrix form as

For a $K$-layer GNN, the outputs of the final layer, i.e., $h_{i}^{(K)}$,$i=1,\cdots,n$, will be used for prediction. For graph classification tasks, we can apply a READOUT function, e.g., summation, to aggregate node features $h_{i}^{(K)}$ to obtain the entire graph’s representation h_{G}={\rm READOUT}\big{(}\big{\{}h_{i}^{(K)}\ \big{|}\ v_{i}\in V\big{\}}\big{)}. A classifier can be applied upon $h_{G}$ to predict the labels.

Normalization. Generally, given a set of values $\left\{x_{1},x_{2},\cdots,x_{m}\right\}$, a normalization operation first shifts each $x_{i}$ by the mean $\mu$, and then scales them down by standard deviation $\sigma$: $x_{i}\rightarrow\gamma\frac{x_{i}-\mu}{\sigma}+\beta$, where $\gamma$ and $\beta$ are learnable parameters, $\mu=\frac{1}{m}\sum_{i=1}^{m}x_{i}$ and $\sigma^{2}=\frac{1}{m}\sum_{i=1}^{m}\left(x_{i}-\mu\right)^{2}$. The major difference among different existing normalization methods is which set of feature values the normalization is applied to. For example, in computer vision, BatchNorm normalizes the feature values in the same channel across different samples in a batch. In NLP, LayerNorm normalizes the feature values at each position in a sequence separately.

Evaluating and Understanding Normalization for GNNs

In this section, we first adapt and evaluate existing normalization methods to GNNs. Then we give an explanation of the effectiveness of the variant of InstanceNorm, and show why the widely used BatchNorm fails to have such effectiveness. The understanding inspires us to develop better normalization methods, e.g., GraphNorm.

To investigate what normalization methods are effective for GNNs, we first adapt three typical normalization methods, i.e., BatchNorm, LayerNorm, and InstanceNorm, developed in other domain to GNNs. We apply the normalization after the linear transformation as in previous works (Ioffe & Szegedy, 2015; Xiong et al., 2020; Xu et al., 2019). The general GNN structure equipped with a normalization layer can be represented as:

We then describe the concrete operations of the adaptations of the normalization methods. Consider a batch of graphs $\left\{G_{1},\cdots,G_{b}\right\}$ where $b$ is the batch size. Let $n_{g}$ be the number of nodes in graph $G_{g}$. We generally denote $\hat{h}_{i,j,g}$ as the inputs to the normalization module, e.g., the $j$-th feature value of node $v_{i}$ of graph $G_{g}$, $i=1,\cdots,n_{g},j=1,\cdots,d,g=1,\cdots,b$. The adaptations take the general form:

where the scopes of mean $\mu$, standard deviation $\sigma$, and affine parameters $\gamma,\beta$ differ for different normalization methods. For BatchNorm, normalization and the computation of $\mu$ and $\sigma$ are applied to all values in the same feature dimension across the nodes of all graphs in the batch as in Xu et al. (2019), i.e., over dimensions $g,i$ of $\hat{h}_{i,j,g}$. To adapt LayerNorm to GNNs, we view each node as a basic component, resembling words in a sentence, and apply normalization to all feature values across different dimensions of each node, i.e., over dimension $j$ of $\hat{h}_{i,j,g}$. For InstanceNorm, we regard each graph as an instance. The normalization is then applied to the feature values across all nodes for each individual graph, i.e., over dimension $i$ of $\hat{h}_{i,j,g}$.

In Figure 2 we show training curves of different normalization methods in graph classification tasks. We find that LayerNorm hardly improves the training process in most tasks, while our adaptation of InstanceNorm can largely boost the training speed compared to other normalization methods. The test performances have similar trends. We summarize the final test accuracies in Table 1. In the following subsections, we provide an explanation for the success of InstanceNorm and its benefits compared to BatchNorm, which is currently adapted in many GNNs.

2 Shift in InstanceNorm as a Preconditioner

As mentioned in Section 1.1, the scale-invariant property of the normalization has been investigated and considered as one of the ingredients that make the optimization efficient. In our analysis of normalizations for GNNs, we instead take a closer look at the shift operation in the normalization. Compared to the image and sequential data, the graph is explicitly structured, and the neural networks exploit the structural information directly in the aggregation of the neighbors, see Eq. (1). Such uniqueness of GNNs makes it possible to study how the shift operation interplays with the graph data in detail.

We show that the shift operation in our adaptation of InstanceNorm serves as a preconditioner of the aggregation in GNNs and hypothesize this preconditioning effect can boost the training of GNNs. Though the current theory of deep learning has not been able to prove and compare the convergence rate in the real settings, we calculate the convergence rate of GNNs on a simple but fully characterizable setting to give insights on the benefit of the shift operation.

Let $Q,N$ be defined as in Eq. (6), $0\leq\lambda_{1}\leq\cdots\leq\lambda_{n}$ be the singular values of $Q$. We have $\mu_{n}=0$ is one of the singular values of $QN$, and let other singular values of $QN$ be $0\leq\mu_{1}\leq\mu_{2}\leq\cdots\leq\mu_{n-1}$. Then we have

where $\lambda_{i}=\mu_{i}$ or $\lambda_{i}=\mu_{i-1}$ only if there exists one of the right singular vectors $\alpha_{i}$ of $Q$ associated with $\lambda_{i}$ satisfying $\mathbf{1}^{\top}\alpha_{i}=0$.

We hypothesize that precoditioning $Q$ can help the optimization. In the case of optimizing the weight matrix $W^{(k)}$, we can see from Eq. (6) that after applying normalization, the term $Q$ in the gradient of $W^{(k)}$ will become $QN$ which makes the optimization curvature of $W^{(k)}$ smoother, see Appendix A.5 for more discussions. Similar preconditioning effects are believed to improve the training of deep learning models (Duchi et al., 2011; Kingma & Ba, 2015), and classic wisdom in optimization has also shown that preconditioning can accelerate the convergence of iterative methods (Axelsson, 1985; Demmel, 1997). Unfortunately, current theoretical toolbox only has a limited power on the optimization of deep learning models. Global convergence rates have only been proved for either simple models, e.g., linear models (Arora et al., 2018a), or extremely overparameterized models (Du et al., 2018; Allen-Zhu et al., 2019; Du et al., 2019a; Cai et al., 2019; Du et al., 2019b; Zou et al., 2020). To support our hypothesis that preconditioning may suggest better training, we investigate a simple but characterizable setting of training a linear GNN using gradient descent in Appendix A.2. In this setting, we prove that:

where $\rho_{1}$ is the convergence rate.

which indicates that the model with shift converges faster than the vanilla model.

To check how much the matrix $N$ improves the distribution of the spectrum of matrix $Q$ in real practice, we sample graphs from different datasets for illustration, as showed in Figure 3 (more visualizations for different types of graph can be found in Appendix D.1). We can see that the singular value distribution of $QN$ is much smoother, and the condition number is improved. Note that for a multi-layer GNN, the normalization will be applied in each layer. Therefore, the overall improvement of such preconditioning can be more significant.

3 Heavy Batch Noise in Graphs Makes BatchNorm Less Effective

The above analysis shows the adaptation of InstanceNorm has the effect of preconditioning the aggregation of GNNs. Then a natural question is whether a batch-level normalization for GNNs (Xu et al., 2019) has similar advantages. We show that BatchNorm is less effective in GNNs due to heavy batch noise on graph data.

In BatchNorm, the mean $\mu_{B}$ and standard deviation $\sigma_{B}$ are calculated in a sampled batch during training, which can be viewed as random variables by the randomness of sampling. During testing, the estimated dataset-level statistics (running mean $\mu_{D}$ and standard deviation $\sigma_{D}$) are used instead of the batch-level statistics (Ioffe & Szegedy, 2015). To apply Theorem 3.1 to BatchNorm for the preconditioning effect, one could potentially view all graphs in a dataset as subgraphs in a super graph. Hence, Theorem 3.1 applies to BatchNorm if the batch-level statistics are well-concentrated around dataset-level statistics, i.e., $\mu_{B}\approx\mu_{D}$ and $\sigma_{B}\approx\sigma_{D}$. However, the concentration of batch-level statistics is heavily domain-specific. While Shen et al. (2020) find the variation of batch-level statistics in typical networks is small for computer vision, the concentration of batch-level statistics is still unknown for GNNs.

We study how the batch-level statistics $\mu_{B},\sigma_{B}$ deviate from the dataset-level statistics $\mu_{D},\sigma_{D}$. For comparison, we train a 5-layer GIN with BatchNorm on the PROTEINS dataset and train a ResNet18 (He et al., 2016) on the CIFAR10 dataset. We set batch size to 128. For each epoch, we record the batch-level max/min mean and standard deviation for the first and the last BatchNorm layer on a randomly selected dimension across batches. In Figure 4, pink line denotes the dataset-level statistics, and green/blue line denotes the max/min value of the batch-level statistics. We observe that for image tasks, the maximal deviation of the batch-level statistics from the dataset-level statistics is negligible (Figure 4) after a few epochs. In contrast, for the graph tasks, the variation of batch-level statistics stays large during training. Intuitively, the graph structure can be quite diverse and the a single batch cannot well represent the entire dataset. Hence, the preconditioning property also may not hold for BatchNorm. In fact, the heavy batch noise may bring instabilities to the training. More results may be found in Appendix D.2.

Graph Normalization

Although we provide evidence on the indispensability and advantages of our adaptation of InstanceNorm, simply normalizing the values in each feature dimension within a graph does not consistently lead to improvement. We show that in some situations, e.g., for regular graphs, the standard shift (e.g., shifting by subtracting the mean) may cause information loss on graph structures.

We consider $r$-regular graphs, i.e., each node has a degree $r$. We first look into the case that there are no available node features, then $X_{i}$ is set to be the one-hot encoding of the node degree (Xu et al., 2019). In a $r$-regular graph, all nodes have the same encoding, and thus the columns of $H^{(0)}$ are the same. We study the output of the standard shift operation in the first layer, i.e., $k=1$ in Eq. (6). From the following proposition, we can see that when the standard shift operation is applied to GIN for a $r$-regular graph described above, the information of degree is lost:

Such information loss not only happens when there are no node features. For complete graphs, we can further show that even each node has different features, the graph structural information, i.e., adjacency matrix $A$, will always be ignored after the standard shift operation in GIN:

This analysis inspires us to modify the current normalization method with a learnable parameter to automatically control how much the mean to preserve in the shift operation. Combined with the graph-wise normalization, we name our new method Graph Normalization, i.e., GraphNorm. For each graph $G$, we generally denote value $\hat{h}_{i,j}$ as the inputs to GraphNorm, e.g., the $j$-th feature value of node $v_{i}$, $i=1,\cdots,n$, $j=1,\cdots,d$. GraphNorm takes the following form:

where $\mu_{j}=\frac{\sum_{i=1}^{n}\hat{h}_{i,j}}{n},\hat{\sigma}_{j}^{2}=\frac{\sum_{i=1}^{n}\left(\hat{h}_{i,j}-\alpha_{j}\cdot\mu_{j}\right)^{2}}{n}$, and $\gamma_{j},\beta_{j}$ are the affine parameters as in other normalization methods. By introducing the learnable parameter $\alpha_{j}$ for each feature dimension $j$, we are able to learn how much the information we need to keep in the mean. It is easy to see that GraphNorm has stronger expressive power than InstanceNorm. Formally, we have the following fact:

To validate our theory and the proposed GraphNorm in real-world data, we conduct an ablation study on two typical datasets, PROTEINS and IMDB-BINARY. As shown in Figure 5, the graphs from PROTEINS and IMDB-BINARY exhibit irregular-type and regular-type graphs, respectively. We train GIN/GCN using our adaptation of InstanceNorm and GraphNorm under the same setting in Section 5. The training curves are presented in Figure 5. The curves show that using a learnable $\alpha$ slightly improves the convergence on PROTEINS, while significantly boost the training on IMDB-BINARY. This observation verify that shifting the feature values by subtracting the mean may lose information, especially for regular graphs. And the introduction of learnable shift in GraphNorm can effectively mitigate the expressive degradation.

Experiments

In this section, we evaluate and compare both the training and test performance of GraphNorm with other normalization methods on graph classification benchmarks.

We use eight popularly used benchmark datasets of different scales in the experiments (Yanardag & Vishwanathan, 2015; Xu et al., 2019), including four medium-scale bioinformatics datasets (MUTAG, PTC, PROTEINS, NCI1), three medium-scale social network datasets (IMDB-BINARY, COLLAB, REDDIT-BINARY), and one large-scale bioinformatics dataset ogbg-molhiv, which is recently released on Open Graph Benchmark (OGB) (Hu et al., 2020). Dataset statistics are summarized in Table 1. We use two typical graph neural networks GIN (Xu et al., 2019) and GCN (Kipf & Welling, 2017) for our evaluations. Specifically, we use a five-layer GCN/GIN. For GIN, the number of sub-layers in MLP is set to 2. Normalization is applied to each layer. To aggregate global features on top of the network, we use SUM readout for MUTAG, PTC, PROTEINS and NCI1 datasets, and use MEAN readout for other datasets, as in Xu et al. (2019). Details of the experimental settings are presented in Appendix C.

We plot the training curves of GIN with GraphNorm and other normalization methodsThe graph size normalization in the preliminary version of Dwivedi et al. (2020) does not show significant improvement on the training and test performance, so we do not report it. on different tasks in Figure 2. The results on GCN show similar trends, and are provided in Appendix D.3. As shown in Figure 2, GraphNorm enjoys the fastest convergence on all tasks. Compared to BatchNorm used in Xu et al. (2019), GraphNorm converges in roughly 5000/500 iterations on NCI1 and PTC datasets, while the model using BatchNorm does not even converge in 10000/1000 iterations. Remarkably, though InstanceNorm does not outperform other normalization methods on IMDB-BINARY, GraphNorm with learnable shift significantly boosts the training upon InstanceNorm and achieves the fastest convergence. We also validate the test performance and report the test accuracy in Table 1,2. The results show that GraphNorm also improves the generalization on most benchmarks.

For reference, we explain the possible reasons of higher test accuracy in two folds. First, as shown in Figure 2, using proper normalization helps the model find a minimum with a higher training accuracy. Second, as suggested by Hardt et al. (2016), faster training leads to smaller generalization gap. Since the test accuracy equals the training accuracy plus the generalization, these two views together suggest better normalization leads to better test performance.

1 Ablation Study

In this subsection, we summarize the results of some ablation studies, including BatchNorm with learnable shift, BatchNorm with running statistics and the effect of batch size. Due to the space limitation, the detailed results can be found in Appendix D.

We conduct experiments on BatchNorm to investigate whether simply introducing a learnable shift can already improve the existing normalization methods without concrete motivation of overcoming expressiveness degradation. Specifically, we equip BatchNorm with a similar learnable shift as GraphNorm and evaluate its performance. We find that the learnable shift cannot further improve upon BatchNorm (See Appendix D), which suggests the introduction of learnable shift in GraphNorm is critical.

We study the variant of BatchNorm which uses running statistics to replace the batch-level mean and standard deviation (Similar idea is also proposed in Yan et al. (2019)). At first glance, this method may seem to be able to mitigate the problem of large batch noise. However, the running statistics change a lot during training, and using running statistics disables the model to back-propagate the gradients through mean and standard deviation. Results in Appendix D show this variant has even worse performance than BatchNorm.

We further compare the GraphNorm with BatchNorm with different batch sizes (8, 16, 32, 64). As shown in Appendix D, our GraphNorm consistently outperforms the BatchNorm on all the settings.

Conclusion and Future Work

In this paper, we adapt and evaluate three well-used normalization methods, i.e., BatchNorm, LayerNorm, and InstanceNorm to GNNs. We give explanations for the successes and failures of these adaptations. Based on our understanding of the strengths and limitations of existing adaptations, we propose Graph Normalization, that builds upon the adaptation of InstanceNorm with a learnable shift to overcome the expressive degradation of the original InstanceNorm. Experimental results show GNNs with GraphNorm not only converge faster, but also achieve better generalization performance on several benchmark datasets.

Though seeking theoretical understanding of normalization methods in deep learning is challenging (Arora et al., 2018b) due to limited understanding on the optimization of deep learning models and characterization of real world data, we take an initial step towards finding effective normalization methods for GNNs with theoretical guidance in this paper. The proposed theories and hypotheses are motivated by several simple models. And we are not able to give concrete theoretical results to problems such as: the convergence rate of general GNNs with normalization, the spectrum of $Q$ normalized by learnable shift, etc. We believe the analyses of more realistic but complicated settings, e.g., the dynamics of GraphNorm on deep GNNs, are good future directions.

Acknowledgements

We thank Mozhi Zhang and Ruosong Wang for helpful suggestions on the paper; Zhiyuan Li and Kaifeng Lyu for helpful discussion on the literature of normalization methods; and Prof. Yang Yuan for support of computational resources. This work was supported by National Key R&D Program of China (2018YFB1402600), Key-Area Research and Development Program of Guangdong Province (No. 2019B121204008), BJNSF (L172037), Beijing Academy of Artificial Intelligence, Project 2020BD006 supported by PKU-Baidu Fund, NSF CAREER award (1553284) and NSF III (1900933).

References

Appendix A Proofs

We first introduce the Cauchy interlace theorem:

Using Lemma A.1, the theorem can be proved as below.

where $\mathbf{0}=\left[\underbrace{0,\cdots,0}_{\times n-1}\right]^{\top}$.

Using Lemma A.1 and taking $R=U^{\top}\bar{C}U$ and $S=U_{1}^{\top}\bar{C}U_{1}$, we have the eigenvalues of $U_{1}^{\top}\bar{C}U_{1}$ are interlacing between the eigenvalues of $U^{\top}\bar{C}U$. Note that the eigenvalues of $DU^{\top}\bar{C}UD$ are $\mu_{1}^{2}\leq\mu_{2}^{2}\leq\cdots\leq\mu_{n-1}^{2}$ and $\mu_{n}^{2}=0$, and by Eq. (18), the eigenvalues of $DU^{\top}\bar{C}UD$ contain the eigenvalues of $U_{1}^{\top}\bar{C}U_{1}$ and . Since the eigenvalues of $U^{\top}\bar{C}U$ are $\lambda_{1}^{2}\leq\lambda_{2}^{2}\leq\cdots\leq\lambda_{n}^{2}$ (By similarity of $U^{\top}\bar{C}U$ and $\bar{C}$), we then have

Since $U_{1}^{\top}U_{1}=I_{n-1}$, the equation above is equivalent to

i.e., $U_{1}z$ is the eigenvector of $\bar{C}$ associated with $\mu$. By noticing $U_{1}z$ lies in the orthogonal complement space of $\mathbf{1}$ and the eigenvector of $\bar{C}$ is right singular vector of ${Q}$, we complete the proof. ∎

A.2 Concrete example of the acceleration

To get more intuition on how the preconditioning effect of the shift can accelerate the training of GNNs, we provide a concrete example showing that shift indeed improves the convergence rate. Note that the global convergence rate of widely-used deep GNNs on general data remains highly unexplored, and the existing works mainly focus on some simplified case, e.g., GNTK (Du et al., 2019b). To make things clear without loss of intuition, we focus on a simple linear GNN applied to a well-specified task where we are able to explicitly compare the convergence rates.

We describe each sample, i.e., graph, with $n$ nodes by a tuple $G=\left\{X,Q,\mathbf{p},y\right\}$, where

The whole dataset $S=\left\{G_{1},\cdots,G_{m}\right\}$ consists of $m$ graphs where $G_{i}=\left\{X_{i},Q_{i},\mathbf{p}_{i},y_{i}\right\}$. We make the following assumptions on the data generation process:

We assume $X_{i}$, $Q_{i}$, $\mathbf{p}_{i}$ are drawn from three independent distributions in an i.i.d. manner, e.g., $X_{1},\cdots,X_{m}$ are i.i.d..

A few remarks are in order, firstly, the assumption that each graph has the same number of nodes and the number $n$ is equal to feature dimension $d$ can be achieved by “padding”, i.e., adding dummy points or features to the graph or the feature matrix. The assumption that $\mathbf{1}^{\top}Y^{-1}XQ=0$ is used to guarantee that there is no information loss caused by shift ($\mathbf{1}^{\top}Y^{-1}YNY^{\top}=0$). Though we make this strong assumption to ensure no information loss in theoretical part, we introduce “learnable shift” to mitigate this problem in the practical setting. The theory taking learnable shift into account is an interesting future direction.

We make the technical assumption that there is a constant $b$ such that the distributions of $X_{i},Q_{i},\mathbf{p}_{i}$ ensures

Then, the model with shift can be represented as:

where $N=I_{n}-\frac{1}{n}\mathbf{1}\mathbf{1}^{\top}$.

We consider using square loss as training objective, i.e.,

We consider using gradient descent to optimize the objective function. Let the initial parameter $\mathbf{w}_{0}=0$. The update rule of $w$ from step $t$ to $t+1$ can be described as:

which indicates the shifted model has a faster convergence rate.

We firstly reformulate the optimization problem in matrix form.

Then the gradient descent update can be explicitly writen as:

which converges to $\mathbf{w}_{*}=\left(ZZ^{\top}\right)^{\dagger}Z\mathbf{y}$ according to classic theory of least square problem (Horn & Johnson, 2012), where $\left(ZZ^{\top}\right)^{\dagger}$ is the Moore–Penrose inverse of $ZZ^{\top}$.

By simultaneously subtracting $\mathbf{w}_{*}$ in the update rule, we have

Let $\sigma_{\max}(A)$ and $\sigma_{\min}(A)$ be the maximal and mininal positive eigenvalues of $A$, respectively. Then the optimial learning rate (the largest learning rate that ensures $I_{d}-\eta ZZ^{\top}$ is positive semidefinite) is $\eta=\frac{1}{\sigma_{\max}(ZZ^{\top})}$. Under this learning rate we have the convergence rate following Eq. (39):

which is the empirical estimation of the covariance matrix of the combined feature. By concentration inequality, we know this quantity is concentrated to the covariance matrix, i.e.,

Let $\mathbf{z}_{1},\cdots,\mathbf{z}_{m}$ be i.i.d. zero-mean random vectors with covariance matrix $\Sigma$ such that $\left\|\mathbf{z}\right\|_{2}\leq\sqrt{b}$ almost surely. Then for all $\delta>0$, the sample covariance matrix $\hat{\Sigma}=\frac{1}{m}\sum_{i=1}^{m}\mathbf{z}_{i}\mathbf{z}_{i}^{\top}$ satisfies

Let $\mathbf{z}_{1},\cdots,\mathbf{z}_{m}$ be i.i.d. zero-mean random vectors with covariance matrix $\Sigma$ such that $\left\|\mathbf{z}\right\|_{2}\leq\sqrt{b}$ almost surely. Then with probability $1-\epsilon$, the sample covariance matrix $\hat{\Sigma}=\frac{1}{m}\sum_{i=1}^{m}\mathbf{z}_{i}\mathbf{z}_{i}^{\top}$ satisfies

where we hide constants $b,\left\|\Sigma\right\|_{2},d$ in the big-O notation and highlight the dependence on the number of samples $m$.

Combining with previous results, we conclude that:

where $\rho_{1},\rho_{2}$ are the constants in the statement of the theorem. This inequality means the shifted model has better convergence speed by Eq. (41). ∎

A.3 Proof of Proposition 4.1

A.4 Proof of Proposition 4.2

where $\otimes$ represents the Kronecker product, and thus $\left(H^{(k-1)}QN\right)^{\top}\otimes S$ is an operator on matrices.

Analogously, the gradient of $W^{(k)}$ without normalization consists a $\left(H^{(k-1)}Q\right)^{\top}\otimes I_{n}$ term. As suggested by Theorem 3.1, $QN$ has a smoother distribution of spectrum than $Q$, so that the gradient of $W^{(k)}$ with normalization enjoys better optimization curvature than that without normalizaiton.

Appendix B Datasets

Detailed of the datasets used in our experiments are presented in this section. Brief statistics of the datasets are summarized in Table 3. Those information can be also found in Xu et al. (2019) and Hu et al. (2020).

PROTEINS is a dataset where nodes are secondary structure elements (SSEs) and there is an edge between two nodes if they are neighbors in the amino-acid sequence or in 3D space. It has 3 discrete labels, representing helix, sheet or turn. NCI1 is a dataset made publicly available by the National Cancer Institute (NCI) and is a subset of balanced datasets of chemical compounds screened for ability to suppress or inhibit the growth of a panel of human tumor cell lines, having 37 discrete labels. MUTAG is a dataset of 188 mutagenic aromatic and heteroaromatic nitro compounds with 7 discrete labels. PTC is a dataset of 344 chemical compounds that reports the carcinogenicity for male and female rats and it has 19 discrete labels.

IMDB-BINARY is a movie collaboration dataset. Each graph corresponds to an ego-network for each actor/actress, where nodes correspond to actors/actresses and an edge is drawn betwen two actors/actresses if they appear in the same movie. Each graph is derived from a pre-specified genre of movies, and the task is to classify the genre graph it is derived from. REDDIT-BINARY is a balanced dataset where each graph corresponds to an online discussion thread and nodes correspond to users. An edge was drawn between two nodes if at least one of them responded to another’s comment. The task is to classify each graph to a community or a subreddit it belongs to. COLLAB is a scientific collaboration dataset, derived from 3 public collaboration datasets, namely, High Energy Physics, Condensed Matter Physics and Astro Physics. Each graph corresponds to an ego-network of different researchers from each field. The task is to classify each graph to a field the corresponding researcher belongs to.

Ogbg-molhiv is a molecular property prediction dataset, which is adopted from the the MOLECULENET (Wu et al., 2017). Each graph represents a molecule, where nodes are atoms and edges are chemical bonds. Both nodes and edges have associated diverse features. Node features are 9-dimensional, containing atomic number and chirality, as well as other additional atom features. Edge features are 3-dimensional, containing bond type, stereochemistry as well as an additional bond feature indicating whether the bond is conjugated.

Appendix C The Experimental Setup

For the medium-scale bioinformatics and social network datasets, we use 5-layer GIN/GCN with a linear output head for prediction followed Xu et al. (2019) with residual connection. The hidden dimension of GIN/GCN is set to be 64. For the large-scale ogbg-molhiv dataset, we also use 5-layer GIN/GCN(Xu et al., 2019) architecture with residual connection. Following Hu et al. (2020), we set the hidden dimension as 300.

For the medium-scale bioinformatics and social network datasets, we compare several competitive baselines as in Xu et al. (2019), including the WL subtree kernel model (Shervashidze et al., 2011), diffusion-convolutional neural networks (DCNN) (Atwood & Towsley, 2016), Deep Graph CNN (DGCNN) (Zhang et al., 2018) and Anonymous Walk Embeddings (AWL) (Ivanov & Burnaev, 2018). We report the accuracies reported in the original paper (Xu et al., 2019). For the large-scale ogbg-molhiv dataset, we use the baselines in Hu et al. (2020), including the Graph-agnostic MLP model, GCN (Kipf & Welling, 2017) and GIN (Xu et al., 2019). We also report the roc-auc values reported in the original paper (Hu et al., 2020).

We use Adam (Kingma & Ba, 2015) optimizer with a linear learning rate decay schedule. We follow previous work Xu et al. (2019) and Hu et al. (2020) to use hyper-parameter search (grid search) to select the best hyper-parameter based on validation performance. In particular, we select the batch size $\in\{64,128\}$, the dropout ratio $\in\{0,0.5\}$, weight decay $\in\{5e-2,5e-3,5e-4,5e-5\}\cup\{0.0\}$, the learning rate $\in\{1e-4,1e-3,1e-2\}$. For the drawing of the training curves in Figure 2, for simplicity, we set batch size to be 128, dropout ratio to be 0.5, weight decay to be 0.0, learning rate to be 1e-2, and train the models for 400 epochs for all settings.

Using the chosen hyper-parameter, we report the averaged test performance over different random seeds (or cross-validation). In detail, for the medium-scale datasets, following Xu et al. (2019), we perform a 10-fold cross-validation as these datasets do not have a clear train-validate-test splitting format. The mean and standard deviation of the validation accuracies across the 10 folds are reported. For the ogbg-molhiv dataset, we follow the official setting (Hu et al., 2020). We repeat the training process with 10 different random seeds.

For all experiments, we select the best model checkpoint with the best validation accuracy and record the corresponding test performance.

Appendix D Additional Experimental Results

As stated in Theorem 3.1, the shift operation $N$ serves as a preconditioner of ${Q}$ which makes the singular value distribution of ${Q}$ smoother. To check the improvements, we sample graphs from 6 median-scale datasets (PROTEINS, NCI1, MUTAG, PTC, IMDB-BINARY, COLLAB) for visualization, as in Figure 7.

D.2 Visualization of noise in the batch statistics

We show the noise of the batch statistics on the PROTEINS task in the main body. Here we provide more experiment details and results.

For graph tasks (PROTEINS, PTC, NCI1, MUTAG, IMDB-BINARY datasets), we train a 5-layer GIN with BatchNorm as in Xu et al. (2019) and the number of sub-layers in MLP is set to 2. For image task (CIFAR10 dataset), we train a ResNet18 (He et al., 2016). Note that for a 5-layer GIN model, it has four graph convolution layers (indexed from 0 to 3) and each graph convolution layer has two BatchNorm layers; for a ResNet18 model, except for the first 3$\times$3 convolution layer and the final linear prediction layer, it has four basic layers (indexed from 0 to 3) and each layer consists of two basic blocks (each block has two BatchNorm layers). For image task, we set the batch size as 128, epoch as 100, learning rate as 0.1 with momentum 0.9 and weight decay as 5e-4. For graph tasks, we follow the setting of Figure 2 (described in Appendix C).

The visualization of the noise in the batch statistics is obtained as follows. We first train the models and dump the model checkpoints at the end of each epoch; Then we randomly sample one feature dimension and fix it. For each model checkpoint, we feed different batches to the model and record the maximum/minimum batch-level statistics (mean and standard deviation) of the feature dimension across different batches. We also calculate dataset-level statistics.

As Figure 4 in the main body, pink line denotes the dataset-level statistics, and green/blue line denotes the maximum/minimum value of the batch-level statistics respectively. First, we provide more results on PTC, NCI1, MUTAG, IMDB-BINARY tasks, as in Figure 8. We visualize the statistics from the first (layer-0) and the last (layer-3) BatchNorm layers in GIN for comparison. Second, we further visualize the statistics from different BatchNorm layers (layer 0 to layer 3) in GIN on PROTEINS and ResNet18 in CIFAR10, as in Figure 9. Third, we conduct experiments to investigate the influence of the batch size. We visualize the statistics from BatchNorm layers under different settings of batch sizes , as in Figure 10. We can see that the observations are consistent and the batch statistics on graph data are noisy, as in Figure 4 in the main body.

D.3 Training Curves on GCN

As Figure 2 in the main body, we train GCNs with different normalization methods (GraphNorm, InstanceNorm, BatchNorm and LayerNorm) and GCN without normalization in graph classification tasks and plot the training curves in Figure 6. It is obvious that the GraphNorm also enjoys the fastest convergence on all tasks. Remarkably, GCN with InstanceNorm even underperforms GCNs with other normalizations, while our GraphNorm with learnable shift significantly boosts the training upon InstanceNorm and achieves the fastest convergence.

D.4 Further Results of Ablation Study

We conduct experiments on BatchNorm to investigate whether simply introducing a learnable shift can already improve the existing normalization methods without concrete motivation of overcoming expressiveness degradation. Specifically, we equip BatchNorm with a similar learnable shift ($\alpha$-BatchNorm for short) as GraphNorm and evaluate its performance. As shown in Figure 12, the $\alpha$-BatchNorm cannot outperform the BatchNorm on the three datasets. Moreover, as shown in Figure 5 in the main body, the learnable shift significantly improve upon GraphNorm on IMDB-BINARY dataset, while it cannot further improve upon BatchNorm, which suggests the introduction of learnable shift in GraphNorm is critical.

We study the variant of BatchNorm which uses running statistics (MS-BatchNorm for short) to replace the batch-level mean and standard deviation (similar idea is also proposed in Yan et al. (2019)). At first glance, this method may seem to be able to mitigate the problem of large batch noise. However, the running statistics change a lot during training, and using running statistics disables the model to back-propagate the gradients through mean and standard deviation. Thus, we also train GIN with BatchNorm which stops the back-propagation of the graidients through mean and standard deviation (DT-BatchNorm for short). As shown in Figure 12, both the MS-BatchNorm and DT-BatchNorm underperform the BatchNorm by a large margin, which shows that the problem of the heavy batch noise cannot be mitigated by simply using the running statistics.

We further compare the GraphNorm and BatchNorm with different batch sizes (8, 16, 32, 64). As shown in Figure 11, our GraphNorm consistently outperforms the BatchNorm on all the settings.

Appendix E Other Related Works

Due to space limitations, we add some more related works on normalization and graph neural networks here. Zou et al. (2019) used normalization to stabilize the training process of GNNs. Zhao & Akoglu (2020) introduced PAIRNORM to prevent node embeddings from over-smoothing on the node classification task. Our GraphNorm focuses on accelerating the training and has faster convergence speed on graph classification tasks. Yang et al. (2020) interpreted the effect of mean subtraction on GCN as approximating the Fiedler vector. We analyze more general aggregation schemes, e.g., those in GIN, and understand the effect of the shift through the distribution of spectrum. Some concurrent and independent works (Li et al., 2020; Chen et al., 2020; Zhou et al., 2020b, a) also seek to incorporate normalization schemes in GNNs, which show the urgency of developing normalization schemes for GNNs. In this paper, we provide several insights on how to design a proper normalization for GNNs. Before the surge of deep learning, there are also many classic architectures of GNNs such as Scarselli et al. (2008); Bruna et al. (2013); Defferrard et al. (2016) that are not mentioned in the main body of the paper. We refer the readers to Zhou et al. (2018); Wu et al. (2020); Zhang et al. (2020) for surveys of graph representation learning.