Source-linked AI summary

GraphNorm: A Principled Approach to Accelerating Graph Neural Network Training

Tianle Cai, Shengjie Luo, Keyulu Xu, Di He, Tie-Yan Liu, Liwei Wang

arXiv:2009.03294v3cs.LGmath.OCstat.ML

TL;DR

GNN training can be unstable and slow, motivating the search for effective graph-specific normalization. The paper adapts existing methods, explains InstanceNorm’s optimization behavior and limitation, and proposes GraphNorm with a learnable shift, which yields faster convergence and better generalization on graph classification benchmarks.

  • Problem

    GNN optimization is less well understood, and training is often unstable and slow; the paper asks which normalization methods are effective for GNNs.

  • Method

    The paper adapts BatchNorm, LayerNorm, and InstanceNorm to GNNs, analyzes their optimization effects and limitations, and proposes GraphNorm with a learnable shift.

  • Results

    GNNs with GraphNorm converge faster and achieve better generalization performance than other normalization methods on graph classification benchmarks.

  • Takeaways & Limitations

    GraphNorm provides a normalization approach for GNNs that combines InstanceNorm’s acceleration effect with a learnable shift intended to avoid expressiveness degradation.

  • Takeaways & Limitations

    The theoretical analyses are motivated by simple models and do not establish convergence rates for general normalized GNNs or characterize the spectrum of Q under a learnable shift.

Abstract

from arXiv · show

Normalization is known to help the optimization of deep neural networks. Curiously, different architectures require specialized normalization methods. In this paper, we study what normalization is effective for Graph Neural Networks (GNNs). First, we adapt and evaluate the existing methods from other domains to GNNs. Faster convergence is achieved with InstanceNorm compared to BatchNorm and LayerNorm. We provide an explanation by showing that InstanceNorm serves as a preconditioner for GNNs, but such preconditioning effect is weaker with BatchNorm due to the heavy batch noise in graph datasets. Second, we show that the shift operation in InstanceNorm results in an expressiveness degradation of GNNs for highly regular graphs. We address this issue by proposing GraphNorm with a learnable shift. Empirically, GNNs with GraphNorm converge faster compared to GNNs using other normalization. GraphNorm also improves the generalization of GNNs, achieving better performance on graph classification benchmarks.

1. Introduction

The paper studies normalization for GNNs, where training is often unstable and slow, and finds that existing methods behave differently after adaptation to graphs. It explains InstanceNorm’s optimization benefit, identifies its expressiveness limitation, and proposes GraphNorm with a learnable shift.

  • GNN training is often unstable and converges slowly, while its optimization is less understood than its expressive power and generalization.
  • The paper asks which normalization methods are effective for GNNs because different neural architectures and domains favor different normalization methods.
  • InstanceNorm adapts per-graph node-representation normalization and converges faster than BatchNorm and LayerNorm on most datasets.
  • InstanceNorm’s shift acts as a preconditioner for graph aggregation, smoothing optimization curvature, whereas BatchNorm’s effect is weaker because graph batch statistics are noisy.
  • InstanceNorm can reduce expressiveness on highly regular graphs by removing mean statistics that contain graph structural information.
  • GraphNorm adds a learnable shift to InstanceNorm, and experiments report faster convergence and better generalization than competing normalization methods.

2. Preliminaries

The preliminaries introduce graph and node-feature notation, neighborhood-aggregation GNNs, graph-level readout, and the general shift-and-scale form of normalization. They also identify GCN and GIN as concrete GNN examples.

  • A graph is represented as G = (V, E), with node features Xi, adjacency matrix A, and degree matrix D.
  • GNNs iteratively update each node representation by aggregating representations from its neighboring nodes.
  • Different aggregation functions yield different architectures, including GCN and GIN.
  • For graph classification, a READOUT function aggregates final-layer node features into a graph representation used by a classifier.
  • Normalization shifts each feature value by its mean and scales it by its standard deviation, with learnable parameters γ and β.
  • Normalization methods mainly differ in which set of feature values supplies the statistics, such as batches for BatchNorm or positions for LayerNorm.

