Original: On the Interpretability of Regularisation for Neural Networks Through Model Gradient Similarity
Revised: On the Interpretability of Regularisation for Neural Networks Through Model Gradient Similarity

Original: Anonymous Author(s)

Original: Af ﬁ liation

Original: Address

Original: Abstract
Revised: Abstract

Original: Most complex machine learning and modelling techniques are prone to over- ﬁ ttingand may subsequently generalise poorly to future data.
Revised: Most complex machine learning and modelling techniques are prone to over-ﬁtting and may subsequently generalise poorly to future data.

Original: Arti ﬁ cial neural networksare no different in this regard and, despite having a level of implicit regularisationwhen trained with gradient descent, often require the aid of explicit regularisers.
Revised: Artiﬁcial neural networks are no different in this regard and,despite having a level of implicit regularisation when trained with gradient descent, often require theaid of explicit regularisers.

Original: We introduce a new framework, Model Gradient Similarity (MGS), that (1) servesas a metric of regularisation, which can be used to monitor neural network training,(2) adds insight into how explicit regularisers, while derived from widely differentprinciples, operate via the same mechanism underneath by increasing MGS, and(3) provides the basis for a new regularisation scheme which exhibits excellentperformance, especially in challenging settings such as high levels of label noise orlimited sample sizes
Revised: We introduce a new framework, Model Gradient Similarity (MGS), that(1) serves as a metric of regularisation, which can be used to monitor neural network training, (2) addsinsight into how explicit regularisers, while derived from widely different principles, operate via thesame mechanism underneath by increasing MGS, and (3) provides the basis for a new regularisationscheme which exhibits excellent performance, especially in challenging settings such as high levelsof label noise or limited sample sizes

Original:  Introduction
Revised:  Introduction

Original: Since the inception of modern neural network architecture and training, many efforts have been madeto understand their generalisation properties.
Revised: Since the inception of modern neural network architecture and training, many efforts have been madeto understand their generalisation properties.

Original: When neural networks are trained with Gradient Descent(GD) this induces implicit regularisation in the network, causing it to attain surprising generalisationperformance with no extra intervention.
Revised: When neural networks are trained with Gradient Descent(GD) this induces implicit regularisation in the network, causing it to attain surprising generalisationperformance with no extra intervention.

Original: Despite this, neural networks will in many instances over ﬁ tin the presence of noise and have been shown to have the capacity to completely memorise data[Zhang et al., 2016].
Revised: Despite this, neural networks will in many instances overﬁtin the presence of noise and have been shown to have the capacity to completely memorise data[Zhang et al., 2016].

Original: Therefore, alongside the desire to understand how neural networks are implicitlyregularised, explicit regularisation has been an ongoing pursuit to improve generalisation even further.
Revised: Therefore, alongside the desire to understand how neural networks are implicitlyregularised, explicit regularisation has been an ongoing pursuit to improve generalisation even further.

Original: A broad range of explicit regularisers now exist that attack the problem from many different angles.
Revised: A broad range of explicit regularisers now exist that attack the problem from many different angles.

Original: Below, we provide a short overview of some of the most commonly used methods.
Revised: Below, we provide a short overview of some of the most commonly used methods.

Original: Weight decay[Krogh and Hertz, 1991] places a L2-penalty on the parameters to encourage a minimum-normsolution.
Revised: Weight decay[Krogh and Hertz, 1991] places a L2-penalty on the parameters to encourage a minimum-normsolution.

Original: Normalisation schemes , that normalise the data as it ﬂ ows through the network, havebeen shown to act as explicit regularisers [Ioffe and Szegedy, 2015].
Revised: Normalisation schemes , that normalise the data as it ﬂows through the network, havebeen shown to act as explicit regularisers [Ioffe and Szegedy, 2015].

Original: Double back-propagationincorporates the gradients with respect to the loss itself as part of the loss, hence the name, as thegradient of the gradient will be evaluated in the complete back-propagation step [Drucker and Le Cun,1991].
Revised: Double back-propagationincorporates the gradients with respect to the loss itself as part of the loss, hence the name, as thegradient of the gradient will be evaluated in the complete back-propagation step [Drucker and Le Cun,1991].

Original: Although these methods were initially used as a part of a energy minimisation principle,they are now believed to help ﬁ nding " ﬂ at minima", which have been shown to produce solutionsthat generalise better [Zhao et al., 2022].
Revised: Although these methods were initially used as a part of a energy minimisation principle,they are now believed to help ﬁnding "ﬂat minima", which have been shown to produce solutionsthat generalise better [Zhao et al., 2022].

Original: Dropout turns off connections between neurons duringtraining and thus forces the network not to rely on any one given connection [Srivastava et al., 2014].
Revised: Dropout turns off connections between neurons duringtraining and thus forces the network not to rely on any one given connection [Srivastava et al., 2014].

Original: Another way to view the learning problem is to ignore the network itself and instead focus on theoptimisation scheme.
Revised: Another way to view the learning problem is to ignore the network itself and instead focus on theoptimisation scheme.

Original: Here, a multitude of different learning rate and optimisers exist that offerexplicit regularisation.
Revised: Here, a multitude of different learning rate and optimisers exist that offerexplicit regularisation.

Original: Some standout examples include Cyclical Learning Rate [Li and Yang, 2020]and the Adam optimisation algorithm
Revised: Some standout examples include Cyclical Learning Rate [Li and Yang, 2020]and the Adam optimisation algorithm

Original: [Kingma and Ba, 2015].
Revised: [Kingma and Ba, 2015].

Original: More recently, especially in worksrelated to the Neural Tangent Kernel (NTK)
Revised: More recently, especially in worksrelated to the Neural Tangent Kernel (NTK)

Original: [Jacot et al., 2018], focus has been put onto functionalregularisation.
Revised: [Jacot et al., 2018], focus has been put onto functionalregularisation.

Original: The core idea is to view the neural network architecture as being an opaque function approximator and instead regularise it as it is viewed from function space rather than the parameterspace.
Revised: The core idea is to view the neural network architecture as being an opaque functionapproximator and instead regularise it as it is viewed from function space rather than the parameter space.

