Title: Are GATs Out of Balance?
Venue: NeurIPS 2023
Accepted: yes
Keywords: graph attention networks

gradient flow

conservation law
Abstract: While the expressive power and computational capabilities of graph neural networks (GNNs) have been theoretically studied, their optimization and learning dynamics, in general, remain largely unexplored. Our study undertakes the Graph Attention Network (GAT), a popular GNN architecture in which a node's neighborhood aggregation is weighted by parameterized attention coefficients. We derive a conservation law of GAT gradient flow dynamics, which explains why a high portion of parameters in GATs with standard initialization struggle to change during training. This effect is amplified in deeper GATs, which perform significantly worse than their shallow counterparts. To alleviate this problem, we devise an initialization scheme that balances the GAT network. Our approach i) allows more effective propagation of gradients and in turn enables trainability of deeper networks, and ii) attains a considerable speedup in training and convergence time in comparison to the standard initialization. Our main theorem serves as a stepping stone to studying the learning dynamics of positive homogeneous models with attention mechanisms.

Are GATs Out of Balance?

Abstract

While the expressive power and computational capabilities of graph neural networks (GNNs) have been theoretically studied, their optimization and learning dynamics, in general, remain largely unexplored. Our study undertakes the Graph Attention Network (GAT), a popular GNN architecture in which a node's neighborhood aggregation is weighted by parameterized attention coefficients. We derive a conservation law of GAT gradient flow dynamics, which explains why a high portion of parameters in GATs with standard initialization struggle to change during training. This effect is amplified in deeper GATs, which perform significantly worse than their shallow counterparts. To alleviate this problem, we devise an initialization scheme that balances the GAT network. Our approach i) allows more effective propagation of gradients and in turn enables trainability of deeper networks, and ii) attains a considerable speedup in training and convergence time in comparison to the standard initialization. Our main theorem serves as a stepping stone to studying the learning dynamics of positive homogeneous models with attention mechanisms.

1 Introduction

A rapidly growing class of model architectures for graph representation learning are Graph Neural Networks (GNNs) [20] which have achieved strong empirical performance across various applications such as social network analysis [8], drug discovery [27], recommendation systems [61], and traffic forecasting [50]. This has driven research focused on constructing and assessing specific architectural designs [59] tailored to various tasks.

On the theoretical front, mostly the expressive power [51, 26] and computational capabilities [45] of GNNs have been studied. However, there is limited understanding of the underlying learning mechanics. We undertake a classical variant of GNNs, the Graph Attention Network (GAT) [48, 9]. It overcomes the limitation of standard neighborhood representation averaging in GCNs [29, 22] by employing a self-attention mechanism [47] on nodes. Attention performs a weighted aggregation over a node's neighbors to learn their relative importance. This increases the expressiveness of the neural network. GAT is a popular model that continues to serve as strong baseline.

However, a major limitation of message-passing GNNs in general, and therefore GATs as well, is severe performance degradation when the depth of the network is even slightly increased. SOTA results are reportedly achieved by models of 2 or 3 layers. This problem has largely been attributed to over-smoothing [35], a phenomenon in which node representations become indistinguishable when multiple layers are stacked. To relieve GNNs of over-smoothing, practical techniques inspired by classical deep neural networks tweak the training process, e.g. by normalization [11, 64, 65, 66] or regularization [38, 42, 55, 67], and impose architectural changes such as skip connections (dense

and residual) [34, 53] or offer other engineering solutions [33] or their combinations [12]. Other approaches to overcome over-smoothing include learning CNN-like spatial operators from random paths in the graph instead of the point-wise graph Laplacian operator [17] and learning adaptive receptive fields (different 'effective' neighborhoods for different nodes) [63, 36, 52]. Another problem associated with loss of performance in deeper networks is over-squashing [1], but we do not discuss it since non-attentive models are more susceptible to it, which is not our focus.

Alternative causes for the challenged trainability of deeper GNNs beyond over-smoothing is hampered signal propagation during training (i.e. backpropagation) [25, 62, 13]. To alleviate this problem for GCNs, a dynamic addition of skip connections to vanilla-GCNs, which is guided by gradient flow, and a topology-aware isometric initialization are proposed in [25]. A combination of orthogonal initialization and regularization facilitates gradient flow in [21]. To the best of our knowledge, these are the only works that discuss the trainability issue of deep GNNs from the perspective of signal propagation, which studies a randomly initialized network at initialization. Here, we offer an insight into a mechanism that allows these approaches to improve the trainability of deeper networks that is related to the entire gradient flow dynamics. We focus in particular on GAT architectures and the specific challenges that are introduced by attention.

Our work translates the concept of neuronwise balancedness from traditional deep neural networks to GNNs, where a deeper understanding of gradient dynamics has been developed and norm balancedness is a critical assumption that induces successful training conditions. One reason is that the degree of initial balancedness is conserved throughout the training dynamics as described by gradient flow (i.e. a model of gradient descent with infinitesimally small learning rate or step size). Concretely, for fully connected feed-forward networks and deep convolutional neural networks with continuous homogenous activation functions such as ReLUs, the difference between the squared l 2 -norms of incoming and outgoing parameters to a neuron stays approximately constant during training [15, 32]. Realistic optimization elements such as finite learning rates, batch stochasticity, momentum, and weight decay break the symmetries induced by these laws. Yet, gradient flow equations can be extended to take these practicalities into account [31].

Inspired by these advances, we derive a conservation law for GATs with positive homogeneous activation functions such as ReLUs. As GATs are a generalization of GCNs, most aspects of our insights can also be transferred to other architectures. Our consideration of the attention mechanism would also be novel in the context of classic feed-forward architectures and could have intriguing implications for transformers that are of independent interest. In this work, however, we focus on GATs and derive a relationship between model parameters and their gradients, which induces a conservation law under gradient flow. Based on this insight, we identify a reason for the lack of trainability in deeper GATs and GCNs, which can be alleviated with a balanced parameter initialization. Experiments on multiple benchmark datasets demonstrate that our proposal is effective in mitigating the highlighted trainability issues, as it leads to considerable training speed-ups and enables significant parameter changes across all layers. This establishes a causal relationship between trainability and parameter balancedness also empirically, as our theory predicts. The main theorem presented in this work serves as a stepping stone to studying the learning dynamics of positive homogeneous models with attention mechanisms such as transformers [47] and in particular, vision transformers which require depth more than NLP transformers[49]. Our contributions are as follows.

We derive a conservation law of gradient flow dynamics for GATs, including its variations such as shared feature weights and multiple attention heads.

This law offers an explanation for the lack of trainability of deeper GATs with standard initialization.

To demonstrate empirically that our theory has established a causal link between balancedness and trainability, we propose a balanced initialization scheme that improves the trainability of deeper networks and attains considerable speedup in training and convergence time, as our theory predicts.

2 Theory: Conservation laws in GATs

Preliminaries For a graph G = ( V, E ) with node set V and edge set E ⊆ V × V , where the neighborhood of a node v is given as N ( v ) = { u | ( u, v ) ∈ E } , a GNN layer computes each node's representation by aggregating over its neighbors' representation. In GATs, this aggregation is weighted by parameterized attention coefficients α uv , which indicate the importance of node u for v .

Given input representations h l -1 v for all nodes v ∈ V , a GAT 1 layer l transforms those to:

We consider ϕ to be a positively homogeneous activation functions (i.e ϕ ( x ) = xϕ ′ ( x ) and consequently, ϕ ( ax ) = aϕ ( x ) for positive scalars a ), such as a ReLU ϕ ( x ) = max { x, 0 } or LeakyReLU ϕ ( x ) = max { x, 0 } + -α max {-x, 0 } . The feature transformation weights W s and W t for source and target nodes, respectively, may also be shared such that W = W s = W t .

