Dynamic Rescaling for Training GNNs

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.

1. 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:

• 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.

• 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.

2. 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.

3. 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.

4. 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.

We thank the reviewer for acknowledging the novelty of our work and their interest in our insights. We address their concerns as follows.

We 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.

1. 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.
2. 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.
3. We derive a synthetic learning task that supports our insights into learning dynamics and provides novel insights into Grokking.
4. Our methodological contributions lead to a couple of novel insights, which are discussed next.

As requested by the reviewer, we further strengthen our insights and deepen our empirical analysis with the support of additional experiments, as follows.

1. 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.

2. 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.

3. 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].

Due 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.

[1] Lobacheva et al. Large Learning Rates Improve Generalization: But How Large Are We Talking About? (2023).

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.

1. 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.

• 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.

• 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.

• 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).

2. 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.

3. Finally, we include an elaborate discussion of the limitations of this work also in the global response.

[1] Lobacheva et al. Large Learning Rates Improve Generalization: But How Large Are We Talking About? (2023).

[2] Zhao et al. Penalizing gradient norm for efficiently improving generalization in deep learning (2022).

[3] Lewkowycz et al. The large learning rate phase of deep learning: the catapult mechanism. (2020).

[4] Seong et al. Towards Flatter Loss Surface via Nonmonotonic Learning Rate Scheduling. (2018). [5] Dinh et al. Sharp minima can generalize for deep nets. (2017).

[6] Mustafa et al. Are GATs Out of Balance? (2023).

[7] Kunin et al. Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics. (2021)

[8] Jaiswal et al. Old can be Gold: Better Gradient Flow can Make Vanilla-GCNs Great Again. (2022).

We thank the reviewer for the constructive feedback and address their concerns and questions as follows.

1. 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.

2. 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.

Table 1: Mean test accuracy (%) over 10 runs for 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.

3. We include an elaborate discussion of the limitations of this work in the global response.

[1] Lobacheva et. al. Large Learning Rates Improve Generalization: But How Large Are We Talking About? (2023)

I thank the authors for their prompt response on my concerns. I especially appreciate them redoing their experiments by performing the rebalancing every 10 epochs, which gives a better idea of the improvement that comes from rebalancing vs the choice of the gradient stabilizer, i.e., gradient clipping (GC) and / or gradient penalty (GP).

However, my concern that the authors do not adequately establish a relationship between their rescale mechanism and the choice of gradient stabilization approach (for which they have found GC to be optimal) is still concerning. The two must be very intimately intertwined since they both directly affect the scales of the gradient updates to the weights, which is then main premise that this work is meant to study. Although the authors' claim that "a deeper study of gradient regularization is beyond the scope of this paper", developing a more complete understanding (beyond it being presented as just a hyperparameter) of how the choice of the gradient stabilization mechanism, be it GC or GP, affects the proposed dynamic rebalancing process, feels imperative.

I thank the authors for their response and the set of experiments that they have conducted. Although some of my concerns have been clarified, I will need some additional deliberation on my end to decide upon my final score.

Meanwhile, I still have one major concern about the experiment on the prevention of exploding gradients presented in Table 1 of the rebuttal. In the (second) "with GC" row, if I understand correctly, the authors apply both gradient penalty and gradient clipping. Although with only gradient penalty, i.e., the first row, there is some drop in performance from the results reported in Table 1 of the main paper, of around 2% from 54% to 52%), applying both gradient penalty and clipping in conjunction leads to a significant boost of around 5% (from 54% to 59%) relative to the authors' original results. Is my understanding correct? If so, then it may be unfair to state that "the performance is adversely impacted" with the use of gradient penalty, since in conjunction with GC, it gives a significant boost.

Since the approach involves dynamic rescaling wrt relative gradient norms, I believe developing an understanding of the interactions between the rescale operations and the regularization is key for the soundness of the central claims, something that is lacking from the work and the rebuttal in their current state.