Original: This is usually achieved via a form of kernel ridge regression or an approximation thereof[Bietti et al., 2019, Hoffman et al., 2019].
Revised: This is usually achieved via a form of kernel ridge regression or an approximation thereof[Bietti et al., 2019, Hoffman et al., 2019].

Original: Our contribution .
Revised: Our contribution .

Revised: The fact that so many different explicit regularisers exist shows that their precisemechanisms are not well understood.

Revised: Ultimately, this leads back to the insufﬁcient understanding ofwhat causes neural networks to generalise in the ﬁrst place.

Original: In this article, a new framework called Model Gradient Similarity (MGS) isintroduced.
Revised: In this article, a new framework called Model Gradient Similarity (MGS) is introduced.

Original: MGS builds on the idea that regularisation and, to a large extent, generalisation, is tightlyconnected to how the model evolves during training.
Revised: MGS builds on the idea that regularisation and, to alarge extent, generalisation, is tightly connected to how the model evolves during training.

Original: We show that MGS can be used to analyseneural network training, monitoring how model gradients evolve together.
Revised: We showthat MGS can be used to analyse neural network training, monitoring how model gradients evolvetogether.

Original: It is revealed that explicitregularisers, despite their fundamentally different construction, operate via the same underlyingmechanisms in terms of increasing model gradient similarity.
Revised: It is revealed that explicit regularisers, despite their fundamentally different construction,operate via the same underlying mechanisms in terms of increasing model gradient similarity.

Original: This insight lays the foundation for anew class of regularisers aimed at directly boosting model gradient similarity.
Revised: Thisinsight lays the foundation for a new class of regularisers aimed at directly boosting model gradientsimilarity.

Original: We derive two simple MGS regularisation techniques from ﬁ rst principles and compare performance to a wide range ofregularisation methods.
Revised: We derive two simple MGS regularisation techniques from ﬁrst principles and compareperformance to a wide range of regularisation methods.

Original: In a majority of the cases, the MGS regularisation achieves top performanceand exhibits robustness qualities.
Revised: In a majority of the cases, the MGSregularisation achieves top performance and exhibits robustness qualities.

Original: This suggests that the MGS framework opens the door for thefurther development of both novel regularisers and the improvement of existing ones.
Revised: This suggests that the MGSframework opens the door for the further development of both novel regularisers and the improvementof existing ones.

Original:  Model Gradient Similarity
Revised:  Model Gradient Similarity

Original: Let x ∈ X denote an input observation in a data batch and y ∈ Y the target.
Revised: Let x ∈ X denote an input observation in a data batch and y ∈ Y the target.

Original: We will train themodel f θ , parameterised by θ , using loss function L via gradient descent (GD).
Revised: We will train themodel f θ , parameterised by θ , using loss function L via gradient descent (GD).

Original: Thus, the standardsingle-observation batch update rule using GD at time i is θ i +

Revised: Thus, the standardsingle-observation batch update rule using GD at time i is θ i +

Original: =
Revised: =

Original: θ i − η ∇ θ L ( f ( x ) , y ) , where η isthe learning rate.
Revised: θ i − η ∇ θ L ( f ( x ) , y ) , where η isthe learning rate.

Original: When training a neural network, f θ shall represent the raw output from the networkbefore any additional normalisation is performed (e.g. softmax as in the case of classi ﬁ cation).
Revised: When training a neural network, f θ shall represent the raw output from the networkbefore any additional normalisation is performed (e.g. softmax as in the case of classiﬁcation).

Original: The loss gradients with respect to the parameters ∇ θ L ( loss-parameter gradient ) are thus themain contributing factor to the update of the model.
Revised: The loss gradients with respect to the parameters ∇ θ L ( loss-parameter gradient ) are thus themain contributing factor to the update of the model.

Original: In the general case for common losses, suchas mean-squared error or cross-entropy, we apply the chain rule to the loss-parameter gradient:∇ [Equation]
Revised: In the general case for common losses, suchas mean-squared error or cross-entropy, we apply the chain rule to the loss-parameter gradient: 61 ∇ θ L = ∇ θ f θ · ∇ f L .

Original: We denote the two components as: the loss gradients with respect to the model 60 output ∇ f L ( loss-model gradients ) and the model gradients with respect to the model parameters∇ θ f θ ( model-parameter gradients ).
Revised: We denote the two components as: the loss gradients with respect to the modeloutput ∇ f L ( loss-model gradients ) and the model gradients with respect to the model parameters 63 ∇ θ f θ ( model-parameter gradients ).

Original: Studying the in ﬂ uence that certain input data has on other data points during training is a commontechnique from robust statistics.

Original: The use of in ﬂ uence functions to study model behaviour has beenexplored for neural networks to, for example, help identify important training samples and detectmislabelled data [Koh and Liang, 2017].

Original: Recently this ideahas been broadened under the generalconcept of gradient similarity/coherenceand has  been suggested as an explanation for a neuralnetwork’s ability to generalise [Faghri et al., 2020, Chatterjee and Zielinski, 2022].
Revised: The idea of gradient similarity/coherence has recently been suggested as an explanation for a neuralnetwork’s ability to generalise [Faghri et al., 2020, Chatterjee and Zielinski, 2022].

Original: It is here, however,mainly considered from the perspective of the loss-parameter gradient ∇ θ L and not the modelparameter gradients ∇ θ f θ .
Revised: It is here, however,mainly considered from the perspective of the loss-parameter gradient ∇ θ L and not the modelparameter gradients ∇ θ f θ .

Original: Contrary to the loss-parameter gradients, which diminish as any minimais approached, the model-parameter gradients display different behaviour and describe the evolutionof the function realised by the model.
Revised: Contrary to the loss-parameter gradients, which diminish as any minimais approached, the model-parameter gradients display different behaviour and describe the evolutionof the function realised by the model.