Definition 2.1. Given training data { ( x m , y m ) } M m =1 ⊂ R d × R p for M ≤ V , let f : R d → R p be the function represented by a network constructed by stacking L GAT layers as defined in Eq. (1) and (2) with W = W s = W t and h 0 m = g ( x m ) . Each layer l ∈ [ L ] of size n l has associated parameters: a feature weight matrix W l ∈ R n l × n l -1 and an attention vector a l ∈ R n l , where n 0 = d and n L = p . Given a differentiable loss function ℓ : R d × R p → R , the loss L = 1 / M ∑ M i =1 ℓ ( f ( x m ) , y m ) is used to update model parameters w ∈ { W l , a l } L l =1 with learning rate γ by gradient descent, i.e., w t +1 = w t -γ ∇ w L , where ∇ w L = [ ∂ L / ∂w 1 , . . . , ∂ L / ∂w | w | ] and w 0 is set by the initialization scheme. For an infinitesimal γ → 0 , the dynamics of gradient descent behave similarly to gradient flow defined by d w / d t = -∇ w L , where t is the continuous time index.

We use W [ i, :] , W [: , i ] , and a [ i ] to denote the i th row of W , column of W , and entry of a , respectively. Note that W l [ i, :] is the vector of weights incoming to neuron i ∈ [ n l ] and W l +1 [: , i ] is the vector of weights outgoing from the same neuron. For the purpose of discussion, let i ∈ [ d l ] , j ∈ [ d l -1 ] , and k ∈ [ d l +1 ] , as depicted in Fig. 1. ⟨· , ·⟩ denotes the scalar product.

Conservation Laws We focus the following exposition on weight-sharing GATs, as this variant has the practical benefit that it requires fewer parameters. Yet, similar laws also hold for the non-weightsharing case, as we state and discuss in the supplement. For brevity, we also defer the discussion of laws for the multi-headed attention mechanism to the supplement.

Theorem 2.2 (Structure of gradients) . Given the feature weight and attention parameters W l and a l of a layer l in a GAT network as described in Def. 2.1, the structure of the gradients for layer l ∈ [ L -1] in the network is conserved according to the following law:

Corollary 2.3 (Norm conservation of weights incoming and outgoing at every neuron) . Given the gradient structure of Theorem 2.2 and positive homogeneity of the activation ϕ in Eq. (1), gradient flow in the network adheres to the following invariance for l ∈ [ L -1] :

(a)

δ

with GD

(b)

δ

(c)

c

with Xavier

(d)

c

with Bal

O

Layer 1

Layer 2

Layer 1

Layer 2

Layer 1

Layer 9

Layer 1

Layer 9

1e-8

1e-8

1e-7

1e-7

1

2.5

1.0

1

0.05

0.005

0

0

0.0

0.0

0.00

-2

-1

2.5

2.5

0.5

0

0.000

0

1e4

0

1e4

0

100

0

100

0

5000

0

5000

0

5000

0

5000

Epoch

Epoch

Epoch

Epoch

Epoch

Epoch

Epoch

Epoch

̸

It then directly follows by integration over time that the l 2 -norms of feature and attention weights incoming to and outgoing from a neuron are preserved such that

where c = d d t ( ∥ ∥ W l [ i, :] ∥ ∥ 2 -∥ ∥ a l [ i ] ∥ ∥ 2 ) at initialization (i.e. t = 0 ).

We denote the 'degree of balancedness' by c and call an initialization balanced if c = 0 .

Corollary 2.4 (Norm conservation across layers) . Given Eq. (4), the invariance of gradient flow at the layer level for l ∈ [ L -1] by summing over i ∈ [ n l ] is:

Remark 2.5 . Similar conservation laws also hold for the original less expressive GAT version [48] as well as for the vanilla GCN [29] yielded by fixing a l = 0 .

Insights We first verify our theory and then discuss its implications in the context of four different initializations for GATs, which are summarized in Table 1. Xavier (Xav.) initialization [19] is the standard default for GATs. We describe the Looks-Linear-Orthogonal (LLortho) initialization later.

In order to empirically validate Theorem 2.2, we rewrite Equation 3 to define δ as:

We observe how the value of δ varies during training in Fig. 2 for both GD and Adam optimizers with γ = 0 . 1 . Although the theory assumes infinitesimal learning rates due to gradient flow, it still holds sufficiently ( δ ≈ 0 ) in practice for normally used learning rates as we validate in Fig. 2. Furthermore, although the theory assumes vanilla gradient descent optimization, it still holds for the Adam optimizer (Fig. 2b). This is so because the derived law on the level of scalar products between gradients and parameters (see Theorem 2.2) still holds, as it states a relationship between gradients and parameters that is independent of the learning rate or precise gradient updates.

Having verified Theorem 2.2, we use it to deduce an explanation for a major trainability issue of GATs with Xavier initialization that is amplified with increased network depth.

The main reason is that the last layer has a comparatively small average weight norm, as E ∥ ∥ W L [: , i ] ∥ ∥ 2 = 2 n L / ( n L + n L -1 ) << 1 , where the number of outputs is smaller than the layer width n L << n L -1 . In contrast, E ∥ ∥ W L -1 [ i, :] ∥ ∥ 2 = 2 n L -1 / ( n L -1 + n L -2 ) = 1 and E a L -1 [ i ] 2 = 2 / (1 + n L -1 ) . In consequence, the right-hand side of our derived conservation law

0.000

0.005

0.010

0.015

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Feat.Weights

Att.Weights

Xavier

LLortho+Bal

Frac. of Sig. Params

1.0

1.0

2

2

Layer

2

1

Loss

Layer

2

Loss

0.5

0.5

3

3

4

4

5

0

5

0

0

150

00

5

4600

0

350

1.5

1500

0.0

1

2

3

4

5

0.0

1

2

3

4

5

1

3

Train

11

Layer

Epoch

Valid

Epoch

Layer

Figure 3: For L = 5 , a high fraction of parameters change, and gradients propagate to earlier layers for both the unbalanced and balanced initialization Xav. and Bal O initialization achieve 75 . 5% and 79 . 9% test accuracy, respectively. Note that the balanced initialization in (b) is able to train faster.

0.00

0.02

0.04

0.06

0.08

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Feat.Weights

Att.Weights

Xav

LLortho+Bal

Frac. of Sig. Params

1.0

1.0

Layer

4

2

2

Loss

Layer

4

2

2

Loss

0.5

0.5

6

6

8

8

10

0

10

0

0.0

0.0

0

1250

00

25

3

5

0000

Train

0

30 00 00

1200

12345678910

12345678910

Epoch

Valid

Epoch

Layer

Layer

Figure 4: For L = 10 , unbalanced initialization is unable to train the model as relative gradients are extremely small, and the parameters (and loss) do not change. However, the balanced initialization is able to propagate gradients, change parameters, and thus attain training loss → 0 . Xav. and Bal O achieve 39 . 3% and 80 . 2% test accuracy, respectively. Note that the balanced initialization in (b) is also able to train faster than the 5 layer network with Xavier initialization in Fig. 3b.

is relatively small, as ∑ n L k =1 ( W L ki ) 2 << 1 , on average, while the weights on the left-hand side are orders of magnitude larger. This implies that the relative gradients on the left-hand side are likely relatively small to meet the equality unless all the signs of change are set in such a way that they satisfy the equality, which is highly unlikely during training. Therefore, relative gradients in layer L -1 and accordingly also the previous layers of the same dimension can only change in minor ways.

The amplification of this trainability issue with depth can also be explained by the recursive substitution of Theorem 2.2 in Eq. (8) that results in a telescoping series yielding:

Generally, 2 n 1 < n 0 and thus E ∥ ∥ W 1 [ j :] ∥ ∥ 2 = 2 n 1 / ( n 1 + n 0 ) < 1 . Since the weights in the first layer and the gradients propagated to the first layer are both small, gradients of attention parameters of the intermediate hidden layers must also be very small in order to satisfy Eq. (9). Evidently, the problem aggravates with depth where the same value must now be distributed over the parameters and gradients of more layers. Note that the norms of the first and last layers are usually not arbitrary but rather determined by the input and output scale of the network.

Analogously, we can see see how a balanced initialization would mitigate the problem. Equal weight norms ∥ ∥ W L -1 [ i, :] ∥ ∥ 2 and ∥ ∥ W L [: , i ] ∥ ∥ 2 in Eq. (8) (as the attention parameters are set to 0 during balancing, see Procedure 2.6) would allow larger relative gradients in layer L -1 , as compared to the imbalanced case, that can enhance gradient flow in the network to earlier layers. In other words, gradients on both sides of the equation have equal room to drive parameter change.

