Are GATs Out of Balance?
Authors: Nimrah Mustafa
Aleksandar Bojchevski
Rebekka Burkholz
Venue: NeurIPS 2023 poster (2023)
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?
Nimrah Mustafa" nimrah. www.mmmfccsissa...ommm
Aleksandar Bojchevskil a. bojchheski wuweeiueennnnenneeeeeeeeeeeeeeee
Rebekka Burkholz" burkholz@cispa.com de
*CISPA Helmholtz Center for Information Security, 66123 Saarbrücken, Germany Unniversity of Cologne, 50923 Köln, Germany
Abstract
While the expressive power and computational capabilities of graph neural net- works (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 which node's neighbor- hood 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 dur- ing training. This effect amplified deeper GATs, which perform significantly worse than their shallow counterparts. To alleviate this problem, devise initialization scheme that balances the GAT network. Our approach i) allows effective propagation of gradients and in turn enables trainability of deeper networks, and fi) attains a considerable speedup in training and convergence time in comparison to the standard initialization. Our main theorem serves stepping stone to studying the learning dynamics of positive homogeneous models with attention mechanisms.
Introduction
rapidly growing class of model architectures for graph representation learning Graph Neural Networks (GNNs) 20 which have achieved strong empirical performance across various applications such social network analysis [3) 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.
the theoretical front, mostly the expressive power 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 GNNs, the Graph Attention Network (GAT) 48 91.
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 node's 's neighbors 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, major limitation of message-passing GNNs in general, and therefore GATs well. is severe performance degradation when the depth of network is even slightly increased. SOTA results are reportedly achieved by models of 2 or layers. This problem has largely been attributed to over-smoothing [35], a phenomenon which node representations become indistigguishable 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 [T]] 64 65| 666 or regularization
38 55 67]. and impose architectural changes such skip connections (dense
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

and residual) [34 53 or offer other engineering solutions [33 or their combinations [12. Other approaches to overcome over-smoothing include learning CNN- spatial operators from random paths in the graph instead of the point-wise graph Laplacian operator [7 and learning adaptive
receptive fields (different 'effective' neighhorhoods 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-awas -wwre 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 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 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 deeper understanding of gradient dynamics has been developed and norm balanced- ness is 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. model of gradient descent with infinitesimally small learning rate or step size). Concretely, for fully connected ferdffowwrd networks and deep convolutional neural networks with continuous homogenous activation functions such as ReLUS,, the difference between the squared 12-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 (3].
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 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 GAT' and GCNs, which can be alleviated with 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 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 transformers [47 and in particular, vision transformers which require depth more than NLP transformers 49) Our contributions are as follows.
We derive 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 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 and edge set E C V V, where the neighborhood of node e is given as NVvv) = (u)(u,v) E E). GNN layer computes each node's 's representation by aggregating over its neighbors' representation. In GATs, this aggregation weighted by parameterized attention coefficients Cuv. which indicate the importance of node for
2

---

1+1
Table I: Attributes of init, schemes
W
Init.
Feat.
Attn.
Bal.?
Xav
Xavier
Xavier
No
Xavz
Xavier
Zero
No
11-1
W
im
a
W "112k
nist
Balx
Xavier
Zero
Yes
Balo
LLortho
Zero
Yes
Figure 1: Params, balance for neuron [nx]
Given input representations h for all nodes E V, GAT
layer/ transforms those to:
nt = o( Σ d'u W' where
vEN(n)
exp((a')T LeakyReLU # +W/hh E E'EN(()) exp((a') (Wag + with
(2)
We consider & to be a positively homogeneous activation functions (i.e o(x) = xo' (x) and conse- quently, ((rr) aδ(x)) for positive scalars a), such ReLU ((x) max .r, 0) LeakyReI.U
<(x) = max xr 0) + +-a max(-=x 0}. The feature transformation weights W, and W2 for source and target nodes, respectively, may also be shared such that W W, = W..
Definition 2.1. Given training data {mmm 3m 1M C R2 x EP for M S V. f: R R' be the function represented by a network constructed by stacking GAT layers as defined in Eq. [ and (2) with W = W. We and ho. g(zm). ). Each layer I E [L] of size nı has associated parameters: and n,, Given a diffcrentiable loss function é Rd x IP + R, the loss L 1/M Di=1 ((f(mm) vm) is used to update model parameters
a feature weight matrix W E IR and an attention vector a' E I where
{W' at R==1 with learning rate y by gradient descent, i.e., w = w² VV. where V [ ac/owwi]] and w is set by the initialization scheme, For an infinitesimal 7 O, dynamics of gradient descent behave similarly gradient flow defined by dw/at V C. where t is the continuous time index.
We use W W[: i]. and ali] to denote the 8th row of W. column of W. and entry of a, respectively. Note that M [ :] is the vector of weights incoming to neuron E [ru] and W++1 [= [ is the vector of weights outgoing from the same neuron. For the purpose of discussion, let E [d:].j [d:-1], and k E [d, as depicted in Fig. (..-) 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-weight- sharing case, as we state and discuss 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 and a' of layer in a GAT network described in Def. 2.1 the structure of the gradients for layer I E IL- 1] in the network is conserved according to the following law:
(w [. VWIka [i], V Va (W'+1 [1 Vwinaas
Corollary 2.3 (Norm conservation of weights incoming and outgoing at every neuron). Given the gradient structure of Theorem 2.2 positive homogeneity of the activation in Eq. gradient flow in the network adheres to the following invariance for [L 1]:
d
[|W'
||a' (I =11?)
dt
TTis definition of GAT is termed GATv2
9 Hereafter, refer to GATv2 GAT.

---

Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
25
066
0005
2.00
-55
5000
6000
5000
Fpezch
Epesh
Fjuch
Fpoch
Fpuch
Fpachh
Fpuzhh
Fpoch
(a) 5 with GD
(b) with Adam
(c) c with Xavier
(d) c with Balo
Figure 2: Visualizing conservation laws: (a),(b) show å = 0 (can not empirically be exactly zero due
to finite y = 0.1) from Eq. ( for hidden layers when L = 3. (c).(d) show the value of c from Eq. (5) which is determined by the mitiaaliztiinn and should be 0 when the network is balanced, for L = 10 trained with GD. With Xav. the network unbalanced and hence c of 0 for all neurons. With Balo, = 0 exactly initialization for all neurons by design but during training. it varies slightly and e 0 as the network becomes slightly unbalanced due to the finite = 0.1. As both ē and e approximately meet their required value, we conclude that the derived law holds for practical purposes as well.
then directly follows by integration over time that the 12- 12-norms of feature and attention weights incoming to and outgoing from a neuron are preserved such that
|| [r, ||a' - | ill? -.
(5)
where c = dt (IIw [x.1]= ||a'[(] ||a(( at initialization (i.e. t t= 0).
We denote the 'degree of balancedness' by e and call an initialization balanced if c O.
Corollary 2.4 (Norm conservation across layers), Given Eq. (4) the invariance of gradient flow at the layer level for E [L 1] by summing over E [mi] is:
d dt
dt (1W4+)11))
d
Remark 2.5. Similar conservation laws also hold for the original less expressive GAT version well as for the vanilla GCN 29 yielded by fixing 0.
48
as
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 Xavier (Xav.) initialization [19 is the standard default for GATs. We describe the Looks-Linear OotheaonO (LLortho) initialization later.
In order to empirically validate Theorem 2.2 we rewrite Equation to define â as:
(W' [1. WWTJ:: L) (a', V (W'+1 [ i). (w++ =
(7)
We observe how the value of 8 varies during training in Fig. for both GD and Adam optimizers with y 0.1. Although the theory assumes infinitesimal learning rates due to gradient flow, it still holds sufficiently (5 0) 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 so because the derived law on the level of scalar products between gradients and parameters (see Theorem 2.2) still holds, as it states relationship between gradients and parameters that independent of the learning rate or precise gradient updates.
Having verified Theorem 2.2 we use it to deduce an explanation for major trainability issue of GATs with Xavier initialization that amplified with increased network depth.
The main reason that the last layer has a comparatively small average weight norm, as E @| [.
2nz/(nu + NL-1) << 1, where the number of outputs is smaller than the layer width ny. << n1-11 In contrast, E ||[ [i.:]] 2nz-1/(mL + nT.-2) = and 1 2/(1 + n)-1). In consequence, the right-hand side of our derived conservation law n
Ea
V V
Vwz C
(WL-1)²
(a 10
Σ
(W
(8)
wh-1
W
a

---

0.000
0.005
0.010
0.015
Xav
XavoZar
Xav+Bal
LLortho+Bal
Feat. Weights
Weights
Xavier
LLortho+Bal
Frac. of Params
Layer
2
Loss
Layer
3 2 1
3
0.5
4499
0.0
1150
2300
3450
Train
375
750
1125
1500
2
3
Epoch
Valid
Epoch
(a) Frac. of sig. params with relative change > 0.05 (b) Rel. grad. norms feat. wghts. & training curve
Figure 3: For I 5, high fraction of parameters change, and gradients propagate earlier layers for both the unbalanced and balanced initialization Xav. and Balo initialization achieve 75.5% and 79.9% test accuracy, respectively. Note that the balanced initialization (b) is able train faster.
0.00
0.02
0.04
0.06
0.08
ZeroAAtt
Feat. Weights
Weights
Xav
LLortho+Bal
Frac. of Parama
1.0
2
4 2
Layer
4
Loss
Layer
0.5
6
6
10 8
8
10
1250
2500
3750
5000
300
800
900
1200
0.0
Train
2
5 Layer
Layer
Epoch
Valid
Epoch
Frac. of sig- params with relative change > 0.5 (b) Rel. grad. norms of feat. wghts. & training
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 able to propagate gradients, change parameters, and thus attain training loss -> O. Xav. and Balo achieve 39.3% and 80.2% test accuracy, respectively. Note that the balanced initialization in (b) also able to train faster than the 5 layer network with Xavier initialization in Fig. 3b)
is relatively small, as ERe (W << 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 highly unlikely during training. Therefore, relative gradients in layer L Ø-1 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 substitu- tion of Theorem 2.2 in Eq. 8) that results in a telescoping series yielding:
V (1112
WO
C
L-1
R
C
RL-1 1
VW (1)2
Jm
(9)
Jm
W
ki
W
m=
jm
1=1
a=1
0
Ki
Generally, 2n, no and thus E [ 2 2n1 / + no) < 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,
how balanced initialization would mitigate the problem. Equal weight norms | =Il and [c, in Eq. (8) (as the attention parameters are set to 0 during balancing, see Procedure 2.6) would allow larger relative gradients in layer I. - 1, as compared to
the imbalanced case, that can enhance gradient flow in the network to earlier layers. other words,
gradients on both sides of the equation have equal room drive parameter change.
5

---

Todl,
loat H-
I-1,
F-
1-10
ZauAA 12.8%
Xav-Bal 72.5% .3%
LLatho+Ball 246
Nomm "0
10"1
Rd Gldd Mami AM
2x
70"
10"
2000 4030
2000
4100
500
1000
1590
500
1013
Epp:3
Eboch
Epoch
Eboch
Figure 5: Relative gradient norms of feature (left axis, solid) and of attention (right axis, stylized) parameters for E [1,5, 10j and L = 10, sampled every 25 epochs. Test accuracy is at the top. Both attention and feature gradients at the first, middle, and last layer of the network with both balanced initializations are much larger than with unbalanced initialization (note axis scales).
To visualize the trainability issue, we study the relative change [<w'-wo\/w"] of trained network parameters (w") w.r.t their initial value. In order to observe meaningful relative change, we only con- sider parameters with a significant contribution to the model output (w" 2 10-4). We also plot the rel- ative gradient norms of the feature weights across all layers to visualize their propagation. We display these values for L = 5 and = 10 in Fig. and respectively, and observe a stark contrast, similar to the trends in relative gradient norms of selected layers of 10 layer GAT in Fig. 5 In all experiments, GAT was trained with the specified initialization on Cora using GD with y 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. Ba] and Mal 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 unable 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 E2, a' 0 leads to Ohr 1/)N())),, 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 modes[30) Although this helps generalization in some cases (see Fig. (). it alone does not help the trainability issue as seen in Fig. 4a may in fact worsen it (see Fig. which is explained by Eq. g These observations underline the further need for parameter balancing.
we note that in order to achieve balancedness, (i.e. set c = 0 in 5), the randomly initialized parameters W and a' must satisfy the following equality for 1 E [A]:
Procedure 2.6 (Balancing). Based on Eq. 5 from the norm preservation law 2.3
|| =I] ||a" [i]]
|| il|
This can be achieved by scaling the randomly initialized weights as follows:
1. Set a' = 0 for / E [L].
2. Set Wi. =
[[]] v
Bi. for i E [ni] where 8, is hyperparameter
3. Set w' [1,
w [.]] [t:]]] for i E [nv] and I E [L - 1]
[w
In step 2, 3g 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 normal or uniform distribution. We explain why we set B, for e [ni] in the context of an orthogonal initialization.
Remark 2.7. This procedure balances the network starting from / towards L. Yet, could also be carried out in reverse order by first defining [W i]]? = 33 for i E [nz-1].
Balanced Orthogonal Initialization The feature weights W for 1 E [L] need to be initialized randomly before balancing, which can simply done by using Xavier initialization. However, existing works have pointed out benefits of orthogonal initialization to achieve initial dynamical

---

Mith
000
500
(a) Test accuracy (higher is better).
(b) Epochs to best model (lower is better).
Figure GAT trained on Cora using SGD
enables training very deep architectures. In line with these findings, we initialize W 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 (balanced) LL-oothogonal 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.
isometry for DNNs and CNNs 37 44
23
We outline the initialization procedure as follows: Let || and } denote row-wise and column- wise matrix concatenation, respectively. Let us draw a random orthonormal submatrix U and define W = [U- -U'] where U E R # x W = [ || -U' + U² U']l where U I 1 for (2, ...., L 13, and W [U_ - -U'1 where UL R " Since UL is an orthonormal matrix, [ :]]|² 2 by definition, and by recursive application of Eq. 2.3
(with a' = 0), balancedness requires that |W''t [ ..]l1² Therefore, we set B, = 2 for i E nı-
3 Experiments
The main purpose of our experiments is to verify the validity of our theoretical insights and deduce an explanation for 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 (..)). 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, study the different GAT initializations listed in Table gen- cralization (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 experiments 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 reported for accuracy and epochs to the best model. All reported results use ReL.U 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
either Nvidia T4 Tensor Core
com/ wwt/eeoattiioooeeeeee htt//eaaataaalM/////a.
Baceeddddnaa...........ccc.ac... Initializatinn
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 learning rate of 0, 0.05 and 0.005 for [2, 5). L = [10, 20), and L = 40, respectively, allows for reasonably stable training Cora, Citeseer, and
SGD
7

---

Table 2: Mean accuracy (%)
95% CI over five runs of GAT with width= 64 trained using SGD.
Init.
= 5
L = 10
L = 20
L = 40
Citeseer
Xav
71.82 +: 2.73
58.40 + 2.25
24.70 + 8,90
19.23 + 1.54
18.72 + 1.15
Balx
71.62 = 0.80
68.83 = 1.62
64.13 1.57
54.88 = 7.95
42.63 = 17.47
Balo
72.02 + 0.63
70.63 0.60
68.83 1.68
68.70 1.18
63.40 + 1.43
Pubmed
Xav
77.26 + 1.39
70.68 2.16
67.32 10.70
36.52 11.50
27.20 13.99
Balx
78.02 + 0.73
75.66 + 1.81
77.60 + 1.56
76.44 + 1.70
75.74 + 2.94
Balo
77.68 + 0.45
76.62 + 1.59
77.04 + 2.14
78.20 + 0.61
77.80 + 1.41
Xav
27.32 +: 0.59
24.60 + 0.93
24.08 + 0.80
22.29 + 3.26
19.46 + 5.75
Actor
Balx
26.00 + 0.59
23.93 + 1.42
24.21 + 0.78
24.74 1.14
23.88 + 0.97
Balo
26.59 + 1.03
24.61 = 0.77
24.17 = 0.62
24.24 1.05
23.93 + 1.53
Chameleon
Xav
52.81 + 1.37
54.21 + 1.05
30.31 5.96
22.19 + 2.04
22.28 - 3.15
Balx
51.18 + 1.94
54.21 0.82
52.11 3.72
51.89 1.89
38.64 + 10.31
Balo
50.00 + 3.07
53.95 1.81
51.84 3.21
52.72 0.13
44.30 1.61
Cornell
Xav
42.70 +: 2.51
41.08 + 2.51
42.70 + 1.34
25.41 + 14.64
22.70 + 13.69
Balx
41.08 + 6.84
35.14 11.82
41.08 2.51
40.00 + 4.93
37.84 + 5.62
Balo
42.16 + 1.64
43.24 2.12
36.76 5.02
35.68 3.29
36.22 3.42
Squirrel
Xav
35.20 +: 0.44
40.96 + 0.92
21.65 1.52
20.23 + 1.69
19.67 = 0.29
Balx
35.95 + 1.69
40.98 + 0.87
38.98 + 1.49
38.35 + 1.07
25.38 + 4.62
Balo
35.83 + 0.92
42.52 + 1.19
38.85 + 1.36
39.15 + 0.44
26.57 + 1.83
Xav
60.00 + 1 1.34
60.54 +: 3.42
58.92 + 1.34
49.73 + 20.97
17.84 + 26.98
Texas
Balx
60.54 + 1.64
61.62 + 2.51
61.62 + 2.51
58.92 + 2.51
56.22 + 1.34
Balo
60.00 + 1.34
57.30 1.34
56.76 0.00
58.38 + 4.55
57.30 3.91
Wisconsin
Xav
51.37 + 6.04
5.93
LL.76
43.14 25.07
31.76 = 31.50
Balx
49.80 + 8.79
54.51 4.19
47.84 7.16
50.59 10.49
41.18 1.54
Balo
49.80 + 4.24
55.69 + 3.64
51.76 + 5.68
49.02 + 4.36
48 48.24 + 4.52
Pubmed. For the remaining datasets, set the learning rate to 0.05, 0.01, 0.005 and 0.0005 for = [2, 5), L = 10, = 20, and L 40, respectively. Note that we do not perform fine-tuning of the learning rate other hyperparameters with respect to performance and use the same settings for all initializations to allow fair comparison.
Figure Ø highliighs that models initialized with balanced parameters, Balo, 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 Interestingly, the default Xavier (Xav) initialized deeper models tend to improve performance with width but cannot compete with balanced initialization schemes. For example, the accuracy achieved by Xav for L = 10 increases from 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 higher number (almost random) features supports better models. This would also be in line with studies of overparameterizztion in vanilla feed-forward architectures, where higher width also aids random feature models
The observation that width overparameterized models enable training deeper models may be of independent interest the context of how overparameterization 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 Balo initialized model for
is already able to attain 79. 79.7% even with =idth= 64 and improves to 80.9% with width= 512. Primarily the

---

Xey
LLotho UMM iraoredd
X0+
Xon-Eemoid
Lorthoeea Pltreed
Coma
3500
2500
Fourst 750
2000
1020
1000
30
10
8
70
10
20
5
90
20
10
20
Daph
Dkpth
Doth
Dapth
3epth
Daph
(a) Test accuracy.
(b) Epochs to convergence.
Figure 7: GAT with 64 hidden dimensions trained using Adam.
parameter balancing must be responsible for the improved performance, as it is not sufficient to simply initialize the attention weights to zero, shown in Figure 4a and
For the remaining datasets, report only the performance of networks with 64 hidden dimension in 2 brevity. Since we have considered the Xavz initialization only an ablation study and find it to be ineffective (see Figure 6aD. we do not discuss it any further.
Note that for some datasets (c.g. Pubmed, Wisconsin), deeper models (e.g. L = 40) are unable to train 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 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 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. in appendix). In fact, the improved performance and training speed-up models with balanced initialization as opposed to models with standard initialization upheld even more 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-ar performance on these datasets is achieved by specialized models. We do not compare with them for two reasons: i) they do not employ 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 a balanced initialization in training deep GATs in comparison to LipschitzNorm [14), normalization technique proposed specifically for training deep GNNs with self-attention layers such as GAT (see Table 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 Balo initialized models achieve higher accuracy than 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 : 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. 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
cases, the results follow a similar tendency as

---

in 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 wGAT [8] (see Fig. [T]] in appendix). Note that the law also holds for the non-attentive GCNs (see Table in appendix) which 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.
Discussion
GATs are powerful graph neural network models that form cornerstone of learning from graph- based 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 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
[4], where
only the last linear layer of a classic feed-forward neural network architecture can adapt to a task. do not require infinite width
Conservation laws for basic feed-forward architectures [15
3 32
assumption like NTK-related theory and highlight more nuanced issues for trainability. Furthermore, they intrinsically linked to implicit regularization effects during gradient descent [1]. 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 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 (TO) and thus enables signal to pass 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.
5
Acknowledgements
We gratefully acknowledge funding from the European Research Council (ERC) under the Horizon Europe Framework Programme (HORIZON) for proposal number 101116395 SPARSE-ML.
References
[1] Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. In International Conference on Learning Representations, 2021.
[2] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. International Conference on Learning Representations, 2019.
[3] Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang Stronger generalization bounds for deep nets via compression approach, 2018.
10

---

[4] 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.
[5] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
[6] David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams, The shattered gradients problem: If resnets the answer, then what is the question? In International Conference on Machine Learning, 2018.
[7] 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.
[8] T. Bian, X. Xiao, T. Xu, P. Zhao, W. Huang, Y. Rong, and Huang. Rumor detection on social media with bi-directional graph convolutional networks. In AAAI Conference on Artificial Intelligence, 2020.
[9] Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? In International Conference on Learning Representations, 2022.
[10] Rebekka Burkholz and Alina Dubatovka. Initialization of ReLUs for dynamical isometry. In Advances in Neural Information Processing Systems, volume 32, 2019.
[11] Tianle Cai, Shengjie Luo, Keyulu Xu, Di He, Tie Yan Liu, and Liwei Wang. Graphnorm: principled approach to accelerating graph neural network training. In International Conference on Machine Learning. 2021.
[12] 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,
[13] 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.
[14] George Dasoulas, Kevin Scaman, and Aladin Virmaux. Lipschitz normalization for self-attention layers with application to graph neural networks. In International Conference Machine Learning, 2021.
[15] 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.
[16] Moshe Eliasof, Eldad Haber, and Eran Treister. Pde-gen: Novel architectures for graph neural networks motivated by partial differential equations. In Advances in Neural Information Processing Systems, 2021.
[17] Moshe Eliasof, Eldad Haber, and Eran Treister. pathgen: Learning general graph spatial operators from paths. In International Conference Machine Learning, 2022.
[18] Moshe Eliasof, Lars Ruthotto, and Eran Treister. Improving graph neural networks with learnable propagation operators. In International Conference on Machine Learning. 2023.
[19] 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.
[20] M. Gori, G. Monfardini, and F. Scarselli. A new model for learnig in graph domains. In IEEE International Joint Conference Neural Networks, 2005.
[21] Kai Guo, Kaixiong Zhou, Xia Hu, Yu Li, Yi Chang, and Xin Wang. Orthogonal graph neural networks. In AAAI Conference Artificial Intelligence, 2022.
11

---

[22] William L. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in Neural Information Processing Systems. 2018.
[23] Wei Hu, Lechao Xiao, and Jeffrey Pennington. Provable benefit of orthogonal initialization in optimizing deep linear networks. In International Conference on Machine Learning, 2020.
[24] 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.
[25] Ajay Jaiswal, Peihao Wang, Tianlong Chen, Justin F. Rousseau, Ying Ding, and Zhangyang Wang. Old can be gold: Better gradient flow can make vanilla-gens great again. In Advances in Neural Information Processing Systems, 2022.
[26] Stefanie Jegelka. Theory of graph neural networks: Representation and learning. In The International Congress of Mathematicians, 2022.
[27] S. Kearnes, K. McCloskey, M. Berndl, V. Pande, and P. Riley. Molecular graph convolutions: Moving beyond fingerprints. In Journal of Computer- Molecular Design, 2016.
[28] Dongkwan Kim and Alice Oh. How find your friendly neighborhood: Graph attention design with self-supervision. In International Conference on Learming Representations, 2021.
[29] Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
[30] 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.
[31] D. Kunin, J. Sagastuy-Brent S. Ganguli, D.L.K. Yamins, and Tanaka. Neural mechanies: Symmetry and broken conservation laws in deep learning dynamics. In International Conference on Learning Representations, 2021.
[32] Thien Le and Stefanie Jegelka. Training invariances and the low-rank phenomenon: beyond linear networks. International Conference on Learning Representations, 2022.
[33] Guohao Li, Matthias Müller, Bernard Ghanem, and Vladlen Koltun. Training graph neural networks with 1000 layers. In International Conference on Machine Learning, 2022.
[34] Guohao Li, Matthias Mülerr, Ali Thabet, and Bernard Ghanem. Deepgens: Can gens go as deep
cnns? In International Conference on Computer Vision, 2019.
[35] Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. In AAAI Conference Artificial Intelligence, 2018.
[36] Meng Liu, Hongyang Gao, and Shuiwang Ji. Towards deeper graph neural networks. In SIGKDD International Conference on Knowledge Discovery and Data Mining, 2020.
[37] Dmytro Mishkin and Jiri Matas. All you need is good init. In International Conference on Learning Representations, 2016.
[38] Pál András Papp, Karolis Martinkus, Lukas Faber, and Roger Wattenhofer. Dropgnn: Ran- dom dropouts increase the expressiveness of graph neural networks. In Advances in Neural Information Processing Systems, 2021.
[39] Hongbin Pei, Bingzhe Wei, Kevin metric graph convolutional networks. In International Conference on Learning Representations, 2020.
Chang, Yu Lei, and Bo Yang. Geom-gen: Geo-
[40] Oleg Platonov, Denis Kuznedelev, Michael Diskin, Artem Babenko, and Liudmila Prokhorenkova. critical look at the evaluation of gnns under heterophily: are we really making progress? In International Conference on Learning Representations, 2023.
12

---

[41] Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. In Advances in Neural Information Processing Systems. 2020.
[42] 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.
[43] Karthik Abinav Sankararaman, Soham De, Zheng Xu, W. Ronny Huang, and Tom Goldstein. The impact of neural network overparameteriiation on gradient confusion and stochastic gradient descent. In International Conference on Machine Learning, 2020.
[44] 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.
[45] F. Scarselli, M. Gori, A.C Tsoi, M. Hagenbuchner, and G. Monfardini. Computational capabili- ties of graph neural networks. In IEEE Transactions on Neural Networks, 2009.
[46] Wenling Shang, Kihyuk Sohn, Diogo Almeida, and Honglak Lee. Understanding and im- proving convolutional neural networks via concatenated rectified linear units. In International Conference on Machine Learning, 2016.
[47] 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 Informa- tion Processing Systems, 2017.
[48] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations, 2018.
[49] 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.
[50] Z. Wu, S. Pan, G. Long, Jiang, and C. Zhang. Graph wavenet for deep spatial- graph modeling. In International Joint Conference on Artificial Intelligence, 2019.
[51] K. Xu, W. Hu, J. Leskovec. and S. Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
[52] 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
[53] 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.
[54] 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.
[55] Han Yang, Kaili Ma, and James Cheng. Rethinking graph regularization for graph neural networks. In Advances in Neural Information Processing Systems, 2021.
[56] Zhilin Yang, William W. Cohen, and Ruslan Salakhutdinov. Revisiting semi- supervipe learning with graph embeddings. In International Conference on Machine Learning, 2016.
[57] Gilad Ychudai and Ohad Shamir. On the power and limitations of random features for under- standing neural networks. In Advances in Neural Information Processing Systems, 2019.
[58] 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.
13

---

[59] Jiaxuan You, Rex Ying, and Jure Leskovec. Design space for graph neural networks. In Advances Neural Information Processing Systems, 2021.
[60] Seongjun Yun, Minbyul Jeong, Rachyun Kim, Jaewoo Kang, and Hyunwoo J. Kim. Graph transformer networks. In Advances in Neural Information Processing Systems, 2020.
[61] M. Zhang and Y. Chen. Inductive matrix completion based on graph neural networks. In International Conference on Learning Representations, 2020.
[62] 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.
[63] 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.
[64] Lingxiao Zhao and Leman Akoglu. Paimorm: Tackling oversmoothing in gnns. In International Conference Learning Representations, 2020.
[65] 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.
[66] Kuangqi Zhou, Yanfci 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.
[67] Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layer- dependent importance sampling for training deep and large graph convolutional networks. In Advances in Neural Information Processing Systems, 2019.
Proofs of Theorems
Notation. Scalar and element-wis 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 casier follow.
Definition 6.1 (Rescale invariance). The loss L(0) is rescale invariant with respect to disjoint subsets of the parameters 0 and 0g if for every À >> 0 we have C(O) _((20), 1- 0s)), where (01, 0z. 0g).
frequently utilize the following relationship.
Lemma 6.2 (Gradient structure due to rescale invariance [3I). The rescale invariance of C enforces the following geometric constraint the gradients of the loss with respect to its parameters:
(01, Vo.C) - (02, C) = O.
(10)
We first use this rescale invariance to prove our main theorem for GAT with shared feature transformation weights and single attention head, as described in Def. 2. The underlying principle generalizes to other GAT versions 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.
14

---

Our main proof strategy relies on identifying a multiplicative rescale invariance in GATs that allows to apply Lemma 6.2 We identify rescale invariances for every neuron i at layer I that induce the stated gradient structure.
Proof of Theorem 2.2
Specifically, define the components of in-coming weights to the neuron 01 {wo w E [ I} and the union of all out-going edges (regarding features and attention) as 02 {w w W [i]]. It is left to show that these parameters are invariant under rescaling.
B U [a"(i]]
Let us, therefore, evaluate the GAT loss X01 and X-102 and show that it remains invariant for any choice of À > O. Note that the only components the network that potentially change under rescaling are h2 [i]. h [j], and an We denote the scaled network components with tilde resulting in h, h+ [k], and d'v As we show, parameters of upper layers remain unaffected, as h [R]
coincides with its original non-scaled variant h [k] = h [k].
Let us start with the attention coefficients. Note that a' = x a'aa'1 and [i. \W") This implies that
exp(e
á uv Ene.N(M) exp(el a'u because
(11)
e = (dd) (W + 1)) Euvz
(12)
which follows from the positive homogeneity of Φ2 that allows
Ri-1
ni-1
Env =
Wi N DD) +
1 +
(13)
-x-'Aa' Σ w j (h 'D1 + n + Σ
l Φ2 Σ
+ h [1))
(14)
= -w
(15)
Since āyu = yo it follows that
Mz-1
hig
[1]
= Φ1
Σ
L
Xw
(i 1igh--h--0
=EN(u)
Σ
a Σ
[il
(EN(u)
[
In the next layer, we therefore have
n2
[++1[]] [k]
Φ1
Σ uEN(v)
a
L-1
Σ
W
+1 [ks i1
Σ
I+1
w ++ [k: ilth111
auv
uEN(v)
J+1 auv
Σ w
[k, []
uEN(()
[k]
Thus, the output node representations of the network remain unchanged, and according to Def. the loss L is rescale-invariant.
6.
15

---

Consequently, as per Lemma 6.2
the constraint in Eq. 10) can be written as:
[1. V Wi'l..) f)
(a V) (Wt++ [. V (wi- = 0.
which can be rearranged to Eq. (.2
(W V..iill..wicc Vwika - (d'lt), V Valle = (W"+) [3, d). V Vw
thus proving Theorem
2.2
Proof Corollary
2.3
Given that gradient flow applied on loss & is captured by the differential
equation
dw vww
(16)
which implies:
d
|| [1. -I| [i,)) d (ii,;)) [i. Vw(R=f)
(17)
dt
dt
substituting in Eq. Theorem 2.2
similarly for gradient flow of w [:, and a' a'] as done in Eq. 17 yields
(IWW :.:11 -Il'' d (1 di 2((W (i.), [w) - (-2) V 2(W |; i). Vwi+
d
Therefore, Eq. and consequently Eq. in Corollary 2.3 hold.
Summing over i E [mi] on both sides of Eq.4 yields Corollary 2.4
Theorem 6.3 (Structure of gradients for GAT without weight-sharing). Let GAT network as defined by Def. 2.1 consisting of L layers be given. The feature transformation parameters W and W of the source and target nodes of an edge and attention parameters a' of a GAT layer are defined according Fq.47) and (2).
Then the gradients for layer E [T. - 1] in the network are governed by the following law: :.),,, VWggaa + (i. :], W) (a' [1), Voll C) -WW++6 [ W C) (18)
Proof. The proof is analogous to the derivation of Theorem 2.2 We follow the same principles and define the disjoint subsets 0 and Og of the parameter set 6, associated with neuron layer I accordingly, follows:
0 = ww/w E :] U fwow E :]
0₂ [w/w E u
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 [1] h++1()) and
a We denote the scaled network components with tilde resulting [i] [i. [[), and a
As we show, parameters of upper layers remain unaffected, as [ coincides with its original non-scaled variant [ [K)]
Let start with the attention coefficients. Note that a'|i x a W. |. = AW and W| = J.... This implies that
exp(ein)
a EvENTO) (xp(ettr) aye because
(19)
an = -aa²T 222(1T (2( W!b = ciuu:
(20)
16

---

which follows from the positive homogeneity of @2 that allows
02)
(21)
n--1
+
12
+ 1/0.11
(22)
A
'''' Adl
Whi w
(23)
ni
+
02
Σ
WHi + w/b..o44
(24)
(25)
Since ā'y anuos follows that
heg
[i] = 01
Σ
a
=EN(u)
Ri-1
= X
Σ
Σ
zu
REN(u)
Ah (1]
In the next layer, therefore have
T2
hp
[k]] = Φ1
Σ
Oue
I+1
W
[k: 2)
(EN(w)
m2
Φ1
Σ
Que
I+1
wi++
uEA((e
Φ1
Σ
I+1 Ouv
W
[k: 111
(EN(e)
[k]]
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 and attention parameters of of an attention head k E [K] in a GAT layer of a L layer network. Then the gradients of layer E [L 1) respect the law:
K
K
K
Σ (.: C)
Σ v
= > k
(w (1 V (
(26)
Proof. Each attention head is independently defined by Eq [ and Eq 2 and thus Theorem holds for each head, separately. The aggregation of multiple heads in layer is done over node representations of each head in either one of two ways:
Concatenation: h4 =IIk Φ o(EuEN(O) Olkeo' -W/h(=)) or
Average: K, = k Σ + (( Okav' W
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. D
17

---

For the most general case of a GAT layer without shaght-shhaing and with multi-header 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.
Training Dynamies
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 7 0.1, 7 0.05 and y 0.01 for L = 5, L 10 and 20, respectively. For Adam, y = 0.005 and 7 = 0.0001 for L = 5 and L = 10, respectively.
Relative change of feature transformation parameter we is defined as [(="-=o)/"r" where wh is the initialized value and w" is the value for the model with maximum validation accuracy during training. Absolute change |a" -aal aul is used for attention parameters a since attention parameters are initialized with 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 without considering the unchanged parameters.
Frae Deama
Wnights
10
Weights
10
Fost Weights
10
Weights
Fox. Puaan
Feat. Weights
Al.. Weights
0.5
0.5
06
0.6
0.5
Xas ZancAA
LLotho+BB
0.0
5 Layer
10
8 -2y0
10
0 12 Laywr
15
20
12 Layar
10
20
Layor
Layar
(a) T=
(b) 10
(c) = 20
Figure 8: Layer-wise fraction of feature transformation parameters W and attention parameters a' 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 5 layer unbalanced network initialized with Xavier, but this fraction becomes significantly large when depth is increased
to L 10. Note that W 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' 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' and all a E a! 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 / as ||V w w*ll, ///l////// ///[///|, and || [V_.a"l//|]aa)) respectively. Figures @bjand[bb and Figure [2 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 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 message- passing GNNs (MPGNNss.
To understand the impact of different architectural variations, common regular- ization strategies such as dropout and weight decay, and different activation functions, we conducted additional experiments. For a 10-layer network with width 64, Table report our results for SGD and Adam, respectively. The values of hyperparameters were taken from [8].
Training variations
18

---

Table Mean test accuracy(%)
95% CI over five runs of GAT trained on Cora using SGD
Variation
L
Xav
Balx
Balo
5
73.00 + 3.02
76.96 + 2.21
79.48 + 0.43
None (Vanilla GAT)
10
25.48 + 18.13
77.72 + 1.49
79.46 1.34
5
73.56 + 2.71
77.44 1.54
79.58 + 0.53
attention heads = 8
10
25.50 + 18.18
77.02 2.76
79.06 0.73
5
75.68 1.80
79.20 = 1.07
79.64 0.36
activation = ELU
10
73.02 + 2.27
78.64 + 1.72
47.76 + 7.39
5
42.14 = 15.97
79.18 + 1.17
81.00 = 0.62
dropout 0.6
10
24.90 + 9.50
30.94 + 1.04
44.40 + 1.84
5
67.26 + 6.30
77.36 + 1.74
79.56 + 0.48
weight decay 0.0005
10
18.78 + 11.96
76.56 + 2.91
79.40 + 1.15
5
70.04 + 3.45
77.46 + 1.19
79.66 + 0.72
weight sharing False
10
19.54 15.64
76.16 + 2.28
80.06 1.12
Table 4: Mean accuracy(%)) + 95% CI over five of GAT trained on Cora using Adam
Variation
L
Xav
Balx
Balo
5
76.18 + 1.61
79.38 + 2.24
80.20 + 0.57
None (Vanilla GAT)
10
70.86
3.99
75.72 + 2.35
79.62 1.27
75.62 1.74
78.54 + 1.06
79.56 0.85
attention heads = 8
10
70.94 + 2.76
75.48 + 2.48
79.74 + 1.10
5
76.56 + 1.72
78.72 1.02
79.56 + 0.64
activation = ELU
10
75.30 = 1.42
78.48 + 1.79
76.52 0.97
5
76.74+1.44 1.44
+ +
79.92 + 0.48
dropout 0.6
10
32.48 + 6.99
31.76 + 1.33
76.34 + 0.95
5
75.10 + 2.05
78.52 1.41
79.80 + 0.63
weight decay 0.0005
10
28.74 12.04
74.68 + 3.06
79.70 + 1.14
5
76.90 = 0.92
78.90 + 1.07
80.62 + 1.36
weight sharing False
10
74.46 1.90
75.50 3.49
79.78 0.70
We find that the balanced initialization performs similar in most variations, outperforming the unbalanced standard initialization by substantial margin. Dropout generally does not seem to be helpful to the deeper network (I. 10). regardless of the initialization. From Table 3 and 4 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 positively homogeneous activation function, which is an assumption in out theory. In practice, however, it does not impact the Xavier-balanced initialization (Balx). However, the Looks-Linear orthogonal structure is specific to ReLUs. Therefore, the orthogonality and balancedness of the Balo initialization are negatively impacted by ELU, although the Adam optimizer seems compensate for it to extent.
In addition increasing the trainability of deeper networks, the balanced initialization matches the state-of-ther 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 scripfoof
19

---

Cora that follows the original (AT(v1) paper and added our balanced initialization to the code. As evident from Table the balanced initialization matches SOTA performance of GAT(v1) on Cora (83-84%) (on 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-ard performance of GAT(v2) on Pubmed, as shown in Table
Table 5: GAT performance on Cora and Pubmed: mean test accuracy (%) + 95% CI over five runs. Balo
Xav
Balx
Cora
84.50 + 0.52
84.58 + 0.65
84.55 = 0.47
Pubmed
78.38 + 0.77
78.52 = 0.54
78.56 = 0.20
Comparison with Lipshitz Normalization A feasible comparison can be carried out with [4] 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, 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 and layer-norm (5)
Table 6: Comparing balanced LL-orthogonal initialization to Lipschitz normalization applied with a standard (imbalanced) Xavier initialization: the balanced initialization results in a much higher accuracy as the depth of the network increases.
Dataset
Cora
Citeseer
Pubmed
Layers
Lip. Norm.
Balo Init.
Lip. Norm.
Balo Init.
Lip. Norm.
Balo Init.
82.1
79.5
65.4
67.7
74.8
76.0
77.1
80.2
63.0
67.7
73.7
75.7
10
78.0
79.6
43.6
67.4
52.8
76.9
20
72.2
77.3
18.2
66.3
23.3
77.3
40
12.9
75.9
18.1
63.2
36.6
77.5
Impact of orthogonal initialization In the main paper, we have advocated and presented results for using balanced LL-oorhogonal 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 [6], and matrices with looks-linear structure using an identity submatrix (LLidentity) since a looks-linear structure would be more effective for ReLUs [T].
In line with [8] 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 identity matrices have row-wise and column-wii 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 general- ization ability of the network in most cases and speeds up training, particularly for deeper networks. However, note that Lothooall 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- Northogonal 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 I.Lrtho+Bal initialization also speeds up training over the L.Oohhho initialization even in cases in which the generalization performance of the model is at par for both initializations.
20

---

Xavier
Kavier+Bal
Iddenitty
Identity-Bal
LLidentiy
Widentity@Bal
LLotho
LLortho-Bal
90
80
70
40
Text
80
Citeseer
Pubmed
Number Layers
Figure 9: GAT network with varying initialization trained using SGD.
Note that identity initializations have also been explored 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. due to the reasons regarding feature diversity and ReLU activation discussed above, the balanced looks-lincar random orthogonal initialization (LLortho+ Bal.) outperforms initialization with identity matrices (see Fig. g 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.
counter-example [7].. However, potentially
Applicability to other MPGNNs As GAT is 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 reported for GAT. As expected, the trainability benefits of balanced initialization also evident in deeper GCN networks.
Table GCN with varying initialization trained using SGD,
Dataset
Depth
Xavier
Xavier+Bal
LLortho
LLortho+Bal
2
77.8 + 0.9
80.5 + 0.5
78.0 + 0.3
80.9 + 0.4
5
78.3 0.8
80.3 + 0.8
79.6 + 1.0
10
24.1 4.5
77.6
80.0
80.0 + 0.9
Cora
20
14.4 11.2
62.8 3.6
0.5
78.8 1.5
40
13.4 0.9
65.9
28.3 16.2
77.1 0.9
64
0.8
33.0 13.4
27.3 12.7
76.7 1.3
80
12.4 19.3
33.8 12.9
21.3
77.1 1.3
2
66.6 + 20.0
71.3 1.8
3.2
72.3 + 0.9
5
60.9 12.3
66.9 + 15.0
6.4
70.1 1.8
10
23.8 36.8
66.0
0.9
69.8 + 10.9
Citeseer
20
16.4 + 18.2
10.0
6..0+8.6
69.7 + 4.5
40
13.9 + 56.8
9
4..8
64.7 + 13.6
64
13.8 + 41.4
29.5 + 15.0
37.3 + 79.6
66.3 + 0.5
80
12.4
42.7
25.8 =
30.1 + 21.8
64.1 + 3.2
The underlying principle of a balanced network initialization holds in general. However, adapting
the balanced initialization to different methods entails modification for GAT's to correspond to the specific architecture of the other method. For example, the proposed
conservation law derived
balanced initialization can be applied directly to a more recent variant of GAT, =GAT [8) for which the same conservation law holds and also benefits from balanced initialization. We conduct
21

---

additional experiments to verify this. The results in Fig, TO follow a similar pattern as for GATs: balanced orthogonal initialization with looks linear structure (LLortho+Bal) of wGAT performs the best, particularly by wide margin at much higher depths (64 and 80 layers).
Xavier
Kavier-
Identity
Ieentity4Ba
LLidentity
Lideentit-BBll
LLortho
LLortho-Bal
70
50
30
10
20
64
SGD
Adam
Number of Layers
Figure 10: wGAT network with varying trained Cora.
In this work, we focus exposition on GATs and take the first step in modeling the training dynamics of attention-based models for graph learning. An intriguing direction future work is to derive modifications in the conservation law for other iiteettoon-bsse models utilizing the dot-product self-attention mechanism such as SuperGAT (28) and other architectures based on Transformers [47), that widely used in Large Language Models (LLMs) and have recently also been adapted for graph learning [58 60
Datasets Summary
We used nine benchmark datasets for the transductive node classification task in our experiments, as summarized in Table
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 Cora dataset, but the PyG dataset has 10356 < (2 x 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 [4), where duplicate edges are removed.
All results in this paper are according the dataset statistics reported in Table
Table 8: Summary datasets used in experiments.
Dataset
Cora
Cite.
Pubmed
Cham.
Squi.
Actor
Corn.
Texas
Wisc.
Nodes
2708
3327
19717
2277
5201
7600
183
183
251
Edges
10556
9104
88648
31371
198353
26659
277
279
450
Features
1433
3703
500
2325
2089
932
1703
1703
1703
Classes
5
5
5
Train
140
120
60
1092
2496
3648
87
87
120
# Validation
500
500
500
729
1664
2432
59
59
80
Test
1000
1000
1000
456
1041
1520
37
37
51
22

---

Xav
Xav ZaroZeroAtt
Xav+Bal
LBottho+Bal
Change of Feat. Weights
10
10"
10
Layer
ZaroZeroAAt
Xav+daal
LLortho+taal
Abs. Change of Waights
10
10
10
1
10
2
4
6
9
10
Layer
(a) L = 10: The balanced initialization allows larger changes in a higher number of parameters in W across all layers / with the highest margin in = 1. The change distribution for the parameters in a' is missing for and E [7, 10] because these parameters remain unchanged (see Fig.
Xav
Xav=ZeroAlt
Xav+ Bal
Loottho+Bal
Rel. Change of Fast, Weights
10
15
10
10
11
12
13
17
18
Layer
Xav
Zev+ZeroAt
Xav+Bal
LLortho+Ball
Abe Change All. Weights
10
b
10
10
10
10 -77
4
5
8
9
10
12
13
14
15
18
17
18
19
20
Layer
(b) L. = 20: The same pattern as in L. = seen. However, difference between the trainability models with unbalanced and balanced initialization becomes even more prominent. No attention parameters change at all with standard Xavier initialization (and hence their change distribution is not visible in the plot).
Figure 11: Distributions of the relative and absolute change in the values of feature transformation parameters W' and attention parameters a', respectively, when trained using SGD with unbalanced (Xav. and Xav+ZeroAtt) and balanced (Xav+Bal and LLortho+Bal) initialization. The black markers in each standard box- whisker plot denote the mean. In general, larger changes occur in attention parameters at later layers closer to the output of the network, whereas feature parameters change more in the earlier layers at the input of the network. We observe this also from the perspective of relative gradients, as higher gradients cause higher parameter changes.
23

---

feature parameters
0.00
0.02
0.04
0.06
Xav
ZevoZeroAtt
Xav+Bal
LLortho+Bal
Layer
4 3 2
Loss
Layer
4 3 2
Loss
Layer
4 3 2 1
Loes
Layer
3 2
4
0
14
28
42
56
Epoch
Epoch
Epoch
Epoch
attention parameters
0.00
0,02
0.04
Xav
Xav+ZeroAtt
Xav+Bal
LLortho+Bal
Layer
2 1
Loss
Layer
2
Loss
Layer
3 2
Loss
Layer
2
0
14
28
42
56
24
N Epoch
Epoch
Epoch
Epoch
(a) L 5: test accuracy (left to right) is 71.6%, 76.2%, 77.3%, and 80.5%.
feature parameters
0.000
0.025
0.050
0.075
Xav
Xav+ZeroAtt
Xav+Bal
LLortho+Bal
Layer
808 6 2
Loss
Layer
8 6 2
Loss
Layer
108 08642 2
Loss
Layer
808 6 2
150
300
600
175
350
700
125
250
500
100
200
300
Epoch
Epoch
Epoch
Epoch
attention parameters
0.025
0.050
0.075
Xav
Zevo
Xav+Bal
LLortho+Bal
0 6 2
8 6 2
0086 6 2
10 8 6 2
Layer
Loss
Layer
Loss
Layer
Loes
Layer
Loes
150
600
175
350
700
125
100
200
400
Epoch
Epoch
Epoch
Epoch
(b) test accuracy (left right) 63.2%, 68.8%, 71.8%, and 78.0%.
Figure 12: Layer-wise relative gradient norms of GAT parameters trained with Adam, which in itself improves trainability over SGD, especially for = (see Fig even with standard initialization. However, with balanced initialization, the training is faster and the model also generalizes better, as evident from the significant margin in test accuracy. A general noticeable trend, particularly with balanced initialization, is that the attention parameters have higher gradients the end of the network (and thus change more the later layers (see Fig. IT) whereas the feature parameters have high gradients at the start of the network. Also, it can be seen that the feature parameters begin to change few epochs before the attention parameters change noticeably, while the attention parameters continue to change for few a representational benefit could drive this behaviour, i.e. with depth, the learned representation of neighbors becomes increasingly informative for node and thus leads higher activity of the attention mechanism.
epochs after the gradients of feature parameters have dropped. We speculate
24


Reviews:
[{"review_id": "5bhhHYU4ik", "reviewer": "Reviewer_cXS5", "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.\n", "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.\n2. The trainability issue is explained intuitively in pages 4-5:\n(a) What is the reason that the trainability issue is amplified with depth? This is not explained. \n(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.\n(c) Suggestion: It would be helpful to explain why the balanced initialization solves the trainability issue (it is not explained).\n", "limitations": "Limitations are not discussed. Maybe it would be helpful to add a section on this.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel conservation law for GAT architectures.\n2. The paper and technical details are well written.\n", "weaknesses": "There are no major weaknesses.\n1. The motivation for using the balanced initialization and why it should improve only for deep networks is not so clear.\n2. The experimental results are not complete.\n\nSee Questions for more details.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank reviewer cXS5 for their insightful review and constructive questions, which are answered below. We would be happy to include the additional experiments and clarifications in the main paper as well. \n\n**1. Effect of orthogonal initialization:** We include results to compare the LLortho and LLortho+Bal initializations (See **Figure 1(a)** in additional PDF). We observe that balancing the LLortho 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 the orthogonal initialization itself also has a positive effect on trainability in particular of deep networks. This follows from the comparison of Xavier+Bal and LLOrtho+Bal (or LLOrtho). Contributing to this fact is the mostly-balanced nature of the LLortho initialization i.e. given hidden layers of equal dimensions the network is balanced in all layers except the first and last layers, 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. \n\n**2a): Amplification by depth:** \nThe reason the trainability issue is amplified with depth can also be explained by our main theorem. Recursive substitution of Theorem 2.2 on the first term of the left side of Equation 8 (Page 4) results in a telescoping series yielding:\n$$ \\sum_{j=1}^{n_{1}} \\sum_{m=1}^{n_{0}} {W_{jm}^{(1)}}^2 \\frac{\\nabla_{W_{jm}^{(1)}} \\mathcal{L}}{W_{jm}^{(1)}}  - \\sum_{l=1}^{L-1} \\sum_{o=1}^{n_{l}} {a_{o}^{(l)}}^2 \\frac{\\nabla_{a_{o}^{(l)}} \\mathcal{L}}{a_{o}^{(l)}} = \n\\sum_{i=1}^{n_{L-1}} \\sum_{k=1}^{n_L} {W_{ki}^{(L)}}^2 \\frac{\\nabla{W_{ki}^{(L)}} \\mathcal{L} }{W_{ki}^{(L)}} $$\nNote that $n_0$ is the input feature dimension. Generally, $2n_1 < n_0$ and thus $\\mathbb{E} \\left\\lVert W^1[j:] \\right\\rVert ^2 = 2n_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 balance the equation. As the scalar product of attention parameters and gradients summed over the first and all intermediate layers must be less than the scalar product of weights and gradients of the first layer, the attention parameters and gradients in each intermediate layer do not have much room to change which hampers trainability. Evidently, the problem aggravates with depth where the same value must now be distributed over the parameters and gradients of more layers. \n\n**2b) Connection to fully-connected networks:** \nWhile it is true that a similar explanation as given in the paper holds also for fully connected networks (FCs), where the attention parameters are missing in Equation 8, the problem is more severe for GATs for two reasons:  i) Firstly, as explained above in (a), the summation of all intermediate layers of attention parameters is not a concern in FCs. \nii) Secondly, FCs can achieve perfect dynamical isometry by employing an orthogonal looks-linear structure of feature weights, which enables signals to pass through very deep architectures [2]. However, due to the peculiarity of neighborhood aggregation in GATs (or GNNs), the same technique does not induce perfect dynamical isometry. Exploring how dynamical isometry can be achieved or approximated in general GNNs is an interesting follow-up direction of this work. \n\n**2c) Effect of balanced initialization:** \nOn pages 4 and 5, we use Equation 8 to explain how norm imbalance is a cause of the trainability issue. We revisit this to analogously see how the balanced initialization mitigates the problem. A balanced initialization implies the weight norm of the last and second to the last layer are equal on both sides of the equation (as the attention parameters $a$ are set to 0). This allows larger relative gradients of the second to last layer (left side of the equation) (as compared to when the weights on the left were much larger than the weights on the right) which 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. \n\n**3) Limitations:** \nWe mention inline in the paper (line 159) that our theory defines a coarse-level conservation law, and thus cannot completely explain fine-grained training dynamics such as why the attention parameters change most in the first layer.  Secondly, the conservation law applies only to the self-attention defined in the original GAT and GATv2 models, and their architectural variations such as $\\omega$GAT [3]. Note that the conservation law also holds for the non-attentive GCNs which are a special case of GATs (where the attention parameters $a$ are simply zero). Modeling different kinds of self-attention such as the dot-product self-attention in [4] entails modification of the conservation law, which has been left for future work. Following the reviewer's suggestion, we will dedicate an independent section to discuss these limitations. \n\n[1] Brody et al. How attentive are graph attention networks? In International Conference on Learning Representations, 2022.\n\n[2] Burkholz et al. Initialization of ReLUs for dynamical isometry. In Advances in Neural Information Processing Systems, volume 32, 2019.\n\n[3] Eliasof, et al. Improving Graph Neural Networks with Learnable Propagation Operators. In International Conference on Machine Learning, 2023.\n\n[4] Kim et al. How to find your friendly neighborhood: Graph attention design with self-supervision. In International Conference on Learning Representations, 2021."}, {"review_id": "dKiMai80RZ", "reviewer": "Reviewer_yMby", "summary": "The 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", "questions": "1. 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", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The 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\n", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer yMby for recognizing the novelty and relevance of our work and appreciate their suggestions, for which we will be taking the following actions: \ni) introduce more subsections and break down larger paragraphs to improve the readability and comprehension of the paper;\nii) cite the proposed relevant literature [1-5];\niii) add experiments in line with the suggestions.\n\n**Comparison with other methods:** \nWe would like to highlight that the main contribution of our work is the derivation of a conservation law for GATs and the insight into how this law can explain trainability issues for standard initialization schemes.\nIn line with this reasoning, it is not feasible to directly compare with [1,2,4], as they rely on different GNN architectures with potentially different inherent conservation laws.\nSince [3] proposes an architectural variant of GAT, the $\\omega$GAT, our proposed balanced initialization would apply to this setting and could potentially improve the trainability of deep $\\omega$GAT. Since the code for this very recent work has not been shared, we cannot provide extensive experiments in this short time frame.\nHowever, we do implement $\\omega$GAT in our setup and the results follow a similar pattern to GAT (See **Figure 1(b)** of the attached PDF). A balanced orthogonal initialization with looks linear structure (LLortho+Bal) of $\\omega$GAT performs the best, particularly by a wide margin at much higher depths (64 and 80 layers).\n\nYet, a feasible comparison can be carried out with [6] which proposes a Lipschitz normalization technique aimed to improve the performance of deeper GATs in particular. \nWe use the code provided by [6] 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 our results of LLortho+Bal initialization (Bal$_O$ Init.), reported in the main paper, as follows:\n\n| Layers | Cora          |                  | Citeseer   |                  | Pubmed     |                  |\n| :-----: | :-----------: | :--------------: | :--------: | :--------------: | :--------: | :--------------: |\n|          | Lip. Norm.    | Bal$_O$ Init. | Lip. Norm. | Bal$_O$ Init. | Lip. Norm. | Bal$_O$ Init. |\n| 2   | **82\\.1** | 79\\.5        | 65\\.4  | **67\\.7**    | 74\\.8  | **76\\.0**    |\n| 5   | 77\\.1     | **80\\.2**    | 63     | **67\\.7**    | 73\\.7  | **75\\.7**    |\n| 10  | 78\\.0     | **79\\.6**    | 43\\.6  | **67\\.4**    | 52\\.8  | **76\\.9**    |\n| 20  | 72\\.2     | **77\\.3**    | 18\\.2  | **66\\.3**    | 23\\.3  | **77\\.3**    |\n| 40  | 12\\.9     | **75\\.9**    | 18\\.1  | **63\\.2**    | 36\\.6  | **77\\.5**    |\n\nAs evident from the above table, the balanced initialization results in a much higher accuracy as the depth of the network increases than the application of the Lipschitz normalization to a standard-initialized network. \nNote that Lipschitz normalization has been shown to outperform other previous normalization techniques for GNNs such as pair-norm and layer-norm. \n\n**Higher depth:** \nAdditional results for 64 and 80-layer networks are presented in **Figure 1** in the additionally submitted PDF. The improved performance of models with balanced initialization as opposed to models with standard initialization is upheld even more so for very deep networks.\n\n**Initialization with identity matrix:** \nRegarding the initialization of GNNs with identity matrices, the identity matrix can be considered a special case of an orthogonal initialization.\nGiven that the hidden layers are all of the same dimensions, a network initialized with identity weight matrices for the intermediate layers and zero attention parameters would be balanced for all the hidden layers (but not with respect to the first and last layers). \nHowever, for ReLUs a looks-linear structure where the submatrix is initialized as identity matrix would be more effective.\nNote that identity initializations have also been explored in the context of standard feed forward neural networks.\nWhile 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]).\nWe conduct experiments using identity matrices to initialize the hidden layers and Xavier initialization for the first (input) and last (output) layers. \nWe compare this with a balanced version by adjusting the weights of the first and last layer to have norm $1$ (as identity matrices have row-wise and column-wise norm $1$). \n\nHowever, 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 **Figure 1** in additional PDF). In most cases, the balanced versions outperform the imbalanced version of the base initialization.\n\n**SuperGATs**: \nThe proposed initialization cannot be directly applied to the self-attention layer used in the SuperGAT architecture [5]. \nSuperGAT combines the attention layer of the original GAT architecture with the dot-product self-attention similar to the Transformer architecture. \nThis requires a modification of the derived conservation law, which would be an intriguing direction for future investigations. \n\n[6] Dasoulas et al. Lipschitz Normalization for Self-Attention Layers with Application to Graph Neural Networks. In International Conference on Machine Learning, 2021.\n\n[7] Bartlett et al. Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks. In International Conference on Machine Learning, 2018.\n\n[8] Chen et al. Measuring and relieving the over-smoothing problem for\ngraph neural networks from the topological view. In Proceedings of the AAAI Conference on Artificial Intelligence, 2020."}, {"review_id": "0BtyBvE8qp", "reviewer": "Reviewer_d9X8", "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.", "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. ", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The problem is well identified.\n- The proposed theories are solid, and well serves the problem.\n- 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.\n- Some formats could be improved for better illustration, such as larger figure size and fitted table style.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank reviewer d9X8 for their constructive feedback. In line with their suggestions, we will i) include the details of equations used in lines 140-142 arising from the definition of Xavier initialization, and ii) increase figure sizes and improve the table style for better illustration. \n\nThe 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. \nFor example, the proposed balanced initialization can be applied directly to a more recent variant of GAT, **$\\omega$GAT** [1], for which the same conservation law holds. We conduct additional experiments on $\\omega$GAT to verify this. The results follow a similar pattern as for GATs. (See **Figure 1(b)** of the attached PDF). A balanced orthogonal initialization with looks linear structure (LLortho+Bal) of $\\omega$GAT performs the best, particularly by a wide margin at much higher depths (64 and 80 layers).\n\nHowever, the derived conservation law and consequent balanced initialization can not be directly applied to SuperGAT [2], which employs a different kind of self-attention, namely the dot-product attention similar to the Transformer architecture.\nAn intriguing direction for future work is to derive modifications in the conservation law for such other attention-based models. Currently, as suggested by the reviewer, we will mention this limitation in the main paper.\n \n[1] Dongkwan Kim and Alice Oh. How to find your friendly neighborhood: Graph attention design with self-supervision. In International Conference on Learning Representations, 2021.\n\n[2] Moshe Eliasof, Lars Ruthotto, and Eran Treister. Improving Graph Neural Networks with Learnable Propagation Operators. In International Conference on Machine Learning, 2023."}, {"review_id": "270vK3Nu0b", "reviewer": "Reviewer_hDWR", "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.", "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. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer hDWR for their interesting comments on our work and the constructive feedback. \nIn response to their question, we discuss the adaptability of the balanced initialization scheme to other MPGNNs, which we would be happy to add to the manuscript. \n\nWhile in principle the argument of imbalanced-ness is applicable to nearly all types of MPGNNs, initializing the model in a balanced manner requires deriving conservation laws inherent to the specific architecture. \nNevertheless, we outline a number of cases in which our derived conservation law (and consequent balanced initialization) is directly applicable.\n\nFirstly, we note that GAT is a generalization of the GCN, and therefore all theorems are also applicable to GCNs (where the attention parameters are simply 0).\nWe provide additional experiments in the following table, 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 that we reported in the paper.\n\nTable 1: **GCN** trained on **Cora** using SGD: \n\n| Depth  | Xavier                 | Xavier+Bal            | LLortho               | LLortho+Bal           |\n| :----: | :--------------------: | :-------------------: | :-------------------: | :-------------------: |\n| 2  | 77\\.8 &pm; 0.9   | 80\\.5 &pm; 0.5  | 78\\.0 &pm; 0.3  | **80\\.9 &pm; 0.4**  |\n| 5  | 73\\.2 &pm; 3.4   | 78\\.3 &pm; 0.8  | **80\\.3 &pm; 0.8**  | 79\\.6 &pm; 1.0  |\n| 10 | 24\\.1 &pm; 4.5   | 77\\.6 &pm; 2.0  | 80\\.0 &pm; 1.1  | **80\\.0 &pm; 0.9**  |\n| 20 | 14\\.4 &pm; 11.2  | 62\\.8 &pm; 3.6  | 78\\.7 &pm; 0.5  | **78\\.8 &pm; 1.5** |\n| 40 | 13\\.4 &pm; 0.9   | 65\\.9 &pm; 7.3  | 28\\.3 &pm; 16.2 | **77\\.1 &pm; 0.9**  |\n| 64 | 9\\.8 &pm; 5.4    | 33\\.0 &pm; 13.4 | 27\\.3 &pm; 12.7 | **76\\.7 &pm; 1.3**  |\n| 80 | 12\\.4 &pm; 19.3  | 33\\.8 &pm; 12.9 | 38\\.9 &pm; 21.3 | **77\\.1 &pm; 1.3**  |\n\nTable 2: **GCN** trained on **Citeseer** using SGD: \n\n| Depth  | Xavier                 | Xavier+Bal            | LLortho               | LLortho+Bal           |\n| :----: | :--------------------: | :-------------------: | :-------------------: | :-------------------: |\n| 2  | 66\\.6 &pm; 20.0  | 71\\.3 &pm; 1.8  | 66\\.0 &pm; 3.2  | **72\\.3 &pm; 0.9**  |\n| 5  | 60\\.9 &pm; 12.3  | 66\\.9 &pm; 15.0 | 69\\.0 &pm; 6.4  | **70\\.1 &pm; 1.8**  |\n| 10 | 23\\.8 &pm; 36.8  | 66\\.0 &pm; 5.9  | **70\\.6 &pm; 0.9**  | 69\\.8 &pm; 10.9 |\n| 20 | 16\\.4 &pm; 18.2  | 47\\.9 &pm; 10.0 | 67\\.0 &pm; 8.6  | **69\\.7 &pm; 4.5**  |\n| 40 | 13\\.9 &pm; 56.8  | 37\\.3 &pm; 92.8 | 44\\.8 &pm; 6.8  | **64\\.7 &pm; 13.6** |\n| 64 | 13\\.8 &pm; 41.4  | 29\\.5 &pm; 15.0 | 37\\.3 &pm; 79.6 | **66\\.3 &pm; 0.5**  |\n| 80 |  12\\.4 &pm; 42.7 | 25\\.8 &pm; 3.6  | 30\\.1 &pm; 21.8 | **64\\.1 &pm; 3.2**  |\n\nSecondly, we recall that we cater to several architectural variations within the GAT architecture itself. For example, both the GATv1 [1] and GATv2 [2] have the same conservation law and we derive more general versions of the law for two GAT variants: \ni) unshared weights for feature transformation of source and target nodes, and \nii) multiple attention heads. Furthermore, residual skip connections between layers are also supported in a balanced initialization provided their parameters are initialized with zero. \n\nLastly, our derived conservation law also holds for a more recent architectural variation of GAT, the **$\\omega$GAT** [3], which also seems to benefit from a balanced initialization. \nWe conduct additional experiments to verify this. The results follow a similar pattern as for GATs (See **Figure 1(b)** of the attached PDF). \nA balanced orthogonal initialization with looks linear structure (LLortho+Bal) of $\\omega$GAT performs the best, particularly by a wide margin at much higher depths (64 and 80 layers).\n\nIn 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 such as SuperGAT [4] and Transformers which both utilize the dot-product self-attention mechanism.  \n\n[1] Petar Veličković et al. Graph Attention Networks. In International Conference on Learning Representations, 2018.\n\n[2] Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? In International Conference on Learning Representations, 2022.\n\n[3] Dongkwan Kim and Alice Oh. How to find your friendly neighborhood: Graph attention design with self-supervision. In International Conference on Learning Representations, 2021.\n\n[4] Moshe Eliasof, Lars Ruthotto, and Eran Treister. Improving Graph Neural Networks with Learnable Propagation Operators. In International Conference on Machine Learning, 2023."}]

Author rebuttal:
We thank the reviewers for acknowledging the importance and relevance of our work and appreciate their constructive feedback, insightful questions, and encouraging comments. We take several actions to improve the paper in line with the reviewers' suggestions, summarized as follows:

i) We provide a more detailed explanation of some key ideas in the paper such as why the increasing depth of the network amplifies the trainability issue, how balancing the network mitigates the problem, and the effects of an orthogonal initialization.

ii) We conduct additional experiments as requested by the reviewers. Firstly, we include comparisons between imbalanced and balanced versions of Looks-Linear-Orthogonal, Identity, and Looks-Linear-Identity initializations. We also show the benefits of a balanced initialization on two other GNN models: the standard GCN and an architectural variation of GAT, the $\omega$GAT. We present these results in an additionally attached PDF and inline within individual responses to reviewers. 

iii) We discuss the limitations of our work in more detail such as the need for modifications of the derived conservation law for other attention mechanisms such as self-attention in SuperGAT and Transformers. 

We are grateful to the reviewers for their suggestions that enhance the presentation of our work. We will update our manuscript accordingly and would be happy to engage in discussions with the reviewers and answer any further questions. 

Meta-review: The reviewers appreciate the novel and solid technical contributions, and unanimously recommend acceptance and weak acceptance. The questions from the initial reviews were satisfactorily addressed.
Decision: Accept (poster)