Original: To help understand why they can be useful in understandingneural network generalisation and regularisation thereof, we will borrow some inspiration aboutgradient similarity from Charpiat et al.
Revised: To help understand why they can be useful in understandingneural network generalisation and regularisation thereof, we will borrow some inspiration aboutgradient similarity from Charpiat et al.

Original: [2019] and use it to de ﬁ ne the core ideas behind MGS.
Revised: [2019] and use it to deﬁne the core ideas behind MGS.

Original: 2.1 De ﬁ ning MGS and its Relation to Generalisation
Revised: 3.1 Deﬁning MGS and its Relation to Generalisation

Original: When judging the similarity of two points x and x ′ from the model’s point of view, one may beinclined to simply compare the output f θ ( x ) to the other f θ ( x ′ ) .
Revised: When judging the similarity of two points x and x (cid:2) from the model’s point of view, one may beinclined to simply compare the output f θ ( x ) to the other f θ ( x (cid:2) ) .

Original: Unfortunately, due to the non-linearnature of neural networks, the same output might be produced for a number of different reasons andit is therefore dif ﬁ cult to draw conclusions about similarity directly from the output.
Revised: Unfortunately, due to the non-linearnature of neural networks, the same output might be produced for a number of different reasons andit is therefore difﬁcult to draw conclusions about similarity directly from the output.

Original: An alternative notion of similarity between x and x ′ can be de ﬁ ned by how much changing the modeloutput for x (by training on this sample) changes the output for x ′ .
Revised: An alternative notion of similarity between x and x (cid:2) can be deﬁned by how much changing the modeloutput for x (by training on this sample) changes the output for x (cid:2) .

Original: If the points are dissimilar fromthe model’s standpoint, then changing f θ ( x ) should have little in ﬂ uence over f θ ( x ′ ) and vice-versa.
Revised: If the points are dissimilar fromthe model’s standpoint, then changing f θ ( x ) should have little inﬂuence over f θ ( x (cid:2) ) and vice-versa.

Original: That is, after one step using GD for a single input sample x the parameters are changed by:
Revised: That is, after one step using GD for a single input sample x the parameters are changed by:

Original: [Equation]
Revised: [Equation]

Original: This induces the following change in the output of the model for that speci ﬁ c x :
Revised: This induces the following change in the output of the model for that speciﬁc x :

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: On the other hand, this update will also yield a change in the model for another x ′ given by:
Revised: On the other hand, this update will also yield a change in the model for another x (cid:2) given by:

Original: [Equation]
Revised: [Equation]

Original: Therefore, the kernel k θ ( x, x ′ )
Revised: Therefore, the kernel k θ ( x, x (cid:2) )

Original: = ∇ θ f ( x ) · ∇ θ f ( x ′ ) (derived from the model-parameter gradients)is crucial to understanding how the model evolves after one step of GD.
Revised: = ∇ θ f ( x ) · ∇ θ f ( x (cid:2) ) (derived from the model-parameter gradients)is crucial to understanding how the model evolves after one step of GD.

Original: It determines how much ofthe loss-model gradient ∇ f L ( f ( x ) , y ) is used to update the model, up to a scaling by the learningrate.
Revised: It determines how much ofthe loss-model gradient ∇ f L ( f ( x ) , y ) is used to update the model, up to a scaling by the learningrate.

Original: Another way to look at it is that k θ ( x, x ′ ) describes the in ﬂ uence learning x has over x ′ .
Revised: Another way to look at it is that k θ ( x, x (cid:2) ) describes the inﬂuence learning x has over x (cid:2) .

Original: Ifk θ
Revised: Ifk θ

Original: ( x, x ′ ) is large, which occurs when the model gradients are similar, then an update for f ( x ) willalso move f ( x ′ ) in the same direction.
Revised: ( x, x (cid:2) ) is large, which occurs when the model gradients are similar, then an update for f ( x ) willalso move f ( x (cid:2) ) in the same direction.

Original: Likewise, if k ( x, x ′ ) is small, meaning the model gradientsare dissimilar, then the update for f ( x ) will have little affect on f ( x ′ ) .
Revised: Likewise, if k ( x, x (cid:2) ) is small, meaning the model gradientsare dissimilar, then the update for f ( x ) will have little affect on f ( x (cid:2) ) .

Original: The cartoon example in ﬁ gure1 illustrates how this property can be useful to understand how a model generalises better if it exhibitshigher MGS.
Revised: The cartoon example in ﬁgure1 illustrates how this property can be useful to understand how a model generalises better if it exhibitshigher MGS.

Original: 2.2 The contribution of MGS to Gradient Descent
Revised: 3.2 The contribution of MGS to Gradient Descent

Original: If equations 1 and 2 above are combined, the complete update for a single step of gradient descentwith training data X , is:[Equation] 
Revised: If equations 1 and 2 above are combined, the complete update for a single step of gradient descentwith training data X , is:

Original: [Equation]
Revised: [Equation]

Original: Here K θ ( X ) =
Revised: Here K θ ( X ) =

Original: k ( X, X ) is the kernel matrix for data X .
Revised: k ( X, X ) is the kernel matrix for data X .

Original: The update for each model output f θ ( x k )can thus be seen as a weighted average of the loss-model gradients ∇ f L ( f ( x j ) , y j ) :
Revised: The update for each model output f θ ( x k )can thus be seen as a weighted average of the loss-model gradients ∇ f L ( f ( x j ) , y j ) :

Original: [Equation]
Revised: [Equation]

Original: From equation 3, it follows that the kernel ultimately controls the weighting of the loss-modelgradients: the greater the similarity exhibited between the model gradients, the greater the averagingeffect will be.
Revised: From equation 3, it follows that the kernel ultimately controls the weighting of the loss-modelgradients: the greater the similarity exhibited between the model gradients, the greater the averagingeffect will be.

Original: This partially explains how GD induces implicit regularisation as, unless K θ ( X ) is dominated by its diagonal, a model update for observation x k will also utilise the loss gradients ofother observations.
Revised: This partially explains how GD induces implicit regularisation as, unless K θ ( X ) is dominated by its diagonal, a model update for observation x k will also utilise the loss gradients ofother observations.