feat, I=1

feat, I=5

feat, I=10

attn,I=1

attn,I=5

attn, I=10

Rel. Grad. Norm of Feat

2×10

Xav:39.3%

Xav+ZeroAtt:12.9%

10

Xav+Bal:72.5%

10°

0

10

LLortho+Bal:80.2%

Rel. Grad. Norm of Attn

3×10-7

10

10

10

2×10~710

10

-2

10

10

2

10

3

10

0

2000

4000

0

2000

4000

0

500

1000

1500

0

500

1000

Epoch

Epoch

Epoch

Epoch

To visualize the trainability issue, we study the relative change | ( w ∗ -w 0 ) / w ∗ | . of trained network parameters ( w ∗ ) w.r.t. their initial value. In order to observe meaningful relative change, we only consider parameters with a significant contribution to the model output ( w ∗ ≥ 10 -4 ). We also plot the relative gradient norms of the feature weights across all layers to visualize their propagation. We display these values for L = 5 and L = 10 in Fig. 3 and 4, respectively, and observe a stark contrast, similar to the trends in relative gradient norms of selected layers of a 10 layer GAT in Fig. 5. In all experiments, GAT was trained with the specified initialization on Cora using GD with γ = 0 . 1 for 5000 epochs.

Interestingly, the attention parameters change most in the first layer and increasingly more towards the last layer, as seen in Fig. 3a and 4a. This is consistently observed in both the 5 and 10 layer networks if the initialization is balanced, but with unbalanced initialization, the 10 layer network is unable to produce the same effect. As our theory defines a coarse-level conservation law, such fine-grained training dynamics cannot be completely explained by it.

Since the attention parameters are set to 0 in the balanced initializations (see Procedure 2.6), as ablation, we also observe the effect of initializing attention and feature parameters with zero and Xavier, respectively. From Eq.(2), a l = 0 leads to α uv = 1 / |N ( v ) | , implying that all neighbors u of the node v are equally important at initialization. Intuitively, this allows the network to learn the importance of neighbors without any initially introduced bias and thus avoids the 'chicken and egg' problem that arises in the initialization of attention over nodes[30]. Although this helps generalization in some cases (see Fig. 6), it alone does not help the trainability issue as seen in Fig. 4a. It may in fact worsen it (see Fig. 5), which is explained by Eq. 8. These observations underline the further need for parameter balancing.

Procedure 2.6 (Balancing) . Based on Eq. (5) from the norm preservation law 2.3, we note that in order to achieve balancedness, (i.e. set c = 0 in Eq.(5)), the randomly initialized parameters W l and a l must satisfy the following equality for l ∈ [ L ] :

This can be achieved by scaling the randomly initialized weights as follows:

Set a l = 0 for l ∈ [ L ] .

Set W 1 [ i, :] = W 1 [ i, :] ∥ W 1 [ i, :] ∥ √ β i , for i ∈ [ n 1 ] where β i is a hyperparameter

In step 2, β i determines the initial squared norm of the incoming weights to neuron i in the first layer. It can be any constant thus also randomly sampled from a normal or uniform distribution. We explain why we set β i = 2 for i ∈ [ n 1 ] in the context of an orthogonal initialization.

Remark 2.7 . This procedure balances the network starting from l = 1 towards l = L . Yet, it could also be carried out in reverse order by first defining ∥ ∥ W L [: , i ] ∥ ∥ 2 = β i for i ∈ [ n L -1 ] .

Balanced Orthogonal Initialization The feature weights W l for l ∈ [ L ] need to be initialized randomly before balancing, which can simply be done by using Xavier initialization. However, existing works have pointed out benefits of orthogonal initialization to achieve initial dynamical

(a) Test accuracy (higher is better).

(b) Epochs to best model (lower is better).

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

AEX

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Width = 512

Width = 512

80

0009

70

5000

Accuracy (%)

88

4000

0

3000

30

2000

20

1000

10

2

5

10

20

2

5

10

20

2

5

10

20

2

5

10

20

D

2

5

10

20

2

5

10

20

2

5

10

20

2

5

10

20

Depth

Depth

Depth

udeg

Depth

Depth

Depth

Depth

isometry for DNNs and CNNs [37, 44, 23], as it enables training very deep architectures. In line with these findings, we initialize W l , before balancing, to be an orthogonal matrix with a looks-linear (LL) mirrored block structure [46, 6], that achieves perfect dynamical isometry in perceptrons with ReLU activation functions. Due to the peculiarity of neighborhood aggregation, the same technique does not induce perfect dynamical isometry in GATs (or GNNs). Exploring how dynamical isometry can be achieved or approximated in general GNNs could be an interesting direction for future work. Nevertheless, using a (balanced) LL-orthogonal initialization enhances trainaibility, particularly of deeper models. We defer further discussion on the effects of orthogonal initialization, such as with identity matrices, to the appendix.