3. Evaluating and Understanding Normalization for GNNs

The paper adapts normalization methods to GNNs, finding that InstanceNorm accelerates training through preconditioning while BatchNorm suffers from noisy graph-batch statistics. GraphNorm addresses InstanceNorm’s expressiveness limitation on highly regular graphs with a learnable shift and achieves faster convergence.

  • 3.1. Adapting and Evaluating Normalization for GNNs: GraphNorm further accelerates training over InstanceNorm, especially on highly regular graphs such as IMDB-BINARY.Figure 2 reports GraphNorm as converging faster than all other evaluated methods.
  • 3.1. Adapting and Evaluating Normalization for GNNs: InstanceNorm normalizes node hidden representations separately within each graph and converges faster than BatchNorm and LayerNorm in most graph-classification tasks.Normalization is applied after the linear transformation, with InstanceNorm computing statistics across nodes of each individual graph.
  • 3.2. Shift in InstanceNorm as a Preconditioner: In a simple fully characterizable setting, the shifted model converges faster than the vanilla model.The paper presents this convergence comparison as theoretical insight rather than a proof of convergence-rate comparisons in real settings.
  • 3.2. Shift in InstanceNorm as a Preconditioner: The InstanceNorm shift acts as a preconditioner by replacing the aggregation matrix Q with QN, producing a smoother singular-value distribution and improved condition number.The analysis connects this spectral change to smoother optimization curvature and potentially faster convergence.
  • 3.3. Heavy Batch Noise in Graphs Makes BatchNorm Less Effective: BatchNorm is less effective because graph batch statistics deviate substantially from dataset-level statistics, introducing noise that may destabilize optimization.This deviation is large for PROTEINS graph classification but negligible for CIFAR10 image classification after a few epochs.

4. Graph Normalization

The section identifies expressiveness loss from InstanceNorm's standard shift on regular graphs and introduces GraphNorm, which uses a learnable shift to preserve mean information.

  • Expressiveness limitation: For regular graphs, subtracting mean statistics can remove graph-structural information and degrade GNN expressiveness.With one-hot degree features, normalization can produce a zero matrix without graph-structure information.
  • Expressiveness limitation: For complete graphs, the standard shift can remove structural information from GIN after multiplying by the normalization operator.The adjacency-derived structural information in Q is ignored after the shift.
  • Expressiveness limitation: Mean statistics after aggregation can contain graph-structural information, so discarding them may reduce neural-network expressiveness.The paper contrasts this with image data, where removing global mean information need not alter object semantics.
  • GraphNorm design: GraphNorm adds a learnable parameter that controls how much mean information to preserve during the shift operation.It combines graph-wise normalization with a learnable shift and is designed to address InstanceNorm's expressive degradation.
  • Empirical validation: On PROTEINS and IMDB-BINARY, the learnable shift slightly improves convergence on irregular graphs and significantly improves training on regular graphs.The experiments compare GIN and GCN using InstanceNorm and GraphNorm under the same settings.

5. Experiments

Experiments compare normalization methods across graph-classification benchmarks and ablations. GraphNorm shows the fastest convergence and generally better generalization, while BatchNorm variants perform worse or remain behind it.

  • Experimental setup: The evaluation uses eight graph-classification benchmarks spanning bioinformatics, social-network, and OGB datasets, with GIN and GCN.The datasets include MUTAG, PTC, PROTEINS, NCI1, IMDB-BINARY, COLLAB, REDDIT-BINARY, and ogbg-molhiv.
  • Main results: GraphNorm achieves the fastest convergence across all reported tasks.The training curves compare GraphNorm with other normalization methods, with similar trends observed for GCN.
  • Main results: On NCI1 and PTC, GraphNorm converges in roughly 5000/500 iterations, whereas BatchNorm does not converge within 10000/1000 iterations.The figures correspond to the respective datasets and reported training comparisons.
  • Main results: GraphNorm improves generalization on most benchmarks according to the reported test accuracies.The paper attributes higher test performance to both higher training accuracy and a potentially smaller generalization gap from faster training.
  • Ablation study: BatchNorm with a learnable shift does not outperform BatchNorm, while the running-statistics variant performs even worse than BatchNorm.The running-statistics variant also prevents back-propagation through the mean and standard deviation.
  • Ablation study: GraphNorm consistently outperforms BatchNorm across batch sizes 8, 16, 32, and 64.This comparison is reported in the ablation results.