Original: Furthermore, it has been observed from the Neural Tangent Theory perspective that K θ ( X ) will oftenalign with the ideal kernel Y Y T in classi ﬁ cation problems, which perfectly discriminates targets.
Revised: Furthermore, it has been observed from the Neural Tangent Theory perspective that K θ ( X ) will oftenalign with the ideal kernel Y Y T in classiﬁcation problems, which perfectly discriminates targets.

Original: This phenomenon has already been investigated as one of the reasons behind implicit generalisationfor neural networks when trained with SGD [Kopitkov and Indelman, 2020, Baratin et al., 2021].
Revised: This phenomenon has already been investigated as one of the reasons behind implicit generalisationfor neural networks when trained with SGD [Kopitkov and Indelman, 2020, Baratin et al., 2021].

Original: Here, we also see that the alignment will naturally engage the averaging effect, grouping gradientsaccording to their target classes.
Revised: Here, we also see that the alignment will naturally engage the averaging effect, grouping gradientsaccording to their target classes.

Original: It is still unclear what the cause of the alignment is.
Revised: It is still unclear what the cause of the alignment is.

Original: However, MGSprovides an explanation as to why such an alignment is useful.
Revised: However, MGSprovides an explanation as to why such an alignment is useful.

Original: More detail on the connection to
Revised: More detail on the connection to

Original: Neural Tangent Theory is presented in the supplementary material.
Revised: Neural Tangent Theory is presented in the supplementary material.

Original: To summarise, apart from capturing to what extent the model gradients are aligned with each other,
Revised: To summarise, apart from capturing to what extent the model gradients are aligned with each other,

Original: MGS also determines how loss-model gradients for observations that have high MGS with each otherare averaged in the GD step.
Revised: MGS also determines how loss-model gradients for observations that have high MGS with each otherare averaged in the GD step.

Original: 2.3 Metrics to Quantify MGS
Revised: 3.3 Metrics to Quantify MGS

Original: Since the kernel K θ ( X ) captures coordinated learning between similar observations (equation 3), itis desirable to identify metrics which can summarise the overall model gradient similarity in a batchof data.
Revised: Since the kernel K θ ( X ) captures coordinated learning between similar observations (equation 3), itis desirable to identify metrics which can summarise the overall model gradient similarity in a batchof data.

Original: To ﬁ nd relevant metrics of K θ we note that the spectrum of K θ ( X ) =
Revised: To ﬁnd relevant metrics of K θ we note that the spectrum of K θ ( X ) =

Original: ( ∇ θ f ( X ))
Revised: ( ∇ θ f ( X ))

Original: T ∇ θ f
Revised: T ∇ θ f

Original: ( X ) is thesame as that of ˜ C
Revised: ( X ) is thesame as that of ˜ C

Original: = ∇ θ f ( X )( ∇ θ f ( X ))
Revised: = ∇ θ f ( X )( ∇ θ f ( X ))

Original: T ; the non-centered version of the sample covariance matrix
Revised: T ; the non-centered version of the sample covariance matrix

Original: C = ( ∇

Original: [Equation]

Original:  f ( X ) − ∇ θ f ( X ))( ∇ [Equation] f ( X ) − ∇ θ f ( X )) T .
Revised: C = ( ∇ θ f ( X ) − ∇ θ f ( X ))( ∇ θ f ( X ) − ∇ θ f ( X )) T .

Original: The spectrum of C and ˜ C are interlaced astheir corresponding eigenvalues 
Revised: The spectrum of C and ˜ C are interlaced astheir corresponding eigenvalues λ 1 ≤ ˜ λ 1 ≤ . . .

Original: [Equation] [Equation] n .
Revised: ≤ λ n ≤ ˜ λ n .

Original: We also note that the trace and determinant of C are commonly used measures of the overall varianceand covariance for a given sample of data.
Revised: We also note that the trace and determinant of C are commonly used measures of the overall varianceand covariance for a given sample of data.

Original: Here, the model-parameter gradients ∇ f θ constitute thedata.
Revised: Here, the model-parameter gradients ∇ f θ constitute thedata.

Original: Moreover, the trace and determinant of ˆ C give a rough approximation of the same metrics of C .
Revised: Moreover, the trace and determinant of ˆ C give a rough approximation of the same metrics of C .

Original: Proofs of these two properties are given in the supplementary material.
Revised: Proofs of these two properties are given in the supplementary material.

Original: Thus, the summary statistics tr K θ and det K θ can be seen as approximations of the model gradientvariance and covariance.
Revised: Thus, the summary statistics tr K θ and det K θ can be seen as approximations of the model gradientvariance and covariance.

Original: Smaller values for the trace and determinant of the kernel, K θ , re ﬂ ect ahigher degree of model gradient similarity and thus a larger averaging effect in equation 3.
Revised: Smaller values for the trace and determinant of the kernel, K θ , reﬂect ahigher degree of model gradient similarity and thus a larger averaging effect in equation 3.

Original: These metrics also serve as proxies for measuring the magnitude of the diagonal elements and relativemagnitude of the diagonal-to-off-diagonal elements of K θ .
Revised: These metrics also serve as proxies for measuring the magnitude of the diagonal elements and relativemagnitude of the diagonal-to-off-diagonal elements of K θ .

Original: From this perspective, smaller values forthe trace and determinant of the K θ re ﬂ ect how well the kernel can be summarised by a low-rankrepresentation.
Revised: From this perspective, smaller values forthe trace and determinant of the K θ reﬂect how well the kernel can be summarised by a low-rankrepresentation.

Original: Low-rank K θ can be viewed as indicative of coordinated learning in that it restrictsthe number of independent directions in which the functions can evolve in equation 3.
Revised: Low-rank K θ can be viewed as indicative of coordinated learning in that it restrictsthe number of independent directions in which the functions can evolve in equation 3.

Original: From a practical standpoint, the exact values of these metrics are not important.
Revised: From a practical standpoint, the exact values of these metrics are not important.

Original: Rather, we wish tomonitor how they evolve during training.
Revised: Rather, we wish tomonitor how they evolve during training.