We outline the initialization procedure as follows: Let ∥ and | = denote row-wise and columnwise matrix concatenation, respectively. Let us draw a random orthonormal submatrix U l and define W 1 = [ U 1 ∥ -U 1 ] where U 1 ∈ R n 1 2 × n 0 , W l = [[ U l ∥ -U l ] | = [ -U l ∥ U l ]] where U l ∈ R n l 2 × n l -1 2 for l = { 2 , . . . , L -1 } , and W L = [ U L | = -U L ] where U L ∈ R n L × n l -1 2 . Since U L is an orthonormal matrix, ∥ ∥ W L [: , i ] ∥ ∥ 2 = 2 by definition, and by recursive application of Eq. (2.3 (with a l = 0 ), balancedness requires that ∥ ∥ W 1 [ i : , ] ∥ ∥ 2 = 2 . Therefore, we set β i = 2 for i ∈ n 1 .

3 Experiments

The main purpose of our experiments is to verify the validity of our theoretical insights and deduce an explanation for a major trainability issue of GATs that is amplified with increased network depth. Based on our theory, we understand the striking observation in Figure 4a that the parameters of default Xavier initialized GATs struggle to change during training. We demonstrate the validity of our theory with the nature of our solution. As we balance the parameter initialization (see Procedure 2.6), we allow the gradient signal to pass through the network layers, which enables parameter changes during training in all GAT layers. In consequence, balanced initialization schemes achieve higher generalization performance and significant training speed ups in comparison with the default Xavier initialization. It also facilitates training deeper GAT models.

To analyze this effect systematically, we study the different GAT initializations listed in Table 1 on generalization (in % accuracy) and training speedup (in epochs) using nine common benchmark datasets for semi-supervised node classification tasks. We defer dataset details to the supplement. We use the standard provided train/validation/test splits and have removed the isolated nodes from Citeseer. We use the Pytorch Geometric framework and run our experiments on either Nvidia T4 Tensor Core GPU with 15 GB RAM or Nvidia GeForce RTX 3060 Laptop GPU with 6 GB RAM. We allow each network to train, both with SGD and Adam, for 5000 epochs (unless it converges earlier, i.e. achieves training loss ≤ 10 -4 ) and select the model state with the highest validation accuracy. For each experiment, the mean ± 95% confidence interval over five runs is reported for accuracy and epochs to the best model. All reported results use ReLU activation, weight sharing and no biases, unless stated otherwise. No weight sharing leads to similar trends, which we therefore omit. Our experimental code is available at https://github.com/RelationalML/GAT_Balanced_Initialization .

SGD We first evaluate the performance of the different initialization schemes of GAT models using vanilla gradient descent. Since no results have been previously reported using SGD on these datasets, to the best of our knowledge, we find that setting a learning rate of 0 . 1 , 0 . 05 and 0 . 005 for L = [2 , 5] , L = [10 , 20] , and L = 40 , respectively, allows for reasonably stable training on Cora, Citeseer, and

Pubmed. For the remaining datasets, we set the learning rate to 0 . 05 , 0 . 01 , 0 . 005 and 0 . 0005 for L = [2 , 5] , L = 10 , L = 20 , and L = 40 , respectively. Note that we do not perform fine-tuning of the learning rate or other hyperparameters with respect to performance and use the same settings for all initializations to allow fair comparison.

Figure 6 highlights that models initialized with balanced parameters, Bal O , consistently achieve the best accuracy and significantly speed up training, even on shallow models of 2 layers. In line with our theory, the default Xavier (Xav) initialization hampers model training (see also Figure 4a). Interestingly, the default Xavier (Xav) initialized deeper models tend to improve in performance with width but cannot compete with balanced initialization schemes. For example, the accuracy achieved by Xav for L = 10 increases from 24 . 7% to 68 . 7% when the width is increased from 64 to 512 . We hypothesize that the reason why width aids generalization performance is that a higher number of (almost random) features supports better models. This would also be in line with studies of overparameterization in vanilla feed-forward architectures, where higher width also aids random feature models [2, 4, 57, 43].

The observation that width overparameterized models enable training deeper models may be of independent interest in the context of how overparameterization in GNNs may be helpful. However, training wider and deeper (hence overall larger) models is computationally inefficient, even for datasets with the magnitude of Cora. In contrast, the Bal O initialized model for L = 10 is already able to attain 79 . 7% even with width = 64 and improves to 80 . 9% with width = 512 . Primarily the

(a) Test accuracy.

(b) Epochs to convergence.

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

80

Citeseer

Pubmed

3000

Citeseer

Pubmed

75

2500

Epochs To Convergence

Accuracy (%)

70

2000

65

1500

60

1000

55

500

50

0

2

5

10

20

2

5

10

20

2

5

10

20

2

5

10

20

2

5

10

20

2

5

10

20

Depth

Depth

Depth

Depth

Depth

Depth

parameter balancing must be responsible for the improved performance, as it is not sufficient to simply initialize the attention weights to zero, as shown in Figure 4a and 5.

For the remaining datasets, we report only the performance of networks with 64 hidden dimension in Table 2 for brevity. Since we have considered the Xav Z initialization only as an ablation study and find it to be ineffective (see Figure 6a), we do not discuss it any further.

Note that for some datasets (e.g.Pubmed, Wisconsin), deeper models (e.g. L = 40 ) are unable to train at all with Xavier initialization, which explains the high variation in test accuracy across multiple runs as, in each run, the model essentially remains a random network. The reduced performance even with balanced initialization at higher depth may be due to a lack of convergence within 5000 epochs. Nevertheless, the fact that the drop in performance is not drastic, but rather gradual as depth increases, indicates that the network is being trained, i.e. trainability is not an issue.

As a stress test, we validate that GAT networks with balanced initialization maintain their trainability at even larger depths of 64 and 80 layers (see Fig. 9 in appendix). In fact, the improved performance and training speed-up of models with balanced initialization as opposed to models with standard initialization is upheld even more so for very deep networks.

Apart from Cora, Citeseer and Pubmed, the other datasets are considered heterophilic in nature and standard GNN models do not perform very well on them. State-of-the-art performance on these datasets is achieved by specialized models. We do not compare with them for two reasons: i) they do not employ an attention mechanism, which is the focus of our investigations, but comprise various special architectural elements and ii) our aim is not to outperform the SOTA, but rather highlight a mechanism that underlies the learning dynamics and can be exploited to improve trainability. We also demonstrate the effectiveness of a balanced initialization in training deep GATs in comparison to LipschitzNorm [14], a normalization technique proposed specifically for training deep GNNs with self-attention layers such as GAT (see Table 6 in appendix).

Adam Adam, the most commonly used optimizer for GNNs, stabilizes the training dynamics and can compensate partially for problematic initializations. However, the drop in performance with depth, though smaller, is inevitable with unbalanced initialization. As evident from Figure 7, Bal O initialized models achieve higher accuracy than or at par with Xavier initialization and converge in fewer epochs. As we observe with SGD, this difference becomes more prominent with depth, despite the fact that Adam itself significantly improves the trainability of deeper networks over SGD. Our argument of small relative gradients (see Eq.(8)) also applies to Adam. We have used the initial learning rates reported in the literature [9]: 0 . 005 for Cora and Citeseer, and 0 . 01 for Pubmed for the 2 and 5 layer networks. To allow stable training of deeper networks, we reduce the initial learning rate by a factor 0 . 1 for the 10 and 20 layer networks on all three datasets.

Architectural variations We also consider other architectural variations such as employing ELUs instead of ReLUs, using multiple attention heads, turning off weight sharing, and addition of standard elements such as weight decay and dropout. In all cases, the results follow a similar tendency as already reported. Due to a lack of space, we defer the exact results to the appendix. These variations further increase the performance of networks with our initialization proposals, which can therefore be regarded as complementary. Note that residual skip connections between layers are also supported

in a balanced initialization provided their parameters are initialized with zero. However, to isolate the contribution of the initialization scheme and validate our theoretical insights, we have focused our exposition on the vanilla GAT version.

Limitations The derived conservation law only applies to the self-attention defined in the original GAT and GATv2 models, and their architectural variations such as ω GAT[18] (see Fig. 10 in appendix). Note that the law also holds for the non-attentive GCNs (see Table 7 in appendix) which are a special case of GATs (where the attention parameters are simply zero). Modeling different kinds of self-attention such as the dot-product self-attention in [28] entails modification of the conservation law, which has been left for future work.

4 Discussion

GATs [9] are powerful graph neural network models that form a cornerstone of learning from graphbased data. The dynamic attention mechanism provides them with high functional expressiveness, as they can flexibly assign different importance to neighbors based on their features. However, as we slightly increase the network depth, the attention and feature weights face difficulty changing during training, which prevents us from learning deeper and more complex models.

We have derived an explanation of this issue in the form of a structural relationship between the gradients and parameters that are associated with a feature. This relationship implies a conservation law that preserves a sum of the respective squared weight and attention norms during gradient flow. Accordingly, if weight and attention norms are highly unbalanced as is the case in standard GAT initialization schemes, relative gradients for larger parameters do not have sufficient room to increase.

This phenomenon is similar in nature to the neural tangent kernel (NTK) regime [24, 54], where only the last linear layer of a classic feed-forward neural network architecture can adapt to a task. Conservation laws for basic feed-forward architectures [15, 3, 32] do not require an infinite width assumption like NTK-related theory and highlight more nuanced issues for trainability. Furthermore, they are intrinsically linked to implicit regularization effects during gradient descent [41]. Our results are of independent interest also in this context, as we incorporate the attention mechanism into the analysis, which has implications for sequence models and transformers as well.

One of these implications is that an unbalanced initialization hampers the trainability of the involved parameters. Yet, the identification of the cause already contains an outline for its solution. Balancing the initial norms of feature and attention weights leads to more effective parameter changes and significantly faster convergence during training, even in shallow architectures. To further increase the trainability of feature weights, we endow them with a balanced orthogonal looks-linear structure, which induces perfect dynamical isometry in perceptrons [10] and thus enables signal to pass even through very deep architectures. Experiments on multiple benchmark datasets have verified the validity of our theoretical insights and isolated the effect of different modifications of the initial parameters on the trainability of GATs.

References

Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. In International Conference on Learning Representations , 2021.

Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations , 2019.

Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach, 2018.

Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Generalization of two-layer neural networks: An asymptotic viewpoint. In International Conference on Learning Representations , 2020.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.

David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question? In International Conference on Machine Learning , 2018.

Peter L. Bartlett, David P. Helmbold, and Philip M. Long. Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks. In International Conference on Machine Learning , 2018.

T. Bian, X. Xiao, T. Xu, P. Zhao, W. Huang, Y. Rong, and J. Huang. Rumor detection on social media with bi-directional graph convolutional networks. In AAAI Conference on Artificial Intelligence , 2020.

Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? In International Conference on Learning Representations , 2022.

Rebekka Burkholz and Alina Dubatovka. Initialization of ReLUs for dynamical isometry. In Advances in Neural Information Processing Systems , volume 32, 2019.

Tianle Cai, Shengjie Luo, Keyulu Xu, Di He, Tie-Yan Liu, and Liwei Wang. Graphnorm: A principled approach to accelerating graph neural network training. In International Conference on Machine Learning , 2021.