6. Conclusion and Future Work

The paper adapts three established normalization methods to GNNs and proposes GraphNorm with a learnable shift. Experiments report faster convergence and better generalization, while the theoretical analysis remains limited in scope.

  • Conclusion: The paper adapts BatchNorm, LayerNorm, and InstanceNorm to GNNs and analyzes their successes and failures.The methods are evaluated as normalization choices for graph neural networks.
  • Conclusion: GraphNorm builds on InstanceNorm by adding a learnable shift to address expressive degradation.The design is motivated by the limitations identified for standard InstanceNorm on graph data.
  • Conclusion: GraphNorm-equipped GNNs converge faster and achieve better generalization on several benchmark datasets.These are the paper's reported experimental outcomes.
  • Future work: The theoretical analyses are motivated by simple models and do not provide concrete convergence rates for general normalized GNNs.The paper also leaves the spectrum of Q normalized by learnable shift and realistic deep-GNN dynamics for future work.

A.1. Proof of Theorem 3.1

The proof uses the Cauchy interlace theorem and similarity and projection properties to relate eigenvalues of normalized matrix products to those of the underlying matrices.

  • Interlacing argument: The proof begins by invoking the Cauchy interlace theorem for a symmetric block matrix and its principal submatrix.The theorem characterizes how the eigenvalues of the two matrices interlace.
  • Spectral reduction: Similarity preserves eigenvalues, allowing the proof to analyze NQ^⊤QN through a similar matrix representation.The proof also uses singular values as square roots of eigenvalues of P^⊤P.
  • Projection structure: N is a projection onto the orthogonal complement of the subspace spanned by 1, with 1 as an eigenvector associated with eigenvalue 0.An orthogonal matrix U provides the basis used in the reduction.
  • Interlacing argument: The proof applies eigenvalue interlacing to transformed matrices and tracks equality conditions through orthogonality and eigenvector relationships.The intermediate steps use R = U^⊤C̄U and S = U_1^⊤C̄U_1 before identifying the relevant eigenvectors.
  • Proof conclusion: The final step identifies U_1z as an eigenvector of C̄ and relates it to a right singular vector of Q.This completes the spectral proof after substituting the orthogonality condition.
  • Interpretation: A later example uses a simple linear GNN to illustrate how the shift's preconditioning effect can improve convergence.The example is motivated by the limited understanding of global convergence for deep GNNs on general data.

A.2.1. SETTINGS

This section specifies graph data, a linear GNN analysis, and gradient-descent convergence under assumptions on graph features, aggregation, and noise. It shows that shifting representations improves the effective spectral conditioning and convergence rate, while smoothing the aggregation spectrum.

  • Model and data: Each graph is represented as G = {X, Q, p, y}, combining node features, neighbor aggregation, node-importance weights, and a label.X is the feature matrix, Q represents neighbor aggregation, p supports READOUT, and y is the label.
  • Assumptions: The analysis assumes independently sampled graph components, equal node counts with n = d, full-rank expected features, bounded noise, and spectral conditions on the feature covariance.Padding can enforce equal sizes, while learnable shift is introduced practically to mitigate possible information loss from the theoretical shift assumption.
  • Optimization: The shifted linear model uses a graph-level shift before aggregation and is optimized with square loss by gradient descent from w0 = 0.The update uses learning rate η, and the shifted representation is encoded through a transformed feature vector.
  • Theoretical result: The shifted model has a faster convergence rate than the vanilla model under the stated assumptions and sufficiently favorable sample and noise conditions.The theorem gives the comparison with high probability for small enough δ1 and sufficiently large m.
  • Optimization: The shift operation acts as a preconditioner by producing a smoother singular-value distribution for the aggregation matrix, improving optimization curvature.The analysis characterizes convergence through the effective condition number σ_min/σ_max of the relevant covariance-like matrices.