Original: Utilising the kernel also constitutes a pragmatic solutionsince calculating C , of size p × p , is near infeasible for larger networks, whereas K θ is only of sizen × n , where p and n are the number of parameters and data batch respectively, and generally n ≪ p .
Revised: Utilising the kernel also constitutes a pragmatic solutionsince calculating C , of size p × p , is near infeasible for larger networks, whereas K θ is only of sizen × n , where p and n are the number of parameters and data batch respectively, and generally n (cid:6) p .

Original:  Tracking MGS Metrics During Training
Revised:  Tracking MGS Metrics During Training

Original: Let us now investigate how the metrics tr K θ and det K θ evolve during training.
Revised: Let us now investigate how the metrics tr K θ and det K θ evolve during training.

Original: Using these metrics,a wide range of regularisers, representing the most commonly used ones, are compared.
Revised: Using these metrics,a wide range of regularisers, representing the most commonly used ones, are compared.

Original: Additionally,two new regularisation methods based on MGS are plotted alongside as a reference.
Revised: Additionally,two new regularisation methods based on MGS are plotted alongside as a reference.

Original: These areintroduced in the next section.
Revised: These areintroduced in the next section.

Original: The results are shown in ﬁ gures 2, 3, and 4.
Revised: The results are shown in ﬁgures 2 and 3.

Original: In each example, all methods have their hyper-parameterstuned, using grid search and cross-validation.

Original: Additionally the results are aggregated after usingdifferent training splits and network initialisations.

Original: Figure 2 presents results from a simple fullyconnected network (FCN) architecture trained on a toy problem generated from two concentric cirlesperturbed with noise.
Revised: Figure 2 presents results from a simple FCN architecturetrained on a toy problem generated from two concentric cirles perturbed with noise.

Original: Figure 3 stems from training a large convolution network (AlexNet) on acorrupted version of MNIST with label noise and restricted training size.
Revised: Figure 3stem from training a large convolution network (AlexNet) on a corrupted version of MNIST withlabel noise.

Original: Figure 4 is created ina similar way but instead uses the more dif ﬁ cult Fashion MNIST dataset [Xiao et al., 2017] anda network with residual connections (ResNet20) to verify if the phenomenon still occurs in more complex scenarios.

Original: In each example, despite being different problems and different architectures,all regularisation methods exhibit the same behaviour: when regularisation is applied the rate of MGS metric growth is decreased, indicating higher model gradient similarity.
Revised: In each example, despite being different problems and different architectures, allregularisation methods exhibit the same behaviour: when regularisation is applied the rate of MGSmetric growth is decreased, indicating higher model gradient similarity.

Original: Furthermore, the testaccuracy or generalisation performance ( ﬁ gures 3 & 4) is strongly correlated with the evolution ofthe MGS metrics. 
Revised: Furthermore, the test accuracyor generalisation performance (ﬁgure 3) is strongly correlated with the evolution of the MGS metrics. Once any explicit regularisation is used, MGS metric growth is slowed. Coupled with this, the decision methods that constrain the MGS metrics better. (Complete set of decision boundaries are provided in the supplementary material.)

Original: The slower the rate of MGS growth, the better the ﬁ nal generalisation, and whenthe growth plateaus so does the test accuracy.
Revised: The slower the rate of MGS growth, the better the ﬁnal generalisation, and when the growth plateausso does the test accuracy.

Original: We list some interesting observations from the MNIST example in ﬁ gure 3:
Revised: We list some interesting observations from the MNIST example in ﬁgure 3:

Original: • The unregularised network (gray) goes through a rapid boost in accuracy and then quickly
Revised: • The unregularised network (gray) goes through a rapid boost in accuracy and then quickly

Original: over ﬁ ts.
Revised: overﬁts.

Original: When test accuracy peaks, there is a small plateau in the MGS metrics.
Revised: When test accuracy peaks, there is a small plateau in the MGS metrics.

Original: However,
Revised: However,

Original: once it starts to decline in accuracy, the MGS metrics again grow, indicating model gradient
Revised: once it starts to decline in accuracy, the MGS metrics again grow, indicating model gradient

Original: similarity is decreased.
Revised: similarity is decreased.

Original: • Weight penalty (blue) is able to regularise the network initially, but also eventually over ﬁ ts.
Revised: • Weight penalty (blue) is able to regularise the network initially, but also eventually overﬁts.

Original: This is re ﬂ ected in the MGS metrics as they plateau when test accuracy peaks and then start
Revised: This is reﬂected in the MGS metrics as they plateau when test accuracy peaks and then start

Original: to increase at a seemingly proportional rate to the decline in performance.
Revised: to increase at a seemingly proportional rate to the decline in performance.

Original: • All methods that achieve high test accuracy control the growth of the MGS metrics.
Revised: • All methods that achieve high test accuracy control the growth of the MGS metrics.

Original: Still,
Revised: Still,

Original: many of them begin to over ﬁ t towards the end of training which coincides with the MGS metrics also gradually increasing.
Revised: many of them begin to overﬁt towards the end of training which coincides with the MGSmetrics also gradually increasing.

Original: • Only the MGS penalty (to be introduced in the next section) is able to maintain a stable test
Revised: • Only the MGS penalty (to be introduced in the next section) is able to maintain a stable test

Original: accuracy
Revised: accuracy

Original: • Dropout ( brown) has gaps in its determinant metric.
Revised: • Dropout (brown) has gaps in its determinant metric.

Original: This means that it is causing K θ to besingular which could have rami ﬁ cations on the stability of the training.
Revised: This means that it is causing K θ to be singular which could have ramiﬁcations on the stability of the training.

Original:  bench experiments  section5.
Revised: More extensive test bench experiments are provided in section

Original:  MGS regularisation
Revised:  MGS regularisation

Original: In the previous section, we saw a clear connection between boosting MGS (re ﬂ ected by lower MGSmetrics) and the ability of the network to generalise.
Revised: In the previous section, we saw a clear connection between boosting MGS (reﬂected by lower MGSmetrics) and the ability of the network to generalise.