Tianlong Chen, Kaixiong Zhou, Keyu Duan, Wenqing Zheng, Peihao Wang, Xia Hu, and Zhangyang Wang. Bag of tricks for training deeper graph neural networks: A comprehensive benchmark study. In IEEE Transactions on Pattern Analysis and Machine Intelligence , 2023.

Weilin Cong, Morteza Ramezani, and Mehrdad Mahdavi. On provable benefits of depth in training graph convolutional networks. In Advances in Neural Information Processing Systems , 2021.

George Dasoulas, Kevin Scaman, and Aladin Virmaux. Lipschitz normalization for self-attention layers with application to graph neural networks. In International Conference on Machine Learning , 2021.

S. S. Du, W. Hu, and J. D. Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In Advances in Neural Information Processing Systems , 2018.

Moshe Eliasof, Eldad Haber, and Eran Treister. Pde-gcn: Novel architectures for graph neural networks motivated by partial differential equations. In Advances in Neural Information Processing Systems , 2021.

Moshe Eliasof, Eldad Haber, and Eran Treister. pathgcn: Learning general graph spatial operators from paths. In International Conference on Machine Learning , 2022.

Moshe Eliasof, Lars Ruthotto, and Eran Treister. Improving graph neural networks with learnable propagation operators. In International Conference on Machine Learning , 2023.

Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In International Conference on Artificial Intelligence and Statistics , volume 9, pages 249-256, May 2010.

M. Gori, G. Monfardini, and F. Scarselli. A new model for learnig in graph domains. In IEEE International Joint Conference on Neural Networks , 2005.

Kai Guo, Kaixiong Zhou, Xia Hu, Yu Li, Yi Chang, and Xin Wang. Orthogonal graph neural networks. In AAAI Conference on Artificial Intelligence , 2022.

William L. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in Neural Information Processing Systems , 2018.

Wei Hu, Lechao Xiao, and Jeffrey Pennington. Provable benefit of orthogonal initialization in optimizing deep linear networks. In International Conference on Machine Learning , 2020.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems , 2018.

Ajay Jaiswal, Peihao Wang, Tianlong Chen, Justin F. Rousseau, Ying Ding, and Zhangyang Wang. Old can be gold: Better gradient flow can make vanilla-gcns great again. In Advances in Neural Information Processing Systems , 2022.

Stefanie Jegelka. Theory of graph neural networks: Representation and learning. In The International Congress of Mathematicians , 2022.

S. Kearnes, K. McCloskey, M. Berndl, V. Pande, and P. Riley. Molecular graph convolutions: Moving beyond fingerprints. In Journal of Computer-Aided Molecular Design , 2016.

Dongkwan Kim and Alice Oh. How to find your friendly neighborhood: Graph attention design with self-supervision. In International Conference on Learning Representations , 2021.

Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations , 2017.

Boris Knyazev, Graham W. Taylor, and Mohamed R. Amer. Understanding attention and generalization in graph neural networks. In Advances in Neural Information Processing Systems , 2019.

D. Kunin, J. Sagastuy-Brena, S. Ganguli, D.L.K. Yamins, and H. Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. In International Conference on Learning Representations , 2021.

Thien Le and Stefanie Jegelka. Training invariances and the low-rank phenomenon: beyond linear networks. In International Conference on Learning Representations , 2022.

Guohao Li, Matthias Müller, Bernard Ghanem, and Vladlen Koltun. Training graph neural networks with 1000 layers. In International Conference on Machine Learning , 2022.

Guohao Li, Matthias Müller, Ali Thabet, and Bernard Ghanem. Deepgcns: Can gcns go as deep as cnns? In International Conference on Computer Vision , 2019.

Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. In AAAI Conference on Artificial Intelligence , 2018.

Meng Liu, Hongyang Gao, and Shuiwang Ji. Towards deeper graph neural networks. In SIGKDD International Conference on Knowledge Discovery and Data Mining , 2020.

Dmytro Mishkin and Jiri Matas. All you need is a good init. In International Conference on Learning Representations , 2016.

Pál András Papp, Karolis Martinkus, Lukas Faber, and Roger Wattenhofer. Dropgnn: Random dropouts increase the expressiveness of graph neural networks. In Advances in Neural Information Processing Systems , 2021.

Hongbin Pei, Bingzhe Wei, Kevin Chen-Chuan Chang, Yu Lei, and Bo Yang. Geom-gcn: Geometric graph convolutional networks. In International Conference on Learning Representations , 2020.

Oleg Platonov, Denis Kuznedelev, Michael Diskin, Artem Babenko, and Liudmila Prokhorenkova. A critical look at the evaluation of gnns under heterophily: are we really making progress? In International Conference on Learning Representations , 2023.

Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. In Advances in Neural Information Processing Systems , 2020.

Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convolutional networks on node classification. In International Conference on Learning Representations , 2020.

Karthik Abinav Sankararaman, Soham De, Zheng Xu, W. Ronny Huang, and Tom Goldstein. The impact of neural network overparameterization on gradient confusion and stochastic gradient descent. In International Conference on Machine Learning , 2020.

Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations , 2014.

F. Scarselli, M. Gori, A.C Tsoi, M. Hagenbuchner, and G. Monfardini. Computational capabilities of graph neural networks. In IEEE Transactions on Neural Networks , 2009.

Wenling Shang, Kihyuk Sohn, Diogo Almeida, and Honglak Lee. Understanding and improving convolutional neural networks via concatenated rectified linear units. In International Conference on Machine Learning , 2016.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , 2017.

Petar Veliˇ ckovi´ c, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations , 2018.

Peihao Wang, Wenqing Zheng, Tianlong Chen, and Zhangyang Wang. Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory to practice. In International Conference on Learning Representations , 2022.

Z. Wu, S. Pan, G. Long, J. Jiang, and C. Zhang. Graph wavenet for deep spatial-temporal graph modeling. In International Joint Conference on Artificial Intelligence , 2019.

K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations , 2019.

Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning on graphs with jumping knowledge networks. In International Conference on Machine Learning , 2018.

Keyulu Xu, Mozhi Zhang, Stefanie Jegelka, and Kenji Kawaguchi. Optimization of graph neural networks: Implicit acceleration by skip connections and more depth. In International Conference on Machine Learning , 2021.

Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In Marina Meila and Tong Zhang, editors, International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 11727-11737, 2021.

Han Yang, Kaili Ma, and James Cheng. Rethinking graph regularization for graph neural networks. In Advances in Neural Information Processing Systems , 2021.

Zhilin Yang, William W. Cohen, and Ruslan Salakhutdinov. Revisiting semi-supervised learning with graph embeddings. In International Conference on Machine Learning , 2016.

Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. In Advances in Neural Information Processing Systems , 2019.

Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do transformers really perform bad for graph representation? In Advances in Neural Information Processing Systems , 2021.

Jiaxuan You, Rex Ying, and Jure Leskovec. Design space for graph neural networks. In Advances in Neural Information Processing Systems , 2021.

Seongjun Yun, Minbyul Jeong, Raehyun Kim, Jaewoo Kang, and Hyunwoo J. Kim. Graph transformer networks. In Advances in Neural Information Processing Systems , 2020.

M. Zhang and Y. Chen. Inductive matrix completion based on graph neural networks. In International Conference on Learning Representations , 2020.

Wentao Zhang, Zeang Sheng, Ziqi Yin, Yuezihan Jiang, Yikuan Xia, Jun Gao, Zhi Yang, and Bin Cui. Model degradation hinders deep graph neural networks. In ACM SIGKDD Conference on Knowledge Discovery and Data Mining , 2022.

Wentao Zhang, Mingyu Yang, Zeang Sheng, Yang Li, Wen Ouyang, Yangyu Tao, Zhi Yang, and Bin Cui. Node dependent local smoothing for scalable graph learning. In Advances in Neural Information Processing Systems , 2021.

Lingxiao Zhao and Leman Akoglu. Pairnorm: Tackling oversmoothing in gnns. In International Conference on Learning Representations , 2020.

Kaixiong Zhou, Xiao Huang, Yuening Li, Daochen Zha, Rui Chen, and Xia Hu. Towards deeper graph neural networks with differentiable group normalization. In Advances in Neural Information Processing Systems , 2020.