B. Datasets

The paper evaluates graph neural networks on benchmark datasets spanning bioinformatics, social networks, and molecular property prediction. The datasets differ in scale, graph semantics, and node or edge features.

  • Dataset coverage: The benchmark collection includes medium-scale bioinformatics and social-network datasets plus the large-scale ogbg-molhiv molecular dataset.The section refers readers to dataset statistics in Table 3 and prior dataset descriptions.
  • Dataset semantics: Bioinformatics graphs represent structured biological objects, with PROTEINS nodes corresponding to secondary-structure elements and molecular datasets encoding chemical entities.ogbg-molhiv represents molecules using atom nodes and chemical-bond edges with associated features.

C. The Experimental Setup

The experiments compare GNN architectures and normalization methods across eight graph-classification benchmarks using tuned optimization settings and validation-based evaluation. Additional analyses examine conditioning and batch-statistics noise.

  • Network architecture: Experiments use 5-layer GIN or GCN models with residual connections, 64-dimensional hidden states for medium-scale data, and 300-dimensional states for ogbg-molhiv.Both settings use a linear output head for prediction.
  • Comparisons: Normalization comparisons include GraphNorm, InstanceNorm, BatchNorm, LayerNorm, and models without normalization, alongside established graph-learning baselines.The baseline list includes WL subtree kernels, DCNN, DGCNN, and Anonymous Walk Embeddings.
  • Training protocol: Hyperparameters are selected by grid search using validation performance, with Adam, linear learning-rate decay, and specified batch-size, dropout, weight-decay, and learning-rate ranges.The candidate batch sizes are 64 and 128; candidate learning rates range from 1e-4 to 1e-2.
  • Evaluation: Medium-scale datasets use 10-fold cross-validation, while ogbg-molhiv follows its official split and reports test performance from the best validation checkpoint.Reported performance is averaged over random seeds or cross-validation folds.
  • Diagnostic analyses: Visualization studies inspect the spectrum of aggregation matrices and the variability of BatchNorm statistics across datasets, layers, and batch sizes.The batch-statistics analysis records extrema across batches and compares them with dataset-level statistics.

D.3. Training Curves on GCN

Training-curve experiments on GCN compare GraphNorm, InstanceNorm, BatchNorm, LayerNorm, and no normalization. GraphNorm shows the fastest convergence across the reported tasks, while InstanceNorm can underperform the alternatives.

  • Training curves: GraphNorm achieves the fastest convergence on all tasks in the reported GCN training curves.The comparison includes GraphNorm, InstanceNorm, BatchNorm, LayerNorm, and GCN without normalization.
  • Training curves: InstanceNorm can underperform other normalization methods in GCNs, whereas GraphNorm’s learnable shift substantially improves training over InstanceNorm.The passage attributes the improvement specifically to GraphNorm’s learnable shift.

D.4. Further Results of Ablation Study

Additional ablations show that GraphNorm’s advantage persists across batch sizes, while straightforward BatchNorm variants do not resolve the batch-noise problem. The accompanying figures examine normalization performance, singular values, and batch-statistic noise across datasets, depths, and batch sizes.

  • BatchNorm with learnable shift: α-BatchNorm cannot outperform BatchNorm on the three evaluated datasets, indicating that adding a learnable shift alone does not reproduce GraphNorm’s benefit.
  • BatchNorm with running statistics: MS-BatchNorm and DT-BatchNorm underperform BatchNorm by a large margin, so running statistics and stopped gradients do not mitigate heavy batch noise.
  • The effect of batch size: GraphNorm consistently outperforms BatchNorm across batch sizes 8, 16, 32, and 64.
  • Additional analyses: The ablation figures compare normalization training performance, singular-value distributions, and batch-statistic noise across graph datasets, network depths, and batch sizes.
Loading 2009.03294v3…