Original: The effect on MGS by the explicit regulariserswas common to all the methods despite their widely differing design.
Revised: The effect on MGS by the explicit regularisers was common to all the methods despite their widely differing design.

Original: Thus, we may think of eachregularisation method as acting as a proxy for enforcing model gradient similarity.
Revised: Thus, we may think of eachregularisation method as acting as a proxy for enforcing model gradient similarity.

Original: This motivates theconstruction of a new regularisation scheme that directly optimises MGS.
Revised: This motivates theconstruction of a new regularisation scheme that directly optimises MGS.

Original: We thus modify the original loss to include a new penalty term g ( K θ ( X ))) which acts on the MGSkernel:
Revised: We thus modify the original loss to include a new penalty term g ( K θ ( X ))) which acts on the MGSkernel:

Original: [Equation]
Revised: [Equation]

Original: We already have two prime candidates for g , namely the previously investigated MGS metrics:
Revised: We already have two prime candidates for g , namely the previously investigated MGS metrics:

Original: [Equation]
Revised: [Equation]

Original: where α is a penalty factor and λ i are the eigenvalues from the spectral decomposition of K θ .
Revised: where α is a penalty factor and λ i are the eigenvalues from the spectral decomposition of K θ .

Original: Aspreviously mentioned in section 2.3, smaller values for these penalty metrics re ﬂ ect higher 
Revised: Aspreviously mentioned in section 3.3, smaller values for these penalty metrics reﬂect higher modelgradient similarity.

Revised: Adding an explicit penalty on these metrics thus forces the network to learn datain groups and preventing it from learning points individually (over-ﬁtting or memorization).

Original: The evolution of the two were shown in ﬁ gures 2, 3 and 4.
Revised: The evolution of the two were shown in ﬁgures 2 and 3.

Original: The metrics seem to behave similarly.
Revised: The metrics seem to behave similarly.

Original: However, g tr ( K θ ) is more ef ﬁ cient to compute as it does not require either the complete spectraldecomposition nor the full gradient similarity kernel to be computed.
Revised: However, g tr ( K θ ) is more efﬁcient to compute as it does not require either the complete spectraldecomposition nor the full gradient similarity kernel to be computed.

Original: However, since both metricsare gradient based, they are, of course, computationally more burdensome than commonly usedregularisers such as weight decay.
Revised: However, since both metricsare gradient based, they are, of course, computationally more burdensome than commonly usedregularisers such as weight decay.

Original: For calculations of the penalties we make use of Neural Tangents[Novak et al., 2019] and Fast Finite Width NTK
Revised: For calculations of the penalties we make use of Neural Tangents[Novak et al., 2019] and Fast Finite Width NTK