Kuangqi Zhou, Yanfei Dong, Kaixin Wang, Wee Sun Lee, Bryan Hooi, Huan Xu, and Jiashi Feng. Understanding and resolving performance degradation in graph convolutional networks. In Conference on Information and Knowledge Management , 2021.

Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layerdependent importance sampling for training deep and large graph convolutional networks. In Advances in Neural Information Processing Systems , 2019.

6 Proofs of Theorems

Notation. Scalar and element-wise products are denoted ⟨ , ⟩ and ⊙ , respectively. Boldface lowercase and uppercase symbols represent vectors and matrices, respectively.

Our proof of Theorem 2.2 utilizes a rescale invariance that follows from Noether's theorem, as stated by [31]. Note that we could also derive the gradient structure directly from the derivatives, but the rescale invariance is easier to follow.

Definition 6.1 (Rescale invariance) . The loss L ( θ ) is rescale invariant with respect to disjoint subsets of the parameters θ 1 and θ 2 if for every λ > 0 we have L ( θ ) = L (( λθ 1 , λ -1 θ 2 , θ d )) , where θ = ( θ 1 , θ 2 , θ d ) .

We frequently utilize the following relationship.

Lemma 6.2 (Gradient structure due to rescale invariance [31]) . The rescale invariance of L enforces the following geometric constraint on the gradients of the loss with respect to its parameters:

We first use this rescale invariance to prove our main theorem for a GAT with shared feature transformation weights and a single attention head, as described in Def. 2.1. The underlying principle generalizes to other GAT versions as well, as we exemplify with two other variants. Firstly, we study the case of unshared weights for feature transformation of source and target nodes. Secondly, we discuss multiple attention heads.

Proof of Theorem 2.2 Our main proof strategy relies on identifying a multiplicative rescale invariance in GATs that allows us to apply Lemma 6.2. We identify rescale invariances for every neuron i at layer l that induce the stated gradient structure.

Specifically, we define the components of in-coming weights to the neuron as θ 1 = { w | w ∈ W l [ i, : ] } and the union of all out-going edges (regarding features and attention) as θ 2 = { w | w ∈ W l +1 [: , i ] } ∪ { a l [ i ] } . It is left to show that these parameters are invariant under rescaling.

Let us, therefore, evaluate the GAT loss at λθ 1 and λ -1 θ 2 and show that it remains invariant for any choice of λ > 0 . Note that the only components of the network that potentially change under rescaling are h l u [ i ] , h l +1 v [ j ] , and α l uv . We denote the scaled network components with a tilde resulting in ˜ h l u [ i ] , ˜ h l +1 v [ k ] , and ˜ α l uv As we show, parameters of upper layers remain unaffected, as ˜ h l +1 v [ k ] coincides with its original non-scaled variant ˜ h l +1 v [ k ] = h l +1 v [ k ] .

Let us start with the attention coefficients. Note that ˜ a l [ i ] = λ -1 a l [ i ] and ˜ W l [ i, j ] = λW l [ i, j ] . This implies that

which follows from the positive homogeneity of ϕ 2 that allows

̸

̸

Since ˜ α l uv = α l uv , it follows that

In the next layer, we therefore have

Thus, the output node representations of the network remain unchanged, and according to Def. 6.1, the loss L is rescale-invariant.

Consequently, as per Lemma 6.2, the constraint in Eq.(10) can be written as:

