Dynamic Rescaling for Training GNNs
Authors: Nimrah Mustafa
Rebekka Burkholz
Venue: NeurIPS 2024 poster (2024)
Keywords: graph neural network
rescale invariance
generalization
network balance
Primary area: graph_neural_networks

Abstract: Graph neural networks (GNNs) with a rescale invariance, such as GATs, can be re-parameterized during optimization through dynamic rescaling of network parameters and gradients while keeping the loss invariant. In this work, we explore dynamic rescaling as a tool to influence GNN training dynamics in two key ways: i) balancing the network with respect to various criteria, and ii) controlling the relative learning speeds of different layers. We gain novel insights, unique to GNNs, that reveal distinct training modes for different tasks. For heterophilic graphs, achieving balance based on relative gradients leads to faster training and better generalization. In contrast, homophilic graphs benefit from delaying the learning of later layers. Additionally, we show that training in balance supports larger learning rates, which can improve generalization. Moreover, controlling layer-wise training speeds is linked to grokking-like phenomena, which may be of independent interest.

Dynamic Rescaling for Training GNNs
Nimrah Mustafa CISPA
Rebekka Burkholz CISPA
66123 Saarbrücken, Germany nimrah., wwsttaffdcissa.de
66123 Saarbrücken, Germany burkholz@cispa.de
Abstract
Graph neural networks (GNNs) with rescale invariance, such as GATs, can be re-parameterized during optimization through dynamic rescaling of network parameters and gradients while keeping the loss invariant. In this work, we explore dynamic rescaling as tool to influence GNN training dynamics in two key ways: i) balancing the network with respect to various criteria, and ii) controlling the relative learning speeds of different layers. We gain novel insights, unique to GNNs, that reveal distinct training modes for different tasks. For heterophilic graphs, achieving balance based on relative gradients leads to faster training and better generalization. contrast, homophilic graphs benefit from delaying the learning of later layers. Additionally, show that training in balance supports larger learning rates, which can improve generalization. Moreover, controlling layer-wise training speeds is linked to grokking llike phenomena, which may be of independent interest.
Deep neural networks (DNNs) with positively homogeneous non-linear activation functions such as ReLU exhibit the rescale invariance property (31, i.e. scaling down incoming weights to neuron and scaling up the outgoing weights from the neuron by the same factor does not alter the function represented by the network.
Under gradient flow that assumes infinitesimally small learning rate, this rescale symmetry induces the conservation of a relationship between network parameters and gradients, determined by the initial state g More specifically, it is known that for traditional DNNs and convolutional neural networks (CNNs) with homogenous activation functions such as ReLUs, the difference between the squared L2-norms of incoming and outgoing parameters to a neuron stays constant (and is thus conserved). When this conserved quantity is (nearly) zero, the network is said to be in a balanced state. Training DNNs generally considered well-conditioned by balanced state of this conservation law.
In the context of Graph Neural Networks (GNNs), the first work presenting insights regarding norm balance derives the conservation law for GATs (2] and demonstrates how a balanced initialization enhances the trainability of GATs, particularly deeper networks (28). This balanced state that would hold under gradient flow could potentially be beneficial throughout training. However, in practice, factors that drive optimization such as finite learning rates, momentum, weight decay, and batch stochasticity break the rescale (and other) symmetries [55] and consequent conserved quantities, causing the network to topple out of balance due to which training conditions may deteriorate.
Motivated by the positive outcomes of training model balanced initialization, we utilize the identified rescale invariance of GAT's (88 to further investigate the effects of maintaining this balance throughout training rather than only initialization, by dynamic rescaling, i.e. scaling network parameters at the neuron level during training in loss-invariant manner. To this end, we derive general procedure balance GAT network. not necessarily only w.r.t. parameter weight norms, but any criterion that is function of network parameters and gradients. More specifically, we propose a novel criterion based on relative gradients of network parameters and demonstrate that this criterion, or combination of both criteria, often offers practical gains in terms of training speedup or better generalization.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

The implications of the core concept of dynamic rescaling extends beyond training in balance. For example, it enables arbitrary control of the order in which network layers learn during training. This can be viewed as inducing an imbalanced state in the network that may be desirable in some cases [16] Based on the experimental exploration of our methodological ideas, we draw various novel insights into how dynamic rescaling can potentially be leveraged. Firstly, specifically regarding graph learning, we discover potential trends of optimal learning dynamics for homophily and heterophily. We observe that, in terms of training speed and generalization, training all network layers in balance tends to be more beneficial for heterophilic than homophilic graphs. On the contrary, homophilic graphs tend to hypothesize that larger
benefit from more focused learning in the first (earlier) layers. Secondly, learning rates that tend to increasingly disrupt network balance can be better supported by dynamic rebalancing of the network in improving generalization. Thirdly, we encounter other interesting phenomena similar to grokking (33 that we relate to our insight on layer learning order.
This work is also motivated by earlier studies on deep feedforward non-linear neural networks that exploit rescale invariance, using transformations respecting loss- loss-invariant& symmetry to teleport parameters to another point in the loss landscape with steeper gradients to improve optimization [46 and/or convergence rate guarantees [47.. To the best of our knowledge, we are the first to conduct an GNNs, where identifying and exploiting the rescale invariance is not as straightforward, due to their peculiar architectural design elements such as node-wise neighborhood aggregation that make identifying the corresponding rescale invariances challenging. Thus, this largely remains unexplored territory for GNNs and our understanding of the underlying gradient dynamics in GNNs lags in comparison to DNNs. Rather than proposing one-size-fite fits-all solution or achieving state-of-the-an performance, this paper aims to contribute insights into GNN learning dynamics.
initial exploration of these ideas
In summary, our contributions are as follows:
1. We prove that, given a GNN exhibits rescale invariance (like GCNs GATs), we can manipulate gradient norms while leaving the GNN and the loss invariant and thus influence the learning dynamics during optimization.
2. We derive the procedure to balance a GAT network w.r.t. any criterion that is function of network parameters and gradients by dynamic rescaling.
3. We suggest a novel criterion for balancing based on relative gradients and find it to be promising for improving generalization and training speed in practical settings.
4. We explore our conceptual ideas empirically and find promising directions to utilize dynamic rescaling for more practical benefits, by training in balance or controlling order of learning among network layers.
5. We discuss novel insights regarding i) trends in training dynamics for homophilic and heterophilic graphs ii) larger learning rates and iii) interesting grokking-like phenomena.
Related work
There has recently been an increased interest in studying training dynamics
43 and generalization
[39] Approaches to address trainability issues in GNNs include initialization 20
28
13
12]
normalization 3 5
50
48 skip-connections [3, regularization 33 44), their combinations 25
architectural variations 29 and insights from graph signal processing based on spectral properties
45.. While our approach of dynamic rescaling allows us to use gradients the network to control the rate at which GNN layers (and potentially neurons) learn non-linear transformations of their features, draws an interesting parallel by modulating message passing updates based on gradients to control the rate at which nodes learn. The closest work to ours is (28) which proposes using balanced norms at initialization. Our work differs in mainly three ways as we: i) derive how a balanced state can be achieved not only at initialization but also during training by dynamic rescaling, ii) propose a different criterion for balance based on relative gradients, and iii) present insights on training in and out of balance in light of input graph homophily and heterophily.
For traditional DNNs, there is a deeper understanding of loss invariant symmetries and their im- pact on gradient dynamics 8 17 40 Several studies exploit rescale symmetry and corresponding conservation laws in ffedd-fowwrrd and convolutional) neural networks in various ways to aid opti- mization [1 26 46), regularization [37], and compression[38] 338. [7] introduce a set of nonlinear,

---