Original: [Novak et al., 2022] libraries which are built on JAX[Bradbury et al.
Revised: [Novak et al., 2022] libraries which are built on JAX[Bradbury et al.

Original: We note that the penalties are equivalent to penalising the arithmetic and geometric mean of the eigenvalues of the kernel, respectively.
Revised: We note that the penalties are equivalent to penalising the arithmetic and geometric mean of the eigenvalues of the kernel, respectively.

Original: The former will mainly be affected by large leading eigenvalues.
Revised: The former will mainly be affected by large leading eigenvalues.

Original: The latter essentially penalises the eigenvalues of the kernel on a log-scale and may thus take more ofthe overall structure of the kernel into account.
Revised: The latter essentially penalises the eigenvalues of the kernel on a log-scale and may thus take more ofthe overall structure of the kernel into account.

Original: Supplemental ﬁ gure A.1 depicts results on the toytwo-circle data from using both penalties.
Revised: Supplemental ﬁgure A.1 depicts results on the toytwo-circle data from using both penalties.

Original: The differences are subtle.
Revised: The differences are subtle.

Original: For pragmatic considerations,on real data we therefore utilise only the trace penalty for explicit regularisation, but both metrics canbe used for monitoring neural network training.
Revised: For pragmatic considerations,on real data we therefore utilise only the trace penalty for explicit regularisation, but both metrics canbe used for monitoring neural network training.

Original: Details on how K θ is calculated with regards to mini-batches, vector-outputs and large data sets areprovided in the supplementary material.
Revised: Details on how K θ is calculated with regards to mini-batches, vector-outputs and large data sets areprovided in the supplementary material.

Original:  Experiments
Revised:  Experiments

Original: We compare the performance of optimising MGS directly versus the most common, explicit regularisers.
Revised: We compare the performance of optimising MGS directly versus the most common, explicit regularisers.

Original: Although there exists a plethora of explicit regularisers and modi ﬁ cations thereof, the oneschosen are a representative of the main types of regularisers with the most widespread use in practice.
Revised: Although there exists a plethora of explicit regularisers and modiﬁcations thereof, the ones [Table] chosen are a representative of the main types of regularisers with the most widespread use in practice [Table]

Original: As the aim is to compare performance between regularisers, a systematic investigation is done wherethe regularisers are rigorously tested against common over ﬁ tting scenarios caused by target noise andtraining size.
Revised: As the aim is to compare performance between regularisers, a systematic investigation is done wherethe regularisers are rigorously tested against common overﬁtting scenarios caused by target noise andtraining size.

Original: Two main tests are done, one on a classi ﬁ cation problem and the other on a regressionproblem.
Revised: Two main tests are done, one on a classiﬁcation problem and the other on a regressionproblem.

Original: For each test two different types of architectures are chosen.
Revised: For each test two different types of architectures are chosen.

Original: This is to test how universaleach method is in handling vastly different problems and networks.
Revised: This is to test how universaleach method is in handling vastly different problems and networks.

Original: Finally, a test of robustnessto variability found in common training setups is performed.
Revised: Finally, a test of robustnessto variability found in common training setups is performed.

Original: Complete details on the experiments,complete results and code to run them can be found in the supplementary material.
Revised: Complete details on the experiments,complete results and code to run them can be found in the supplementary material.

Original: 5.1 Classi ﬁ cation: corrupted MNIST
Revised: 6.1 Classiﬁcation: corrupted MNIST

Original: We generate a corrupted version of the popular MNIST dataset by applying a motion blur to eachtraining image.
Revised: We generate a corrupted version of the popular MNIST dataset by applying a motion blur to eachimage.

Original: The test set is kept uncorrupted, which ascertains how well the network generalises:if the network over ﬁ ts the training set, it will learn the corrupted version and perform poorly on theclean data.

Original: A testbench of challenging scenarios are created by varying the amount of label noise andtraining size.
Revised: A testbench of challenging scenarios are created by varying the amount of label noise andtraining size.

Original: Each regularisation method is tuned in an identical setting, of medium dif ﬁ culty, andthen retains the chosen hyper-parameters for the entire testbench.
Revised: Each regularisation method is tuned in an identical setting and then retains the chosenparameters for the entire testbench.

Original: This ensures an even playing ﬁ eldand that each method is given the same starting point.
Revised: This ensures an even playing ﬁeld and that each method is giventhe same starting point.

Original: For each testing scenario, multiple runs areperformed using a new network initialisation and resampling of the training data.
Revised: For each testing scenario, multiple runs are performed using a new networkinitialisation and resampling of the training data.

Original: No early stoppingis used as we purely want to test how well each regularisation method can prevent over-

Original: ﬁ tting on itsown.

Original: Table 1 shows the results for 4 scenarios from the 143 scenario large testbench (see supplemental).
Revised: Table 1 shows the results for 4 scenarios from the 143 scenario large testbench (see supplemental).

Original: Inthis scenario 3000 samples are used for training at 4 noise levels.
Revised: Inthis scenario 3000 samples are used for training at 4 noise levels.

Original: It is immediately clear that MGSachieves the best performance overall in many aspects.
Revised: It’s immediately clear that MGSachieves the best performance overall in many aspects.

Original: Not only does it reach higher accuracy levelsthan the other methods, it is also robust in its performance.
Revised: Not only does it reach higher accuracy levelsthan the other methods, it’s also robust in its performance.

Original: This can be seen both by the standarddeviation of its ﬁ nal test accuracy, but also when comparing the max test accuracy to its ﬁ nal one.
Revised: This can be seen both by the standarddeviation of its ﬁnal test accuracy, but also when comparing the max test accuracy to its ﬁnal one.

Original: When comparing the maximum values, it is also evident that early stopping would not have resultedin better performance for any of the methods compared to the ﬁ nal MGS result. Additionally insupplemental ﬁ gures A.2 and A.3, the evolution of test accuracy during training is depicted, visuallyshowing that early stopping techniques do not exhibit improved performance over MGS. Indeed,
Revised: For

Original: MGS max and ﬁ nal test accuracy essentially coincide, indicating that MGS regularised networks donot over ﬁ t. The
Revised: MGS max and ﬁnal test accuracy essentially coincide, indicating that MGS regularised networks donot overﬁt.

Original: [Table] Only Dropout and MGS achieve an acceptable level of ﬁ nal test loss. provide a marginal increase in performance compared to the unregularised network. In the supplemental results on a which only MGS is able to attain good performance.
Revised: The same conclusions hold true for a fully connected network (see supplemental) whereperformance is overall worse for all methods but MGS outperforms the others.

Original: [Table]

Original: 5.2 Regression: Facial Keypoints
Revised: 6.2 Regression: Facial Keypoints

Original: A similar setup to the classi ﬁ cation is created using corrupted version of the Facial Keypoints dataset was used as a regression benchmark.
Revised: A similarly corrupted version of the Facial Keypoints data set was used as a regression benchmark.

Original: Based on images of human faces, the problem is to predictthe coordinates of 15 key facial features.
Revised: Based on images of human faces, the problem is to predict the coordinates of 15 key facial features.

Original: Each image is ﬁ rst corrupted using a motion blur.
Revised: Each image is ﬁrst corrupted using a motion blur.

Original: Noise isthen introduced by adding normally distributed values, using different scale parameters, to the targetvariables.
Revised: Noise is then introduced by adding normallydistributed values, using different scale parameters, to the target variables.

Original: Here, we ﬁ xed the number of training samples while the amount of target noise was varied.
Revised: Here, we ﬁxed thenumber of training samples while the amount of target noise was varied.

Original: Each method is tuned using the same scenario ﬁ rst.
Revised: Each method is tunedusing the same scenario ﬁrst.

Original: Once again no early stopping is used.

Original: Similarconclusions can be drawn as from the classi ﬁ cation results.
Revised: Similar conclusions can be drawn as from the classiﬁcation results.

Original: The results are visible in table 2 wherethe noise column corresponds to the scale parameter of the normally distributed noise.
Revised: The results are visible in table 2 where the noise column corresponds to the scale parameter of thenormally distributed noise.

Original: Interestingly,only Dropout and MGS were actually able to converge to a satisfactory performance level.
Revised: Interestingly, only Dropout and MGS were actually able to convergeto a satisfactory performance level.

Original: However,when results are compared for a FCN architecture (table A.2, supplemental), Dropout falls into thesame category as the other methods while MGS performance remains high.
Revised: However, when results are compared for a FCN architecture(table A.2, supplemental), Dropout falls into the same category as the other methods while MGSperformance remains high.

Original: This shows a weaknessin a method such as Dropout: that it is inevitably architecture dependent. 
Revised: This shows a weakness in a method such as Dropout: that it is inevitablyarchitecture dependent. [Table]

Revised: penalty to a lesser degree.

Original: 5.3 Training parameter robustness
Revised: 6.3 Training parameter robustness

Original: Finally, we test the robustness of each regularisation method by changing: amount of training size,label noise, batch size, learning rate, and epochs.
Revised: Finally, we test the robustness of each regularisation method by changing: amount of training size,label noise, batch size, learning rate, and epochs.

Original: A middle-ground scenario was chosen from this
Revised: A middle-ground scenario was chosen from this

Original: MNIST testbench as a starting point to tune each method.
Revised: MNIST testbench as a starting point to tune each method.

Original: Then, the ﬁ ve training variables werechanged individually to both higher and lower values, tracking the performance of each method aftertraining.
Revised: Then, the ﬁve training variables werechanged individually to both higher and lower values, tracking the performance of each method aftertraining.

Original: Consistent with the current ﬁ ndings, MGS outperforms each method substantially.
Revised: Consistent with the current ﬁndings, MGS outperforms each method substantially.

Original: In ﬁ gure5 we use radar charts to summarise the test bench results.
Revised: In ﬁgure4 we use radar charts to summarise the test bench results.

Original: We note that MGS is the least affected by a change in all but one of the test parameters or conditions,and exhibits superior performance compared to the other regularisers.
Revised: We note that MGS is the least affected by a change in all but one of the test parameters or conditions,and exhibits superior performance compared to the other regularisers.

Original: Performance is only affectedby learning rate to some extent.
Revised: Performance is only affectedby learning rate to some extent.

Original: This is not unexpected as the same can be seen for the other methodsapart from weight penalty.
Revised: This is not unexpected as the same can be seen for the other methodsapart from weight penalty.

Original: Also, as is evident from the GD update step (equation 1), the one directlycontributing training parameter in MGS is the learning rate.
Revised: Also, as is evident from the GD update step (equation 1), the one directlycontributing training parameter in MGS is the learning rate.

Original:  Conclusion
Revised:  Conclusion

Original: We introduced the concept of Model Gradient Similarity (MGS) and discussed its connections toregularisation for models trained with gradient descent and the generalisation properties of neuralnetworks.
Revised: We introduced the concept of Model Gradient Similarity (MGS) and discussed its connections toregularisation for models trained with gradient descent and the generalisation properties of neuralnetworks.

Original: We proposed metrics that can be used to summarise MGS for a model and track the trainingof different neural network architectures for various learning problems.
Revised: We proposed metrics that can be used to summarise MGS for a model and track the trainingof different neural network architectures for various learning problems.

Original: It was shown that a wide range of explicit regularisers all appeared to attempt to enforce higher modelgradient similarity, i.e. lower MGS metrics.
Revised: It was shown that a wide range of explicit regularisers all appeared to attempt to enforce higher modelgradient similarity, i.e. lower MGS metrics.

Original: Moreover, higher test accuracy performance was shownto be re ﬂ ected in lower MGS metric values.
Revised: Moreover, higher test accuracy performance was shownto be reﬂected in lower MGS metric values.

Original: Based on these ﬁ ndings, a new type of regulariser, geared toward direct control of MGS, wasdesigned and found to achieve top performance in several rigorous test bench experiments.
Revised: Based on these ﬁndings, a new type of regulariser, geared toward direct control of MGS, wasdesigned and found to achieve top performance in several rigorous test bench experiments.

Original: Itsoverall robustness to label noise and training parameter settings was also an indication that directlyoptimising MGS comes closer to a more holistic approach to regularisation.
Revised: Itsoverall robustness to label noise and training parameter settings was also an indication that directlyoptimising MGS comes closer to a more holistic approach to regularisation.

Original: Taken together, these results provided insight into the underlying mechanisms of neural networkregularisation.
Revised: Taken together, these results provided insight into the underlying mechanisms of neural networkregularisation.

Original: Due to the higher computational costs for gradient based regularisers, such as the
Revised: Due to the higher computational costs for gradient based regularisers, such as the

Original: MGS metric penalities introduced here or loss-gradient penalties, their use for direct optimisation isnot always ef ﬁ cient.
Revised: MGS metric penalities introduced here or loss-gradient penalties, their use for direct optimisation isnot always efﬁcient.

Original: To scale for use in larger networks and in more complex settings, additional workis needed to obtain more ef ﬁ cient ways to compute or approximate the MGS kernel or metrics thereof.
Revised: To scale for use in larger networks and in more complex settings, additional workis needed to obtain more efﬁcient ways to compute or approximate the MGS kernel or metrics thereof.

Original: Future work could thus focus both on how MGS can be used to design new regularisers as well asto improve upon existing ones.
Revised: Future work could thus focus both on how MGS can be used to design new regularisers as well asto improve upon existing ones.

Original: The MGS metrics can be useful as KPI’s for measuring a network’scurrent capacity for under/
Revised: The MGS metrics can be useful as KPI’s for measuring a network’scurrent capacity for under/over-ﬁtting.

Revised: Finally, the grouping effect of regularised neural networktraining, where model gradient similarity encourages coordinated learning across observations,suggests that MGS regularisation can be explored for joint prediction modeling and clustering.

Revised: Checklist

Revised: For all

Revised: authors

Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s

Original: over-
Revised: contributions and scope? [Yes] (b) Did you describe the limitations of your work? [Yes]

Revised: (c) Did you discuss any potential negative societal impacts of your work?

Revised: [N/A]

Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Revised: them?

Revised: [Yes]

Revised: If you are including theoretical results...

Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: ﬁ tting.
Revised: [Yes] (b) Did you include complete proofs of all theoretical results? [Yes] If you ran experiments...

Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Revised: mental results (either in the supplemental material or as a URL)?

Revised: [Yes]

Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Revised: how they

Revised: were chosen)?

Revised: [Yes]

Revised: (c) Did you report error bars (e.g., with respect to the random seed after running ex-

Revised: periments multiple times)?

Revised: [Yes] Standard errors are provided in the tables but for

Revised: visualisation purposes are omitted from some of the ﬁgures.

Revised: (d) Did you include the total amount of compute and the type of resources used (e

Revised: (a) If your work uses existing assets, did you cite the creators?

Revised: [Yes] In supplementarymaterial.

Revised: (b) Did you mention the license of the assets?

Revised: [No]

Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Revised: [No]

Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?

Revised: [No]

Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Revised: information or offensive content?

Revised: [No]

Original: Finally, the grouping effect of regularised neural networktraining, where model gradient similarity encourages coordinated learning across observations,suggests that MGS regularisation can be explored for joint prediction modeling and clustering.
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review

Revised: Board (IRB) approvals, if applicable?

Revised: [N/A]

Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Revised: spent on participant compensation?

Revised: [N/A]