SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Remark 2.7. This procedure balances the network starting from  $l = 1$  towards  $l = L$ . Yet, it could also be carried out in reverse order by first defining  $\left\| W^L [.;i]\right\| ^2 = \beta_i$  for  $i\in [n_{L - 1}]$ .
Balanced Orthogonal Initialization The feature weights  $W^{l}$  for  $l \in [L]$  need to be initialized randomly before balancing, which can simply be done by using Xavier initialization. However, existing works have pointed out benefits of orthogonal initialization to achieve initial dynamical
(a) Test accuracy (higher is better).
Figure 6: GAT trained on Cora using SGD
(b) Epochs to best model (lower is better).
isometry for DNNs and CNNs [37, 44, 23], as it enables training very deep architectures. In line with these findings, we initialize  $W^l$ , before balancing, to be an orthogonal matrix with a looks-linear (LL) mirrored block structure [46, 6], that achieves perfect dynamical isometry in perceptrons with ReLU activation functions. Due to the peculiarity of neighborhood aggregation, the same technique does not induce perfect dynamical isometry in GATs (or GNNs). Exploring how dynamical isometry can be achieved or approximated in general GNNs could be an interesting direction for future work. Nevertheless, using a (balanced) LL-orthogonal initialization enhances trainability, particularly of deeper models. We defer further discussion on the effects of orthogonal initialization, such as with identity matrices, to the appendix.
We outline the initialization procedure as follows: Let  $\parallel$  and  $\models$  denote row-wise and columnwise matrix concatenation, respectively. Let us draw a random orthonormal submatrix  $U^l$  and define  $W^{1} = [U^{1}\parallel -U^{1}]$  where  $U^{1}\in \mathbb{R}^{\frac{n_{1}}{2}\times n_{0}}$ ,  $W^{l} = [[U^{l}\parallel -U^{l}]]\models [-U^{l}\parallel U^{l}]$  where  $U^{l}\in \mathbb{R}^{\frac{n_{l}}{2}\times \frac{n_{l - 1}}{2}}$  for  $l = \{2,\ldots ,L - 1\}$ , and  $W^{L} = [U^{L}\models -U^{L}]$  where  $U^{L}\in \mathbb{R}^{n_{L}\times \frac{n_{l - 1}}{2}}$ . Since  $U^{L}$  is an orthonormal matrix,  $\| W^L [,:i]\| ^2 = 2$  by definition, and by recursive application of Eq. (2.3) (with  $a^l = 0$ ), balancedness requires that  $\| W^{1}[i:,]\|^{2} = 2$ . Therefore, we set  $\beta_{i} = 2$  for  $i\in n_1$ .
The main purpose of our experiments is to verify the validity of our theoretical insights and deduce an explanation for a major trainability issue of GATs that is amplified with increased network depth. Based on our theory, we understand the striking observation in Figure 4a that the parameters of default Xavier initialized GATs struggle to change during training. We demonstrate the validity of our theory with the nature of our solution. As we balance the parameter initialization (see Procedure 2.6), we allow the gradient signal to pass through the network layers, which enables parameter changes during training in all GAT layers. In consequence, balanced initialization schemes achieve higher generalization performance and significant training speed ups in comparison with the default Xavier initialization. It also facilitates training deeper GAT models.
To analyze this effect systematically, we study the different GAT initializations listed in Table 1 on generalization (in % accuracy) and training speedup (in epochs) using nine common benchmark datasets for semi-supervised node classification tasks. We defer dataset details to the supplement. We use the standard provided train/validation/test splits and have removed the isolated nodes from CiteSeer. We use the Pytorch Geometric framework and run our experiments on either Nvidia T4 Tensor Core GPU with 15 GB RAM or Nvidia GeForce RTX 3060 Laptop GPU with 6 GB RAM. We allow each network to train, both with SGD and Adam, for 5000 epochs (unless it converges earlier, i.e. achieves training loss  $\leq 10^{-4}$ ) and select the model state with the highest validation accuracy. For each experiment, the mean  $\pm 95\%$  confidence interval over five runs is reported for accuracy and epochs to the best model. All reported results use ReLU activation, weight sharing and no biases, unless stated otherwise. No weight sharing leads to similar trends, which we therefore omit. Our experimental code is available at https://github.com/RelationalML/GAT_Balanced_Initialization.
SGD We first evaluate the performance of the different initialization schemes of GAT models using vanilla gradient descent. Since no results have been previously reported using SGD on these datasets, to the best of our knowledge, we find that setting a learning rate of 0.1, 0.05 and 0.005 for  $L = [2,5]$ ,  $L = [10,20]$ , and  $L = 40$ , respectively, allows for reasonably stable training on Cora, CiteSeer, and
Table 2: Mean accuracy(%) ± 95% CI over five runs of GAT with width  $= {64}$  trained using SGD.
Init.	L=2	L=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
Actor	Xav	27.32 ± 0.59	24.60 ± 0.93	24.08 ± 0.80	22.29 ± 3.26	19.46 ± 5.75
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
Texas	Xav	60.00 ± 1.34	60.54 ± 3.42	58.92 ± 1.34	49.73 ± 20.97	17.84 ± 26.98
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	51.37 ± 8.90	51.76 ± 3.64	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.24 ± 4.52
Pubmed. For the remaining datasets, we set the learning rate to 0.05, 0.01, 0.005 and 0.0005 for  $L = [2,5], L = 10, L = 20$ , and  $L = 40$ , respectively. Note that we do not perform fine-tuning of the learning rate or other hyperparameters with respect to performance and use the same settings for all initializations to allow fair comparison.
Figure 6 highlights that models initialized with balanced parameters,  $\mathrm{Bal}_O$ , consistently achieve the best accuracy and significantly speed up training, even on shallow models of 2 layers. In line with our theory, the default Xavier (Xav) initialization hampers model training (see also Figure 4a). Interestingly, the default Xavier (Xav) initialized deeper models tend to improve in performance with width but cannot compete with balanced initialization schemes. For example, the accuracy achieved by Xav for  $L = 10$  increases from  $24.7\%$  to  $68.7\%$  when the width is increased from 64 to 512. We hypothesize that the reason why width aids generalization performance is that a higher number of (almost random) features supports better models. This would also be in line with studies of overparameterization in vanilla feed-forward architectures, where higher width also aids random feature models [2, 4, 57, 43].
The observation that width overparameterized models enable training deeper models may be of independent interest in the context of how overparameterization in GNNs may be helpful. However, training wider and deeper (hence overall larger) models is computationally inefficient, even for datasets with the magnitude of Cora. In contrast, the  $\mathrm{Bal}_O$  initialized model for  $L = 10$  is already able to attain  $79.7\%$  even with width=64 and improves to  $80.9\%$  with width=512. Primarily the
(a) Test accuracy.
(b) Epochs to convergence.
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, as shown in Figure 4a and 5.
For the remaining datasets, we report only the performance of networks with 64 hidden dimension in Table 2 for brevity. Since we have considered the  $\mathrm{Xav}_Z$  initialization only as an ablation study and find it to be ineffective (see Figure 6a), we do not discuss it any further.
Note that for some datasets (e.g. PubMed, Wisconsin), deeper models (e.g.  $L = 40$ ) are unable to train at all with Xavier initialization, which explains the high variation in test accuracy across multiple runs as, in each run, the model essentially remains a random network. The reduced performance even with balanced initialization at higher depth may be due to a lack of convergence within 5000 epochs. Nevertheless, the fact that the drop in performance is not drastic, but rather gradual as depth increases, indicates that the network is being trained, i.e. trainability is not an issue.
As a stress test, we validate that GAT networks with balanced initialization maintain their trainability at even larger depths of 64 and 80 layers (see Fig. 9 in appendix). In fact, the improved performance and training speed-up of models with balanced initialization as opposed to models with standard initialization is upheld even more so for very deep networks.
Apart from Cora, CiteSeer and Pubmed, the other datasets are considered heterophilic in nature and standard GNN models do not perform very well on them. State-of-the-art performance on these datasets is achieved by specialized models. We do not compare with them for two reasons: i) they do not employ an attention mechanism, which is the focus of our investigations, but comprise various special architectural elements and ii) our aim is not to outperform the SOTA, but rather highlight a mechanism that underlies the learning dynamics and can be exploited to improve trainability. We also demonstrate the effectiveness of a balanced initialization in training deep GATs in comparison to LipschitzNorm [14], a normalization technique proposed specifically for training deep GNNs with self-attention layers such as GAT (see Table 6 in appendix).
Adam Adam, the most commonly used optimizer for GNNs, stabilizes the training dynamics and can compensate partially for problematic initializations. However, the drop in performance with depth, though smaller, is inevitable with unbalanced initialization. As evident from Figure 7,  $\mathrm{Bal}_O$  initialized models achieve higher accuracy than or at par with Xavier initialization and converge in fewer epochs. As we observe with SGD, this difference becomes more prominent with depth, despite the fact that Adam itself significantly improves the trainability of deeper networks over SGD. Our argument of small relative gradients (see Eq.(8)) also applies to Adam. We have used the initial learning rates reported in the literature [9]: 0.005 for Cora and CiteSeer, and 0.01 for Pubmed for the 2 and 5 layer networks. To allow stable training of deeper networks, we reduce the initial learning rate by a factor 0.1 for the 10 and 20 layer networks on all three datasets.
Architectural variations We also consider other architectural variations such as employing ELUs instead of ReLUs, using multiple attention heads, turning off weight sharing, and addition of standard elements such as weight decay and dropout. In all cases, the results follow a similar tendency as already reported. Due to a lack of space, we defer the exact results to the appendix. These variations further increase the performance of networks with our initialization proposals, which can therefore be regarded as complementary. Note that residual skip connections between layers are also supported
in a balanced initialization provided their parameters are initialized with zero. However, to isolate the contribution of the initialization scheme and validate our theoretical insights, we have focused our exposition on the vanilla GAT version.
Limitations The derived conservation law only applies to the self-attention defined in the original GAT and GATv2 models, and their architectural variations such as  $\omega$ GAT[18] (see Fig. 10 in appendix). Note that the law also holds for the non-attentive GCNs (see Table 7 in appendix) which are a special case of GATs (where the attention parameters are simply zero). Modeling different kinds of self-attention such as the dot-product self-attention in [28] entails modification of the conservation law, which has been left for future work.
GATs [9] are powerful graph neural network models that form a 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 a conservation law that preserves a sum of the respective squared weight and attention norms during gradient flow. Accordingly, if weight and attention norms are highly unbalanced as is the case in standard GAT initialization schemes, relative gradients for larger parameters do not have sufficient room to increase.
This phenomenon is similar in nature to the neural tangent kernel (NTK) regime [24, 54], where only the last linear layer of a classic feed-forward neural network architecture can adapt to a task. Conservation laws for basic feed-forward architectures [15, 3, 32] do not require an infinite width assumption like NTK-related theory and highlight more nuanced issues for trainability. Furthermore, they are intrinsically linked to implicit regularization effects during gradient descent [41]. Our results are of independent interest also in this context, as we incorporate the attention mechanism into the analysis, which has implications for sequence models and transformers as well.
One of these implications is that an unbalanced initialization hampers the trainability of the involved parameters. Yet, the identification of the cause already contains an outline for its solution. Balancing the initial norms of feature and attention weights leads to more effective parameter changes and significantly faster convergence during training, even in shallow architectures. To further increase the trainability of feature weights, we endow them with a balanced orthogonal looks-linear structure, which induces perfect dynamical isometry in perceptrons [10] and thus enables signal to pass even through very deep architectures. Experiments on multiple benchmark datasets have verified the validity of our theoretical insights and isolated the effect of different modifications of the initial parameters on the trainability of GATs.
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.
[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. In International Conference on Learning Representations, 2019.
[3] Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach, 2018.
[59] Jiaxuan You, Rex Ying, and Jure Leskovec. Design space for graph neural networks. In Advances in Neural Information Processing Systems, 2021.
[60] Seongjun Yun, Minbyul Jeong, Raehyun 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. Pairnorm: Tackling oversmoothing in gnns. In International Conference on 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, Yanfei Dong, Kaixin Wang, Wee Sun Lee, Bryan Hooi, Huan Xu, and Jiashi Feng. Understanding and resolving performance degradation in graph convolutional networks. In Conference on Information and Knowledge Management, 2021.
[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.
6 Proofs of Theorems
Notation. Scalar and element-wise products are denoted  $\langle ,\rangle$  and  $\odot$ , respectively. Boldface lowercase and uppercase symbols represent vectors and matrices, respectively.
Our proof of Theorem 2.2 utilizes a rescale invariance that follows from Noether's theorem, as stated by [31]. Note that we could also derive the gradient structure directly from the derivatives, but the rescale invariance is easier to follow.
Definition 6.1 (Rescale invariance). The loss  $\mathcal{L}(\theta)$  is rescale invariant with respect to disjoint subsets of the parameters  $\theta_{1}$  and  $\theta_{2}$  if for every  $\lambda >0$  we have  $\mathcal{L}(\theta) = \mathcal{L}((\lambda \theta_1,\lambda^{-1}\theta_2,\theta_d))$ , where  $\theta = (\theta_{1},\theta_{2},\theta_{d})$ .
We frequently utilize the following relationship.
Lemma 6.2 (Gradient structure due to rescale invariance [31]). The rescale invariance of  $\mathcal{L}$  enforces the following geometric constraint on the gradients of the loss with respect to its parameters:
$$
\left\langle \theta_ {1}, \nabla_ {\theta_ {1}} \mathcal {L} \right\rangle - \left\langle \theta_ {2}, \nabla_ {\theta_ {2}} \mathcal {L} \right\rangle = 0. \tag {10}
$$
We first use this rescale invariance to prove our main theorem for a GAT with shared feature transformation weights and a single attention head, as described in Def. 2.1. The underlying principle generalizes to other GAT versions as well, as we exemplify with two other variants. Firstly, we study the case of unshared weights for feature transformation of source and target nodes. Secondly, we discuss multiple attention heads.
Proof of Theorem 2.2 Our main proof strategy relies on identifying a multiplicative rescale invariance in GATs that allows us to apply Lemma 6.2. We identify rescale invariances for every neuron  $i$  at layer  $l$  that induce the stated gradient structure.
Specifically, we define the components of in-coming weights to the neuron as  $\theta_{1} = \{w\mid w\in W^{l}[i,:]\}$  and the union of all out-going edges (regarding features and attention) as  $\theta_{2} = \{w\mid w\in W^{l + 1}[:$ $,i]\} \cup \{a^l [i]\}$ . It is left to show that these parameters are invariant under rescaling.
Let us, therefore, evaluate the GAT loss at  $\lambda \theta_{1}$  and  $\lambda^{-1}\theta_{2}$  and show that it remains invariant for any choice of  $\lambda > 0$ . Note that the only components of the network that potentially change under rescaling are  $h_u^l[i]$ ,  $h_v^{l+1}[j]$ , and  $\alpha_{uv}^l$ . We denote the scaled network components with a tilde resulting in  $\tilde{h}_u^l[i]$ ,  $\tilde{h}_v^{l+1}[k]$ , and  $\tilde{\alpha}_{uv}^l$ . As we show, parameters of upper layers remain unaffected, as  $\tilde{h}_v^{l+1}[k]$  coincides with its original non-scaled variant  $\tilde{h}_v^{l+1}[k] = h_v^{l+1}[k]$ .
Let us start with the attention coefficients. Note that  $\tilde{a}^l [i] = \lambda^{-1}a^l [i]$  and  $\tilde{W}^l [i,j] = \lambda W^l [i,j]$ . This implies that
$$
\tilde {\alpha} _ {u v} ^ {l} = \frac {\exp \left(e _ {u v} ^ {l}\right)}{\sum_ {u ^ {\prime} \in \mathcal {N} (v)} \exp \left(e _ {u v} ^ {l}\right)} = \alpha_ {u v} ^ {l}, \text {b e c a u s e} \tag {11}
$$
$$
\tilde {e} _ {u v} ^ {l} = \left(a ^ {l}\right) ^ {\top} \cdot \phi_ {2} \left(W ^ {l} \left(h _ {u} ^ {l - 1} + h _ {v} ^ {l - 1}\right)\right) = e _ {u v} ^ {l}, \tag {12}
$$
which follows from the positive homogeneity of  $\phi_{2}$  that allows
$$
\begin{array}{l} \tilde {e} _ {u v} ^ {l} = \lambda^ {- 1} a ^ {l} [ i ] \phi_ {2} \left(\sum_ {j} ^ {n _ {l - 1}} \lambda W ^ {l} [ i, j ] \left(h _ {u} ^ {l - 1} [ j ] + h _ {v} ^ {l - 1} [ j ]\right) + \sum_ {i ^ {\prime} \neq i} ^ {n _ {l}} a ^ {l} [ i ] \phi_ {2} \left(\sum_ {j} ^ {n _ {l - 1}} W ^ {l} [ i, j ] \left(h _ {u} ^ {l - 1} [ j ] + h _ {v} ^ {l - 1} [ j ]\right) \right. \right. (13) \\ = \lambda^ {- 1} \lambda a ^ {l} [ i ] \phi_ {2} \left(\sum_ {j} ^ {n _ {l - 1}} W ^ {l} [ i, j ] \left(h _ {u} ^ {l - 1} [ j ] + h _ {v} ^ {l - 1} [ j ]\right) + \sum_ {i ^ {\prime} \neq i} ^ {n _ {l}} a ^ {l} [ i ] \phi_ {2} \left(\sum_ {j} ^ {n _ {l - 1}} W ^ {l} [ i, j ] \left(h _ {u} ^ {l - 1} [ j ] + h _ {v} ^ {l - 1} [ j ]\right) \right. \right. (14) \\ = e _ {u v} ^ {l}. (15) \\ \end{array}
$$
Since  $\tilde{\alpha}_{uv}^{l} = \alpha_{uv}^{l}$ , it follows that
$$
\begin{array}{l} \tilde {h _ {u}} ^ {l} [ i ] = \phi_ {1} \left(\sum_ {z \in \mathcal {N} (u)} \alpha_ {z u} ^ {l} \sum_ {j} ^ {n _ {l - 1}} \lambda W ^ {l} [ i, j ] h _ {z} ^ {l - 1} [ j ]\right) \\ = \lambda \phi_ {1} \left(\sum_ {z \in \mathcal {N} (u)} \alpha_ {z u} ^ {l} \sum_ {j} ^ {n _ {l - 1}} W ^ {l} [ i, j ] h _ {z} ^ {l - 1} [ j ]\right) \\ = \lambda h _ {u} ^ {l} [ i ] \\ \end{array}
$$
In the next layer, we therefore have
$$
\begin{array}{l} \tilde {h} _ {v} ^ {l + 1} [ k ] = \phi_ {1} \left(\sum_ {u \in \mathcal {N} (v)} \alpha_ {u v} ^ {l + 1} \sum_ {i} ^ {n _ {l}} \lambda^ {- 1} W ^ {l + 1} [ k, i ] \tilde {h} _ {u} ^ {l} [ i ]\right) \\ = \phi_ {1} \left(\sum_ {u \in \mathcal {N} (v)} \alpha_ {u v} ^ {l + 1} \sum_ {i} ^ {n _ {l}} \lambda^ {- 1} W ^ {l + 1} [ k, i ] \lambda h _ {u} ^ {l} [ i ]\right) \\ = \phi_ {1} \left(\sum_ {u \in \mathcal {N} (v)} \alpha_ {u v} ^ {l + 1} \sum_ {i} ^ {n _ {l}} W ^ {l + 1} [ k, i ] h _ {u} ^ {l} [ i ]\right) \\ = h _ {v} ^ {l + 1} [ k ] \\ \end{array}
$$
Thus, the output node representations of the network remain unchanged, and according to Def. 6.1, the loss  $\mathcal{L}$  is rescale-invariant.
Consequently, as per Lemma 6.2, the constraint in Eq.(10) can be written as:
$$
\langle W ^ {l} [ i,: ], \nabla_ {W ^ {l} [ i,: ]} \mathcal {L} \rangle - \langle a ^ {l} [ i ], \nabla_ {a ^ {l} [ i ]} \mathcal {L} \rangle - \langle W ^ {l + 1} [:,: i ], \nabla_ {W ^ {l + 1} [:,: i ]} \mathcal {L} \rangle = 0.
$$
which can be rearranged to Eq.(2.2):
$$
\langle W ^ {l} [ i,: ], \nabla_ {W ^ {l} [ i,: ]} \mathcal {L} \rangle - \langle a ^ {l} [ i ], \nabla_ {a ^ {l} [ i ]} \mathcal {L} \rangle = \langle W ^ {l + 1} [:,: i ], \nabla_ {W ^ {l + 1} [:,: i ]} \mathcal {L} \rangle .
$$
thus proving Theorem 2.2.
Proof of Corollary 2.3 Given that gradient flow applied on loss  $\mathcal{L}$  is captured by the differential equation