data-dependent symmetries, relate conserved quantities to the convergence rate and sharpness of the optima, and provide insights into how initialization impacts convergence and generalizability. Gradi- ent flow equations for neural networks have also been extended to account for realistic optimization elements such as finite learning rates, momentum, weight decay, etc. [15]
2 Dynamic rescaling
Preliminaries Consider a L layer GAT network f with positively homogeneous activation ø
(i. 6(x) zo (z)) and consequently, @(ax) a0(x) for positive scalars a) such as ReLU
6(x) = m 0} or LeakyReLU o(x) = max z 0} + -a max{ -2, 0}. Then, as shown by [28]. [:, i], and
the parameters [i. [
associated with a hidden unit i the network layer 1, [[. [ = A w [t, i] and a' [i] = A a where À > O such that f = f i.e. the rescaling respects the network symmetry.
may be respectively scaled to [;, = A AW [. :]. W
Note that W :] and [;+ denote weights incoming to and outgoing from neuron E [di]. respectively, where di is the width of layer I. Given a network parameter and its gradient VoL w.r.t. the network loss L, the relative gradient AO of parameter 0 is defined as:
A = VoL/O for 0 #0 or A = 0 for A =
(I)
This rescale property is more powerful than it might appear at first, as it provides us the means to significantly influence the training dynamics. It suggests that have high number of degrees of freedom to pick a parameterization without changing the function of a GAT. Concretely, we can use any scaling factors X!" > 0 that are associated with features in the middle layers and define another parameterization W j] = x 126-11 1
Wi j] and a'[: = a' [ ((= that will induce the same function. Yet, according to the following lemma, the gradients of the parameters depend on the rescale factors and can therefore be controlled correspondingly.
Lemma 2.1 (Gradient scaling). Under the rescale invariance of GATs, if a parameter is scaled by i = X0, then its gradient is scaled as V VgC- X VoL.
We defer the proof to the appendix A. It generally implies that we have the freedom to pick any
Wwaj and V = X ((") V [g C induce favorable learning dynamics. Considering the flexibility, in move during gradient descent, the choice likely has a significant
positively homogeneous constants so that our resulting gradients V C x x-11)/1 />mh
which gradient direction influence on our learning success. We aim to exploit this fact in this work and discover conceptual insights into what criteria could constitute choices. Thus, follow the following procedure.
Balancing criteria and procedure The rescale invariance property allows to rescale parameters to fulfill the desired criterion not only at initialization but also during training without changing the network output while potentially improving the training dynamics.
We could indeed choose a relatively general criterion g R. - R that depends on our rescaled parameters and gradients and determines our choice of scaling factors. As long as we can solve
g (x ['..),, W
[t. A (1) (1) al aajj A V V [w: C. A Vweng C, V C) = 0.
(2)
For unique nonzero scaling parameters, 9 can act our guide during gradient descent. Importantly not all reasonable criteria determine the scaling factors.
As has been derived recently for GAT's 28), the rescale invariance also induces a conservation law that holds throughout training and is characterized by the fact that the scaling factors cancel out. In fact, any such law that remains invariant under specific parameter transformations is also linked to an invariance like the rescale invariance. according Noether's theorem. Specifically, in the case of GATs, the following equations
(3) hold regardless of the scaling factors. This law implies that a specific sum of L2-norms of the corresponding parameters stays conserved throughout gradient descent if the learning rates are sufficiently small.
(w' :]. Cw (a' [.. V) (W'++ Vw "(.aI =0
3

---

Recently, it has been shown that balancing these squared parameter 12 1 norms at initialization such [: a]]l [||
that w (.:i] | in GAT's [28]. For larger learning rates, this balance might get disturbed during training. Yet, we could use our rescaling degrees of freedom to bring the parameters back in balance. To induce even better trainability, we propose another criterion to rescale a GAT network. which balances the norms of gradients relative to the corresponding parameters, ie., A VoL/0. Intuitively, this should allow the parameters in different layers to move at similar speeds and ensure good trainability in all parts of the network.
0, vi E di, E [L-1], induces good initial trainability
Aw J. [.:]]] || |Aw [=]]=== =
(4)
Yet, as this criterion is not naturally preserved during gradient descent or gradient flow like the weight norms, fulfilling the equation above requires frequent rescaling during training. Thus, based on Eq. balancing a neuron i in layer w.r.t. relative gradients requires fulfilling::
W S
0.
(5)
[i.:]
W
a
[
Balancing the entire network requires fulfilling Eq.15) Vi E [d), E [L 1]. In this case, note that every weight [i] is eventually scaled by /2 Thus, balancing the entire network requires iterative rescaling until convergence of all rescaling factors to
For each iteration E [1, T] and given ((1)(0) 1, Vi [d]. [L 1], the scaling factor X (0)()) given by:
is
A Aw"li [i, (1-1))
10000)
; where
(6)
(a 1| + AWi++| (1 2
xmi VoL if E w 16 ;1> 0
;1>0
A0(1)
0
sort vol 101001
if (w 00 a : 0.
Ideally, this process is repeated to convergence until X (1) = 1, Vi [d], [L 1]. The number of required iterations depends on the frequency of rebalancing during training as well as the network find that this is not too computationally expensive process
parameters and gradients. In practice, and a few iterations (< 10) are sufficient to (mostly if not completely) balance the network. In terms of time complexity, this only affects the training time linearly depending on the frequency of rebalancing, which is controllable hyperparameter.
Implications Dynamic rescaling opens up the opportunity to control training dynamics in several ways for which we lay out two key ideas as follows.
Firstly, dynamic rescaling can be used to train networks in balance w.r.t certain criteria. We propose that balancing based on relative gradients may be one such good candidate. Our intuition is that balanced relative gradients allow all layers (and neurons) in the network relatively equal opportunity to learn by propagating gradients to drive parameter change throughout the network, thereby enhancing trainability. We observe that this novel insight shows the promising potential of being translated into practical gains such as faster or better generalization on real-world data, particularly for heterophilic tasks.
Secondly, dynamic rescaling allows us to control of relative training speed the level or neuron level. In principle, by rescaling layers, we can configure the relative order in which they learn arbitrarily, at any time during training. This is direct consequence of the conservation law that parameters and their gradients in the network adhere to. The underlying insight is that by controlling the parameter weight or relative gradient norms, the gradients that drive parameter change can be influenced. As result, layer receiving relatively larger gradients (due to relatively scaled-down

---

Table 1: Results of training a 5-layer GAT network with various dynamic rescaling (DR) settings. The mean =95% CI test metric at the epoch of the best validation metric across 10 splits is re- ported using the best learning rate from (0.01, 0.001, 0.005]. The evaluation metric accuracy for roman-empire and amazon-ratings. and ROC AUC for the remaining three datasets.
roman -mmpee
amazon-ratings
tolokers
questions
minesweeper
w/o DR
.478 + .2099
.545 = .0043
.6493 + .008
.5829 + .0172
.5057
.0058
DRw
.3307 = .670
4547 + .0041
.6451 = .0135
5
+ .0145
50588 = .0051
DRaG
.5422 = .0234
.4400 = .0029
.6637 = 00888
5869 + .0145
5065 + .0076
DRc
.6731 + .0149"
.526 = .0042
6642 + .0064
5699 + .0132
5080 + .0053
weight norms) thus learns 'more' or 'faster' than other layers. While the possibilities are numerous, in this work, we limit our investigation simple but core case: allowing the network to (initially) focus learning on specific layer by scaling down its weight norm at initialization. We observe that initially allowing more focused learning in the earlier layers of the network can improve the convergence time substantially while retaining even improving the generalization for homophilic tasks.
Our findings suggest that the order which layers learn influences convergence time and generaliza- tion. This building block may be leveraged devise more sophisticated learning sequences such as training layers cyclically or in a task-dependent manner rather than in predefined order. We elaborate on the current limitations of dynamic rescaling in the appendix.
While these concepts apply to any GNN, provided its rescale invariance has been identified, we focus investigation the GAT architecture in this work as they GCN architecture that is the building block of more complex GNNs. Furthermore, GAT serves as a strong basis for graph learning with attention mechanism, in which there has been increased interest recently 18
generalization of the more basic
3
Experiments
We divide our exploration of the ideas discussed $ into three parts: 1) practical gains on real-world data of training in balance by dynamic rescaling, 2) empirical insights into the layer- level order of learning, and 3) observation of grokking-like phenomenon, which related to 2). Hereafter, we use the notation DRw, and DR to denote dynamic rescaling w.r.t. weight norms and relative gradients, respectively, every 10th epoch. DRc denotes a combination of the two by rescaling w.r.t. weight norms every 10th epoch and w.r.t. relative gradients in all other epochs. A maximum 10 iterations for the rebalancing procedure outlined in Eq. (6) and (7) used. All experiments looks-linear orthogonal structure [36 I unless specified otherwise. Experiments were run on an NVIDIA RTX A6000 GPU with 50GB RAM. Our experimental code is available at https //github. wtt..o//eeeaaaiaaaaaaaaaiaaaaaea
the Adam optimizer and networks are randomly initialized with
3.1 Training in balance
We primarily study the effect of training GAT in balanced state based on the relative gradients criterion (see Eq. ()), by dynamic rescaling on five real-world heterophilic benchmark datasets (32
We find that rebalancing the network w.r.t. relative gradient norms is more effective than the criterion based on parameter weight norms, as shown Table I This aligns with observations on CNNs where rebalancing w.r.t. parameter norms was also not effective [37]. Therefore, our insight on the impact of rebalancing w.r.t. relative gradients may also be of independent interest outside the context of GNNs.
In addition to improved generalization in most cases, dynamic rescaling may also provide the benefit of fewer training epochs to attain comparable even slightly better generalization, as shown in Fig. We make an interesting observation that balanced state during training together with larger learning rates results in better generalization than when either of the two components is individually employed.
5

---

0.000,, wio
0.001, DP
0.001, DR
0.01, veoo
Ir 0.01, DRoc
2 tr 0.01, 0.01, DRg
Test 41100 vitSiiio miitc
0p4211 10000 00000000111 m4200
6.80
0.90
3800
Maus
(a) Higher is better.
(b) Lower better.
Figure I: Performance of a 5-layer GAT with various dynamic rescaling (DR) settings using learning rates (Ir) 0.001 and 0.01. Across 10 splits, the mean accuracy is reported for roman-empixe and amazon-rat inns while ROC AUC is reported for minesveeper, questions ionss and tolokers. The case annotated by indicates training for more than 10k epochs.
Layen
Layen2
Layen3
Layend
every Dispocte
every spach
23ah
29"
18
to th-
50
1B00
10000
1300
1X00
19p
3332 LoP
20
20
10
12
14
15
20
Lesining
18305
10:000
10000
LO0CE
Figure 3: A narrow range of larger learn- ing rates results in better generalization for roman - empire dataset on GAT, with further im- provement brought by dynamic rescaling (DR) to balance relative gradients.
Figure 2: The degree of imbalance, i.e. the R.H.S. quantity of Eq. (5). before and after rebalancing every 10 epochs when training roman- -emaree on GAT. A value of 0 indicates complete balance.
This synergetic effect can be attributed to two factors. Firstly, several works report empirical evidence that larger learning rates are usually associated with flatter minima that have been linked to better generalization [4 7 19] On the contrary, larger learning rates push the network out of its balanced state faster and more severely (see Fig. 2). which may impede trainability. However, this can now be addressed by rebalancing the network during training to facilitate trainability. Thus, balancing by dynamic rescaling supports higher learning rates in improving generalization. However, for given task, only a narrow range of these 'large enough' learning rates can produce optimal results [24]. We show in Fig. 3 that training in balance can further improve the performance for this range of larger learning rates.
We use gradient clipping in combination with dynamic rescaling to accommodate exploding or vanishing gradients as result of training with a larger learning rate and any numerical instabilities that may arise due to direct manipulation of parameter weights and gradients. We defer the ablation study of gradient clipping to the appendix.
3.2 Learning layers in order
As discussed in 2 one potential opportunity with dynamic rescaling to improve training dynamics to control the order in which layers (or even neurons) in the network learn. We investigate this experimentally on both synthetic and real-world data, gaining novel insights into potential trends of optimal training dynamics for different tasks based on their homophily and heterophily. We defer the description of the synthetic data generation to the appendix.
The training curves and generalization performance achieved on the synthetic task for various train settings are summarized in Fig. 4 For this task, allowing the network to initially concentrate learning on the first layer results faster training, lower minimum test loss, and better generalization than training the network in the standard setting or with dynamic rescaling. Interestingly, as more initial

---

Standard
L=1
L-2
LF3
L=4
LR5
DR- Ir.0.001
DR- Ir.0.11
1.000
0.98
0.8
0.96
0.6
Test Loss
0.6
Train Ace
Test Acc
0.950
0.94
0.4
0925
0.92
0.2
0.000
0.90
500
Epoch
Epoch
Epoch
Epoch
Figure Performance of five-layer GAT network on synthetic data under varying settings. Standard implies regular training with no constant or dynamic rescaling. L = for / E [5] denotes scaling down the parameters of by constant (A 0.002) at initialization followed by regular training. DR denotes dynamic rescaling to balance relative gradients during training with the specified learning rate (Ir). Note that the train and test accuracy axis have been zoomed in for clarity and the initial (train test) accuracy is lower than 0.9 (but rises sharply in the first few epochs). The best strategy (among considered cases) for this task to (initially) focus the learning more on the first layer.
Cora
Pubmed
Physics
Photo
amazon-ratings
questions
Citeseer
CS
Computers
roman-empire
tolokers
minesweeper
Cervergrroe Gooe
Conveggree Sposh
4000
066
2000
8800
025
20
4m
0
0125
18
10
&
F2tc F2 Fin Traiing
Fihh F2 Ftr
Thaning
(a) parameters W of first layer are initially scaled by À loss invariant manner (i.c., 2 second layer scaled by A 1). implics standard training. BRG and Be denote training in balance hy DR and DRc, respectively. Entries annotated with indicate the best (highest) test metric and (lowest) convergence epoch. Models were trained for maximum 5000 epochs.
(b) X indicates the best from (a). Fl and denote the case where first and second layers the model are frozen, respectively, during training whereas +c indicates the presence of an additional linear classifier layer.
Figure 5: Impact of training layers of two-layer GAT network in and out balance for different tasks. The tasks -ranzonnraains,,, questions, roman- heterophilic and the remaining homophilic. Homophilic tasks tend perform better and converge much faster with learning concentrated in the first layer initially (lower weight norms imply larger relative gradients), whereas heterophilic tasks perform better when layers trained in balance. Interestingly, even freezing the initial values of parameters in the second layer (i.e. only allowing the second layer learn) does not significantly reduce the performance for homophilic tasks, even without an additional classifier layer. On the contrary, freezing the first layer results in a severe drop in performance for all tasks.
tolokers, minesweeper' )
focus is placed on each subsequent layer, the training slows down and generalization worsens, with I 1= being (marginally) the lowest. Dynamic rescaling with learning rate of 0.001, as used in all cases, is not effective. However, using larger learning rate of 0.01 allows the fastest training with the second-best generalization. We expand on this as we analyze the evolution of relative gradients for interesting cases in Fig. g in the appendix.
Note that this synthetic task, which benefits most from concentrated learning in the first layer, is designed to be homophilic and thus differs from the heterophilic tasks that benefit more from training in balance. Prompted by this observation, we next study the impact of ordered learning of layers for real-world homophilic and heterophilic graphs.
As shown Fig. 5 we find that real-world graphs indeed exhibit different trends regarding generaliza- tion and convergence time when trained in balance (w.r.t. relative gradients) and when trained out of balance (by controlling learning order of layers) depending on their homophilic or heterophilic nature. Training layers in balance is generally more effective for heterophilic tasks whereas allowing learning
7

---

Cross Entropy Loss
Cross Entropy Loss
2
10
05
0.0
800
1200
1800
2800
5000
3800
4000
88
180
210
240
220
300
Valid.
Train
Epoch
Train Test
Epoch
(a) Initial focus placed on second layer
Cross Entropy Loss
1.5
1.0
1.0
Cross Entropy
0.5
-3
100
210
270
300
180
210
270
300
Train
Epoch
Train
Epoch
(b) Initial focus placed on third (middle) layer
Figure 6: Evolution of gradient norms in five-layer GCN network trained on synthetic data and evaluated on two validation sets from the same input distribution as the train set (left and right) with initial learning focused the second and third layers in (a) and (b), respectively. The plot design is similar to that described Fig. S except that in this case, the heatmap represents gradient norms.
focus on the first layer tends benefit homophilic tasks, particularly terms of convergence time. We discuss this interesting observation more broadly in light of GNN training dynamics.
Generally, the optimal performance achieved by model is, a large extent, dependent how well the inductive bias the model architecture aligns with the task and its underlying graph structure. For example, widely known that general GNNs, without peeciall--introdun such as GCN perform better on homophilic than on heterophilic tasks. Intuitively, we hypothesize that homophilic tasks rely more on the neighborhood aggregation functionality GNNs rather than feature learning. In this case, aggregation over random transformation of similar features may still be sufficient for good generalization.
architectural elements,
Our insight in line with a recent analysis of training dynamics 43 which shows that the NTK that controls the evolution of the learned GNN function tends to align with the message passing matrix (i.e. the adjacency matrix in most cases). Furthermore, for homophilic graphs, the adjacency matrix also aligns well with the optimal kernel matrix that represents nodes with the same label. As result, on homophilic graphs, the alignment of the underlying structure with the optimal kernel matrix allows parameter-free methods similar to label propagation perform par with GNNs. However, the generalization of GNNs on heterophilic tasks, where the graph structure does not align with the optimal kernel, is mostly adversely impacted by the NTK aligning with the adjacency matrix. In other words, the structural information in the graph is not very relevant for node's label heterophilic settings and thus the node relies more on learning in the feature space rather than neighborhood aggregation. This is also supported by results showing that embedding additional MLP layers the network significantly improves the performance of basic GNNs such as GATs on these heterophilic tasks (32). Thus, we conclude that training in balance to potentially learn better feature transformations all layers (and potentially neighbors farther away in deeper models) more effective in heterophilic cases.
3.3 Grokking-like phenomena
Grokking is defined as a phenomenon where the validation loss reduces, long after a near-zero train loss has been achieved, towards perfect generalization. It has been observed primarily in the context of algorithmic and synthetic datasets 30 problem of fundamental interest, not only because the phenomenon appears be puzzling at first,
33
42 22 IT. Yet, it still regarded as

---

10
Train
Validation
Test
Cross Entropy Loss
Layer
10
Accuracy
0.5
8
0.0
Train
500
1500
2000
3000
4000
1600
5000
2000 Epoch
4000
Valid.
Epoch
6
Train
Validation
Test
Cross Entropy Loss
Layer
Accuracy
0.5
2
Train
800
1200
1600
2000
2400
2800
3200
3600
4000
0.0
Valid.
Booch
1000
2000
3000
4000
Test
Epoch
Figure 7: Left: Layer-wise relative gradient norm (logıo scale) and loss curves similar to Figure in the paper. Right: Corresponding accuracy of the same run. Grokking-like phenomenon can be induced On a 5-layer GAT using real-world roman- roman-empire empire dataset by placing initial training focus on layers 4 (top) and (bottom) by scaling down initial parameter norms, followed by rebalancing w.r.t. relative gradients every 10 epochs staring only at epoch 1000. Note the sharp drop in validation/test loss immediately after rehalancing which also translates more rapid improvement in test accuracy.
but also because it gives rise to nuanced insights into how neural networks learn, as accentuates different training phases.
Grokking also touches upon our study in which order network layers learn best, as it has been attributed to network layers getting out of balance due to weight decay and weight norm decrease More precisely, [4] argues that it is the result of delayed feature learning. According to their theory and observations for 3-layer neural network, primarily the last layer learns first linear combination of random features, as the parameter norms are scaled in such a way that it renders the first layer effectively untrainable, still reaching nearly zero training loss. With increase of the last layer norm, however, feature learning by the first layer begins and starts to drive down the generalization enror.
[2 42.
We conjecture that, depending on the learning task, grokking might be induced also by delayed learning of other layers, not only the last one. Our insights into rescaling based on relative gradients give us the means test this hypothesis.
synthetic node classification task is constructed, similarly as for the experiments in $3.2 except with N 100, p 0.01, and using GCN instead of GAT as the target network. Using 50 nodes as the train set, we train five-layer GCN and evaluate on the remaining nodes as two disjoint test sets each of size 25.
Under the condition that the network is allowed to focus learning on the second or third hidden layer by scaling the weight parameters of the respective layer by > 0.02 at initialization, the resulting training curves in Fig. 6 exxhibii a trend similar to grokking.
While one test set eventually achieves zero validation loss, the other's 's loss is reduced only to certain extent before increasing again to give rise U-shaped curve that is more common. The latter also similar to the only other case of grokking case detected in GNNs [23], to the best of our knowledge. Note that, the former case, learning focused on the middle (third) layer of the network achieves perfect generalization much faster than when learning is focused on the second layer. This indicates that layer-level control of learning may be leveraged to facilitate faster generalization grokking scenarios.

---

We also induce grokking peekking-related related behavior on real-world dataset roman- empire (see Fig. in two steps. Firstly, we allow only the last (or second to last) layer to learn which allows the training accuracy to increase continually while the test accuracy saturates or begins to drop. At this point, we rescale the network to bring all layers in balance w.r.f. relative gradients, following which, the test accuracy immediately begins to improve more rapidly accompanied by drop in training accuracy. This can be interpreted as the network 'learning' more effectively when trained in balance rather than overfitting to the training data. While this is different from grokking where the training accuracy would generally not drop, it is independently an interesting observation on a real-world dataset.
Our key takeaway is that, as opposed to the general perception that grokking is induced by learning only the last layer, we observe a similar pattern by focusing learning on other (hidden) layers. Strikingly, grokking as a phenomenon can also support learning. While this possibility is novel insight, we acknowledge that it a potentially noisy phenomenon that materializes only in specific settings. We would like to highlight that understanding grokking is a separate area on its own and recent efforts are also limited to synthetic data [27]. Our induction of a similar phenomenon in real-world data by influencing learning dynamics can be of independent interest to develop deeper theoretical understanding of such observations. Further investigation for more conclusive insights related to grokking is an interesting direction for future work.
4 Discussion
We have proposed to exploit the rescale invariance of GNNs such as GAT's to control their training dynamics for improved training speed and generalization. To that end, we propose to rescale the parameters in such way that the function, which is defined by a GNN that is trained, remains unperturbed, while the gradients are rescaled. The partial control of the learning dynamics, which is thus available, offers us the opportunity to teleport in the parameter space.
While the guiding principle of this teleportation is flexible, we propose to balance relative gradient norms such that all layers of the GNN are equally involved in the learning process. Our experiments highlight the potential of this rescaling to support larger learning rates and to balance relative gradients in search of flatter optima generally associated with better generalization. To our knowledge, we are the first to tap into the potential that the rescaling flexibility of GNNs has to offer.
While the goal of balancing relative gradients is involve all network layers equally in learning, have also analyzed the other end of the spectrum and discussed how we can influence the order in which GNN layers are learned. In doing so, we find that while training in balance is effective for heterophilic graphs, homophilic graphs tend to benefit more from training layers out of balance, with initially more focus learning in the first layers. This novel insight into trends regarding training modes for different tasks contributes to understanding GNN learning dynamics under homophily and heterophily.
Furthermore, the ordered learning of layers also has implications for grokking like phenomena, which provide insights into the inner mechanies of learning and are therefore also of fundamental interest. With our novel set of experiments, we could show that grokking does not necessarily result from a delay in feature learning, in which primarily the last layer learns in the first training phase. We could also induce a similar phenomenon by giving the middle layer a headstart in learning. Interestingly, this case provides also an example of scenario in which a grokking- llie phenomenon improves the overall generalization performance of the resulting model. We conjecture that unbalanced learning can also act regularization that fights over-fitting
Thus, our explorations and insights serve as the first stepping stone to developing a more compre- hensive theory or set of guidelines to leverage dynamic rescaling for training not only GNNs, but various other deep learning architectures, to improve training speed, generalization, and, potentially, robustness.
Acknowledgements
We gratefully acknowledge funding from the European Research Council (ERC) under the Horizon Europe Framework Programme (HORIZON) for proposal number 101116339 SPARSE-ML.
10

---

References
[1] 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? International Conference on Machine Learning. 2018.
[2] Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? In International Conference on Learning Representations, 2022.
[3] Tianle Cai, Shengjie Luo, Keyulu Xu, Di Hc, Tic- Yan Liu, and Liwei Wang. Graphnorm: A principled approach to accelerating graph neural network training. In International Conference on Machine Learning, 2021.
[4] 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.
[5] George Dasoulas, Kevin Scaman, and Aladin Virmaux. Lipschitz normalization for self-attention layers with application to graph neural networks.
[6] Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. arXiv preprini arXiv:2303.061 2023.
[7] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets, 2017.
[8] 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.
[9] Kimon Fountoulakis, Amit Levi, Shenghao Yang, Aseem Baranwal, and Aukosh Jagannath. Graph attention retrospective. In Journal of Machine Learning Research, 2023.
[10] 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.
[11] Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02679 2023.
[12] Kai Guo, Kaixiong Zhou, Xia Hu, Li, Yi Chang, and Xin Wang. Orthogonal graph neural networks. In AAAI Conference Artificial Intelligence, 2022.
[13] 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.
[14] Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. In The Twelfth International Conference Learning Representations, 2024.
[15] D. Kunin, J. Sagastuy-Brent S. Ganguli, D.L.K. Yamins, and H. Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. International Conference on Learning Representations, 2021.
[16] Daniel Kunin, Allan Raventós, Clémentine Dominé, Feng Chen, David Klindt, Andrew Saxe, and Surya Ganguli. Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning, 2024.
[17] Thien Le and Stefanie Jegelka. Training invariances and the low-rank phenomenon: beyond linear networks, 2022.
[18] Soo Yong Lee, Fanchen Bu, Jaemin Yoo, and Kijung Shin. Towards deep attention in graph neural networks: Problems and remedies. In International Conference on Machine Learning, 2023.

---

[19] Aitor Lewkowyzz. Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. 2020.
[20]
Jiahang Li, Yakun Song, Xiang Song, and David Paul Wipf. On the initialization of graph neural networks, 2023.
[21] Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651-34663, 2022.
[22] Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. arXiv preprint arXiv:2210.01111 2022.
[23] Ziming Liu, Eric J. Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In International Conference on Learning Representations, 2023.
[24] Ekaterina Lobacheva, Eduard Pockonechnyy. Maxim Kodryan, and Dmitry Vetrov. Large learning rates improve generalization: But how large are we talking about? In Mathematics of Modern Machine Learning Workshop at NeurIPS, 2023.
[25] Sitao Luan, Mingde Zhao, Xiao-Wert Chang, and Doina Precup. Training matters: Unlocking potentials of deeper graph convolutional neural networks, 2023.
[26] Qi Meng, Shuxin Zheng, Huishuai Zhang, Wei Chen, Zhi-Ming Ma, and Tie- Liu. G-sgd: Optimizing relu neural networks in its positively scale-invariant space.
[27] Mohamad Amin Mohamadi, Zhiyuan Li, Lei Wu, and Danica Sutherland. Why do you grok? a theoretical analysis of grokking modular addition. In International Conference Machine Learning, 2024.
[28] Nimrah Mustafa and Rebekka Burkholz. Are GATS out of balance? In Advances in Neural Information Processing Systems, 2023.
[29] Nimrah Mustafa and Rebekka Burkholz. Gate: How to keep out intrusive neighbors. In International Conference on Machine Learning, 2024.
[30] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023.
[31] Behnam Neyshabur, Ruslan Salakhutdinov, and Nathan Srebro. Path-sgd: Path-normalized optimization in deep neural networks. In Advances in Neural Information Processing Systems, 2015.
[32] Oleg Platonov, Denis Kuznedelev, Michael Diskin, Artem Babenko, and Liudmila Prokhorenkova. A critical look at the evaluation of gnns under heterophily: are we really making progress? In International Conference on Leurning Representations, 2023.
[33] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets, 2022.
[34] 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.
[35] T. Konstantin Rusch, Benjamin P. Chamberlain, Michael W. Mahoney, Michael M. Bronstein, and Siddhartha Mishra. Gradient gating for deep multi-rate learning graphs, 2023.
[36] 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.
[37] Pierre Stock, Benjamin Graham, Rémi Gribonval, and Hervé Jégou. Equi-normalization of neural networks. In International Conference Learning Representations, 2019.
12

---

[38] Hidenori Tanaka, Daniel Kunin, Daniel L. Yamins, and Surya Ganguli. Pruning neural networks without any data by iteratively conserving synaptic flow. In Advances in Neural Information Processing Systems, 2020.
[39] Huayi Tang and Yong Liu. Towards understanding the generalization of graph neural networks, 2023.
[40] Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal. Understanding the dynamics of gradient flow in overparamettrizzed linear models. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10153-10161, 2021.
[41] Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206. 04817, 2022.
[42] Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint wrXiv:2309.02390 2023.
[43] Chenxiao Yang, Qitian Wu, David Wipf, Ruoyu Sun, and Junchi Yan. How graph neural networks learn: Lessons from training dynamics. In International Conference on Machine Learning, 2024.
[44] Han Yang, Kaili Ma, and James Cheng. Rethinking graph regularization for graph neural networks. In Advances in Neural Information Processing Systems, 202.
[45] Hanqing Zeng, Muhan Zhang, Yinglong Xia, Ajitesh Srivastava, Andrey Malevich, Rajgopal Kamnan, Viktor Prasanna, Long Jin, and Ren Chen. Decoupling the depth and scope of graph neural networks, 2021.
[46] Bo Zhao, Nima Dehmamy, Robin Walters, and Rose Yu. Symmetry teleportation for accelerated optimization, 2022.
[47] Bo Zhao, lordan Ganev, Robin Walters, Rose Yu, and Nima Dehmamy. Symmetries, flat minima, and the conserved quantities of gradient flow, 2023.
[48] Lingxiao Zhao and Leman Akoglu. Pairnorm: Tackling oversmoothing in gnns. In International Conference on Learning Representations, 2020.
[49] Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning. 2022.
[50] 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.
13

---

A Proofs
A.1 Proof of Lemma 2.1
Lemma 2.1 states that, under the rescale invariance of GATs, if parameter is scaled by 0 = X0, then its gradient scaled V&C= VoC.
Proof. We prove this by deriving the gradients of GAT network parameters.
Firstly, given imput representations from layer - and homogeneous activation functions Φ1 and 02. a GAT layer is defined as:
h2 Φ0
Σ a'
where and
(8)
(EN(()
exp(c(re)
a EWEN(()) exp(e
(9)
eyv (a') (W' +b, 1))
(10)
Given the following layer +
[
I+1 aue
Σ w [k, 2
uEN( (v)
we derive the gradient of parameter W [k,]] w.r.t. to network loss using the product rule:
OC
OC
[h|+1(k
aw
Oh.+1
[k]
ow
[ "
OC
aa I+1 UW (k.
1+1 h + Σ
+1 [k, Th
ah [K)
Oh
uEN(v)
Now, W
1+1 "11 [k.]] --1 j]. then in order to keep L unchanged, which implies that h Ah( (as derived in 28 must hold so that
W G. W..
XW Xw
and a'(i) [i] aa [i]
h, [k] b. [[]] remains unchanged, as:
-i+1
-1+1 h
[K]
= di
Σ
Oue
++1
Σ x w [k: i] Ab
uEN(u)
72
Σ
age
1+1
Σ
-
[1; 1]
uEN(e)
Consequently,
0C
ac
-1+
Oh
[K]]
We first show that:
0at+1
ta
aw 1+1
aw
[K
14

---

For this, note that
e =
[ [6 b.111)
Σ a
my
+ m2
Σ w wh+ [m,
[h k.,itt)
miti
=
w [e (un' + sm((tt)
+
Σ a
++ [mn
Σ w (m, a] (ab + (bllel
matt
n2
Aa 2aa+111 (k/ob
w w' [k. [] (1:1 + m.(11)
Σ
a
|mnecs
> w Im, a] h1.00 + b((lt)
mifi
and,
da
delt1
ow
+1 [k,j]
Devi
1+1
tt mp de
a
[k)
(b. + bll + Σ a
m2
++ mm]oz Σ w
[m,
(n'li + by (1 111)
mati
Then.
0a
delt1
ow
[
Dewn
OW
[k
Y21
a ['+1[k] (bL + h/Dl))
Σ a
4 mm]
Σ
w [m, o| [b. b
axl+1
myght
ow
[k:
Also note that:
Oc
OC
1++ [k] h² [A:)
since h,
0h,
[k]
Th
Then, using all the above facts,
1++
0C
OC
Oh
7c
aw 1+1
[k,]]
Oh [+1
[k]
aw
[k
dc
n²
0a
Σ
Abl
Σ
w [k; i| i|lhh
Oue
oh
++ [K]
ow
[ks 5
NEN(())
X
Oc
Σ
+ h [i] + Σ
w 1l++1p 1 [k. ihh
ao 2+1
Oh
aw +1 [
uEN(u)
DC
aw'+
[k
Therefore, scaling the parameter W [ks, by X-1 scales its gradient by À (A-1)-1.
Next, we derive the gradients for parameter a' [i].
15

---

ac
aC
ah [a
@al[d]
Th
We note that under the rescaling, due- alw If aajl = X [t]. the rescale invariance implies that have also W
[i, = Aw [i, k] then:
ac
OC
0e [i]
0a'
de
0a's [i]
OC
22 (W [i, [1,:[[[] +bb
del
0C
de! 02
[1 + +)² [ [h
OC
[i]
Oa
Therefore, scaling the parameter a' a'[i] by X-1 scales its gradient by A (A-1)-1.
This concludes the proof Lemma 2.1
B Additional Details
Synthetic task generation We construct a synthetic node classification task where we can examine the effects of allowing a layer to learn faster relative to others in controlled environment. The 5000 is generated by the Erdös-Rényi (ER) model with edge probability 0.001. Input node features h E Rd are sampled from standard normal distribution. This input graph G with no self-loops on nodes (i.e. v & N(v)) is fed to random Xavier [TO] initialized GAT network M2 of l- layers each of width d and final layer of width C, the number of classes of the classification task. This serves as the task's 's target network since its output embeddings are used to generate node labels in G. More precisely, for each node the node embedding output by M,. h, is effectively a function f of nodes in the k--hpp neighborhood (₂) of node v. f is represented by random GAT' network on the neighborhood aggregated representation of nodes hM to divide nodes into C' clusters. This clustering assigns node labels for our node classification task such that ye argee[c] (v c). We set d 10, C = 2 and & 5. This synthetic data generation ensures that the ground truth model can be represented by a GAT.
input graph with
series of non-linear transformations. Finally, we run K-means clustering
Synthetic task training Given the input graph G with .75/.25/.25 train/validationntess split, train a = layer GAT network with the same architecture as M2. but initialized with a looks-linear orthogonal structure which ensures that the network must learn the non-linear transformations of the target network. Furthermore, as noted in 28 the LL-orthogonal initialization facilitates trainability.
Synthetic task results discussion We analyze the relative gradients of some interesting cases of training a five layer GAT on the synthetic task under various settings. As also observed in Fig. 4 the network generalizes best on the given task when the first layer (initially) learns faster than the other layers in FigBle c). This implies that the task demands more focus on the first layer, and thus enforcing all layers to learn (relatively) at the same rate in Fig Fig((()) a) is not ideal. However, interestingly, larger learning rate in 8 b) enhances the network's 's ability to break free from its balanced state during training to have greater relative gradients the first layer and decrease relative gradients in the latter layers with the last layer receiving the lowest. This also aligns with the pattern observed in Fig. However, due to the dynamic rebalancing every 10 epochs, the network's 's favorable dynamic still disrupted. This results in smaller increase in test accuracy from Fig 8(a) as compared to the
16

---

2-
0.8
Cross Entroyy Loss
0.8
Cross Entropy Loss
06
0.4
02
882
870
1164
1358
1552
1746
1940
189
216
270
Train Test
Train
Epoch
Test
(a) DR of rel. gradients; Ir=0.001, (94.08 1847)
(b) DR of rel. gradients; 1r=0.01, (95.28@) 173)
08
0.8
Croes Entryy Loss
Croes Entropy Loos
06
0
-2
278
322
368
414
460
136
272
816
962
1080
1224
1380
Train Test
Test
Epoch
(c) Init. scaling to focus first layer; (97.04 07.00447 367)
(d) Init. scaling focus last layer; (94.0@) 1260)
Figure 8: Evolution relative gradients norms in a five-layer GAT network trained on synthetic data under varying settings. Test accuracy (%) @ epoch of maximum validation accuracy is reported in parentheses. The colored heatmap in the background displays log1o of relative gradient /2 12-non 2-norm for each layer (left axis) during training. Darker regions correspond to higher relative gradients. The training curves in the foreground show the train and test loss (right axis) for the epoch.
most favorable dynamics in Fig c( for the task. Nevertheless, this supports our hypothesis that larger learning with dynamic rescaling potentially enables better generalization, as shown in Fig. for real-world data. In comparison, initially restricting learning to only the last layer in Fig is the worst case and comparable to enforcing all layers (including the last) learn simultaneously. Although the test accuracy is high (>> 90%) in all cases, we focus on the relative difference in test accuracy under varying conditions, which are small, yet, significant as the goal of the task is make the effects of different training conditions easily observable in isolation.
Ablation of gradient clipping As evident from Table employing gradient clipping alone does not improve the performance in the absence of dynamic rescaling.
Table 2: Results of training a 5-layer GAT network with and without gradient clipping (GC) in the case of dynamic rescaling. The mean +.55% CI test metric the epoch of the best validation metric across 10 splits reported using the best learning rate from {0.01, 0.001, 0.005). The evaluation metric is accuracy for empire and amazon-ratings, and ROC AUC for the remaining datasets.
roman-empire
amazon-ratings
tolokers
questions
minesweeper
w/o DR
0.4978
0.4545 +
0.6493 +
0.5829 +
0.5057 +
w/o GC
0.0209
0.0043
0.008
0.0172
0.0058
w/o DR
0.3711
0.4535 +
0.6441 +
0.5745 +
0.5019 pm
w/ GC
0.0357
0.0025
0.0116
0.0135
0.0092
Limitations
We elaborate on the current limitations of this work as follows:
The primary limitation of dynamically rescaling model that we require the rescale invariance of the model architecture (if it exhibits one), which may vary widely across different GNN architectures. The dynamic rescaling proposed in the paper applies to GCNs and GATs, two standard baseline GNN models commonly used to study and develop
17

---

insights into the training dynamics of GNNs. While in theory, this is prerequisite to dynamically rescale the network during training without altering function potentially improve training dynamics, we conduct an experiment where this prerequisite is not fulfilled by simply replacing the ReL.U activation with the Tanh activation that is positively homogeneous. Consequently, the network no longer exhibits rescale invariance. However, we find that in practice, the advantages of dynamic rescaling to train in balance can still be observed in terms of better generalization. For example, on the roman-empire dataset, using Tanh activation in a layer GAT trained on roman-empire achieves an average test accuracy (over 3 runs) of 58.14 14 + 4.64% and 30.98 = 2.32% with and without rebalancing, respectively. While this trend aligns with that observed using ReLU, the training may be more noisy each rescaling during training is not loss invariant.
From implementation perspective, directly manipulating model parameter values and gradients during rescaling can result numerical instability issues that we currently regulate using gradient clipping, a commonly used practice in machine learning. Nevertheless, more principled approach to tackle this problem could be beneficial.
Rebalancing repeatedly and frequently during training may incur computational overhead in practice. However, it may be offset by the increased training speed, requiring fewer epochs overall. From a time complexity perspective, rebalancing only adds constant factor of operations in each epoch determined by the number of iterations in one rebalancing step. In practice, we find that only a few iterations (< 10) are necessary to balance the network approximately enough to gain better generalization and/or training speed.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:[Yes)
Justification:
The contributions listed in the introduction can be mapped the section ad subsection headings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We mention that we do not experimentally explore all the conceptual ideas discussed the paper.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
19

---

Answer: [Yes]
Justification: The paper has one theoretical result, for which the complete proof with references is provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We report the setup of our experiments, reference benchmark datasets used, and outline the method of synthetic data generation in detail.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
20

---

reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code setup for the experiments provided as supplementary material and will be made publicly available once anonymity restrictions are lifted.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
Cotee//ddeeeebmmmcia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:[Yed Yes]
Justification: All details are included in the Experiments section and complete executable code is provided.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The 95% confidence interval over 10 runs reported for the main experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
21

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details provided in experiments section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA].
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
22

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Reference to the paper that introduced the real-world datasets used in our experiments provided.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
23

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:[NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "TGGeOwfz26", "reviewer": "Reviewer_2KpP", "summary": "The paper studies the various phenomena in GNN training that prevent the usage of large learning rates for faster convergence and better generalization. Based on the theoretical foundation that large learning rates can be stably used only when the rescale symmetry of the loss function is satisfied, the authors propose a way to systematically perform a dynamic rescaling on the network parameters with relative gradient norms to maintain the rescale symmetry throughout training. This allows the usage of large learning rates leading to better convergence and generalization. Experiments reveal that their proposed approach for dynamic rescaling leads to consistently better empirical generalization across multiple benchmarks, as well as interesting phenomena akin to grokking arising from the arbitrary controllability of the order in which layers are learnt as a consequence of their proposed rescaling mechanism.", "questions": "Please refer to the Weaknesses section.", "limitations": "The paper can benefit from a more elaborate discussion of the limitations, which are not very clear from the current version.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a novel insight into the problem of convergence of GNNs trained with large learning rates. Although large learning rates are known to provide better generalization, convergence while using large learning rates still remains a challenge, hindering its adoption in practice. The authors interestingly observe that certain factors that get introduced while training neural networks in practice, such as momentum, weight decay, batch stochasticity break the rescale symmetry of the loss function. In fact, better generalization from large learning rates can only be achieved under rescale symmetry. However, the breaking of the rescale symmetry is the core factor hindering the usage of large learning rates while still achieving stable training dynamics.\n\n2. To remedy the above problem, the authors introduce a dynamic rescaling mechanism, where they periodically rescale the network parameters with relative gradient norms to explicitly maintain rescale symmetry and consequently, the corresponding conservation law throughout training, thereby allowing the usage of large learning rates, while maintaining trainability and achieving generalization.\n\n3. The dynamics rescaling mechanism allow the authors to arbitrarily determine the order in which network layers are trained, which provides greater flexibility while designing learning schemes, especially for OOD generalization tasks.\n\n4. The experiments show that dynamic rescaling of network parameters with the proposed relative gradient norms achieves consistently improved generalization on multiple benchmarks.\n\n5. The authors present interesting experiments with their methodology that illustrate grokking like phenomenon while explicitly controlling layer learning order. Interestingly, they observe that grokking might also be induced by the learning of other intermediate layers, and not just the last.", "weaknesses": "1. Line 8, Abstract is incomplete: \"on relative gradients, that promotes faster and better.\" I suppose the authors meant convergence.\n\n2. What is exactly meant by \"larger learning rates tend to induce improved generalization but make the training dynamics less robust\"? More specifically, what do the authors mean by \"robustness\" in the context of training dynamics? Throughout the paper, the authors talk about improvement in training speed and generalization, but none of them necessarily imply \"robustness of the training dynamics\".\n\n3. Enforcing Lipschitz continuity (to prevent vanishing / exploding gradients) through gradient clipping introduces a sharp penalty leading to instability in the optimization process, whereas gradient penalty is known to produce smoother training dynamics, at least in the context of training WGAN critics [a]. This result is not restricted to WGANs, since gradient clipping, by definition introduces a sharp cut-off, while gradient penalty is inherently a smoother regularization constraint. Then, why do the authors go for gradient clipping? A discussion is necessary.\n\n4. The paper can benefit from a more elaborate discussion of the limitations, which are not very clear from the current version.\n\n5. Minor: Line 110: Comma at the beginning of the sentence.\n\nReferences:\n\n[a] Gulrajani et al, \"Improved Training of Wasserstein GANs\", NeurIPS 2017.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and address their concerns and questions as follows.\n\n1. Thanks for pointing out this typo. We meant training. We would like to further clarify the statement in the abstract regarding larger learning rates and robustness. Here, by robustness, we mean general training stability. Larger learning rates in neural networks are generally associated with better generalization but more unstable training [1]. We state this already observed behavior from literature. Our finding is that dynamic rebalancing of relative gradients further improves the generalization obtained by using larger learning rates. For reference, please see Figure 1 in the attached PDF in the global response. As the reviewer has rightly pointed out, we do not discuss the robustness (to noise, or adversaries) aspect of dynamic rescaling in this work, and using the term training stability would be a better alternative.\n\n2. We opted for gradient clipping, as it is a widely used standard practice, generally used across various domains of machine learning to avoid exploding gradients. However, upon the reviewer's valuable suggestion, we have also tried to replace gradient clipping with an L2-norm gradient penalty. As evident from Table 1 below, we find that the performance is adversely impacted, and gradient clipping (even in combination with gradient regularization) performs better. Nevertheless, it merits a discussion in the related work that we will include.\n\nTable 1: Mean test accuracy (\\%) over 10 runs for \\texttt{roman-empire} dataset on GAT using gradient l2-norm penalty (of strength $\\lambda$) with and without gradient clipping (GC) while rebalancing w.r.t relative gradients every epoch during training. Rebalancing only in combination with gradient clipping without using gradient penalty results in $59.62\\pm2.21\\%$ mean test accuracy.\n| $\\lambda$ |           0.01              |          0.1               |            2                 |           10                 |\n| :----------: | :--------------------: | :--------------------: | :-------------------: | :--------------------: |\n| w/o GC       | $50.43 \\pm 3.07$ | $52.76 \\pm 2.75$ | $51.4 \\pm 3.42$ | $52.44 \\pm 2.43$ |\n| with GC      | $\\mathbf{59.28 \\pm 2.55}$ | $\\mathbf{59.82 \\pm 2.64}$ | $\\mathbf{58.4 \\pm 1.85}$ | $\\mathbf{59.7 \\pm 2.35}$ |\n\n3. We include an elaborate discussion of the limitations of this work in the global response. \n\n[1] Lobacheva et. al. Large Learning Rates Improve Generalization: But How Large Are We Talking About? (2023)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "v0sYpxR9aw", "reviewer": "Reviewer_TkPU", "summary": "This paper investigates the use of dynamic rescaling to train Graph Attention Networks (GATs), a type of Graph Neural Network (GNN). The method aims to enhance the trainability and generalization of GATs by balancing network parameters and gradients during training. GATs' rescale invariance property enables them to adjust the rate at which network layers learn, resulting in better training dynamics optimization. The proposed dynamic rescaling technique balances network parameters using their relative gradients, resulting in faster and more efficient training. Higher learning rates can be achieved by dynamically rebalancing the network during training, resulting in better generalization on real-world data.\nThe document emphasizes the importance of selecting suitable rescaling criteria and procedures based on network parameters and gradients. It emphasizes the potential benefits of dynamic rescaling in controlling the order of learning among network layers, which can affect the overall performance of GATs. Experimental studies are being conducted to validate dynamic rescaling's effectiveness in improving GAT training speed, generalization, and robustness.\nFurthermore, the paper discusses related research on training dynamics and generalization in GNNs, demonstrating how dynamic rescaling can outperform traditional methods such as initialization, normalization, and regularization. The findings indicate that leveraging the rescale invariance of GNNs, such as GATs, can result in sharper optima, better generalization, and greater robustness during training. The paper concludes by discussing grokking-like phenomena associated with learning patterns in GNNs and highlighting promising future research directions in the use of dynamic rescaling for more practical benefits in GNN optimization.", "questions": "1. Line 5: Does a larger learning rate help generalization? If larger learning rates tend to induce improved generalization but disrupt the \"balanced state,\" what is the effect of the balanced state? Does this mean it is not beneficial for generalization?\n\n2. Line 22: To improve the readability of this work, I suggest the authors assume readers generally do not have extensive knowledge about the \"conservation law\" and \"conservation of a relationship between network parameters and gradients.\" Could you please use one more sentence in paragraph 2 to clarify these concepts?\n\n3. Line 31: Please correct the citation format.\n\n4. Why do we need balance? Is it for the trainability of GNNs? If so, I suggest you accurately describe what trainability means. Otherwise, in the next paragraph, starting with “The common criterion used to balance networks,” it seems like you already assume readers are fully convinced of the need for balance.", "limitations": "I would like to encourage the authors to make a separate Limitations section in their paper.", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The paper focuses on a compelling topic and observes several notable phenomena, including: 1. A balanced state during training combined with larger learning rates results in better generalization compared to when either component is used alone. 2. The Grokking-like phenomena.", "weaknesses": "In general, I appreciate the topic of this work, which focuses on the training dynamics of graph neural networks and aims to improve them. However, the following points indicate that it is not yet ready for publication:\n\n1. Presentation: I have included questions 1 to 4 in the next section. While I am aware of the answers to these questions, I suggest the authors use them as references to clarify the corresponding parts of the draft, thereby improving its readability.\n\n2. Analysis: More interesting results might be obtained by examining the connection between learning, grokking, and the balanced state. I recommend that the authors conduct an in-depth analysis of this connection.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and address their concerns as follows. Note that all references to figures are in the PDF attached to the global response unless stated otherwise.\n\n1.  Firstly, to improve the readability and comprehension of the paper, we answer the questions set forth by the reviewer and provide a holistic background and discussion of conservation laws, trainability, and large learning rates. We will include this in the updated version of the paper.\n\n* **Conservation law and balance:** It is known that, for traditional deep feed-forward networks (DNNs) and convolutional neural networks (CNNs) with homogenous activation functions such as ReLUs, the difference between the squared l2-norms of incoming and outgoing parameters to a neuron stays constant (and is thus conserved) under gradient flow, i.e. gradient descent with an infinitesimal learning rate. Such conservation laws arise for neural networks exhibiting a rescale invariance property [7]. When this conserved quantity is (nearly) zero, the network is said to be in a balanced state. The concept of balance at the neuron level is rooted in an understanding of the training dynamics of deep neural networks where it is generally assumed that norm balance induces conditions for successful training. We have discussed more implications of the conservation law for DNNs and CNNs in the related work. In the context of GNNs, the insight regarding norm balance was presented by [6], which derives the conservation law for GATs and demonstrates how a network initialization fulfilling this balance property, in practice, enhances the trainability that we discuss next.\n    \n* **Trainability:** Trainability can be defined as the ability of model parameters to change during training from their initial state to an optimal position. In general, the (relative) gradient norm of parameters serves as a good proxy for this measure, as it quantifies the relative update of parameters in each training epoch, and has been used in literature to study trainability [6,8]. As mentioned above, a balanced initial state enhances trainability [6]. In theory, under gradient flow due to the conservation law, this balanced state should hold throughout training. However, in practice, using finite (and larger) learning rates disturbs this balance. We propose a novel idea of dynamic rescaling of the network to rebalance the network during training. Additionally, we show that norm balance may not be the ideal criterion for better trainability and propose balancing the network w.r.t. relative gradients of parameters instead. Intuitively, this allows all parameters an equal opportunity to change during training and thus improves trainability. \n        \n* **Larger learning rates:** As we state in the paper, several works report empirical evidence that a larger learning rate biases toward flatter minima which have been linked with better generalization [2-5]. The same idea is endorsed by [1], which also states that only a narrow range of these `large enough' learning rates can produce optimal results. However, as discussed above, larger learning rates cause the network state to become more imbalanced during training and thus may impede learning. We show that rebalancing the network during training by dynamic rescaling further improves the generalization brought about by larger learning rates (see Fig. 1).\n\n2. Secondly, as suggested by the reviewer, we elaborate on our insights relating learning, balanced state, and grokking-like phenomena in further depth with the support of additional experiments. Due to space limitations and overlap with questions from other reviewers, we present these further insights in the global response. We request the reviewer to please also consider the global response a part of our response to them.\n\n3. Finally, we include an elaborate discussion of the limitations of this work also in the global response.\n\n[1] Lobacheva et al. Large Learning Rates Improve Generalization: But How Large Are We Talking About? (2023).\n\n[2] Zhao et al. Penalizing gradient norm for efficiently improving generalization in deep learning (2022).\n\n[3] Lewkowycz et al. The large learning rate phase of deep learning: the catapult mechanism. (2020).\n\n[4] Seong et al. Towards Flatter Loss Surface via Nonmonotonic Learning Rate Scheduling. (2018).\n[5] Dinh et al. Sharp minima can generalize for deep nets. (2017).\n\n[6] Mustafa et al. Are GATs Out of Balance? (2023).\n\n[7] Kunin et al. Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics. (2021)\n\n[8] Jaiswal et al. Old can be Gold: Better Gradient Flow can Make Vanilla-GCNs Great Again. (2022).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uN2y0fRUw3", "reviewer": "Reviewer_uEaK", "summary": "This manuscript attempts to study the rescale invariance of GNNs and suggests to dynamically rescale the network using relative gradient, i.e. g/theta, where g is gradient of theta. Several experiments are conducted to to show how the dynamic rescaling using relative gradient norms affects the learning speed and generalization.", "questions": "See the above", "limitations": "See the above", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- this paper is the first propose to use relative gradient norm to study the invariance of GNNs. \n- several interesting experimental phehomenon have been identified, such as the training speed accross the different layers for the grokking", "weaknesses": "- From algorithimic side, my main concern is that the proposed dynamical scaling does not induce a consistent improvement over the tasks the authors have considered. This brings the skeptism whether this strategy works or not. \n- If we did not treat this paper as algorithm paper, instead, both the theoretical and empirical analysis, as well as potential insights, should be in more depth. \nSo, I think this paper initializes a good idea and point to start with, but less developed. It is still not mature yet to be published.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the novelty of our work and their interest in our insights. We address their concerns as follows.\n\nWe would like to reiterate that our goal is not to propose a one-size-fits-all solution or achieve state-of-the-art performance. Rather, this work aims to conduct an exploratory study of a novel idea and gain insights into GNN learning dynamics and Grokking-like phenomena. In doing so, we highlight our main contributions as follows.\n\n1. We prove that if GNNs exhibit a rescale invariance (like GCNs or GATs), we can manipulate gradient norms while leaving the GNN and the loss invariant and thus influence the learning dynamics. \n2. Utilizing this insight, we propose to dynamically rescale GNNs based on a novel criterion, i.e., relative gradient norms, to foster trainability and enable potentially faster training with larger learning rates.\n3. We derive a synthetic learning task that supports our insights into learning dynamics and provides novel insights into Grokking.\n4. Our methodological contributions lead to a couple of novel insights, which are discussed next.\n\nAs requested by the reviewer, we further strengthen our insights and deepen our empirical analysis with the support of additional experiments, as follows. \n\n1. We discover trends in optimal training dynamics regarding homophily and heterophily, which aligns with a recent analysis of learning dynamics in general GNNs. In summary, training GNNs in and out of balance may be more beneficial for heterophilic and homophilic tasks, respectively. \n\n2. We construct a grokking-like phenomenon on real-world data, which we can induce by a combination of different ways to influence gradient dynamics. Interestingly, a worsening training accuracy can be associated with an improving test accuracy, suggesting that training GNNs in balance can potentially promote better feature learning where required (such as in heterophilic tasks) and mitigate over-fitting.\n\n3. While we show in the main paper that dynamic rescaling allows better generalization with larger learning rates, we extend the experiments to cover a wider range of learning rates and observe the trend to be in line with [1].\n\nDue to space limitations and overlap with questions from other reviewers, we have presented a detailed discussion of these insights in the global response. We request the reviewer to please also consider the global response a part of our response to them.\n\n[1] Lobacheva et al. Large Learning Rates Improve Generalization: But How Large Are We Talking About? (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ntdwzEoHV1", "reviewer": "Reviewer_GhMV", "summary": "This paper proposes the dynamics rescaling approach to improve the trainability of GAT. This paper is motivated by the rescale invariance property of GAT, i.e. the function is invariant w.r.t. some scaling terms on a neuron’s weights, and the corresponding conservation law. To balance the network throughout training, this paper proposes a criterion (inspired by the conservation law) to dynamically rescale weights. This approach enables the network to benefit from larger learning rates while circumventing the imbalance problem, and thus improves generalization. The paper also discusses other use cases of the approach, such as enabling layer-wise training to accelerate grokking.", "questions": "Are there any experiments on steering the model to learn certain layers first on real-world datasets to test if this idea of accelerating grokking is actually useful in practice? Particularly, I wonder if developing methods to induce feature learning (or grokking) is really effective for improving the peak performance beyond synthetic datasets, as there are also results suggesting proper way of learning in the kernel regime (or lazy learning regime) is sufficient for good generalization [2]. It would be interesting to see discussions in this direction.\n\n[2] How Graph Neural Networks Learn: Lessons from Training Dynamics. in ICML 2024", "limitations": "While the analytical framework might be generally useful, but the method seems only applicable to those GNNs that satisfy the rescale invariance property.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper studies training dynamics of GNNs, which is an important yet underexplored topic in graph learning. The proposed method is well-motivated from a symmetry perspective, and can effectively improve the trainability and generalizability of GAT by forcing balance. The empirical observations from controlling the order in which layers are trained are also intriguing and might be of independent interest in other areas.", "weaknesses": "1. As preliminary is integrated into the approach, it is hard to parse which part is novel. It would be beneficial to more clearly compare the contributions of this paper and previous papers on training dynamics, especially [1].\n2. The experiment in sec.4.1 is relatively weak as the considered five datasets have similar properties (heterophily). It is also questionable if the proposed method can work on GNNs that do not satisfy the invariance property.\n3. It is not entirely clear how the proposed criterion is related with the conservation law associated with GAT.\n\n[1] Are GATS out of balance? in NeurIPS 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the importance of our work and their constructive feedback. We address their concern as follows. Note that all references to figures are in the PDF attached to the global response unless stated otherwise.\n\n1. We gladly highlight the novelty of our work in comparison to past work. To the best of our knowledge, we are the first to propose dynamic rescaling as a tool to influence the training dynamics of GNNs that can be used in various ways. For example, ensuring all layers have an equal opportunity to drive change in their parameters by training in balance or allowing a particular layer with higher relative gradients to train `more' than others. Particularly concerning [1], our work differs in two ways: \n\n* While [1] is the first work to evaluate the role of network balance in the generalization and training speed of GATs, their method is restricted to balance at initialization. In theory, this balance would be maintained throughout training under gradient flow (i.e. gradient descent with infinitesimal learning rates). However, in practice, as a result of training with finite learning rates (potentially in a range of larger values associated with better generalization - see Fig. 1), the network balance established by the proposed initialization of [1] is disrupted. We show the degree of imbalance in the network during training before and after the rebalancing of the network every 10 epochs in Fig 2. The balanced initialization of [1] thus cannot ensure that the network will stay in balance throughout training. For this, we provide a method using dynamic rescaling and demonstrate the benefits of training GNNs in balance, specifically for heterophilic tasks. \n\n* We challenge the balancing criterion based on norms of parameter values (used in [1]) and instead propose balancing w.r.t. the norm of relative gradients of parameters. Table 1 in the paper shows that relative gradients are often the better criterion for improved training resulting in better generalization. \n\n2. The reviewer has rightly pointed out that a primary limitation of our approach is that we need to know the rescale invariance of the considered GNN architecture. In theory, this is a prerequisite to dynamically rescale the network during training without altering its function to potentially improve training dynamics. However, we conduct an experiment where this prerequisite is not fulfilled by simply replacing the ReLU activation with the Tanh activation that is not positively homogeneous. Consequently, the network no longer exhibits rescale invariance. However, we find that in practice, the advantages of dynamic rescaling to train in balance can still be observed in terms of better generalization. For example, on the roman-empire dataset, using Tanh activation in a 5 layer GAT trained on roman-empire achieves an average test accuracy (over 3 runs) of $58.14\\pm4.64\\%$ and $30.98\\pm2.32\\%$  with and without rebalancing, respectively. While this trend aligns with that observed using ReLU, the training may be more noisy as each rescaling during training is not loss invariant. \n\n3. The conservation law of GAT describes a relationship between the norms of parameters associated with each neuron that the network adheres to under gradient flow. This law is induced by the rescale invariance of the network parameters. Thus, we can rescale the network parameters to fulfill any criterion involving parameter values or norms. One such criterion is the specific sum of parameter l2-norms defined by the conservation law during training. However, as we prove in Lemma 1 in the paper, the parameter gradients also depend on the scaling factor of the parameter values, and thus can also be controlled directly or indirectly by the scaling factors. This allows us to opt for a criterion that depends both on the parameter values and gradients (such as relative gradients) that can be fulfilled by rescaling without changing the network function. However, unlike the criterion based on parameter L2-norms, the proposed criterion of relative gradients is not naturally fulfilled during training and thus requires us to rebalance the network frequently. \n\n4. As the reviewer has rightly noted, the main advantage of training GATs in balance is observed for heterophilic tasks, while for homophilic tasks such as on the Cora and Citeseer datasets, there is no significant impact of training in balance w.r.t. relative gradients which perform similarly to standard training. Interestingly, we can relate this observation to the reviewer's question regarding results in [2], which suggests that a proper way of learning in the kernel regime (or lazy learning regime) is sufficient for good generalization in the case of homophilic datasets.  A detailed discussion of this insight with empirical verification, as well as an example of inducing a grokking-like phenomenon on a real-world dataset, is presented in the global response. Thus, we request the reviewer to read the global response and consider it a part of our response to them.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their encouraging comments and constructive suggestions that we have incorporated to improve the paper. We jointly answer some key common questions of the reviewers here in two parts regarding i) additional insights and ii) limitations. 

**Further insights and additional experiments:**

1. We discover trends regarding homophily and heterophily in the training dynamics of GATs. 

* *Background:* In the context of GNNs, the optimal performance achieved by a model is, to a large extent, dependent on how well the inductive bias of the model architecture aligns with the task and its underlying graph structure. For example, it is widely known that general GNNs, without specially-introduced architectural elements, such as GCN perform better on homophilic than on heterophilic tasks.

* *Observation:* We find that training in balance is generally more effective for heterophilic tasks (as we report in the main paper) than homophilic tasks (such as Cora and Citeseer) where we observe a similar performance with no significant advantage of training in balance. On the contrary, allowing learning focused on the first layer may be beneficial for homophilic tasks, although the effect is not similar in magnitude to that of balancedness for heterophilic tasks. For example, see the test accuracy reported for different influences on gradient dynamics in Fig. 3. 
 
* *Explanation:* Intuitively, we hypothesize that homophilic tasks rely more on the neighborhood aggregation functionality of GNNs rather than feature learning. In this case, an aggregation over a random transformation of similar features may still be sufficient for good generalization. This is in line with a recent analysis of training dynamics [1] which shows that on homophilic graphs, alignment of the underlying structure with the optimal kernel matrix allows parameter-free methods similar to label propagation to perform at par with GNNs. However, this adversely affects generalization on heterophilic tasks where the graph structure does not align with the optimal kernel. In other words, the neighborhood is not very relevant for a node's label in heterophilic settings and thus the node relies more on feature learning rather than neighborhood aggregation. This is also supported by results showing that embedding additional MLP layers in the network significantly improves the performance of basic GNNs such as GATs on these heterophilic tasks [2]. Thus, we conclude that training in balance to learn better feature transformations in all layers (and potentially neighbors farther away in deeper models) is more effective in heterophilic cases. From the perspective of neighborhood aggregation,  dynamic rescaling allows the model to achieve higher values of $\alpha_{ii}$ (i.e. the weight assigned to the node itself in GATs, see Fig. 3) by potentially allowing better training of the attention parameters whereas an opposite but not equally large effect can be observed for homophilic settings with a larger focus on the neighborhood (and thus smaller $\alpha_{ii}$) may be better. For sufficiently shallow models, however, the difference in regular and layer-focused training could be negligible for homophilic tasks. 

2. We induce grokking-related behavior on a real-world dataset (see Fig. 4) in two steps. Firstly, we allow only the last (or second to last) layer to learn which allows the training accuracy to increase continually while the test accuracy saturates or begins to drop. At this point, we rescale the network to bring all layers in balance w.r.t. relative gradients, following which,  the test accuracy immediately begins to improve more rapidly accompanied by a drop in training accuracy. This can be interpreted as the network `learning' more effectively rather than overfitting to the training data. While this is different from grokking where the training accuracy would generally not drop, it is independently an interesting observation on a real-world dataset. 

3. We show that rebalancing the network during training by dynamic rescaling further improves the generalization brought about by larger learning rates (see Fig. 1).

**Limitations:** As suggested by multiple reviewers, we discuss the limitations elaborately in a dedicated section that we will include in the paper as follows. 

1. The primary limitation of dynamically rescaling a model is that we require the rescale invariance of the model architecture (if it exhibits one), which may vary widely across different GNN architectures. The dynamic rescaling proposed in the paper applies to GCNs and GATs, two standard baseline GNN models commonly used to study and develop insights into the training dynamics of GNNs.

2. From an implementation perspective, directly manipulating model parameter values and gradients during rescaling can result in numerical instability issues that we currently regulate using gradient clipping, a commonly used practice in machine learning. Nevertheless, a more principled approach to tackling this problem could be beneficial. 

3. Rebalancing repeatedly and frequently during training may incur computational overhead in practice. However, it may be offset by the increased training speed, requiring fewer epochs overall. From a time complexity perspective, rebalancing only adds a constant factor of operations in each epoch, determined by the number of iterations in one rebalancing step. In practice, we find that only a few iterations (<10) are necessary to balance the network approximately enough to gain better generalization and/or training speed. 

[1] Yang et al. How Graph Neural Networks Learn: Lessons from Training Dynamics. (2024).

[2] Platonov et al. A critical look at the evaluation of GNNs under heterophily: Are we really making progress? (2023).

Meta-review: This work studies dynamic rescaling during GNN training to improve learning speed and generalization. The reviewers overall appreciated the novel use of relative gradients and interesting experiments conducted, along with widely acknowledged useful and relevant empirical results, though with some concerns about significance. There was a robust discussion during rebuttal period and additional clarifications and content to overcome some weaknesses should be made by the authors for a revised manuscript.

Additional context after discussion with the SAC: Reviewer TkPU wanted more effective clarification of the motivation and details on learning rate and grokking. The authors included the additional detail in two follow up responses to address their concerns and improved the paper. Reviewer 2KpP following a discussion, was satisfied with the rebuttal assuming the authors included data from the discussion in the final paper. Overall, while this was borderline, it fell on the side of acceptance.
Decision: Accept (poster)