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 the aid 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) adds insight into how explicit regularisers, while derived from widely different principles, operate via the same mechanism underneath by increasing MGS, and (3) provides the basis for a new regularisation scheme which exhibits excellent performance, especially in challenging settings such as high levels of 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 made to 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 generalisation performance 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ﬁt in 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 implicitly regularised, 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-norm solution.

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, have been 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-propagation incorporates the gradients with respect to the loss itself as part of the loss, hence the name, as the gradient 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 solutions that 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 during training 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 the optimisation 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 offer explicit 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 works related 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 functional regularisation.

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 function approximator 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 .

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 a large 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 show that MGS can be used to analyse neural network training, monitoring how model gradients evolve together.

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: This insight lays the foundation for a new class of regularisers aimed at directly boosting model gradient similarity.

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 compare performance 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 MGS regularisation 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 MGS framework opens the door for the further development of both novel regularisers and the improvement of 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 the model 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 standard singleobservation 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 η is the 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 network before 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 the main 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, such as mean-squared error or cross-entropy, we apply the chain rule to the loss-parameter gradient: ∇ θ L = ∇ θ f θ · ∇ f L .

Original: We denote the two components as: the loss gradients with respect to the model60 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 model output ∇ f L ( loss-model gradients ) and the model gradients with respect to the model parameters ∇ θ 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.
Revised: Studying the inﬂuence that certain input data has on other data points during training is a common technique 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].
Revised: The use of inﬂuence functions to study model behaviour has been explored for neural networks to, for example, help identify important training samples and detect mislabelled data [Koh and Liang, 2017].

Original: Recently this idea has been broadened under the generalconcept of gradient similarity/coherence and has been suggested as an explanation for a neuralnetwork’s ability to generalise [Faghri et al., 2020, Chatterjee and Zielinski, 2022].
Revised: Recently this idea has been broadened under the general concept of gradient similarity/coherence and has been suggested as an explanation for a neural network’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 model-parameter 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 minima is approached, the model-parameter gradients display different behaviour and describe the evolution of 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 understanding neural network generalisation and regularisation thereof, we will borrow some inspiration about gradient 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. [Figure]

Original: 2.1 De ﬁ ning MGS and its Relation to Generalisation
Revised: 2.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 ′ from the model’s point of view, one may be inclined to simply compare the output f θ ( x ) to the other f θ ( x ′ ) .

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-linear nature of neural networks, the same output might be produced for a number of different reasons and it 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 ′ can be deﬁned by how much changing the model output for x (by training on this sample) changes the output for x ′ .

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 from the model’s standpoint, then changing f θ ( x ) should have little inﬂuence over f θ ( x ′ ) 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]

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 ′ given by:

Original: [Equation]
Revised: [Equation]

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

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 ′ ) (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 of the loss-model gradient ∇ f L ( f ( x ) , y ) is used to update the model, up to a scaling by the learning rate.

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 ′ ) describes the inﬂuence learning x has over x ′ .

Original: Ifk θ ( 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: If k θ ( x, x ′ ) is large, which occurs when the model gradients are similar, then an update for f ( x ) will also move f ( x ′ ) 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 ′ ) is small, meaning the model gradients are dissimilar, then the update for f ( x ) will have little affect on f ( x ′ ) .

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 ﬁgure 1 illustrates how this property can be useful to understand how a model generalises better if it exhibits higher MGS.

Original: 2.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:
Revised: If equations 1 and 2 above are combined, the complete update for a single step of gradient descent with training data X , is:

Original: [Equation]
Revised: [Equation]

Original: 

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-model gradients: the greater the similarity exhibited between the model gradients, the greater the averaging effect 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 of other 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 often align 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 generalisation for 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 gradients according 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, MGS provides an explanation as to why such an alignment is useful.

Original: More detail on the connection to Neural Tangent Theory is presented in the supplementary material.
Revised: More detail on the connection to 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 other are averaged in the GD step.

Original: 2.3 Metrics to Quantify MGS
Revised: 2.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), it is desirable to identify metrics which can summarise the overall model gradient similarity in a batch of 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 the same 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 as their 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 variance and covariance for a given sample of data.

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

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 gradient variance 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 a higher 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 relative magnitude 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 for the trace and determinant of the K θ reﬂect how well the kernel can be summarised by a low-rank representation.

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 restricts the 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 to monitor 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 solution since calculating C , of size p × p , is near infeasible for larger networks, whereas K θ is only of size n × n , where p and n are the number of parameters and data batch respectively, and generally n ≪ 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 are introduced in the next section.

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

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