which can be rearranged to Eq.((2.2):

thus proving Theorem 2.2.

Proof of Corollary 2.3 Given that gradient flow applied on loss L is captured by the differential equation

which implies:

substituting in Eq.(4) similarly for gradient flow of W l +1 [: , i ] and a l [ i ] , as done in Eq.(17) yields Theorem 2.2:

Therefore, Eq.(4) and consequently Eq.(5) in Corollary 2.3 hold.

Summing over i ∈ [ n l ] on both sides of Eq.(4) yields Corollary 2.4.

Theorem 6.3 (Structure of gradients for GAT without weight-sharing) . Let a GAT network as defined by Def. 2.1 consisting of L layers be given. The feature transformation parameters W l s and W l t of the source and target nodes of an edge and attention parameters a l of a GAT layer l are defined according to Eq.(1) and (2).

Then the gradients for layer l ∈ [ L -1] in the network are governed by the following law:

Proof. The proof is analogous to the derivation of Theorem 2.2. We follow the same principles and define the disjoint subsets θ 1 and θ 2 of the parameter set θ , associated with a neuron i in layer l accordingly, as follows:

Then, the invariance of node representations follows similarly to the proof of Theorem 2.2.

The only components of the network that potentially change under rescaling are h l u [ i ] , h l +1 v [ j ] , and α l uv . We denote the scaled network components with a tilde resulting in ˜ h l u [ i ] , ˜ h l +1 v [ k ] , and ˜ α l uv As we show, parameters of upper layers remain unaffected, as ˜ h l +1 v [ k ] coincides with its original non-scaled variant ˜ h l +1 v [ k ] = h l +1 v [ k ] .

Let us start with the attention coefficients. Note that ˜ a l [ i ] = λ -1 a l [ i ] , ˜ W s l [ i, j ] = λW l s [ i, j ] and ˜ W l t [ i, j ] = λW l t [ i, j ] . This implies that

which follows from the positive homogeneity of ϕ 2 that allows

̸

̸

Since ˜ α l uv = α l uv , it follows that

In the next layer, we therefore have

Thus, the application of Lemma 6.2 derives Theorem 6.3.

Theorem 6.4 (Structure of gradients for GAT with multi-headed attention) . Given the feature transformation parameters W l k and attention parameters a l k of an attention head k ∈ [ K ] in a GAT layer of a L layer network. Then the gradients of layer l ∈ [ L -1] respect the law:

Proof. Each attention head is independently defined by Eq.1 and Eq.2, and thus Theorem 2.2 holds for each head, separately. The aggregation of multiple heads in a layer is done over node representations of each head in either one of two ways:

Thus, the rescale invariance and consequent structure of gradients of the parameters in each head are not altered and Eq.(6.2) holds by summation of the conservation law over all heads.

For the most general case of a GAT layer without weight-sharing and with multi-headed attention, each term in Theorem 6.3 is summed over all heads. Further implications analogous to Corollary 2.3, and 2.4 can be derived for these variants by similar principles.

7 Training Dynamics

In the main paper, we have shared observations regarding the learning dynamics of GAT networks of varying depths with standard initialization and or balanced initialization schemes. Here, we report additional figures for more varied depths to complete the picture.

All figures in this section correspond to a L layer GAT network trained on Cora. For SGD, the learning rate is set to γ = 0 . 1 , γ = 0 . 05 and γ = 0 . 01 for L = 5 , L = 10 and L = 20 , respectively. For Adam, γ = 0 . 005 and γ = 0 . 0001 for L = 5 and L = 10 , respectively.

Relative change of a feature transformation parameter w is defined as | ( w ∗ -w 0 ) / w ∗ | where w 0 is the initialized value and w ∗ is the value for the model with maximum validation accuracy during training. Absolute change | a ∗ -a 0 | is used for attention parameters a since attention parameters are initialized with 0 in the balanced initialization. We view the fraction of parameters that remain unchanged during training separately in Figure 8, and examine the layer-wise distribution of change in parameters in Figure 11 without considering the unchanged parameters.

(a)

L

= 5

(b)

L

= 10

(c)

L

= 20

Feat.Weights

Att.Weights

Feat.Weights

Att.Weights

Feat.Weights

Att.Weights

Frac. of Params.

1.0

1.0

ofParams.

1.0

1.0

Frac. of Params.

1.0

1.0

Xav

0.5

0.5

0.5

0.5

0.5

0.5

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

0.0

0.0

0.0

0.0

0.0

0.0

1

234

5

12345

246810

246810

48121620

48121620

Layer

Layer

Layer

Layer

Layer

Layer

Figure 8: Layer-wise fraction of feature transformation parameters W l and attention parameters a l with zero relative and absolute change, respectively, trained with SGD. For L = 2 , no parameters with zero change existed. A small number of feature weights do not change in a 5 layer unbalanced network initialized with Xavier, but this fraction becomes significantly large when depth is increased to L = 10 . Note that W 1 contains a much larger number of parameters compared to the intermediate layers (specifically in this case, though it is generally common). At L = 20 , nearly all w ∈ W l with an unbalanced initialization (Xav. and Xav+ZeroAtt) struggle to change during training, whereas the balanced Xavier and LL-orthogonal initialization are able to drive change in most w ∈ W l and all a ∈ a l parameters, allowing the network to train.

To analyze how gradients vary in the network during training, we define the relative gradient norm of feature transformation and attention parameters for a layer l as ∥ ∇ Wl W l ∥ F / ∥ W l ∥ F and ∥ ∇ a l a l ∥ / ∥ a l ∥ , respectively. Figures 3b and 4b, and Figure 12 depict relative gradient norms for training under SGD and Adam respectively.

8 Additional Results

The results of the main paper have focused on the vanilla GAT having ReLU activation between consecutive layers, a single attention head, and shared weights for feature transformation of the source and target nodes, optimized with vanilla gradient descent (or Adam) without any regularization. Here, we present additional results for training with architectural variations, comparing a balanced initialization with a normalization scheme focused on GATs, and discussing the impact of an orthogonal initialization and the applicability of the derived conservation law to other messagepassing GNNs (MPGNNs).

Training variations To understand the impact of different architectural variations, common regularization strategies such as dropout and weight decay, and different activation functions, we conducted additional experiments. For a 10 -layer network with width 64 , Table 3 and 4 report our results for SGD and Adam, respectively. The values of hyperparameters were taken from [48].

We find that the balanced initialization performs similar in most variations, outperforming the unbalanced standard initialization by a substantial margin. Dropout generally does not seem to be helpful to the deeper network ( L = 10 ), regardless of the initialization. From Table 3 and 4, it is evident that although techniques like dropout and weight decay may aid optimization, they are alone not enough to enable the trainability of deeper network GATs and thus are complementary to balancedness.

Note that ELU is not a positively homogeneous activation function, which is an assumption in out theory. In practice, however, it does not impact the Xavier-balanced initialization (Bal X ). However, the Looks-Linear orthogonal structure is specific to ReLUs. Therefore, the orthogonality and balancedness of the Bal O initialization are negatively impacted by ELU, although the Adam optimizer seems to compensate for it to some extent.

In addition to increasing the trainability of deeper networks, the balanced initialization matches the state-of-the-art performance of Xavier initialized 2 layer GAT, given the architecture and optimization hyperparameters as reported in [48]. We used an existing GAT implementation and training script 2 of

Cora that follows the original GAT(v1)[48] paper and added our balanced initialization to the code. As evident from Table 5, the balanced initialization matches SOTA performance of GAT(v1) on Cora (83 -84%) (on a version of the dataset with duplicate edges). GAT(v2) achieved 78 . 5% accuracy on Pubmed but the corresponding hyperparameter values were not reported. Hence, we transferred them from [48]. This way, we are able to match the state-of-the-art performance of GAT(v2) on Pubmed, as shown in Table 5.

.

Comparison with Lipshitz Normalization Afeasible comparison can be carried out with [14] that proposes a Lipschitz normalization technique aimed at improving the performance of deeper GATs in particular. We use their provided code to reproduce their experiments on Cora, Citeseer, and Pubmed for 2,5,10,20 and 40-layer GATs with Lipschitz normalization and compare them with LLortho+Bal initialization in Table 6. Note that Lipschitz normalization has been shown to outperform other previous normalization techniques for GNNs such as pair-norm [64] and layer-norm [5].

Impact of orthogonal initialization In the main paper, we have advocated and presented results for using a balanced LL-orthogonal initialization. Here, we discuss two special cases of orthogonal initialization (and their balanced versions): identity matrices that have also been used for GNNs in [18, 16], and matrices with looks-linear structure using an identity submatrix (LLidentity) since a looks-linear structure would be more effective for ReLUs [10].

In line with [18], identity and LLidentity are used to initialize the hidden layers while Xavier is used to initialize the first and last layers. A network initialized with identity is balanced by adjusting the weights of the first and last layer to have norm 1 (as identity matrices have row-wise and column-wise norm). A network with LLidentity initialization is balanced to have norm 2 in the first and last layer, similar to LLortho initialization. We compare the performance using these four base initializations (one standard Xavier, and three orthogonal cases) and their balanced counterparts in Fig. 9.

We observe that balancing an (LL-)orthogonal initialization results in an improvement in the generalization ability of the network in most cases and speeds up training, particularly for deeper networks. However, note that an (LL-)orthogonal initialization itself also has a positive effect on trainability in particular of deep networks. Contributing to this fact is the mostly-balanced nature of an (LL)orthogonal initialization i.e. given hidden layers of equal dimensions the network is balanced in all layers except the first and last layers (assuming zero attention parameters), which allows the model to train, as opposed to the more severely imbalanced standard Xavier initialization. This further enforces the key takeaway of our work that norm imbalance at initialization hampers the trainability of GATs. In addition, the LLortho+Bal initialization also speeds up training over the LLortho initialization even in cases in which the generalization performance of the model is at par for both initializations.

Xavier

Xavier+Bal

%Identity

Identity+Bal

LLidentity

LLidentity+Bal

■LLortho+Bal

90

80

70

60

Test Accuracy(

50

40

30

20

10

0

2

5

10

20

40

64

80

2

5

10

20

40

64

80

2

5

10

20

40

64

80

Citeseer

Pubmed

Number of Layers

Note that identity initializations have also been explored in the context of standard feed-forward neural networks. While they tend to work in practice, the lack of induced feature diversity can be problematic from a theoretical point of view (see e.g. for a counter-example [7]). However, potentially due to the reasons regarding feature diversity and ReLU activation discussed above, the balanced looks-linear random orthogonal initialization (LLortho+Bal.) outperforms initialization with identity matrices (see Fig. 9). In most cases, the balanced versions outperform the imbalanced version of the base initialization and the performance gap becomes exceedingly large as network depth increases.

Applicability to other MPGNNs As GAT is a generalization of the GCN, all theorems are also applicable to GCNs (where the attention parameters are simply 0 ). We provide additional experiments in Table 7, comparing the performance of GCN models of varying depths with imbalanced and balanced initializations trained on Cora using the SGD optimizer. We use the same hyperparameters as reported for GAT. As expected, the trainability benefits of a balanced initialization are also evident in deeper GCN networks.

The underlying principle of a balanced network initialization holds in general. However, adapting the balanced initialization to different methods entails modification of the conservation law derived for GATs to correspond to the specific architecture of the other method. For example, the proposed balanced initialization can be applied directly to a more recent variant of GAT, ω GAT [18], for which the same conservation law holds and also benefits from a balanced initialization. We conduct

additional experiments to verify this. The results in Fig. 10 follow a similar pattern as for GATs: a balanced orthogonal initialization with looks linear structure (LLortho+Bal) of ω GAT performs the best, particularly by a wide margin at much higher depths (64 and 80 layers).

Xavier

Xavier+Bal

Identity

■Identity+Bal

LLidentity

LLidentity+Bal

WLLortho

■LLortho+Bal

80

70

60

Test Accuracy

50

40

30

20

10

0

2

5

10

20

40

64

80

2

5

10

20

40

64

80

SGD

Adam

Number of Layers

In this work, we focus our exposition on GATs and take the first step in modeling the training dynamics of attention-based models for graph learning. An intriguing direction for future work is to derive modifications in the conservation law for other attention-based models utilizing the dot-product self-attention mechanism such as SuperGAT [28] and other architectures based on Transformers[47], that are widely used in Large Language Models (LLMs) and have recently also been adapted for graph learning [58, 60].

9 Datasets Summary

We used nine benchmark datasets for the transductive node classification task in our experiments, as summarized in Table 8 .

For the Platenoid datasets (Cora, Citeseer, and Pubmed) [56], we use the graphs provided by Pytorch Geometric (PyG), in which each link is saved as an undirected edge in both directions. However, the number of edges is not exactly double (for example, 5429 edges are reported in the raw Cora dataset, but the PyG dataset has 10556 < (2 × 5429) ) edges as duplicate edges have been removed in preprocessing. We also remove isolated nodes in the Citeseer dataset.

The WebKB (Cornell, Texas, and Wisconsin), Wikipedia (Squirrel and Chameleon) and Actor datasets [39], are used from the replication package provided by [40], where duplicate edges are removed.

All results in this paper are according to the dataset statistics reported in Table 8.

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Rel. Change of Feat. Weights

10°

10

10

2

3

5

6

8

9

10

Layer

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Abs. Change of Att. Weights

10

10

-3

百

卤

百

南

百

卤

10

5

10

-7

10

1

2

3

4

5

6

7

8

9

10

Layer

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Rel. Change of Feat. Weights

10

10

2

3

4

5

6

8

9

10

11

12

13

14

15

16

17

18

19

20

Layer

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Abs.Change of Att.Weights

10

-5

10

9

10

-13

10

-17

2

3

4

5

6

8

9

10

11

12

13

14

16

17

18

19

20

Layer

(a)

L

= 5

: test accuracy (left to right) is

71

.

6%

,

76

.

2%

,

77

.

3%

, and

80

.

5%

.

(b)

L

= 10

: test accuracy (left to right) is

63

.

2%

,

68

.

8%

,

71

.

8%

, and

78

.

0%

.

featureparameters

0.00

0.02

0.04

0.06

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Layer

5 4 3 2 1

2

Loss

Layer

2

Loss

Layer

. 5 4 3 2 1

2

Loss

Layer

.5 4321

2

SS07

0

0

0

0

014284256

0

124368

0

1234

0

163284

Epoch

Epoch

Epoch

Epoch

attentionparameters

0.00

0.02

0.04

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Layer

5 4321

2

Loss

Layer

.5 4 3 21

2

Loss

Layer

.5 4 3 2 1

2

SS07

Layer

5 43 21

SS07

十

0

0

0

0

0

14284256

0

1243648

0

1234

0

163284

Epoch

Epoch

Epoch

Epoch

featureparameters

0.000

0.025

0.075

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Layer

108 6 4 2

2

Loss

Layer

108 6 4 2

2

SS07

Layer

108 6 4 2

SS07

Layer

1086 4 2

2

SS07

十

1

1

0

0

0

0

100, 30 40 00

0

15,230.,5 30

0

15 20 .35 50

0

010 200 00

Epoch

Epoch

Epoch

Epoch

attentionparameters

0.000

0.025

0.050

0.075

Xav

Xav+ZeroAtt

Xav+Bal

LLortho+Bal

Layer

108642

2

Loss

Layer

108 6 4 2

2

Loss

Layer

108 6 4 2

2

Loss

Layer

108642

Loss

0

0

0

0

0

09

0

130 55 0

0

15,20 .35 30

0

00000

Epoch

Epoch

Epoch

Epoch

Peer reviews:
comments: SUMMARY:
The work derives a conservation for the dynamics of the GAT gradient flow during training. The conservation law is used to explain why it is challenging to train deep GAT models and in particular why a large portion of parameters do not change much throughout training. A new initialization scheme is introduced to mitigate these issues.

STRENGTHS:
Training GATs is notoriously problematic and this is especially the case with many layers. The work takes an interesting angle of trying to further explain why this is the case by studying conservation laws induced by the gradient flow during training. The observation that weight gradients must be small in order to satisfy such a conservation law is very interesting and practically useful. The balancing algorithm proposed is also practically useful and seems to perform well especially with deep GATs.

WEAKNESSES:
While the work focuses on GAT, it would be interesting to have a statement on more general classes of MPNNs. There are some faint connections in the paper to GCN for instance, but it would be interesting to have a more general result as well. Experimentally this might also be interesting to see if such an initialisation can be adapted to different types of MPNNs.

QUESTIONS:
Would the proposed initialisation technique work well for other types of MPNNs (with appropriate modifications)? Would you have experiments/theory to back this up?

LIMITATIONS:
The authors address the fact that the derived theory defines a conservation law that does not explain detailed dynamics but still remains relatively coarse. As such it is not able to explain some phenomena as why attention parameters change most on the first layer. I still believe that the work is a good step in helping explain phenomena in this direction.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

comments: SUMMARY:
This paper focuses on the parameter struggle training problem of the well-known GNN structure GAT. The authors propose to alleviate the issue via parameter norm balancedness. A conservation law is derived for GATs with positive homogeneous activation function as the theoretical support for the parameter norm balancedness based initialization method. The authors also conduct extensive experiments to prove its effectiveness and fast convergence property.

STRENGTHS:
- The problem is well identified.
- The proposed theories are solid, and well serves the problem.
- The proposed initialization method is simple yet effective in improving accuracy and speeding convergence.

WEAKNESSES:
- Some explanation skips details and may raise confusions. For example, the equations in line 140-142.
- Some formats could be improved for better illustration, such as larger figure size and fitted table style.

QUESTIONS:
GAT mainly exploits the attention mechanism. Can this initialization be adapted to any other methods that utilize the attention mechanism?

LIMITATIONS:
The author may address the limitation from the generalization of the proposed method beyond the GAT structure.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a theoretical analysis of the initialization of GATs and their impact on the performance of such networks, focusing on the performance vs depth aspect.\n\nThe authors propose an initialization algorithm, that starts from a random initialization and modifies the initial random weights to adhere to the findings of the theoretical analysis.\n\nThe authors then show the impact of their proposed initialization on several node classification datasets. Specifically, they show that by initializing GATs with the proposed algorithm, deep GATs can be trained to achieve better than standard (Xavier) initialized GATs.\n\nSTRENGTHS:\nThe paper addresses a real issue with GATs - the degradation of performance when more layers are added.\n\nThe theoretical analysis seems correct to me and is supported with experimental results as well as the inspection of actual training artifacts and details (e.g., the gradients of the GAT layers).\n\nWEAKNESSES:\nThe paper can be slightly better written, in terms of organization. I think that adding more paragraphs/subsections to better separate between the parts of the paper can help to ease the reader.\n\nThe paper lacks a few relevant citations that also consider graph neural networks as gradient flow (see [1][2][3]). However, their focus is not on the initialization of GATs, and therefore the paper here is novel on its own.\n\nWhile the experimental results are compelling and show the benefit of the proposed method, I think that the authors should also include comparisons with other methods.\n\n[1] GRAND: Graph Neural Diffusion\n\n[2] Understanding convolution on graphs via energies\n\n[3] Improving Graph Neural Networks with Learnable Propagation Operators\n\nQUESTIONS:\n1. Can the authors add results with more layers? For instance 64,80 layers? does the method still work with very deep network?\n\n2. Regarding the initialization of graph neural networks, it is discussed in [3][4] that the GNN weights are initialized with identity matrices. Can the authors comment on this point? is it related to the proposed method?\n\n3. Can the proposed initialization scheme be applied to other graph attention layers, such as superGAT[5] ?\n\n[4] PDE-GCN: Novel Architectures for Graph Neural Networks Motivated by Partial Differential Equations\n\n[5] How to Find Your Friendly Neighborhood: Graph Attention Design with Self-Supervision\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
This work proves a conservation law for GAT architectures, which is similar to conservation laws shown for fully connected networks. The conservation law shows a simple connection between the norms of weights of two consecutive layers. Using this law, an intuitive explanation is given for the trainability issues of deep GAT architectures. To mitigate this issue, a balanced initialization is suggested which shows improvements in generalization and training speed of deep GATs.

STRENGTHS:
1. Novel conservation law for GAT architectures.
2. The paper and technical details are well written.

WEAKNESSES:
There are no major weaknesses.
1. The motivation for using the balanced initialization and why it should improve only for deep networks is not so clear.
2. The experimental results are not complete.

See Questions for more details.

QUESTIONS:
1. Does LLOrtho without balancing work well for deep networks? It is not clear if the cause for improvement is only balancing or also the initialization with orthogonal vectors.
2. The trainability issue is explained intuitively in pages 4-5:
(a) What is the reason that the trainability issue is amplified with depth? This is not explained. 
(b) Why does this trainability issue occur only for GATs and not for fully connected networks? It seems that the same explanation holds for FC nets as well.
(c) Suggestion: It would be helpful to explain why the balanced initialization solves the trainability issue (it is not explained).

LIMITATIONS:
Limitations are not discussed. Maybe it would be helpful to add a section on this.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07