Original: Additionally the results are aggregated after usingdifferent training splits and network initialisations.
Revised: Additionally the results are aggregated after using different 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 fully-connected network (FCN) architecture trained 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 3 stems from training a large convolution network (AlexNet) on a corrupted version of MNIST with label noise and restricted training size.

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.
Revised: Figureis created in a similar way but instead uses the more difﬁcult Fashion MNIST dataset [Xiao et al., 2017] and a 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, all regularisation methods exhibit the same behaviour: when regularisation is applied the rate of MGS metric 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 accuracy or generalisation performance (ﬁgures 3 & 4) is strongly correlated with the evolution of the MGS metrics.

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 plateaus so 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 over ﬁ ts.
Revised: • The unregularised network (gray) goes through a rapid boost in accuracy and then quickly 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, once it starts to decline in accuracy, the MGS metrics again grow, indicating model gradient similarity is decreased.
Revised: However, once it starts to decline in accuracy, the MGS metrics again grow, indicating model gradient 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 to increase at a seemingly proportional rate to the decline in performance.
Revised: This is reﬂected in the MGS metrics as they plateau when test accuracy peaks and then start 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, many of them begin to over ﬁ t towards the end of training which coincides with the MGS metrics also gradually increasing.
Revised: Still, many of them begin to overﬁt towards the end of training which coincides with the MGS metrics also gradually increasing.

Original: • Only the MGS penalty (to be introduced in the next section) is able to maintain a stable test accuracy • Dropout ( brown) has gaps in its determinant metric.
Revised: • Only the MGS penalty (to be introduced in the next section) is able to maintain a stable test accuracy • 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  section 5.
Revised: More extensive test bench experiments are provided in section 5.

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 MGS metrics) 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 each regularisation 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 the construction 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 MGS kernel:

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: As previously mentioned in section 2.3, smaller values for these penalty metrics reﬂect higher model gradient similarity.

Revised: Adding an explicit penalty on these metrics thus forces the network to learn data in 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, 3 and 4.

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 spectral decomposition 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 metrics are gradient based, they are, of course, computationally more burdensome than commonly used regularisers 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 of the 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 toy two-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 can be 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 are provided 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 chosen are a representative of the main types of regularisers with the most widespread use in practice.

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 where the regularisers are rigorously tested against common overﬁtting scenarios caused by target noise and training 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 regression problem.

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 universal each 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 [Table] of robustness to 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: 5.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 each training image.

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.
Revised: 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 the clean 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 and training 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, of medium difﬁculty, and then retains the chosen hyper-parameters 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 given the 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 network initialisation 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- ﬁ tting on itsown.
Revised: No early stopping is used as we purely want to test how well each regularisation method can prevent over-ﬁtting on its own.

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: In this 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 is immediately clear that MGS achieves 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 levels than the other methods, it is 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 standard deviation 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.
Revised: When comparing the maximum values, it is also evident that early stopping would not have resulted in better performance for any of the methods compared to the ﬁnal MGS result.

Original: 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.
Revised: Additionally in supplemental ﬁgures A.2 and A.3, the evolution of test accuracy during training is depicted, visually showing that early stopping techniques do not exhibit improved performance over MGS.

Original: Indeed,
Revised: Indeed,

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 do not overﬁt.

Original: [Table] Only Dropout and MGS achieve an acceptable level of ﬁ nal test loss. provide a marginal increase in performancecompared 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) where performance is overall worse for all methods but MGS outperforms the others.

Original: [Table]

Original: 5.2 Regression: Facial Keypoints
Revised: 5.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 similar setup to the classiﬁcation is created using 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 normally distributed 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 the number 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 tuned using the same scenario ﬁrst.

Original: Once again no early stopping is used.
Revised: 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 the normally 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 converge to a satisfactory performance level. However,

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: [Table] all methods, but weight penalty to a lesser degree. when results are compared for a FCN architecture (table A.2, supplemental), Dropout falls into the same category as the other methods while MGS performance 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 inevitably architecture dependent.

Original: 5.3 Training parameter robustness
Revised: 5.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 were changed individually to both higher and lower values, tracking the performance of each method after training.

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 ﬁgure 5 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 affected by 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 methods apart 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 directly contributing 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 to regularisation for models trained with gradient descent and the generalisation properties of neural networks.

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 training of 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 model gradient 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 shown to 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, was designed 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: Its overall robustness to label noise and training parameter settings was also an indication that directly optimising 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 network regularisation.

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 is not 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 work is 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 as to 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’s current capacity for under/over-ﬁtting.

Original: over-
Revised: Finally, the grouping effect of regularised neural network training, where model gradient similarity encourages coordinated learning across observations, suggests that MGS regularisation can be explored for joint prediction modeling and clustering. Acknowledgments and Disclosure of Funding This research was supported by funding from Centiro Solutions, the Swedish Research Council (VR), the Swedish Foundation for Strategic Research, and the Wallenberg AI, Autonomous Systems and Software Program (WASP). Checklist

Original: ﬁ tting.
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope? [Yes] (b) Did you describe the limitations of your work? [Yes] (c) Did you discuss any potential negative societal impacts of your work? [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Revised: [Yes] 2.

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 are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results? [Yes] (b) Did you include complete proofs of all theoretical results? [Yes] 3. If you ran experiments...

Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Revised: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Revised: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Revised: [Yes] Standard errors are provided in the tables but for 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.g., type of GPUs, internal cluster, or cloud provider)?

Revised: [Yes] 4.

Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Revised: [Yes] In supplementary material.

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

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

Revised: [No] (d) Did you discuss whether

Revised: and how consent was obtained from people whose data you’re using/curating?

Revised: [No] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

Revised: [No] 5.

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, if applicable?

Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Revised: [N/A]