Global curvature for second-order optimization of neural networks
Authors: Alberto Bernacchia
Venue: ICML 2025 poster (2025)
Keywords: Curvature
Second-order optimization
Symmetric neural networks
Primary area: deep_learning->theory
TL;DR: We present a theoretical framework for exact computation of the structure of the global curvature by leveraging the intrinsic symmetries of neural networks.

Abstract: Second-order optimization methods, which leverage the local curvature of the loss function, have the potential to dramatically accelerate the training of machine learning models. 
However, these methods are often hindered by the computational burden of constructing and inverting large curvature matrices with $\mathcal{O}(p^2)$ elements, where $p$ is the number of parameters. 
In this work, we present a theory that predicts the \emph{exact} structure of the global curvature by leveraging the intrinsic symmetries of neural networks, such as invariance under parameter permutations. 
For Multi-Layer Perceptrons (MLPs), our approach reveals that the global curvature can be expressed in terms of $\mathcal{O}(d^2 + L^2)$ independent factors, where $d$ is the number of input/output dimensions and $L$ is the number of layers, significantly reducing the computational burden compared to the $\mathcal{O}(p^2)$ elements of the full matrix. 
These factors can be estimated efficiently, enabling precise curvature computations.
To evaluate the practical implications of our framework, we apply second-order optimization to synthetic data, achieving markedly faster convergence compared to traditional optimization methods.
Our findings pave the way for a better understanding of the loss landscape of neural networks, and for designing more efficient training methodologies in deep learning.
Code: \href{https://github.com/mtkresearch/symo_notebooks}{github.com/mtkresearch/symo\_notebooks}

Global curvature for second-order optimization of neural networks
Alberto Bernacchia
Abstract
first-order approaches (Bottou et al., 2018). These meth- ods face a major challenge: the need to compute and invert curvature matrices of size p p. where p is the number of model parameters. Several studies have explored block- diagonal and Kronecker- vature computation in neural networks (Martens & Grosse, 2015; Eschenhagen et al., 2024). Recent advances have scaled these methods effectively (Ba et al., 2017: Anil et al., 2021; Kasimbeg et al., 2025), even enabling second-order pre-training of large language models (Liiu et al., 2025). However, despite their computational efficiency, these ap- proximations lack strong theoretical guarantees for nonlin- ear problems (Bernacchia et al., 2018; Karakida & Osawa, 2020). Exact second -oddrr optimization, avoiding block- diagonal approximations, remains tractable only for small- scale models (Cai et al., 2019; Arbel et al., 2023; Korbit et al., 2024) or specific architectures, e.g. reversible neural networks (Buffelli al., 2024).
Second-order optimization methods, which lever- age the local curvature of the loss function, have the potential to dramatically accelerate the train- ing of machine learning models. However, these methods are often hindered by the computational burden of constructing and inverting large curva- ture matrices with Opp") elements, where p is the number of parameters. In this work, we present theory that predicts the exact structure of the global curvature by leveraging the intrinsic sym- metries of neural networks, such invariance under parameter permutations. For Multi-Layer Perceptrons (MLPs), our approach reveals that the global curvature can be expressed in terms of (dd + I,2) independent factors, where d the number of input/output dimensions and is the number of layers, significantly reducing the computational burden compared to the Of²²) el- ements of the full matrix. These factors can be estimated efficiently, enabling precise curvature computations. To evaluate the practical implica- tions of our framework, we apply second-order op- timization to synthetic data, achieving markedly faster convergence compared traditional opti- mization methods. Our findings pave the way for better understanding of the loss landscape of neural networks, and for designing more efficient training methodologies in deep learning. Code: github.commmmkkresaacchyymm notebooks ooooooooccccyyoooooooooooos
approximations for cur-
Most neural network architectures are built a layered structure, where the core operations consist of matrix-vector products combined with relatively simple nonlincarities. Due to this structural framework, neural networks exhibit well-known symmetries: for example, permuting the rows of parameter matrix while applying the same permutation to the columns of the matrix in the subsequent layer leaves the overall computation of the network unchanged, as illustrated in Figure 1 (Hecht-Nielsen, 1990; Chen et al., 1993). While all neural networks exhibit symmetries, the specific symme- tries vary by architecture. This raises a natural question: can we harness the unique symmetries of given neural network architecture to design a tailored optimization algorithm? Our key contribution is demonstrating that, by leveraging the precise symmetries of model, we compute the cur-
1. Introduction
Neural network models are commonly trained using adap- tive variants of gradient descent and momentum (Schmidt et al., 2021). Second-order optimization methods, which exploit the curvature of the loss function, have shown the potential for significantly faster convergence compared to
W2
Wi
Wz
W
'MediaTek Research, Cambridge, UK. Correspondence to: Alberto Bernacchia bebbertt.bbrnacciicee www.eeeerrrcc.coom......
Figure 1. Illustration of permutation symmetry. The output of a neural network is invariant for swapping any pair of neurons within layer. That corresponds simultaneously swapping the rows of incoming and the columns of outgoing weight matrices.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Global curvature for second-order optimization of neural networks
vature significantly reduced computational cost, making second-order optimization computationally feasible.
In Section 2.1, we review second-order optimization. In Section 2.2 we describe all known symmetries of neural networks, and we motivate our choice for this study. In Section 2.3 we prove that probability distributions remain invariant upon equivariant maps, under a more general case than previously known. In Section 2.4, we describe how the curvature is constrained by the symmetries of the neural network.
While prior work has leveraged continuous symmetries to enhance optimization (Neyshabur et al., 2015; Meng et al., 2018; Zhao et al., 2022; 2023b) and to impose constraints on the curvature (Kunin et al., 2020), our study is the first to utilize discrete symmetries for computing the full curvature matrix. A key innovation of our approach is the computa- tion of the global curvature, averaged over the parameter space, rather than the local curvature, which depends on specific parameter values. The primary motivation for con- sidering global curvature is its efficient computation, which serves as the main contribution of this work. Future research will explore whether global curvature serves as reliable approximation of local curvature and whether offers inher- ent advantages ( (Yang et al., 2021; Yao et al., 2021; Titsias, 2024). A detailed discussion of the related work is given in Appendix A.
2.1. Second-order optimization
We consider a scalar loss function (()) of parameters @ E R". In machine learning. the loss also depends on either a dataset or data distribution, but do not make this dependence explicit here. Second order optimization corresponds to the following update
0 = (o.MVVV(((0)
(1)
where a is the learning rate and M, is the pre-coniitigning matrix, usually the inverse of the curvature matrix. Vari- ous studies have employed different formulations for the curvature matrix, including the Fisher information matrix (Martens & Grosse, 2015; Bernacchia al., 2018; Garcia et al., 2023), the Gauss-Newton matrix (Botev et al., 2017; Yu et al., 2024; Buffelli et al., 2024), the Hessian matrix (Goldfarb et al., 2020), and the gradient covariance matrix (Duchi et al., 2011). Within our framework, all these matri- ces transform identically under the symmetries of a neural network (see Section 3). While their numerical values dif- fer, they share the same underlying structure, making our approach applicable to any of them.
We contribute the following:
For Multi-Layer Perceptrons (MLPs), derive act expression for the global curvature matrix, reveal- ing that it depends on (d² + I!?) unknown factors, where d is the number of input/output dimensions and L is the number of layers. This significantly reduces the computational complexity compared to the original matrix size of Op²²).
We analyze the impact of the activation function on global curvature and find that less symmetric functions, such as ReLU,, exhibit a more intricate curvature struc- ture compared to more symmetric functions, such as linear or Tanh activations. This increased complexity corresponds larger number of unknown factors in the curvature matrix.
In most previous work, the pre-conditioning matrix is local, it depends on the current value of the parameters, M (0.). Here we take a different approach, similar to Titsias (2024), and we consider global averages of the curvature. In particu- lar, consider a probability distribution Pe of parameters O at training step t, which is induced by distribution of initial conditions on parameters evolving under the same training dynamics, We define the gradient mean and covariance
We propose a simple algorithm for estimating the un- known factors of the curvature and for computing the second-order update, that significantly reduces com- plexity by using a surrogate matrix with (d + L?)) elements instead of the Opp²) elements of the full ma- trix.
H. = E VL(O.) A.
(2)
Σ, E VC(0.)VC(.)) - -uP! 8.
(3)
We demonstrate the effectiveness of our approach by running second-order optimization on a two-layer MLP and synthetic data. Empirical results show that the second-order update substantially accelerate conver- gence.
and we set the preconditioning matrix equal to the inverse square root of the gradient covariance
M, Σ 1
(4)
2. Background
The choice of taking the square root is motivated by other adaptive methods, such as Adam (Kingma & Ba, 2014), Adagrad (Duchi et al., 2011), Shampoo (Gupta et al., 2018) and RMSProp (Tieleman & Hinton, 2012). A debate on whether the curvature matrix should be square rooted
In this section, we recall known facts that are instrumental for stating our main results, and we also provide a new theorem that may be of independent interest, Theorem 2.2.
2

---

Global curvature for second-order optimization of neural networks
currently ongoing (Lin al., 2024; Choudhury al., 2024; Morwani et al., 2025), experience is that the square root significantly stabilizes training, Appendix B argues that the square root is also reasonable choice for the toy case of a quadratic loss function.
the following expression
h = We or(hx-1) + b
8= 1= L
(5)
where og is the actibvation function, a pointwise nonlinearity (a) is the identity), hgg = I E IR:do is the input, hn = y E gd is the output and fig E RR is the latent representation of layer l. The parameters of the neural network are the biases bj E Rd and weights We E Tedexde- We consider the following transformation, applying to all layers { = I,...,L.
We note that previous methods take an average of gradients over training iterations or data points, while we define the covariance by an average over an ensemble of models with different initialization. In practice, we estimate the covari- ance within a single model (see Section 3), we assume that the structure present in the global curvature is also present in the local curvature and captures meaningful variations for optimization purposes.
b Vibe
(6)
We +&W.V
(7)
2.2. Symmetries of neural networks
where the matrix V3 (ddxxd= is assumed orthogonal for each layer, V2" = V We denote by G the corresponding transformation acting on the set of all parameters, which combines the effect of V for all layers. This transformation leaves the output of the neural network invariant if it belongs to a symmetry group that depends on the activation function. We consider three cases:
Neural networks exhibit a rich variety of symmetries. The most general and ubiquitous is permutation symmetry (see Figure 1), whereby permuting the rows of weight matrix and simultancously permuting the columns of the subse- quent layer's weight matrix leaves the network's output unchanged (Hecht-Nielsen, 1990). Permutation symmetry applies to nearly all neural networks and is so fundamen- tal that new architectures often require deliberate design choices to break it (Lim et al., 2024). Beyond permutation symmetry, additional symmetries arise specific architec- tures: sign flip for networks with odd activation functions (Chen et al., 1993), rescaling in ReLU networks, exploiting their homogeneity (Neyshabur et al., 2015), translation in Softmax layers (Kunin et al., 2020), scaling in normalization layers (Kunin et al., 2020), general linear transformations of keys and queries in transformers (Ziyin, 2024), deep linear networks (Zhao et al., 2022) and data-dependent transfor- mations (Zhao et al., 2023a). It is unlikely that any other symmetries exist, at least in common neural network archi- tectures (Grigsby et al., 2023; Chen et al., 1993).
Linear activations. The network is invariant if V belongs to the set of orthogonal matrices (dde). that corresponds to the orthogonal group (rotations and reflections). Across all layers, the transformation G belongs to the product G. O(d2) x... O(dz-1). The group is continuous and compact, of dimension 2/11 d=(d==1)
(Zhao et al., 2022).
2
Odd activations (e.g. Tanh) The network is invariant if Ve belongs to the set of signed permutation matrices B(de), that corresponds to the signed symmetric group. Across all layers, the transformation G' belongs to the product G2 = B(di) x x B(dr.-i). The group is discrete and includes [1/-1 2d de! elements (Chen al., 1993).
In Section 2.3, we introduce a critical assumption underpin- ning our work. Specifically, we require not only that the neural network output be invariant under a given symmetry group acting on the parameters, but also that the probability distribution over the initial parameters exhibits the same invariance. As reviewed in Appendix C, this condition fur- ther necessitates that the transformation must be similar to an orthogonal transformation (Flytzanis, 1977). This additional assumption excludes many of the symmetries discussed earlier, such as scaling, translations and general linear transformations, but retains permutations, sign flips, rotations and reflections. Consequently, we restrict our anal- ysis these key symmetry groups, with a particular focus on orthogonal transformations.
Other activations (e.g. ReLU) The network is in- variant if V2 belongs to the set of permutation matrices S(dz). that corresponds to the symmetric group. Across all layers, the transformation G' belongs to the product C, = S(di ) x... x S(d The group is discrete and includes I[_ de! elements (Hecht-Nielsen. 1990).
The three groups satisfy a decreasing set of constraints, G, D G. 3 Gg. Correspondingly, we show in Section 3 that the structure of the curvature increases in complexity as the group gets smaller. We note that there is no symmetry on the input and output side of the MLP, therefore Vo = Ld and V, = La, are fixed to identity matrices. Additional constraints may be introduced the data has symmetries, but we do not consider that case in this work.
A Multi-Layer Perceptron (MLP) of L layers is defined by
3

---

Global curvature for second-order optimization of neural networks
theory
W
W
7e-25
bi
So+02
De
Ws
W2
D2
W,
Ws
Wi bx
W
Wi bi
Ws
Figure Gradient covariance matrix two-layer MLP with Tanh activations, tiny model with neurons in each layer. W1, W2 are the weights and bj, by are the biases of the two layers. A and B: Covariance is measured by averaging 10000 models with different initialization, before training (A: = 0) and after steps gradient descent (B: = 20). C: Covariance generated by our theory for random draw of the factors. We highlight that theory predicts the overall structure, rather than the specific numerical values.
Throughout this work, we assume that the loss function depends the parameters only through the output of the neural network, therefore the loss itself is also invariant for the group action
We consider different groups depending on the activation function, described in Section 2.2, namely G Go. G. or G.. Appendix D shows that Assumption 2. is satisfied by the most common initialization routines used in deep learning. In Pytorch for example, Assumption 2.1 holds for nn init. normal and nn. init.. orthogonal, with all groups G, Gb, G. Also, it holds when us- ing nn. init uniform and nn. init sparse with Gb, G2 but not with Gg, and it still holds when using layer- dependent parameters, for example n.. init ka hamming and nn. init xavier. See Appendix D for details.
((GO) C(O) E
where G is equal to either Go, Gb or Gy.
2.3. Invariant distribution throughout training
In this section, we assume that the parameter distribution is invariant under the specified symmetry group at initializa- tion. then demonstrate that this invariance preserved throughout training, if the parameter update rule is equiv- ariant with respect to the same symmetry group. Previous studies have established that such invariance holds under the assumption of globally invertible map (Köhler et al., 2019). Here we extend this result by proving that the same invariance is maintained under the less restrictive assump- tion of local invertibility, as stated Theorem 2.2. This generalization broadens the applicability of the result and may be of independent interest.
Theorem 2.2. Assume that the update rule O, (,(0,-1) is differentiable and its Jacobian is non-singular almost everywhere. Furthermore, it is equivariant under a volume- preserving transformation G, namely
Ur Gu,())
(1I)
Then, if the probability distribution of parameters is invari- ant at step I - 1, then it must be invariant also at step
Pc(GO,) = Pr(O,)
(12)
Parameters are initialized t 0 according probabil- ity distribution (0p),, and their evolution at subsequent training steps is governed by the update rule
The proof is provided in Appendix E. A similar result can be found in Theorem of Köhler et al. (2019), however they assume that the mapping u, is globally invertible, which quite restrictive. Theorem 2.2 considers the more general case in which the mapping is just locally invertible.
0, (0,11)
At this stage, we do not impose a specific form on the update rule, as the results in this section are derived under general assumptions. However, these results will be utilized for the update rule of Equation (1).
The assumption of equivariance of u, is satisfied by many optimizers in deep learning, provided that the loss is invari- ant. Appendix F shows that both gradient descent (with or without momentum) and the second-order update of Equa- tion (1) are equivariant for all groups Go Gb Gg. while the Adam optimizer equivariant for Gb, G., but not Go.
Assumption 2.1.
The probability distribution the initial
parameters O is invariant under the action of group G
Po(GO,) (fo) YG E G
(10)

---

which implies the possibility of infinitely many solutions Global curvature for second-order optimization of neural networks
theory
W
W.
2x:05
5
W.
Se+02 Wa
W.
Wy
Wi
[n]
Vic
10u
W Wr.
W9
1b:1
Figure Gradient covariance matrix three-layer MLP with Tanh activations, similar to Figure with three layers. Most of the correlations across layers are lost when adding the third layer. The covariance generated by our theory agrees with the observed structure.
Corollary 2.3. Assume the update is equivariant at all steps. Given Assumption 2.1 and Theorem 2.2, by induction, the distribution is invariant at all steps, Pe (c))) = p.p (6,). Vt.
(see Lemma C.3). In Section 3, we identify a solution space that remains valid for all members of given symmetry group. Since this space is infinite, both the mean μ4, and the covariance Σn may evolve to different values at different time steps, while still satisfying Equations (13) and (14).
Remark. This result does not help in finding the functional form of the distribution, that remains unknown except in very simple cases (see Appendix for toy example where the distribution remains Gaussian all time steps).
Appendix H demonstrates that Equations (14) and (15) hold identically for the Hessian matrix when averaged over the parameter distribution. A Additionally, it provides examples of Hessians observed in experiments (Figures 7, 8, 9). Similar results extend to other types of curvature matrices, including the Fisher Information matrix and the
2.4. Constraints on mean and covariance of gradient
In this section, we demonstrate that the invariance of the loss function, outlined in Section 2.2, together with the in- variance of the probability distribution, discussed in Section 2.3, imposes specific constraints on both the gradient mean and covariance. Furthermore, we note that any transforma- tion belonging to one of the three groups Go Gb, Gg is orthogonal, meaning that all its elements satisfy the property GT =
matrix.
3. Results
In this section, derive exact expressions for the mean and covariance of the gradient in a MLP and compare them with empirical results obtained from experiments on syn- thetic data. Given the layered architecture of MLPs and the distinction between weights and biases, we introduce a flattened vector of parameters encompassing all layers, by concatenating and vectorizing all tensors in column-major order, given by
Lemma 2.4. Assume both the loss function and the proba- bility distribution are invariant for orthogonal transfor- mation G. namely L(GO,) = L(O,) and pe (GO,) p. (O.). Assume that the mean Il and covariance E of the gradient exist and are finite. Then, the mean satisfies the eigenvalue equation
0 Vec (W1, bi,... Wr,, by.)
(16)
F4 = G
(13)
We visualize the gradient covariance matrix by vectorizing
and the covariance matrix invariant upon the congruent transformation
the parameters according to Equation (16). For the Hessian matrix, we present analogous results in Figures 7, 8, and 9 of Appendix
Σ, = GT
(14)
Figure 2 presents the gradient covariance matrix for two- layer MLP with Tanh activations in tiny model with five neurons per layer. The results are shown for both empirical experiments (panels A and B, see Section 4 for details) and our theoretical predictions (panel C, see Section 3.2 for de- tails). The structure observed experiments closely aligns with theoretical predictions, both before training (panel A,
Furthermore, any analytic matrix function f(E) of the co-
variance matrix satisfies the same equation
f(E,) = Gf(E))"
(15)
The proof is provided in Appendix G. It important to note that Equations (13) and (14) are linear and homogeneous,
5

---

Global curvature for second-order optimization of neural networks
L 0) and after 20 steps of gradient descent (panel B, t 20). We highlight that our theoretical framework cap- tures the overall structure of the covariance matrix rather than its exact numerical values. The estimation of specific values is further discussed in Section 3.4.
where denotes Kronecker product. For the mean gradient, notation similar to the parameter vector, Equation (16)
0L ac
ac
dL
1 E Vec
(18)
@W1 0b
OW 0b1
Figure further reveals that input weights exhibit column- wise correlations, likely reflecting dependencies in their in- put data, while output weights show row-wise correlations, capturing similarities in their output. Additionally, correla- tions are observed between the rows of the output weights and the columns of the input weights, potentially indicating input-output dependencies. Figure extends this analysis to a three-layer MLP. The observed gradient covariance structure remains consistent with theoretical predictions. However, the covariance matrix becomes block-diagonal, indicating significant reduction in inter-layer correlations as the number of layers increases from two to three. As demonstrated in Section 3.2, our theoretical framework pre- dicts that this absence of correlations between layers persists for architectures with more than three layers. We speculate that the invariance with respect to sign changes of both the incoming and outgoing weights of neuron may lead to cancellations in the covariance of these weights. However, further studies needed to understand this observation.
= Vec (#1 PI.....
11, (r.))
(19)
DC where we define be = Eo Bw: be = Ee ūb,"
ac
and Wt = ac BC T owz Ow Then,
Vec(We). Similarly, we define Eer Eo
Eee = Eo ab, Owp the covariance of the gradient is written in terms of its con- stituent blocks as
ac ac T and Eee. = Eo ab, ōbe'
ac ac T
En
ΣT
E1
ΣT
E1
Σ11.
==
:
(20)
Σ11
ΣT
ELL ΣL..
ST
Σn
Σr
Then, Equation (13) for the mean gradient becomes
Figure 4 presents the gradient covariance matrix for three- layer MLP with ReI.U activations and excluding biases. Once again, our thcoretical framework accurately captures the observed structure of the covariance matrix, both before training and after few optimization steps. Notably, the reduced symmetry of ReLU activations results in intricate covariance structure, characterized by significant correlations across layers. As demonstrated in Section 3.3, these inter-layer correlations persist regardless of the net- work's depth. In the following sections, we provide the theoretical framework used to derive the gradient covari- ance structures observed in Figures 2, 3, and 4.
H, (Ve-1 @ Vt) Pr
(21)
i4, Ve big
(22)
for all e 1,..., L. Equation (14) for the covariance of the gradient becomes
Eee = (Ve- 8 Vo) Eee (V V2))
(23)
E = V₁ Euv Vi
(24)
Eee Ve Em (V 8 V17)
(25)
for all f, 1,..., solutions Equations (21), (22), (23), (24), (25).
., T. In the next sections, provide
3.1. Theoretical results
In this section, we present solutions to Equations (13) and (14) that remain valid for all possible transformations G within given symmetry group. These solutions are derived for each of the three symmetry groups introduced in Section 2.2. Using the vectorization of parameters of Equation (16), the transformation in Equations (6), (7) can be rewritten in block-diagonal form
3.2. Odd activations (e.g. linear or Tanh)
Here, examine both linear and nonlinear odd activations (e.g., Tanh) together, as they yield the same structure for the mean and covariance of the gradient.
Theorem 3.1. Assume L>> 1. If the loss and distribution
of parameters are invariant for either groups Go Go then the mean gradient is equal to
Vo 8 V1
#g = 0
for = I,..., L
(26)
V1
iu = 0
for l I,...,,
(27)
==
(17)
(28)
Vz-1 3 Vz
Vz
where 2y is a vector of size dr. The covariance of the
6

---

Global curvature for second-order optimization of neural networks
theory
W
4:-23
W1
Se+02 We
18
-1e03
W1 We
W.
V5
Figure Gradient covariance matrix for three-layer MLP with ReLU activations. Similar Figure 3 but without bias averaging over 100000 models. In contrast to Tanh activations, ReLU activations retain correlations across layers. The covariance shows remarkably complex structure that is well predicted by theory.
gradient, given in terms of its constituent blocks, is equal to
constituent blocks, is described by Equations (118)-(142) in Appendix J.
En = 11 8 dd
(29)
Ex = th (Ide-1 8 Tde))
E - 2,.......,
(30)
The proof is also provided in Appendix J. We note that the number of undetermined factors is Old + L²). The number of factors is larger than the case of odd activations, thus revealing a more complex structure. We interpret this result as a consequence of the smaller size of the group Gş with respect G.. Gb (see Section 2.2).
EL = Ldz & 4
(31)
E1 = ( & Ia,)
=
(32)
Ex Φel Ide
=
for E = 1-
(33)
(34)
-T En Φ ldi
(35)
3.4. practical algorithm for second-order optimization
T
Σ1-1,L Idy.-1
11-1
(36)
We observed in Figures 2,3,4, that the covariance is very structured. In contrast to previous approximations (Martens & Grosse, 2015; Eschenhagen al., 2024), the matrix block-diagonal, although it is Kronecker-factorizee This structure may capture the most important variations in the loss landscape, therefore we use it for computing the second- order update in Equation (1). We break down the computa- tion of the second-order update into three steps: 1) Estimate the factors of the covariance; 2) Compute the factors of its inverse square root; Compute the matrix matrix-vector vector product between the inverse square root of the covariance and the gradient. We demonstrate that all three steps be per- formed efficiently, leveraging the convenient properties of the matrix. A detailed description of the complete proce- dure provided in Algorithm in the Appendix, using the simple case two-layer MLP with Tanh activation and no bias. Similar steps can be applied to other cases.
where dt. Φe are positive scalars, Φ1 E Rdh Φ1-1 Rd. are column vectors. 4 E xdo, br. E Rd Rdix o, E dd xdu, are positive-definte matrices and 1 Bdoxdaa is a matrix. All other terms are zero. is the commutation matrix (see Chapter 3.7 of Magnus Neudecker (2019)).
Proof is in Appendix We refer to the set (ăL, de, de, Φ Φ1.1. P1, L2, Pz, 11 as the factors, which undetermined. We provide algorithm for estimating the factors in Section 3,4 The number of undetermined factors is of order Odd + L). We note that the values of those factors should ensure that the covariance positive semi-definite.
3.3. Other activations (ReLU)
The gradient covariance, as defined in Equation (3), requires averaging over models initialized with different parame- ters. However, we observe that the number of factors of the covariance is much smaller than the total number of elements in the covariance matrix. As shown in Figures 2, 3, and 4, many elements of the covariance matrix share identical values. We leverage this observation to estimate
In the case of non-symmetric activations, we have the fol- lowing
Theorem 3.2. Assume L > L. If the loss and distribution of parameters are invariant for the group G. then the mean gradient is described by Equations (113)-(117) in Appendix The covariance of the gradient, given in terms of

---

Global curvature for second-order optimization of neural networks
the covariance efficiently, even from the gradient of a single model, by averaging over the corresponding pairs of param- eters. An example of this estimation process is provided for the case of two-layer MLP with Tanh activations and no bias. Similar equations can be derived for the case with bias and ReLU activations. We estimate the factors using the following equations:
more feasible, leverage the fact that any analytic func- tion of the covariance matrix retains the same structure as the covariance itself, as shown by Lemma 2.4 and Equa- tion (15). Therefore, the inverse square root Σ can be fully described by another set of factors, denoted by the vector e Instead of computing the full matrix > and its inverse square root Σー we compute the factors & from Φ. without ever constructing the large matrices. Appendix L provides detailed procedure for efficiently computing Dett In short, because the factors do not depend the layer ,,,,-,,
1
OL
OC
(1)),t
(37)
d1 1=1
aw,
OW1
ij
a
widths d1,...,d,.-1 we compute a surrogate covariance A where the layer widths are set to the smallest values al- lowed (de = for Tanh activations and de = 2 for ReLU (38) activations). The surrogate matrix A is significantly smaller than the full matrix E. For an MLP, the surrogate matrix has Oa + di + L22) elements, while the full matrix has [x, (de-1 + 1)]² elements. We then compute the inverse square root A 1 of the surrogate and derive the correspond- ing factors & jer Appendix I. provides examples comparing the full and surrogate matrices for small MLPs with Tanh or ReLU activations (Figures 11, 12, 13).
di Σ
0C
ac
(f2) d1 1=1
OW2
2 kj
ij
1
Oc
5x
(39)
d1
OW1
aw2
1=1
These are mattixx products of size equal to the neural network width, that can be computed efficiently using GPU.
To quantify the error in estimating the factors of global covariance by single model, we computed their correlation with a high-precision reference estimate. This reference was obtained by averaging across an ensemble of 10000 models, which serves as our ground truth. Table I shows that the correlation between the single-mode model estimate and the ground truth increases with layer width di, for all factors 21. V2 22. To further improve the estimates and reduce error, we apply momentum to Equations (37), (38), and (39) across training iterations. Appendix K provides additional details, along with an empirical analysis of the estimation error in presence of momentum (see Figure 10).
10a
Symo
KFAC
Sharrpo
GD
10-1
Adam
test loss
10-3
10-1
10-1
20
40
60
80
100
iterations
Table 1. Correlation between single-mods estimates of the factors and ground truth, varying the width of the hidden layer. Single model estimates improve with the layer width.
Figure 5. Optimization of a two-layer MLP with linear acti- vations. Test loss vs training iterations. The Symmetry- Optimizer (SymO) is compared with other first- and second-order optimizers.
Layer widths
1
4
(do, d1, d2)
(100, 10, 100)
0.67+0.05
0.38
0.30+0005
(100, 100, 100)
0.90=0.02
0.61++0.04
0.52+00.3
In the final step, we compute the product between the inverse square root of the covariance matrix Σ- 1 vector, as described in Equation (1). This computation can be carried out efficiently by using the estimated factors der without the need to explicitly compute Σ 1 We observe that each block of the covariance matrix (as well as its inverse square root) consists of Kronecker products, which enables efficient computation of matrix-vector products (see e.g. Martens & Grosse (2015)). For the case of a two-layer MLP with Tanh activations and no bias, the update is equal
(100, 1000, 100)
0.96-10.01
0.64
0.54-10.03
and the gradient
(100, 10000, 100)
0.97+0.01
0.65+0.03
0.56+10.03
We denote the flattened vector of all factors by Φb Af- ter obtaining estimates of the factors & by Equations (37), (38), (39), we could construct the covariance matrix Σ us- ing Theorem 3. and then compute its inverse square root Σ 1. However, this process becomes computationally ex- pensive for high-dimensional models. To make this step
8

---

Global curvature for second-order optimization of neural networks
to (see Equation (206) in Appendix L)
Grosse, 2015), Shampoo (Gupta et al., 2018) and our opti- mizer, which call SymO (Symmetry-basee Optimizer). For all optimizers, learning rate is set by grid search. For second-order optimizers, we additionally set a second hyperparameter by grid search: damping À for KFAC, ini- tialization e for Shampoo and decay parameter 3 for SymO. See Table 2 for the hyperparameter values used and Ap- pendix M for details. Figures 5, show the optimization trajectories of the five optimizers for, respectively, linear and Tanh activations. The symmetry-based optimizer SymO tends to converge faster than all other optimizers with nearly identical time per iteration.
ac
dC 2w2 2
(Wi)++1 - -WWi), - a
+
@
(40)
(W2) - -(W2), (vi (w/or
0L
0C
+ 2
OwT
2W2
(41)
Symo
KFAC
Shampoo
109
GD
5. Discussion
Adam
We introduce a novel theoretical framework for deriving the exact structure of the global curvature of neural networks, and for estimating a second-order optimizer. We provide preliminary evidence suggesting that the structure of the global curvature enhances convergence when applied as a preconditioner. However, a comprehensive analysis of the errors introduced by our approximation will be subject of future investigation. While we focus on the case of MLP and synthetic data, our framework is general and can be extended to other neural network architectures, including residual networks, convolutional networks, recurrent networks, and transformers. Future work will explore the symmetries and corresponding curvature structures for these architectures, and will evaluate the effectiveness of our approach on real- world data and larger models.
test 1o0s
10-1-
10
0
20
40
60
80
100
Iterations
Figure 6. Optimization of two-layer MIP with Tanh activa- tions. Test loss vs training iterations. The Symmetry-basee Op- timizer (SymO) is compared with other first and second order optimizers.
4. Experiments
We note that most prior studies on second-order optimiza- tion have relied on block-diagonal and Kronecker-factored upproximations of the curvature matrix (Martens & Grosse, 2015), which are exact only in models that are either linear in the input (Bernacchia et al., 2018) or in the parameters (Karakida Osawa, 2020). Our work demonstrates that the global curvature matrix is not inherently diaa--digonna,, but we do find that matrix blocks exhibit Kronecker- Kroneccer-factoree structure. We also note that the block-diagona terms of the update in Equations (40), (41) are similar to the update of Shampoo (Gupta et al., 2018). Thus, our findings partially justify the approximations used in previous studies and offer a more nuanced understanding of their validity.
We test second-order optimization on a two-layer MLP with- out bias and synthetic data, using either linear or Tanh ac- tivation. The model has layer width dg = 100, d1 70, d2 = 40, with total of 9800 parameters. The synthetic dataset consists of 5000 training and 5000 testing data points, where the input is sampled from a Gaussian distribu- tion with zero mean. The covariance matrix of the input is generated using random orthogonal eigenvectors (Mezzadri, 2007), and the eigenvalues are set on logarithmic grid between 10-5 and 10 We optimize the square loss, where the target output is provided by neural network (teacher) with identical architecture of the network to be optimized (student). The weights of both the teacher and student net- works are drawn randomly from a Gaussian distribution. We use full-batch optimization, where the gradient is computed over the entire training dataset in each iteration. Although full batch training is uncommon in neural networks, second- order optimizers usually benefit from large batches (Zhang et al., 2019; Anil et al., 2021), therefore we leave the study of mini-batch training to future work.
Our framework could have significant applications in fields that require second-ordes estimates for large models. For example, Bayesian deep learning relies on approximating the posterior over parameters by Gaussian distribution (Blundell et al., 2015; Lin et al., 2019). The covariance of this distribution is usually approximated using a diagonal or block-diagonal structure (Daxberger et al., 2021). Our work offers a method for efficiently computing the full covari- ance, which may lead to more accurate Bayesian posterior We compare five optimization algorithms: gradient descent (GD), Adam (Kingma & Ba, 2014), KFAC (Martens &
estimates.
9

---

Global curvature for second-order optimization of neural networks
Impact statement
Chen, A. M., Lu, H.-m., and Hecht-Nielsen, R. On the geometry of feedforward neural network error surfaces. Neural computation, 5(6):910-927, 1993.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Chen, F,, Kunin, D., Yamamura, A., and Ganguli, S. Stochastic collapse: How gradient noise attracts sgd dy- namics towards simpler subnetworks. Advances in Neural Information Processing Systems, 36, 2024.
References
Ainsworth, S. K., Hayase, J., and Srinivasa, S. Git be-baiin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209. 04886, 2022.
Choudhury, S., Tupitsa, N., Loizou, N., Horváth, S., Takáč, M., and Gorbunov, E. Remove that square root: A new efficient scale-invariant version of adagrad. In The Thirty- eighth Annual Conference on Neural Information Pro- cessing Systems, 2024.
Anil, R., Gupta, V., Koren, T., Regan, K., and Singer, Y. Scalable second order optimization for deep learning. arXiv, 2021.
Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, Laplace redux-efforteess bayesian deep learning. Advances in Neural Informa- tion Processing Systems, 34:20089-20103. 2021.
Arbel, M., Menegaux, R., and Wolinski, P. Rethinking Gauss-Newton for learning wver-parameterized models. Advances in Neural Information Processing Systems, 37, 2023.
Du, S. S., Hu, W., and Lee, J. D. Algorithmic regulariza- tion in learning deep homogeneous models: Layers are automatically balanced. Advances in neural information processing systems, 31, 2018.
Ba, J., Grosse, R., and Martens, J. Distributed second-order optimization using kronecker-factored approximations. In International Conference on Learning Represenaaionss 2017.
Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Bernacchia, A., Lengyel, M., and Hennequin, G. Exact nat- ural gradient in deep linear networks and its application to the nonlinear case. Advances in Neural Information Processing Systems. 31, 2018.
Entezari, R., Sedghi, H., Saukh, O., and Neyshabur, B. The role of permutatioon invariance in linear mode connectivity of neural networks. arXiv preprint arXiv:2 10.06296, 2021.
Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Eschenhagen, R., Immer, A., Turner, R., Schneider, F., and Hennig, P. Kronecker-factored approximate curvature for modern neural network architectures. Advances in Neural Information Processing Systems, 36, 2024.
Botev, A., Ritter, H., and Barber, D. Practical Newton optimisation for deep learning. In International Confer- ence on Machine Learning, pp. 557-565. 2017.
Flytzanis, E. Linear operators as measure preserving trans- formations. Annales scientifiques de I' l'Uiiversiéé de Cler- mont. Mathématiques, 65(15):63-75, 1977.
Bottou, L., Curtis, F. E., and Nocedal, J. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-311 2018.
Fukumizu, K. and Amari, S.-i. Local minima and plateaus in hierarchical structures of multilayer perceptrons. Neural networks, 13(3):317-327, 2000.
Buffelli, D., McGowan, J., Xu, W., Cioba, A.. Shiu, D.- s., Hennequin, G., and Bernacchia, A. Exact, tractable gauss-newton optimization in deep reversible architec- tures reveal poor generalization. In The Thirth-eighth Annual Conference on Neural Information Processing Systems, 2024.
Garcia, R., Freddi, F., Fotiadis, S., Li, M., Vakili, S., Bernacchia, A., and Hennequin, Fisher-Legendre (FishLeg) optimization of deep neural networks. In The Eleventh International Conference Learning Repre- sentations, 2023.
Cai, T., Gao, R., Hou, J., Chen, S., Wang, D., He, D., Zhang, Z., and Wang, L. Gram Gram-gauss-newer geuss-newton method: Learning overparameterized neural networks for regression prob- lems. arXiv, 2019.
Garipov, T., Izmailov, P., Podoprikhin, D., Vetrov, D. P., and Wilson, A. G. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances neural information processing systems, 31, 2018
10

---

Global curvature for second-order optimization of neural networks
Goldfarb, D., Ren, Y., and Bahamou, A. Practical quasi- newton methods for training deep neural networks. Ad- vances in Neural Information Processing Systems, 33: 2386,2396, 2020.
Lim, D., Putterman, T. M., Walters, R., Maron, H., and Jegelka, S. The empirical impact of neural parameter sym- metries, or lack thereof. arXiv preprint arX1v:2405.20231, 2024.
Grigsby, E.,, Lindsey, K., and Rolnick, D. Hidden symme- tries of relu networks. In International Conference Machine Learning, pp. 1 1174--1760. PMLR, 2023.
Lin, W. Khan, M. E., and Schmidt, M. Fast and sim- ple natural-gradicns variational inference with mixture of exponential-fame approximations. In International Con- ference on Machine Learning, pp. 3992-4002. PMLR, 2019.
Grosse, R. and Martens, J. A Kronecker-factoree approx- imate Fisher matrix for convolution layers. In Interna- tional Conference on Machine Learning, pp. 573-582. PMLR, 2016.
Lin, W., Dangel, F., Eschenhagen, R., Bae, J., Turner, R. E., and Makhzani, A. Can we remove the square-root in adaptive gradient methods? a second-order perspective. International Conference on Machine Learning, pp. 11949929973. PMLR, 2024.
Gupta, V., Koren, T., and Singer, Y. Sham- poo: Preconditioned Stochastic Tensor Optimization. arXiv:1802.09568 [cs, math, stat]. March 2018. URL http:/ /ttp:p//aarvv.ooo arg//bs// abs /8002.80.00 09568. arXiv: 1802.09568.
Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., Qin, Y., Xu, W., Lu, E., Yan, J., et al. Muon is scalable for 11m training. arXiv preprint arXiv:2502. 16982, 2025.
Hecht-Nielsen, R. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pp. 129-135. Elsevier, 1990.
Magnus, J. R. and Neudecker, H. Matrix differential calcu- lus with applications in statistics and econometrics. John Wiley & Sons, 2019.
Jordan, K., Sedghi, H., Saukh, O., Entezari, R., and Neyshabur, B. Repair: Renormalizing permuted acti- vations for interpolation repair. In The Eleventh Interna- tional Conference on Learning Representations, 2023.
Martens, and Grosse, R. Optimizing neural networks with Kronecker factorrd approximate curvature. In Interna- tional conference on machine learning, pp. 2218-2417. PMLR, 2015.
Karakida, R. and Osawa, K. Understanding approximate Fisher information for fast convergence of natural gradi- ent descent in wide neural networks. Advances in neural information processing systems, 33: 109100991111 18991119901 2020.
Martens, J., Ba, J., and Johnson, M. cronecker-factorae curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018
Kasimbeg, P., Schneider, F., Eschenhagen, R., Bae, J., Sas- try, C. S., Saroufim, M., FENG, B., Wright, L., Yang, E. Z., Nado, Z., et al. Accelerating neural network train- ing: An analysis of the algoperf competition. In The Thirteenth International Conference on Learning Repre- sentations, 2025.
Meng, Q. Zheng. S., Zhang, H., Chen, W., Ma, Z.-M., and Liu, T.-Y. G-sgd: Optimizing relu neural networks in its positively scale-invariant space. arXiv preprint 2018.
Kingma, D. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6900 2014.
Mezzadri, F. How to generate random matrices from the classical compact groups. Notices of the American Math- ematical Society, 54(5):592 604, 2007.
Köhler, J., Klein, L., and Noé, F. Equivariant flows: sam- pling configurations for multi-body systems with sym- metric energies. arXiv preprint arXiv:1910.00753 2019.
Morwani, D., Shapira, I., Vyas, N., Malach, E., Kakade, S., and Janson, L. A new perspective on shampoo's preconditioner. ICLR, 2025.
Korbit, M., Adeoye, A. D., Bemporad, A., and Zanon, M. Exact gauss-newton optimization for training deep neural networks. arXiv preprint arXiv:2405.1 14402, 2024.
Navon, A., Shamsian, A., Fetaya, E., Chechik, G., Dym, N., and Maron, H. Equivariant deep weight space align- ment. In Forty-first International Conference on Machine Learning, 2024.
Krantz, S. G. and Parks, H. R. Geometric integration theory. Springer Science & Business Media, 2008
Kunin, D., Sagastuy- Breena J., Ganguli, S., Yamins, D. L., and Tanaka, H. Neural mechanics: Symmetry and bro- ken conservation laws in deep learning dynamics. arXiv preprint 212:22200222.222.. 2020.
Neyshabur, B., Salakhutdinov, R. R., and Srebro, N. Path- sgd: Path-normalized optimization in deep neural net- works. Advances in neural information processing sys- tems, 28, 2015.

---

Global curvature for second-order optimization of neural networks
Schmidt, R. M., Schneider, F,, and Hennig, P. Descend- ing through Crowded Valley Benchmarking Deep Learning Optimizers. In Proceedings of the 38th Interna- tional Conference on Machine Learning. pp. 9367-9376. PMLR, July 2021. URL https::////......... http:///occeeedeeeee mlr. www.pressssse.eee saress/v133 htt11///cchmtt22//// 21a html. 2640-3498.
Zhang, G., Li, L.,, Nado, Z.. Martens, J., Sachdeva, S., Dahl, G., Shallue, C., and Grosse, R. B. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
ISSN:
Zhang, Y., Zhang, Z., Luo, T., and Xu, Z. Embedding principle of loss landscape of deep neural networks. Ad- vances in Neural Information Processing Systems, 34: 14848-14859, 2021.
Simoncini, Computational methods for linear matrix equations. siam REVIEW, 58(3):377-441, 2016.
Zhao, B., Dehmamy, N., Walters, R., and Yu, R. Symmetry teleportation for accelerated optimization. Advances in neural information processing systems, 35:16699-16690, 2022.
Simsek, B., Ged, F,, Jacot, A., Spadaro, F. Hongler, C., Gerstner, W., and Brea, Geometry the loss landscape in overparameterizee neural networks: Symmetries and invariances. In International Conference Machine Learning, pp. 9722-9732. PMLR, 2021.
Zhao, B., Ganev, I., Walters, R., Yu, R., and Dehmamy, N. Symmetries, flat minima, and the conserved quantities of gradient flow. In International Conference on Learning Representations, 2023a.
Stuart, J. L. and Weaver, J. R. Matrices that commute with permutation matrix. Linear Algebra and Its Applications, 150:255-265, 1991.
Zhao, B., Gower, R. M.. Walters, R., and Yu, Improving convergence and generalization using parameter symme- tries. arXiv preprint 2rXiv: /3404, 2023b.
Stuart, J. L. and Weaver, J. R. Diagonally scaled permu- tations and circulant matrices. Linear Algebra and Its Applications, 412:397-411. 1994.
Ziyin. Symmetry induces structure and constraint of learning, In Forty-first International Conference on Ma- chine Learning. 2024.
Tatro, N., Chen, P.-Y,, Das, P., Melnyk, I., Sattigeri, P., and Lai, R. Optimizing mode connectivity via neuron alignment. Advances Neural Information Processing Systems, 33:15300-15311, 2020.
Ziyin, L., Wang. M., Li, H., and Wu, L. Parameter symme- try and noise equilibrium of stochastic gradient descent. In The Thirty-eight Annual Conference on Neural Infor- mation Processing Systems, 2024.
Tieleman, Т. and Hinton, G. Rmsprop: Divide the gradient by running average of its recent magnitude. coursera: Neural networks for machine learning. COURSERA Neu- ral Networks Mach. Learn, 17, 2012.
Titsias, M. Optimal preconditioning and fisher adaptive langevin sampling. Advances in Neural Information Pro- cessing Systems, 36, 2024.
Yang. Y., Hodgkinson,
L.. Theisen, R., Zou, J., Gonzalez, E., Ramchandran, K., and Mahoney, M. W. Taxonomiz- ing local versus global structure in neural network loss landscapes. Advances in Neural Information Processing Systems, 34:18722-18733, 2021.
Yao, Z., Gholami, A., Shen, S., Mustafa, M., Keutzer, K., and Mahoney, M. Adahessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAI conference artificial intelligence, volume 35, pp. 10665-10673, 2021.
Yu, Y., Xia, R., Ma, Q. Lengyel, M., and Hennequin, G. Second-orde heward-mode optimization of recurrent neural networks for neuroscience. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
12

---

Global curvature for second-order optimization of neural networks
A. Related work
Numerous efforts have been made to make second-order optimization feasible for deep neural networks. Like our work, prior studies have explored rrchitectrresppcific curvature for various models, such as MLPs (Martens & Grosse, 2015), ConvNets (Grosse & Martens, 2016), recurrent networks (Martens et al., 2018), and transformers (Eschenhagen et al.. 2024). However, these approaches rely block-diagonal and pronecker-factoree approximations, which lack theoretical guarantees and are exact only in limited scenarios, such as when networks are linear either in their input (Bernacchia et al., 2018) or their parameters (Karakida & Osawa, 2020). Recent research has aimed to compute curvature exactly, without relying on block-diagonal approximations, but such methods are restricted to small-scale models (Cai et al., 2019; Arbel et al., 2023; Korbit et al., 2024) or to reversible networks (Buffelli et al., 2024). Crucially, none of these studies leverage the inherent symmetries present in neural network architectures. In contrast, our work introduces a novel framework that derives the exact structure of curvature matrices by utilizing model symmetries for the first time. This approach is not limited by model size and can be scaled to accommodate any large neural network with symmetries, providing robust and widely applicable solution to second-orde optimization.
Neural networks exhibit a range of symmetries that can be broadly categorized into discrete (e.g., permutations and sign flips, Chen et al. (1993)) and continuous (e.g., scaling, translations, and rotations, Kunin et al. (2020); Ziyin (2024)). While continuous symmetries have been extensively studied for their potential to improve optimization (Neyshabur et al., 2015; Meng et al., 2018; Zhao al., 2022; 2023b), discrete symmetries have received comparatively less attention. Our work is the first leverage discrete symmetries specifically to enhance optimization. Prior research has employed discrete symmetries in other contexts, such model merging (Tatro et al., 2020; Entezari et al., 2021; Ainsworth al., 2022: Jordan et al., 2023; Navon et al., 2024) and analyzing the structure and stability of fixed points in the loss landscape (Fukumizu & Amari, 2000; Simsek et al., 2021; Zhang et al., 2021). Despite these applications, their direct use in optimization strategies remains unexplored, and our work fills this critical gap. Additionally, the study of optimization dynamics has incorporated both discrete (Chen et al., 2024; Ziyin, 2024) and continuous symmetries (Du et al., 2018; Kunin et al., 2020; Zhao et al., 2023a; Ziyin al., 2024), highlighting the broader interest in understanding how symmetries influence training. By focusing on how discrete symmetries may improve optimization, we introduce a novel perspective that complements existing approaches.
A key distinction of our approach compared to prior work on second-orde optimization and symmetries is our focus on global rather than local curvature, also termed position-independent versus position-dependent curvature (Titsias, 2024), Existing studies predominantly estimate local curvature, which depends on a specific set of model parameters and captures the loss landscape's 's behavior around a single point. Only two previous works, Yang et al. (2021) and Yao et al. (2021), have explored global curvature, though with different interpretations. Yang et al. (2021) characterize global structure through mode connectivity (Garipov et al., 2018), while Yao al. (2021) define via an exponential moving average of curvature. Instead, we define global curvature by an average of the curvature matrix over an ensemble of model parameters, and we show that it can be estimated efficiently.
Similar to our study, Kunin et al. (2020) explores symmetry-induces constraints on curvature. However, their work is restricted to analyzing local curvature, which limits them to continuous symmetries and yields fewer constraints compared to the stronger and more numerous constraints that emerge when analyzing the effect of discrete symmetries on global curvature. Furthermore, global curvature estimation can offer more robust insights for optimization, as it is not tied to a specific parameter configuration. The only previous study to incorporate global curvature is Titsias (2024), which investigates position-indeeennent preconditioning. However, their approach does not integrate symmetries into the computation, whereas our approach uniquely combines global curvature estimation with symmetry constraints.
B. Quadratic loss
In this section, we study a toy example to build some intuition into the dynamics of parameters when optimizing a loss function using the update introduced in Section 2.1, Equation (I). Given a parameter vector E R..", a constant and positive definite Hessian matrix H E RR*P, we consider the quadratic loss function
((0) - (0- -
((0--')"H(0-0))
(42)
The global minimum of this loss is 0" When optimizing this loss by gradient descent, the path taken in the space of parameters depends strongly on the Hessian matrix H. If H is a scalar matrix, H x I, then the loss is isotoopic and the parameters follow a straight line to the minimum. Instead, if H is badly conditioned, then some directions in parameter
13

---

Global curvature for second-order optimization of neural networks
space much more curved than others, and gradient descent takes convoluted path, which may significantly slow down convergence. Second-order optimization straightens the path and accelerates convergence.
The gradient of the quadratic loss is equal to VL = ( 0*). We are given an ensemble of models following normal distribution p(0) O [4g. Σg). where Ne and Dg are, respectively. the mean and covariance of the parameters. Since
the gradients are a linear function the parameters, gradients also follow a normal distribution, VL ~ NN Eg), where the mean and covariance are given by
Fg H - -0) and Eg = HIgH
(43)
We define the optimization dynamics by the continuous flow
d0
dt - -- -0 -0 (-)
(44)
This is the continuous-timme version of the second-order optimization update considered in Section 2.1, Equation (1), with In our derivation, we keep generic value of y and we later argue in favor of the choice 7
Since the update of Equation (44) is linear in the parameters, the distribution of parameters remains Gaussian at all time steps, and so does the distribution of gradients. However, the mean and covariance of both distributions change in time. The dynamics of the mean obtained by averaging Equation (44), while the dynamics of the covariance is obtained by averaging the same equation multiplied by the parameter vector. We find
dpo - - 0²)
(45)
dt
diee - (Eg + EgH
(46)
dl
Using the identity Eg HEgH, and assuming that H and Eg commute, obtain
dpo
aa 2222 (n 0*)
(47)
dl
džie
dt 2ax H 1+27 D147
(48)
We note that these equations simplify by setting = 2. because the dependence on H disappears. In that case, the equation for the covariance becomes die 2a 1-/2 We further assume an isotropic initialization parameters, Eg x I dt at 0. The isotropic initialization a natural choice in absence of any structured prior. In that case, the distribution remains isotropic all times, Ee sT, where the scalar time-dependent. This also ensures commutation of Eg and H. Then, the optimization dynamics is described by
d1tg dt
a
(Hg 0*)
(49)
V8
dt = 2008
(50)
In this case, the mean of the distribution takes a straight path to the minimum 0" When considering optimization in discrete time, this implies that a large learning rate can be used and convergence is faster.
For constant learning rate a, we show that the continuous dynamics converge in finite time. The solutions of Equations (59) are
14g (1 (1- al
/ --m6) (n0 0")) 0"
(51)
14

---

Global curvature for second-order optimization of neural networks
(v50 od)²
(52)
where Mo and 50 are the initial values at = 0, and the dynamics stops at 1 L when all models have converged the minimum 0". We note that discretization of the dynamics in Equations (49). has large error near convergence, even if the learning rate a is small. Therefore, we expect that discrete parameter updates would differ significantly from the predicted solution in Equations (51).52). To avoid this problem, consider exponentially decaying learning rate
3
(53)
with 0 < B < 1. With this choice of the learning rate, and assuming that on and 3 are chosen to satisfy / solutions of Equations (59),(50) are equal to
og the log( })
M4 = (10 -
o
(54)
Ox 3t (3)
$
(55)
In this case well, all models converge quickly (exponentially fast) to the minimum of the loss. With the choice of an exponentially decaying learning rate, the discretization error of the dynamies in Equations (49).(50) remains small, therefore expect the discrete updates to follow closely the solution in Equations (55). The assumed equality / log( is satisfied, for example, when so 1, ag is small and B - orp- We find the relationship B ~ ox to hold in our experiments, see Table 2.
C. Existence of invariant measures
In this section, we discuss the conditions for a transformation G to accept an invariant probability density and, given a suitable G, what can we say about the mean and covariance of the invariant probability. Those conditions are identical to those of the mean and covariance of the loss gradient, discussed in the main text, therefore the results of this section can be used to analyze the properties of the gradient statistics as well. While in the main text we assume that G' is orthogonal, the results of this section do not use that assumption and are valid also for non-orthogonal G.
necessary condition for the existence of an invariant probability density is that the transformation is preumee-reeeriigg. This can be shown by computing the normalization condition. Denoting the probability density p(0) and the linear transformation G, invariance is given by p(G)) = p(0). The domain of p also invariant. Using the change of variable GO', d0 [de(()) the normalization condition can be written as
= f do p(0) |det(G)) f f p(GO') [dc((G)] / do' p(0') [det(G)]
(56)
Therefore, (det(G')| must be equal to one, which means that the transformation must be volume-preserringg For example, the distribution 8(0). the Dirac delta function, is invariant for any volume- volume-preserving transformation, since 5(G0) [det(G)) 60). However, besides volume preservation, the transformation G' must satisfy stronger requirements to accept non-trivial invariant distributions.
useful example is provided by rescaling transformations, in which a parameter wi is multiplied by a constant À and another parameter W2 is multiplied by its inverse 1/A. Rescaling is volume-preserviee and keeps ReLU neural networks invariant (Neyshabur et al., 2015). However, besides the delta function, is not possible for any other probability distribution of Wi and u2 to be invariant for rescaling, because it would have be constant along the one-dimensional curves with constant Wyl1,, which are unbounded.
It was shown in Flytzanis (1977) that, to accept an invariant distribution whose domain spans the entire space, the linear transformation must be similar to an orthogonal transformation. Technically speaking, that is equivalent to the requirement that all eigenvalues of lie on the unit circle of the complex plane. That condition implies volume- cauume--eeeeration. but is much stronger. However, transformations that are similar to orthogonal transformations are not necessarily orthogonal (see Lemma C.2).
15

---

Global curvature for second-order optimization of neural networks
We provide proof that valid for the case of distributions with finite and positive definite covariance. The proof includes procedure for constructing invariant covariance Y., given the transformation G. First, Lemma C.1 shows that the covariance of an invariant distribution must be invariant upon a congruent transformation, akin to Lyapunov equation (Simoncini, 2016). Lemma C.2 provides conditions on the transformation G' to accept an invariant distribution, and lemma C.3 shows how to construct the covariance an invariant distribution.
Lemma C.1. Assume that a probability density is invariant for a linear volume-presse preserving transformation, =(0)) and it has finite mean and covariance. The mean and covariance E of the distribution must satisfy the following equations
/ = Gp Σ = GYO
(57)
(58)
We note that the equation for μ is eigenvalue equation, while the equation for Σ is congruent transformation, also known as homogeneous discrete-time Lyapunov equation (Simoncini, 2016).
Proof. The result can be obtained by change of variable 0 = GO', d0 [det(G))dO0, using volume preservation, [det G) 1 and the invariance of the probability density:
1 = / d0 0 p(0) = [de((G]] / do' Go' p(GO') = G / do' o' p(0') = Gp
(59)
Σ = d0 T p() mpd = [der(G]] / do' Go T r(G'') Gpp dT
(60)
G /
d6' T p(0') mni" Gf = =GEG
(61)
Lemma C.2. the homogeneous discrete-time Lyapunov equation
=
(62)
has a positive definite solution Σ, then G must be similar to an orthogonal matrix, thus it can be diagonalized and its eigenvalues must lie on the unit circle of the complex plane.
Proof. Since V is positive definite, its square root and inverse exist and are unique. We define the matrix U as
U= -= GE
G F
(63)
Substituting this expression into the Lyapunov equation (62), we obtain
Σ == A -
UUT =I
(64)
Therefore, U must be an orthogonal matrix, and G is similar to U, therefore it must be diagonalizable and have all eigenvalues in the unit circle of the complex plane.
For example, a Gaussian distribution with zero mean and covariance equal to Σ is invariant for G - - for fixed E and any orthogonal U. However, other distributions, not necessarily Gaussian, may be invariant for the same group of transformations. for example any distribution with the same level sets. Many other distributions may be invariant only for small subset of U values, for example permutations, reflections rotations with respect specific axis.
Lemma C.3. Given transformation G that is similar an orthogonal transformation, with real Jordan form given by
G = -BV-
(65)
16

---

Global curvature for second-order optimization of neural networks
where B is orthogonal and block-diagonal, with a +1 or diagonal eniry for each real eigenvalue of G and 2 2 rotation block for each pair of complex conjugate eigenvalues of G. The columns of the real matrix V are the real and imaginary parts of all the right eigenvectors of G.
Then, a solution Σ of the Lyapunov equation (62) is equal to
= VDVT
(66)
G has distinct eigenvalues, the matrix D is diagonal with 1 x block, with an arbitrary positive vulue, for each real eigenvalue of G (either +1 or -1), and a x 2 diagonal scalar matrix block, with an arbitrary positive value, for each complex conjugate pair of eigenvalues G. G.
If G has degenerate eigenvalues, D is block-diagone with arbitrary p μ positive definite matrix block for each real eigenvalue of G, where is its multiplicity, and 2μ x 2p block of the form Dy @ 12 where Dj is an arbitrary positive definite matrix of size P X Ps, for each complex conjugate pair of eigenvalues of G, where p is their multiplicity.
Proof. Using the Jordan form (65), we rewrite the Lyapunov equation (62) as
Σ = EV T BTVT
V-' = -V- BT
BT
(67)
We define the positive definite matrix
D D=V-
E= VDVT
(68)
then the above equation is rewritten as
D BDB
(69)
Since B is orthogonal, this equation implies that D and B commute. Therefore, for either +1 or -1 diagonal entries in B, the matrix D has single scalar positive value the diagonal. For each complex conjugate pair of eigenvalues e=ió the matrix B has 2 x 2 block
cos ( -in( ())
cos
sin ②)
)
(70)
(Ø)
The only symmetric matrix that commutes with 2-dimensional rotation is the identity matrix, therefore D must have 2 x 2 positive scalar matrix block.
The case which has degenerate cigenvalues is straightforward generalization of the case with distinct eigenvalues. In that case: when has an identity block Imo where p4 is the multiplicity of eigenvalue 1, then D has a corresponding block with arbitrary positive definite matrix Ho of size Mo No: when B has a minus identity block -Ip, where frr the multiplicity of eigenvalue 11, then D has corresponding arbitrary positive definite matrix D Dr,, of size urx pzz When B has blocks of size 2ps 2p-
(os( sin(
(in( cos( ( )
Tpo 69
(71)
where ns is the multiplicity of conjugate eigenvalue pairs of angle db then has a corresponding block of size 2 x 2p of the form Dps 12 where Dr., an arbitrary positive definite matrix of size Pa pox
D. Invariant initialization of parameters
In this section, we discuss which of the most common parameter inizialization routines used in deep learning are consistent with Assumption 2.1. In other words, we ask whether the probability distribution used to initialize parameters is invariant for the given symmetry group. As described in Section the transformation in Equations (6), (7) can be rewritten in terms of a
17

---

Global curvature for second-order optimization of neural networks
single vector of parameters encompassing all layers, concatenating and vectorizing all tensors. Then, the transformation can be written as 0 GO, where the vector of all parameters is equal
Vec (W1, bi, ......, WL, Py)
(72)
and the transformation has a block-diagonal structure
Vo 0 V1
V1
G
(73)
VL V
V4.
In the next sections, we discuss whether different initialization routines are invariant for the transformation G, when G belongs the three different groups considered in this study, Go Gb, Gs.
D.1. Gaussian (normal) initialization
In this case, the probability distribution used to initialize parameters is given by the standard multivariate Gaussian
_ 2o
Pu (0) = (2TT²)- & exp
(74)
where o² is the variance of the distribution. This distribution is invariant for any orthogonal transformation 0 -> UO, where U is any orthogonal matrix. It is straightforward to verify that the transformation given by Equation (73) is orthogonal, when all V for 1,..., .....,, L are all orthogonal. Therefore, Gaussian initialization is invariant for all Go, Gb, Gs.
D.2. Orthogonal initialization
In this case, each tensor in (W1, bi....... ......,,, Wz, bz.) orthogonal matrices (Mezzadri, 2007), by a QR decomposition of a random Gaussian matrix (c.g. generated as in Section D.I), and then dropping the excess rows or columns. An additional step of flipping the sign of some rows and columns necessary for technical reasons, to make the distribution uniform (see Mezzadri (2007)).
is generated according to uniform distribution in the compact set of
We recall that Vi are orthogonal for all É z I...., L, thus V V V.V." lác. The transformation G in Equation (73) corresponds to transforming each weight matrix W/ V.WVZ, When de- de. then W, W.TWe Id,_, implies
w/" W: = VV-.WV V.
Ve-1V2 Ida-
(75)
Similarly, when de-1 2 de, then WoWT" = Laz implics
Www V. V/ VVVV = =
(76)
Furthermore, the transformation keeps the uniformity of the group (Mezzadri, 2007), therefore the orthogonal initialization also invariant for all Go Gb, G..
D.3. Uniform initialization
In this case, each parameter initialized by drawing number from a uniform distribution U(a, 6), where a and are, respectively, the lower and upper bound of the distribution.
It is straightforward to verify that this distribution is not invariant for orthogonal transformations, therefore ruling out the group G.. In fact, the domain of the distribution is hypercube, which is not invariant for rotations. However, any permutation of the parameters leaves the distribution invariant, therefore the distribution is invariant for the group G.- If the lower and upper bound are opposite, a = -b, then the distribution also invariant for sign flips and therefore invariant for the group Gb.
18

---

Global curvature for second-order optimization of neural networks
D.4. Sparse initialization
In this case, random fraction of parameters is zero, and the others are initialized following Gaussian distribution. This distribution is not invariant for orthogonal transformation, because rotations do not preserve sparsity in general. Therefore the group Go is again ruled out. However, the distribution is invariant for both permutations and sign flips, and thus for the groups Gb, G2.
D.5. Layer-dependent initialization
In some circumstances, different distributions used to initialize distributions in different layers. For example, layer- dependent variances when using the normal distribution, or layer-dependent lower and upper bounds when using the uniform distribution. All results of previous sections apply to this case as well, because the transformation of Equation (73) is block-diagonal and therefore applies separately to different layers. In other words, parameters of different layers are not mixed Equation (73).
E. Proof of Theorem 2.2
By assumption, the Jacobian matrix of is nonsingular except set of measure zero. Then, use the Coarea formula for computing the probability density after the mapping (Krantz & Parks, 2008). Since the dimension of the image and preimage are equal, then the set of preimages is discrete. Therefore, the probability of 0,+1 at step / + is given by the change of variable formula
Pr+1 (02+1)
Σ Pr(O:) det (Ju, (0.))l
(77)
0, (0.11)
where Ju is the Jacobian matrix of u, and (1(6++1) is the set of pre-images of 0i+1, namely all values O₂ such that u, (0:) = 0.+1. Given the assumption of equivariance, we have that u, (GO.) = G0,+1 and therefore GO, E (GB+.11 for all 0, 14)+ On the other hand, for all θ (G,0G.++1) since G is invertible and using again equivariance, must be u, (G-10) 10) 0c+1, and therefore 0-10 E (c(Or+1). Those two observations together imply that the elements of Ui (G0++1) equal the elements of Ui (0++1) multiplied by G. Therefore, we can compute the transformed probability density as
(G0
Σ
(r)) det (Ju, (0.)))
E
Pc(GO.) det (Ju, (GO.))
(78)
0, (GO
(M4(0))))
use invariance of the probability at step t, namely pr (GO,) = (O:). and we note that equivariance and chain rule imply Ju. (GO,) Ju, (0.)G and therefore det (GO.)) det(Ju., (0, det((G) Finally, since G is volume-preseree
(de((G) 1, then
Pc+1(G0r+1)
Σ
(0.) det (Ju, (0,)) = (1+1(0
(79)
(,
F. Equivariance of optimizers
In this section, we ask whether common optimizers used in deep learning are equivariant for the three groups considered in the main text. All results are derived under the assumption that the loss is invariant. We show that both gradient descent (with without momentum) and the second-order update of Equation (1) are equivariant for all groups Go Gb Gg. while Adam optimizer is equivariant for Gb. Gg, but not Go-
The gradient descent update is
0,
(80)
First, we note that the gradient is equivariant when the linear transformation G is orthogonal. Given the loss L()) and the change of variable 0' = GO, we define the transformed loss Ł(0) = L(0'). Using the chain rule, we have
="
VL(GO) - VV''
(81)
19

---

Global curvature for second-order optimization of neural networks
by assumption, the loss invariant, L. and G orthogonal, GT = G, therefore
VL(GO) = = ( = G (
u(GØ) = Gu(0)
(82)
We note that the gradient is not equivariant in general, however it is equivariant when the transformation G is orthogonal. The extension to momentun is straightforward, although the update depends on all previous values of the parameters. In that case
0: = -u (.o............) = 00-1--0--00
- B) Σ
(83)
t'=0
Similar to the case of gradient descent, the update equivariant for any orthogonal transformation (acting at all steps)
u. (f. G0o}) = Gu, (f0)-1..... Out)
(84)
Since all groups considered this work consists of orthogonal transformations, gradient descent (with or without momentum) equivariant for all groups Go Gb. G..
In the of Adam optimizer, the update is equal to (we summarize time-dependent scalars in a²)
8
0: = ur ((0l-...... Oot) 01-1 ag
(85)
Er t-1 1 ((,))) o VL(0,) +
Because of the denominator, this update is not equivariant for all orthogonal transformations. However, it is equivariant for signed permutations. We note that the numerator changes sign while the denominator does not, therefore the update is equivariant for sign flips. Furthermore, the update of parameter depends only on the gradient of that parameter, therefore the update is equivariant for permutations. Therefore, Adam optimizer is equivariant for Gb, G., but not G.-
In the case of the second-order update of Equation (I), we note that Lemma 2.4 applies at = 0 under assumption 2.1. Since the inverse and square root are analytic functions of positive definite matrices, then Mg Equation (1) satisfies Equation (15), provided that Do is positive definite. By induction, we can iteratively apply Theorem 2.2 to prove that the update is equivariant at all steps t, for all groups Go, Gb. Ga.
G. Proof of Lemma 2.4
Given the loss C(0) and the change of variable 6' GO, we define the transformed loss [(0) = -(0'). Using the chain rule, we have
= =Tv())='
VL(GO) = GT -
(86)
by assumption, the loss is invariant, L = Ê, and G is orthogonal, GT G, therefore
VL(GO) = GT VL(()= - G ()
(87)
We note that the gradient is not equivariant in general, however it is equivariant when the transformation G is orthogonal. Similarly, it follows that
V[(G))V/(G0)"' v G"
(88)
Using the invariance of the probability distribution, p(G0) p(0), and volume-preserving change of variable 0' GO,
with [d)) = 1, it is easy to show that
VC) do p(0)VL(G)) / do p(GB)VL(GG) |det(G))=) / de' p(0') (C
(89)
Using similar steps, obtain
E (2(G))V2(G0)) (C())
(90)
20

---

Global curvature for second-order optimization of neural networks
Averaging Equations (87), it follows that
μ Gp and Σ = GEG"
(91)
Furthermore, since G' is orthogonal (G G = 1), any power of Σ satisfies the same equation
n (GEG*) ... (GIG*) GE"
(92)
and does any analytic function of >, because analytic functions defined by power series.
theory
13:01
1.601
W
11:01
W
1.0011
5.002
1,0
b
00:130
1:0c100
3.0e c2
W3
Wo
1.0000 61
b2
13em
bo
Wi
W2
W,
bi
Ws
ba
2.0
2,00
Figure Global Hessian matrix two-layer MLP with Tanh activations. Same as Figure but shows the global Hessian instead the gradient covariance. We highlight that theory predicts the overall structure, rather than the specific numerical values.
A
B
= 20
theory
Wi
Wi
Vs.
br
Vb
bx
1gg
53
Figure 8. Global Hessian matrix for three-layer MLP with Tanh activations. Same as Figure 3 but shows the global Hessian instead of the gradient covariance. We highlight that theory predicts the overall structure, rather than the specifie numerical values.
H. Hessian matrices
In this section, we show that results similar to Theorem 2.4 and Figures 2, 3, 4, hold not only for the gradient covariance, but also for the Hessian. Given the loss L(0) and the change of variable 0' GO, we define the transformed loss Ł(0) L(0"). Applying the chain rule twice, noting that the transformation linear, we have
2 P 3 G"
(93)
by assumption, the loss invariant, Σ, and G is orthogonal, GT G, therefore
VCC(G)) -
V21(0)=-1= -
(94)
21

---

Global curvature for second-order optimization of neural networks
=
theory
W
Wi
1e-02
00:00 10
0.:00 V4
1-81
22462
W.
W.
W1
We
W.
Vb
V6
Figure Global Hessian matrix three-layer MLP with ReLU activations. Same as Figure but shows the global Hessian instead of the gradient covariance. We highlight that theory predicts the overall structure, rather than the specific numerical values.
We define the averaged (or global) Hessian as
H = E V
(95)
Using steps similar to Equation (89), it is straightforward to show that the global Hessian satisfies Equations (14), (15), namely
H GHG and =(H)= = Gf
(96)
for any analytic function f of the global Hessian matrix. Therefore, the statements of Lemma 2.4 for the gradient covariance hold identically for the global Hessian.
Figures 7, 8, show similar results as Figures 2, 3, 4, but for the global Hessian instead of the gradient covariance. The figures show that the global Hessian has structure similar the gradient covariance, which is predicted equally well by our theory. Again, theory predicts the structure, rather than the specific values of the matrix. The specific values need to be estimated, explained in Section 3.4.
I. Proof of Theorem 3.1
We have either Ve E O(de,, the set of orthogonal matrices, or Vj E B(dc), the set signed permutations, for -= 1,..., L-1,
while Vo = Idg and Vr. = Idz. The gradient mean of the bias given by. Equation (22), equal
Ve B
(97)
First note that, since Vz = Tdg. then و is not zero and remains undetermined. Besides that special case, this equation has non-zero solution for jág only if all transformations within the group have common eigenvector corresponding to an cigenvalue equal to one. However, no such cigenvector exists for either groups: for orthogonal matrices, cigenvectors with eigenvalue equal to one correspond to axes rotation and/or reflection, that can point at any direction for different orthogonal tranformations; For signed permutations, flipping simultaneously the sign one row and the same column of a matrix translates to flipping the sign the corresponding component its eigenvector, Therefore signed permutations also do have a common eigenvector and jag must equal to zero.
The mean gradient of the weights is given by Equation (21), equal
b (V:- Vi) be
(98)
This equation has a non-zero solution for μ, only all possible products V within group have a common eigenvector corresponding to an eigenvalue equal to one. However, for the same reasons described above, no such eigenvector exists for either groups and therefore it must be μ4 =
22

---

Global curvature for second-order optimization of neural networks
Now turn the covariance of the gradient. The bias-bias covariance must satisfy Equation (24), equal
Eee Ve Eee Vi
(99)
Since Eer = DT e, we only need to study the cases 8 S f. First, we note that it must be Σe,e- 0 for all é o f", because Equation (99) must hold for all orthogonal or signed permutation matrices Ve, Ve, but that is impossible. For example, setting V2 = Iap- there no matrix Ee, satisfying Se = Vi Ee,e e for all possible Ve in the group. Next, we consider the case f = L'. We note that, since Vz = Idz, then ΣLL remains undetermined. For < L, we multiply Equation (99) by Vi and use its orthogonality (V² V2= = Laz). we obtain
Ex V2 = Ve Eee
(100)
The only matrix that commutes with all orthogonal or signed permutation matrices is scalar matrix (Stuart & Weaver, 1994), therefore we must have
Lee de Ide
(101)
for some scalar value de. The weight-weight covariance must satisfy Equation (23), equal to
Ee = (V2-1 @ Vo) Eee (V Vi7)
(102)
Since Ege- = e. we only need to study the cases & < f", First, we study the case & = l' and E * 1, L. In that case, using the orthogonality of Ve, Ve,, we rewrite Equation (102) as
Eu (Vz- Vi)= = (V2-1 & Vo) En
(103)
Similar to the case of the bias, the only matrix that commutes with all products of orthogonal signed permutation matrices is scalar matrix, therefore it must be
Eye Øe (Ide-, & Ia,) for some scalar value de. For 1, L, we have the following two equations
(104)
En (La. ② V1) E11 (Lao @ vi)
and EL = (VL-1 2 Laz) ELL (V @ Laz)
(105)
Using similar arguments as above, since these equations must be valid for all orthogonal or signed permutation matrices Vi and Vi-1, we must have
En 21 & Ld, and ELL = -Ide- 8 41
(106)
for some matrices 21, Lz. Next, we study the case = 1. We use the commutation matrix K, which is symmetric permutation, thus K = KT, K²2 = I. The commutation matrix exchanges the factors in a Kronecker product: K(A @ B)K- = B 8 A (see Chapter 3.7 of Magnus& & Neudecker (2019)). We rewrite Equation (102) as
Eee+ = (Ve-1 @ Vi) KeerK K (VoT @ V?)
(107)
and solve for Σe.4+1 K instead of Σe,4+1. This equation must be valid for all orthogonal or signed permutation matrices V (f (f=1) 1...., . I. - 1). We show that, if either f > 1 or é < I - 1, then Ee.+1 K = 0 + Σg,41 = 0. In the former case, we set V/2 = Lic. Ve+1 = Ide.r then there is no finite matrix Σe, K that satisfies Er,+1K = (V--1 & La,) Ee,441 K for all Vi-1 in the group. In the latter case, we set Vi-1 = Ia., and V/ = I then there is no finite matrix Ex,411 K that satisfies
Ee,41KK= = E (VTT @ Ia,) for all Ve41 in the group. If = 2, Equation (107) is equal to
Σ12K (Ja, @ V1) EieK (la, @ V/²)
(108)
23

---

Global curvature for second-order optimization of neural networks
that has the finite solution
E1² = ( @ Ia,)
E - ( X Lt, K
(109)
for some matrix 1/. Lastly, we study the case l' + that case, we can set Ve-1 Idi_.. Ve Iap. and either
Vi Taz or V Tap-, in Equation (102) show that Ee,e" must equal to zero.
The bias-weight covariance must satisfy Equation (25), equal
Eee = Ve Eee (V-18V)
(110)
We note that, in this case, Eee * Ef therefore must consider all pairs of values f, f. For l' be Eee O. In that case, we Ve Lif and there is finite matrix Ee satisfying Eee Σe, (V² @ Laz) for all + and l' < L, we set V2 = ld and there no finite matrix E2,441
1, it must
Vi-1 in the group, thus Eee = O. For l' satisfying Ee,441 = Ee,441 (Ia, @ Ve+1) for all Ve+1 in the group, thus Σe,e+1 O. For the special cases l = l" 1 and = = L, we rewrite Equation (1 (110) as, respectively
Σn Vi Σn (Iz, @ VT)
and =1.11,1=Vi Σ,.1,,r (V & ILr.)
(111)
Since these equations must be valid for all orthogonal or signed permutation matrices V and V2-1 we must have
Σn = Φ1 @ Ia, and Σ1-1L = Idy.- @ ΦL-1
-T
(112)
For some column vectors Φ, and Φ -- For other values of E, f",, we set V₂ = Id, in Equation (110) and note that there is no matrix satisfying Enger Σe, (VT-, @ VT) for all - Ve in the group, therefore it must be Σe,e 0.
J. Non-symmetric activations (ReLU) For the mean of the bias gradients, we have
3. Le
8=
(113)
fr, 2L
(114)
where āة RRd, 2e R remain undetermined, dp + L unknowns. For the weights
11 20 1d; Vec (1d,
(115)
f4g <e( 1dt) Vec (1d, 1 I
2_.....,
(116)
11 Ldz Et. Vec (21.1T1.1))
(117)
where 2g E RR Æy d2 E R remain undetermined, do dr + L - unknowns. In total, for the mean gradient have do + 2d,, 2L 3 unknowns.
For covariance of the biases, since Eeg-- =
3x4 Et we only need to study the cases & < We have
((2)
En = 6t La +
=(1)
1 17
(= L-1
(118)
dt
E = Te
1 17
-1 L- l<b
(119)
Vdede:
Ee
L....,, 1-1
(120)
Vdle
E1
(121)
where LL is a symmetric matrix of size d, x dr, We E Rd der 3(0) 5(2)
.01
R. remain undetermined, a total of
dy(d, + 12/2 + (I. 1)d, + I(I. 1)/2 + 2(T 1) unknowns.
24

---

Global curvature for second-order optimization of neural networks
For the weights, since Eee- Eier only need study the cases S f" In layers é = 2, .,..., L 1, we have
of
Eu 22 (Ide 2 la.) +
de 2 la 12e) + oi (Ide
lde de @ (otee) dede-1
(Ids I 60 Iae l
de
de-1
(122)
For l = 2...., ......,L L- 2, we have
58 de de (Idi.n - @ 1a, 12) K
Ee0+1= = dde
(Id.. I2 @ dde) K +
(123)
For 2 < <1+++ < L, we have
Eee
1d (1d 12p- @ 12
WE,t
(124)
Note that this form depends on 4(L -) 2) + 2(L 3) + (L - -3 - 422/2 unknowns. Terms including the first and last layers have slightly different form
1 (89) 2 la.)
) + (0))) 2 la1 12.))
En
(125)
di
1 Vol
(111 lld K di divil
(1 @ ld, (.. K
E12
(126)
Eie- d
1
(1 dp- lai 1bp))
(127)
E1L =
(1)
Wi de @ la, 1/ + v 1ddd
1e0 1160 @ Ia w?? 2) T
(128)
EEL dde
(Idi 1he- 1.
(129)
1
(Id @ Idli K+ + de-1v
1
(11.-2 8 1 1p )
E1-11 = /dL-2
(130)
EL (Lie.. 8 200 +
1 (Ide- 17, 8 00(2))
(131)
di-1
In the special case L 22, 2, have instead
E1 = (wo()) 3 Tot k+ (w/²) @ 1 17,) K
(132)
For the cross-terms, we note that En + first and last one, we have
ST 2.e therefore we must consider all pairs of values e, f.. In all layers except the
((2)
3(2)
Eu
(Ite. @ (at)) devde de-1
1T 2 1e
11LL
(133)
vdi-1
7(1)
((2))
E0,041
(Iee 8 1dea.) + devdett
1 1J 8 21TT.)
1 L-111
(134)
dde+1
Wel
Ee -dide-ide
feee-1 LLL, E' 11,L
(135)
25

---

Global curvature for second-order optimization of neural networks
Terms involving the first and last layer equal
(1)
&
Ins
1d
1
(136)
di
1
(1,,0% & 12)
8= 2..,...,
(137)
ded
Σ1 =
1
(iLI @ 11%
(138)
Vdi
Eu
1
(Weethe, @ 12
8= 2,.,...,4
(139)
de-11de
E7
Vde-1 (Idi.. 8 @11)
(140)
11-1I-= = (Neems & (107) ) + Idee 1o @ 254 1)
(141)
ΣEL VdL-1de (Leothee & w/1)
t t= 2..,...,
(142)
Proof
We have Vi E S(de), the set of permutation matrices, while Vo- = Ldo and VE Ldz. We start from the gradient mean of the bias, given by Equation (22)
i, Ve Phg
(143)
First we note that, since Vz = Idg, then 3N remains undetermined.
For l < L, we note that all permutations V& have common eigenvector, corresponding to the eigenvalue equal one, given by vector with components all equal. That is the only eigenvector common to all permutations, therefore jtg Çe Le for i < L and some scalar Se.
The gradient mean of the weights is given by Equation (21)
H4 (Vt-1 Vi) p4
(144)
For l * I, the only eigenvector common all possible products V @ Vi within the group, corresponding to an eigenvalue equal one, is again the vector will all components equal. Therefore must be bt Ge(1de @ 1d;) for some scalar Çe. For é = 1, I., by similar arguments find, respectively
M1 = Zo 2 14, and P1
ldg 8 Zr.
(145)
for some column vectors Zo. Zz.-
Now we turn to the covariance of the gradient. The bias-bis covariance must satisfy Equation (24), equal to
Eer = Ve See vē
(146)
Since Ee. Lg we only need to study the cases l < f. First, we study the case = E'. For L, we note that Vz = Lad and therefore Σ7 is undetermined. For < L, we multiply Equation (146) by Vz and use its orthogonality (V. Ve = Id,), we obtain
Eu V2 Ve Exe
(147)
We note that a matrix that commutes with all permutation matrices Ve must be the sum of scalar matrix and a matrix with all elements equal (Stuart & Weaver, 1991). Therefore, for < L we have
(())
E,
7(1) 07 Id. +
0 1 1
(148)
de
26

---

Global curvature for second-order optimization of neural networks
for some scalars 0
((2) For é + f', since the only cigenvector common to all permutations is a vector with all components
equal, then Equation (146) is satisfied by
Eee
Eee
i < L, and Ee 1a, 5s if < L 1a, 12 dddde
(149)
ddg
For some scalar Ter and column vector We.
The covariance must satisfy Equation (23), equal to
Eer = (V--1 @ Vi) Eee (V²-1 6 Vi)
(150)
Since Eeee - Eee- we only need to study the cases < f. First, we study the case & = and & o 1, L. In that case, using the orthogonality of Ve, Ve,, we rewrite Equation (150) as
Eu (Vt-1 & Vi) = (Vt-1 @ Va
(151)
Similar to the case of the bias, matrix that commutes with product of permutation matrices must have the form (Stuart Weaver, 1991)
o (Lam @ ld, 1,,1 + oi
of-²)
Eu =0 (Iar, 2 Id.) +
(16-111 & Ide) + dede-1 (Id 1ā. @ 1a 1a, 1%
di
de-1
(152)
For some scalars respectively
(1) of 02 Pi 0i
For l = 1, L, since Vo = Ido VL-1 = -d the covariance block satisfies,
Σ1 (La, Vi) En (Lao v?) and ELL (VL-1 Ia,.) ELL (V @ Idr.)
(153)
Using similar arguments as above, since these equations must be valid for all permutation matrices Vi and VL-1- have
En (1 ( ② la. Lo..
(
di (0)) @ 1
and ELL = (loz.. 000==(101.,1 8
o - di (Id. -6 @ of)) (154)
for some matrices + ()) (2) (1) $ (2) Next, we study the case l' = e 1. We use the commutation matrix K, which is
symmetric permutation, thus K K² K2
x===1 = The commutation matrix exchanges the factors in Kronecker product: K(A @ K B A (see Chapter 3.7 of Magnus & Neudecker (2019)). We rewrite Equation (150)
E (Vt-1 @ Vi) K (Vē & v22)
(155)
and solve for Σg,4 K instead of Ee,+11 We start with E * 1, T, 1, then this equation must be valid for all permutation matrices V1-1, Vj, Ve+1. Using similar arguments as above, we find
(1) V/
(2) (Idr.1
E24+1
(Idr-1 Ihe.. & Tdd) K+ de de 1
la 1x)
(156)
Vde-idem1
for some scalars 1 2/ For l = 1,1-11. Equation (155) becomes, respectively
(1)
Σ₁² = (Id, @ V1) K (V: & vi) and ΣL-LLK = (VL-2 & VZ-1) (Ia, 8 &VT.)
(157)
These equations must valid for all permutations Vi, 1/2 and Vr.22, Vt-1- We find
(11 IL.) K - di divals 1 17 @ ld, ld 1T. K
(2)
En Voz
(158)
1
(1ee-2Wl 22 @ @lde- K+
1
Ide 2b12)) @ ldd 1hr..) K
ΣL-1L Vdz-2
(159)
de-I/de-z
27

---

Global curvature for second-order optimization of neural networks
For some column vectors 11 11 242 22 - The case is distinct, since Equation (155) is equal to
Σ1²K = (lao @ V1) M12K (la, @ V2?)
(160)
that has the solution
E12 = (vl)) 2 11. k+ (v/3) 2 1a, 12.) K
(161)
for some matrix 1. Lastly, we study the case l > l 1, starting with l and l' L. that case, Equation (150) must be valid for all different possible permutations V1-1. V4, Ve Ve. The solution is
WE,
(Ide.. 17, @ 1d
Eee /de-ide-idede
(162)
If f 1, l' L. Equation (150) must be valid for all different possible permutations Ve, V Ve,, and
Ene de (wi.e 12 @ la, 1
(163)
If 1, L, Equation (150) must be valid for all different possible permutations Ve-1, Ve, VL-1, and
Ee Vdeidd--dd (Idr-, 1 @ ld, w.
(164)
Lastly, for e = L, l' = L (with L > 2), Equation (150) must be valid for all different possible permutations V1, Vz-1, and
En = Vdi-1 (w de @ 1a, 12e)) +
(Iuttee.. @ la, la, (1).7 (2) 7)
(165)
Vdode-1d1
The bias-weight covariance must satisfy Equation (25), equal to
Ee Ve Eee (Ve-1&V;))
(166)
We note that, in this case, Ee * E therefore must consider all pairs of values E, f. We start with l' = l. For l * 1, L, Equation (166) must hold for all possible different permutation matrices Ve-1. Ve, therefore
7(1) (17 @ @lde) + de dev/de11 (17m., @ le.12)
7(2) 61
Eu
(167)
dde-1
for some scalars 06 06 For the special cases = = and l' = L, we rewrite Equation (166) as, respectively
Σn V1 En (Id, 2 Vi") and ΣLL ELL (V Itz))
(168)
Since these equations must be valid for all permutation matrices Vi and Vr.11 we must have
En = 3 (1)TT
(2)
2 Tai
+ di
1 17 and E = Vdr.-1
(1de-) @ xσ1)
(169)
For some column vectors (1 Equation (166) as
Φ1 and some matrix $ Next, we look the case = l +++ For - 1, we rewrite
(2)
S. Ve E4mm (V7 s Vit1)
(170)
This equation must be valid for all permutations Ve. Ve+1, therefore
((2)
E = Vdet1
(Iee & 1dmm)) + de dev/det2 (1 @ 18
(171)
28

---

Global curvature for second-order optimization of neural networks
((2) for some scalars D Vi For the special case '' = + I, we rewrite Equation (166) as
EE-IL VL-1 - (VΣ-1 @ laz)
(172)
This equation must be valid for all permutations V1-1, therefore we find
E1.-111= (Tai.. & 4 + (Id... - 12 @ (1.2)
(173)
(1) for column vectors L2 2 L Lastly, study the case ll # For l + L and ll o L, Equation (166) must hold for all possible different permutation matrices Ve, -e Ve, therefore we have
WeP dddde-ide
Ee
(Lae
122.)
(174)
For some scalar Wg,er For E. = T, and 8 # 1, T., Equation (166) becomes
Ei Eie (V 8 V2)
(175)
This equation must hold for all possible different permutation matrices V Ve, therefore we have
dee, (Wi 17 do @ 17)
(176)
For some column vector We. For l o 1, L and l' 1, Equation (166) becomes
Σn1 Vz Ee (las @ V2²)
(177)
This equation must hold for all possible different permutation matrices Vi, Ve, therefore we have
Σn /ded1
(1a, W7 2 12,)
(178)
For some column vector We. For 11 1, and ll L., Equation (166) becomes
Σn V₁ E. (V @ Tay))
(179)
This equation must hold for possible different permutation matrices Vz, Vt-1, therefore we have
Vde-1de (1 1" dg @ (wFT)
(180)
For some column vector Wez. Lastly, we have the case I,, ll In that case Equation (166) becomes
== (Ido vV))
(181)
This equation must hold for possible different permutation matrices V1, therefore have
Σ11 1 (Ār. @ 17 Vin
(182)
for matrix Ωzz-
K. Estimation of factors
In this section, we provide some details on how to estimate the factors of the gradient covariance. First, we provide a derivation of Equations (37). in the main text. We then describe how to improve the estimates by adding momentum, and provide an empirical analysis of the estimation error (Figure 10).
29

---

Global curvature for second-order optimization of neural networks
0.75
betah 0
betah 0.9
0.70
betah = 0.95
mean pairwise correlation
betah 0.98
0.65
0.60
0.55
0.50
0.45
0.40
200
400
600
800
1000
tteraionn
Figure 10. The factors of the covariance can be estimated by single model. Mean pairwise correlation of estimated factors across 10 independentl initialized models (45 pairs of models) as function of training iteration (gradient descent), for different values of the momentum parameter 8 A correlation value close to would imply that hyperparameter estimates agree across independent models. Without momentum (3t, = 0), the hypeparameters of independent models display correlation of 0.5. With momentum, correlations reach about 0.7.
We consider the example of a two-layer MLP with Tanh activations and no bias, but similar procedure can be followed for other cases. The covariance is described in Theorem 3.1, 1 in this case is equal
Σ
Σ1 Σg2
La. V 2 Ia.)
(L & La, K Id, & $2
= K
(183)
First, we assume that the covariance E is given, and we are asked to obtain the factors { $2 11 Later we consider the case in which we have an estimate of Σ, instead of its true value. Given E, to obtain the factors {41, D2 V. we need to invert the following expressions, describing each block of the covariance
En Φ @ Lt
Σx Ld, 42
Σ = (v) 2 II,)
(184)
is convenient to rewrite these equations including the indices of all matrix elements
(Em))s.a (01),t 8sk
(Ee)) 55z (422)s
(E12) (V1),k dn
(185)
The problem inverting these equations is underparameterrized, therefore we look for the least square solution. That is equal to
1
(11),1
(En)
(186)
di
1 di
(12)sk
(E22)
(187)
d1
5=1
(4)) jkc
Σ (E12)
(188)
practice, the covariance E is not given, we only have access to an estimate from the observation of a gradient vector. note that the mean gradient is zero this case (Theorem 3.1). therefore. substitute the values of Σ11, Σι,, Eea with the observed outer product of the observed gradient, and we obtain Equations (3)) in the main text.
These equations show sums over neurons of the first layer, of dimension di The larger is d1, the lower is the expected етrог in the estimates, because there will be a correspondingly large number of repeated values in the gradient covariance, therefore the value can be estimated by summing over the corresponding pairs of parameters. practice, the expected value the error in the estimate depends on the distribution of the gradients, which is unknown and likely non- aon-Gaussian, therefore it not easy to estimate.
30

---

Global curvature for second-order optimization of neural networks
To improve the estimate of the factors, momentum, a standard method for noise reduction. Equations (37),(38),(3))
provide estimate of the factors for fixed iteration step during training. We rewrite those equations here in matrix form for convenience
0C
T
ac
11
(189)
d1
aw1
aw1
0C
OL
T
Q2
(190)
d1
@W2
0C
ac
V d1
(191)
aw2
Across training iterations, instead of using the equations above, we use the following standard instance of momentum for estimating the factors at step given the estimate at step
1
3r
OC
T
0L
(f1); = Bh (i) (01)e-1 +
(192)
d1
OW,
aw;
1 3n
OC
ac
(92), = BA (12)--1 +
(193)
d1
aw,
aw2
Br.
ac
T
OC
T
(V1), = 3h (11), +
(194)
d1
ow
Ow2
All factors are set zero é and estimates at step & divided by the factor (1-3%, B,)) to remove the bias in the estimate, as in standard practice (Kingma & Ba, 2014).
We study the estimation error by running 10 models, each one with its own random initialization and its own individual esimate of the hyperparametsrs. We training each model separately by gradient descent. At each iteration, we estimate the hyperparamaters of each model and save it into a flattened vector containing all factors. Then, we compare the models by computing correlation between each pair models (a total of 45 pairs) and compute the mean correlation across pairs. large correlation would imply that estimates of different models agree, suggesting that the estimates are good. We run the same model described in Section 4, with layer widths do 100, di 70, d2 40, for a total of 9800 parameters, 5000 training imput data points drawn from Gaussian distribution with badly conditioned covariance and the output given by teacher network with identical architecture and random weights (see Section for details).
Figure 10 shows the mean correlation as a function of training iterations for different values of 3r. Without momentum (3,. = 0), the factors estimated from different individual models display a correlation of about 0.5, showing that individual models estimate similar but not identical factors. The introduction momentum allows increasing the correlation up to about 0.7, therefore it provides significant improvement. However, the correlation drops back to about 0.5 after number of iterations that depends the value of 8n. We note that the number of steps takes drop is equal to about 10 times the intrinsic timescale of momentum, equal to (1 -1
surrogate
surrogate
W1
W/
sigenvalues
4
eigenvalues
w5
ba
20
60
0
5
10
15
Figure II. Comparison of full (A) and surrogate (B) matrix for two-layer MLP with Tanh activations. Panels C and D compare the cigenvalues the two matrices. We note that the cigenvalues are identical, although some of the eigenvalucs of the full matrix are repeated multiple times.
31

---

Global curvature for second-order optimization of neural networks
A
full
B
surrogate
full
surrogate
Wh
51
Wa
E:
eigenvalues
elgenvalues
4-
4-
Wh
3-
3
2-
2-
W5
be
1
W9
Figure Similar Figure 11, for a three-layer with Tanh activations.
A
full
surrogate
tull
surrogate
6
W4
eigenvalues
eigenvalus
Wb
Wa
3-
3-
W5
W)
W
Figure 13. Similar to Figures and 12, for a three-layer MLP with ReLU activations.
L. Surrogate matrix
In this section, describe how to perform operations the covariance matrix using smaller surrogate matrix. In particular, we need to compute its inverse square root, that is required by the second-order update of Equation (1). First recall that, using either Theorem 3.1 (Tanh activation) Theorem 3.2 (ReLU activation), we can describe the covariance matrix by a set of factors. We consider the case of Tanh, but similar argument can be used for the case of ReLU, We define the flattened vector of all factors by
= Vec (102. .I. (112,..,.11)), (ōi............. b 2 1 DI, $r, vi)
(195)
Using Theorem 3.1, We compute the covariance matrix given the factors df and the layer widths {de}, and vice versa.
Φ, (di...., dp-1) +
(196)
The next important observation is that, by Lemma 2.4, Equation (15), any analytic function of the covariance matrix must have the same structure of the covariance itself. Therefore, provided that the covariance positive definite, the inverse square root is also described by vector of factors, identified by the "isr" superscript
g jsr = Vec (lod2)
Φ1 " 2
ist jar
ier L w/r
(197)
Also in this case, we can compute the matrix given the factors and the layer widths, and vice versa
jst (di...., d1-1)
(198)
naive approach for computing the factors & jer the inverse square root would be to compute Σ from Φ, then its inverse square root ター ! and finally isr from -
Φ (d
+ -
jar
(199)
32

---

Global curvature for second-order optimization of neural networks
However, the matrix E may be very large and exceed the available memory. Furthermore, computing its inverse square root may be computationally very expensive. Therefore, would like compute o² directly from Φ, without ever computing and square rooting / inverting the big matrix Σ.
We solve this problem by noting that the factors do not depend on the choice of the layer widths. In other words, any alternative choice (d)) of the layer widths would give the same outcome Gla
Φ df - A - A
isr
(200)
where A is a surrogate matrix with the same factors as E but different layer widths. Therefore, we can set the layer widths very small numbers, corresponding to a surrogate matrix A that is easy store in memory and to square root invert. In particular, the smallest number that retains the factors is dž for case of Tanh and dž 2 for the case of ReLU.
For example, in the case of a two-layer MLP (L = 2) with Tanh activation and no bias, the full matrix has size (dod1 + dzd2) x (dod2 + d1d2) and is equal to
1 ld
( @ la,)
Σ == K (vi @ la.))
(201)
ld, $2
while the surrogate matrix has size (du + d2) (do + d2) and is equal to
A $
V
(202)
22
Therefore, the surrogate matrix has a factor of less elements than the full matrix. The main point of this section is that any analytic function of Σ has the same form of Σ, for some other values of the factors, and that the same operation applies to A.
In the case of the inverse square root, we thus have
Dior ldi
(viTT @ Ia,) K
)
Σー==
K
T @ Iti
Ia, Tb.
(203)
/ isrr
A-
3
1
(004)
We note that the knowledge of the surrogate matrix is sufficient to compute the product between the full matrix and any vector. For example, given the gradient vector
VC Vec (
OC Oc ow, OW,
(205)
Using the properties of Kronecker product, we derive the update for a two-layer Tanh MLP without bias
Oc
OC
Vaint OC
Oc OW2
WーTFL
Vec
isr
Tiat
(206)
Ow,
awe
OWT
Figures 11, 12, 13 compare the full and surrogate matrix for a few example MLPs. They show that the eigenvalues of the full and surrogate matrices are identical, although some of the eigenvalues are repeated multiple times in the full matrix.
M. Hyperparameter settings
Table 2 shows the hyperparameter values for the experiments illustrated in Figures 5, 6. We note that most optimizers are quite robust with respect to small changes of the hyperparameter values. An exception is KFAC, we found that was very sensitive to the combination of learning rate a and damping X. In Table 2, boldface values are set by grid search, while others are set to standard values. The latter include momentum parameters for Adam (3) = 0.9, 82 = 0.999) and KFAC (8 0.9), and the exponent of Shampoo, set to p = 1 as common to most applications. For SymO, we set the momentum parameter to 3r 0.95, as that value seemed to stabilize training in all our experiments (See Figure 10 in Appendix K).
In the case of SymO, we choose exponentially decreasing learning rate, an.3. While this choice is uncommon in neural network optimization, it is motivated by the analytical study of a toy quadratic problem, which is detailed in Appendix B. While neural network optimization is very far from a quadratic problem, may represent a reasonable approximation near a fixed point, and we found that it worked fine in the non-convex optimization problems illustrated in Figures 5, 6. More work is necessary to find out the optimal settings of SymO.
33

---

Global curvature for second-order optimization of neural networks
Algorithm 1 Symmetty- Optimization update (SymO)
Input:
gradient VL. previous factors {#1, D2, W W-h momentum parameter Sr. Compute current factors { 22, Vi using Equations (193))))))))1194)
Compute surrogate covariance matrix A by Equation (202)
Divide A by (3) (bias correction)
Compute inverse square root of surrogate matrix A- by numerical algebra Compute factors {dher disr Viier) of A- à by Equation (204)
Compute SymO update by Equation (206)
Return: SymO update, current factors {L1, 22, 11
Table 2. Hyperparameter values used in experiments, for all optimizers, either in the case of linear activations (Figure 5), Tanh activations (Figure 6). Values boldface are set by grid search. others are set to standard values.
GD
Adam
KFAC
Shampoo
SymO
a= 0.01
a= 1c
a= 0.8
Og= 0.05
Linear activation
a= 0.8
31 = 0.9
x= 1.3e
6= le 10
3= 0.95
82 0.999
0.9
p 1
8p 0.95
Q= 0.01
a= 0.006
Q= 0.8
QØ= 0.09
Tanh activation
Q= 0.8
31 0.9
x= 0.017
<= 1c 8
3_ 0.91
32 0.999
3 = 0.9
p= 1
3, = 0.95
34


Reviews:
[{"review_id": "f9M0qZgwDC", "reviewer": "Reviewer_bCpH", "summary": "The authors' work focuses on getting insights on the second moments of $\\Sigma_t  = \\int \\nabla \\mathcal L(\\theta)\\nabla \\mathcal L(\\theta)^\\top dp_t(\\theta) - \\mu_t \\mu_t^\\top$, by exploiting invariances in representation space. In particular, the authors observe that if $G$ is a symmetry of the loss function $\\mathcal L(G\\theta_t) = \\mathcal L(\\theta_t)$, then mean and covariance satisfy eigenvalue equations $\\mu_t = G \\mu_t$, $\\Sigma_t = (G^\\top \\otimes G) \\Sigma_t$.\nBy exploiting symmetry groups $\\mathbb G$ of common activation functions, the authors use the fact that these eigenvalue equations have to hold simultaneously for all $G \\in \\mathbb G$ to show which structure is imposed by the solution space.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments. We are glad that the reviewer recognises that the key contribution of our work is theoretical and considers our work novel and interesting.\nWe answer here the main concerns raised by the reviewer.\n- *Empirical comparison with other second-order methods*. We implemented KFAC and Shampoo. We used the original version of Shampoo with power 1/4. Using more recent versions with power 1/2 (as in SOAP) did not give better results in our experiments. In KFAC, we optimised learning rate and damping. In Shampoo, we optimised learning rate and epsilon. The following table shows test loss values for a few time steps in the 2-layer MLP with Tanh case showed in the paper. We will include the full plot in the final version of the paper. \n\n| time      |    GD     |     Adam     |     KFAC    |      Shampoo    |      SymO |\n|----------------|-------------|----------|----------------|-------------|----------|\n| t=50\t    |             0.0378 |    0.0255    |     0.0091   |     0.0071       |       0.0058 |\n| t=100 | 0.0232  |  0.0116    |     0.0074     |   0.0060      |        0.0056 |\n|t=150 | 0.0177 |   0.0086    |     0.0067    |    0.0057      |        0.0056  |\n| t=200 | 0.0147 |  0.0072       |  0.0063     |   0.0056       |       0.0056 |\n\nWe believe that SymO has better results because it goes beyond block-diagonal preconditioning, and considers also interactions among the two layers. However, the performance of KFAC and Shampoo is very similar to SymO.\n- *Sparsity of Hessian for Tanh versus ReLU*. Our understanding of the stronger sparsity in the Tanh case with respect to the ReLU case comes from the symmetry groups. In the case of Tanh, in addition to the permutation symmetry, there is also a symmetry for switching the sign of parameters in adjacent layers. Therefore, the symmetry group of Tanh has larger size with respect to ReLU (please see Chen et al 1993 for specific numbers). As a consequence, the number of constraints that need to be satisfied are larger and the Hessian has less degrees of freedom. These constraints result in a higher sparsity of the matrix. In the final version of the paper, we will report the specific groups  sizesfor Tanh and ReLU in Section 2.2.\n- *Difference between local and global (average) gradient outer product*. It is not easy to compare the second moment at a single point with the ensemble average, because the former is a rank-1 matrix while the ensemble average is usually full rank. Instead, we provide an additional analysis showing the quality of the estimation of the curvature with a single model. We computed the correlation of the single-model estimate with the estimate obtained by an average over a large sample of models (N=10,000). We observe the following numbers (d0,d1,d2 are the layer sizes):\n\n|Layer sizes|correlation|\n|-|-|\n|(100,10,100)|0.66$\\pm$0.05|\n|(100,100,100)|\t0.75$\\pm$0.03|\n|(100, 1000,100)|0.54$\\pm$0.03|\n|(100, 10000, 100)|0.56$\\pm$0.03|\n\nWe will include a plot with more details in the final version of the paper.", "claims_and_evidence": "I believe the claims of the authors are supported by sufficient empirical evidence.", "methods_and_evaluation_criteria": "I believe the proposed experimental evaluation suffices for the authors' claims. See \"Experimental Designs Or Analyses\" for the only concern.", "theoretical_claims": "I confirm that I checked the proofs in the appendix of the manuscript and I didn't find any criticality.", "experimental_designs_or_analyses": "In section 3.4, I would include in the comparison also some local preconditioning methods other than Adam, to see in practice if the average structure gives comparable results with respect to a local precondition. In particular, I would compare with [1,2,3]\n\n[1] J. Martens et al., \"Optimizing Neural Networks with Kronecker-factored Approximate Curvature\", NeurIPS 2023.\n[2] N. Vyas et al., \"SOAP: Improving and Stabilizing Shampoo using Adam\",  ArXiv 2024\n[3] V. Gupta et al., \"Shampoo: Preconditioned Stochastic Tensor Optimization\", ICML 2018.", "relation_to_broader_scientific_literature": "The key contribution of this work is theoretical, quantifying a precise relationship between symmetries of the neural network and the averaged Hessian structure. The proposed preconditioning strategy resulting from the theoretical investigation sounds well-developed (no inversions of big matrices), placing also the work in the line of works proposing Quasi-Newton methods for neural network training. As mentioned above, experimental comparison with some of these methods is missing in the manuscript.", "essential_references_not_discussed": "To my knowledge, the idea is novel in this context and the relevant literature has been discussed.", "other_strengths_and_weaknesses": "The work is extremely original and the results are of serious interest both from a theoretical point of view and concerning potential applications. In particular, the idea of exploiting neural network's symmetries to infer the structure of the averaged Hessian is of broad interest to develop cheap preconditioners to accelerate training.", "other_comments_or_suggestions": "I believe a point of strong interest would be to quantify how much the second moment in a point $\\nabla L(\\theta_t)\\nabla L(\\theta_t)^\\top $ can deviate from the average $\\mathbb E_{\\theta_t}[\\nabla L(\\theta_t)\\nabla L(\\theta_t)^\\top]$. While not easy to perform since it depends on the update rule, a result of this kind would quantify the local effectiveness of $\\Sigma_t$ as a preconditioner for optimization.\nEven just experimental investigation in this direction would be of interest.", "questions_for_authors": "1. It is interesting to see how your predictions for the structure of the Hessian are extremely sparse in figure 7 (Tanh activation) and less in figure 8 (ReLU activation). Do the authors have any insights about this observation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "K3JDVKuHqo", "reviewer": "Reviewer_ewk1", "summary": "The work attempts to improve the computations of second order methods by analyzing the covariance matrix of the gradients in small MLP networks. They rely on certain symmetrics expected to be in network parameters and derive theory on the structre, as well as explicit solutions for the covariance matrix. They perform minor experiments, very toyish, where their predictions approximately coincide with the general obtain structure on average (over 10K-100K models). They show on this specific toy example accelaration is obtained, compared to standard GD and Adam optimizers. The experiments are not convincing enough, the setting is highly restrictive, \"the input is sampled from a Gaussian distribution with zero mean. The covariance matrix of the input is generated using random orthogonal eigenvectors\". It is very hard to conclude if this behavior is actually general and portrays well large NNs and true data. They perform experiments on 3-layer MLP's with RelU, which can already perform quite well on real (easy) data, such as MNIST. Unfortunately, only a toy example is shown. Thus, although the theory is of some interest, I cannot recommend publication in this form of the work.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments, which will help improving our work.\n\nWe would like to highlight that the main contribution of our work is theoretical. As acknowledged by the other reviewers, the theory presented in our work is novel and it could be of broad interests for other fields of machine learning. We hope that we can convince the reviewer of the value and novelty of our theory.\n\nSpecific concerns:\n- *Comparison with other second-order methods*. We agree with the reviewer that comparison with other second-order methods is appropriate. We implemented KFAC and Shampoo. In KFAC, we optimised learning rate and damping. In Shampoo, we optimised learning rate and epsilon. The following table shows test loss values for a few time steps in the 2-layer MLP with Tanh case showed in the paper. We will include the full plot in the final version of the paper. \n\n| time      |    GD     |     Adam     |     KFAC    |      Shampoo    |      SymO |\n|----------------|-------------|----------|----------------|-------------|----------|\n| t=50\t    |             0.0378 |    0.0255    |     0.0091   |     0.0071       |       0.0058 |\n| t=100 | 0.0232  |  0.0116    |     0.0074     |   0.0060      |        0.0056 |\n|t=150 | 0.0177 |   0.0086    |     0.0067    |    0.0057      |        0.0056  |\n| t=200 | 0.0147 |  0.0072       |  0.0063     |   0.0056       |       0.0056 |\n\nWe believe that SymO has better results because it goes beyond block-diagonal preconditioning, and considers also interactions among the two layers. However, the performance of KFAC and Shampoo is very similar to SymO.\n\n- *Relation to broader scientific literature*. The reviewer thinks that our work bears little relation to broader scientific literature, which is the opposite of what stated by the other reviewers, who recognised important connections of our work with other fields of machine learning. We kindly ask the reviewer to specify what references are missing in our paper. \n\n- *How the method works with mini-batches*. We agree with the reviewer that this is an important question. Typically, second-order methods work better with large batches, while first-order methods sometimes work better with small batches. However, the advantage of first-order methods on small batches concerns the generalisation performance of the optimum. Instead, here we are concerned with the speed of convergence, for which no advantage of small batches has ever been observed. Therefore we chose to analyse the full batch case, and postpone the study of mini-batches to future work. \n\n- *Toy examples*. We would like to clarify that the theory provided in Section 3, which is our main contribution, is valid for any size of a neural network. While Figures 2,3,4 report the case of a neural network with 5 neurons per layer, the choice of 5 neurons is only for illustrative purposes and the same theory is valid for any number of neurons.\n\n- *Variance of prediction*. We agree with the reviewer that the we should provide the quality of the estimation of the curvature with a single model. We computed the correlation of the single-model estimate with the estimate obtained by an average over a large sample of models (N=10,000). We observe the following numbers (d0,d1,d2 are the layer sizes):\n\n| Layer sizes\t|\t\tcorrelation |\n|----------------|-------------|\n|(100,10,100)\t|\t\t0.66 $\\pm$ 0.05 |\n|(100,100,100)\t|\t0.75 $\\pm$ 0.03 |\n|(100, 1000,100)\t|\t0.54 $\\pm$ 0.03 |\n|(100, 10000, 100) |\t0.56 $\\pm$ 0.03 |\n\nWe will include a plot with more details in the final version of the paper.\n\n- *Assumptions are not validated*. We highlight that the theory provided in Section 3, which is our main contribution, is exact and there are no assumptions involved, other than the invariances that are known to hold exactly for neural networks and initialisation routines. We kindly ask to clarify which assumptions the reviewer is concerned about.\n\n- *Experiments are not realistic* We highlight that the main contribution of our work is the theory provided in Section 3. We consider the empirical evaluation as preliminary and not a major contribution of the work. We believe that a full empirical evaluation on a large scale problem deserves a separate paper.\n\nWe would like to emphasize the novelty of our theory, as it introduces a unique approach to determining the curvature of neural networks by leveraging their inherent symmetries — a perspective not previously explored. The theoretical methods we developed to tackle this problem are not only original but also hold significant potential for advancing other fields. The study of neural network curvature and loss landscapes is a topic of broad interest, and our work offers new avenues for gaining deeper insights into this major unsolved challenge. We sincerely hope the reviewer will acknowledge the substantial value of our theoretical contribution.", "claims_and_evidence": "Claim covariance matrix structure can be predicted, based on certain expected symmetries of the parameters and properties of the input data.", "methods_and_evaluation_criteria": "Evaluation is of low quality, only a very simple toy example assuming Gaussian distribution input.", "theoretical_claims": "Theory on covariance matrix of the gradients. Appear fine.", "experimental_designs_or_analyses": "Very scarce, no real analysis. Do not show instances of a certain model, only an average over 10K or 100K models, do not show the variance of their prediction.", "relation_to_broader_scientific_literature": "little.", "essential_references_not_discussed": "refs ok", "other_strengths_and_weaknesses": "Strength: theory and general idea may have merrit.\n\nWeaknesses:\n* Many assumptions are needed to derive the results.\n* Assumptions are not validated on real examples.\n* Experiments are not realistic.\n* No variance of the results is shown.\n* Solution optimization loss (Fig 5), the authors do not show other 2nd order methods, only GD and Adam. This toy examle perhaps works best for 2nd order methods (as it involves Gaussian inputs). Not clear if the proposed solution is faster or simply 2nd order is better here.\n* The experiments are perfromed with full gradient of the entire training set, how SGD works? Can the proposed algorithm work well in mini-batches?", "other_comments_or_suggestions": "see above", "questions_for_authors": "Please explain in more detail if additional experiments were conducted that I missed, or justify why the chosen setting is reasonable to reflect real world DNNs and data.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s4DNWuIloA", "reviewer": "Reviewer_3sTt", "summary": "The submission studies the structure of the \"global curvature\" of deep networks. The main result is that global matrix quantities such gradient covariance and Hessian (where \"global\" means the expected values of those matrices under some distribution on the weights) has a specific matrix structure with much fewer free parameters. As a potential application, the submission proposes an optimizer based on this structure and shows it outperforms SGD/Adam on a two-layer teacher/student problem.\n\n\n## update after rebuttal\n\nThe response has addressed my main issues from the `Claims And Evidence` section below. I was worried that the phrasing in the submission was over-claiming the benefit of \"global curvature\" and the proposed estimator. The proposed wording and additional evidence for the quality of the claimed estimator would clarify that the main benefit is in the structure of the curvature estimator, which is a more appropriate claim given the supporting experiments. With these changes, I think the paper should be accepted.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detail comments that will improve our paper.\n- *Why global curvature*. We agree on the lack of clarity for why global curvature should be used as a preconditioner. We will revise as follows:\n\nIntroduction: “The primary motivation for considering global curvature is its efficient computation, which serves as the main contribution of this work. Future research will explore whether global curvature serves as a reliable approximation of local curvature and whether it offers inherent advantages for enhancing convergence in second-order optimization.”\n\nSection 2.1: “We operate under the assumption that global curvature captures meaningful variations in local curvature for optimization purposes. \n\nDiscussion: “This work provides preliminary evidence suggesting that global curvature enhances convergence when applied as a preconditioner. During optimization, the distribution of parameters collapses into a set of local solutions, therefore our method is expected to better approximate the local curvature towards the end of training, similar to other methods (e.g. Natural Gradient, Gauss-Newton). A comprehensive analysis of the errors introduced by our approximation will be a subject of future investigation.”\n- *Bayesian inference*. We agree that our method could be valuable for Gaussian posterior approximation in Bayesian inference. We will revise the Discussion: “Bayesian deep learning relies on approximating the posterior over parameters by a Gaussian distribution. The covariance of this distribution is usually approximated using a diagonal or block-diagonal structure. Our work offers a method for efficiently computing the full covariance, which may lead to more accurate Bayesian posterior estimates.”\n- *Use of the word “second-order”*. While “second-order” traditionally refers to the computation of second-order derivatives, several methods classified as second-order use only first-order derivatives, such as Natural Gradient (e.g. KFAC), Gauss-Newton, and Shampoo. Our approach can be applied to any of those methods. Furthermore, the gradient covariance provides second-order moments of the distribution, therefore we believe that it may be also called “second-order”. However, we are open to reconsidering the terminology if the reviewer strongly feels that it constitutes a misuse.\n- *Estimation error of the global curvature and dependence on layer size*. We will replace Fig 9 with a plot of the correlation of the single-model estimate with the average over a large sample of models (N=10,000), and how that depends on layer size d1. As the reviewer suggested, we do not observe a simple increase of correlation with layer size d1. We observe the following numbers (d0,d1,d2 are the layer sizes):\n\n|Layer sizes|correlation|\n|-|-|\n|(100,10,100)|0.66$\\pm$0.05|\n|(100,100,100)|\t0.75$\\pm$0.03|\n|(100, 1000,100)|0.54$\\pm$0.03|\n|(100, 10000, 100)|0.56$\\pm$0.03|\n\n- *Theory predicts the overall structure only*. In Fig.2,3,4, we wanted to highlight that, while panels A and B differ significantly, both follow the same structure. Furthermore, the theory of Section 3 remains valid regardless of the specific values of the unknown factors.\n- *Decreasing step size*. We optimize two hyperparameters for all optimization methods except GD (in addition to Adam and SymO, we added Shampoo and KFAC, see answers to other reviewers), ensuring a fair comparison. The choice of an exponentially decreasing learning rate for SymO is based on a theoretical analysis of a quadratic loss function, which we will include in the Appendix.\n- *Why do ReLU activations satisfy fewer invariances?* The intuitive reason is that the symmetry group is smaller, the constraints that are satisfied by the covariance are fewer and therefore the covariance has more degrees of freedom. We will report the groups sizes in Section 2.2.\n - *Why an extra hidden layer reduces the inter-layer correlations?* Odd activations introduce an invariance with respect to sign changes in both the incoming and outgoing weights of a neuron. We speculate that this invariance may lead to cancellations in the correlations of these weights. Further studies are needed to answer this question.\n- *What is the probability distribution?* We agree that we do not know the distribution of the gradients, other than its symmetries and the structure of its first two moments. We will add an analytic study of a quadratic loss to the Appendix, for which the parameter and gradient distributions remain Gaussian.\n- *Use of the word “curvature”*. Curvature is frequently linked to Fisher Information and Gauss-Newton methods, which do not directly compute the Hessian. The gradient covariance is essentially equivalent to the “empirical” Fisher Information matrix, which is also considered a form of curvature. Nevertheless, we are open to changing the terminology if the reviewer believes it is a misuse.\n- *Use of the word “hyperparameter”*. We agree and we will replace it with “factors”.", "claims_and_evidence": "The main claim of the submission, that the expected gradient covariance has the claimed structure, is well supported. The visualization provided are convincing up to minor comments (see below).\n\nWhat's less clear is why this global quantity matters, whether for the gradient covariance or Hessian. The submission seems to take for granted that \"global\" is better than \"local\" in its introduction, and the only argument in favor of the global approach is that an algorithm is proposed that uses an estimator inspired by the global curvature structure. \n\n> We demonstrate the effectiveness of our approach by running exact second-order optimization on a two-layer MLP and synthetic data\n\nThis appears to assume that the optimizer works better because it estimates the global curvature. I strongly disagree, as the evaluation of the proposed estimator is insufficient and potentially misleading (details in Methods And Evaluation Criteria).\n\nThe one potential application I can see for the theory in the global curvature is in enabling efficient full-Gaussian posterior approximation in Bayesian inference at a reduced cost. This is mentioned tangentially in the discussion section, but should be expanded upon to make the link obvious.\n\n---\n\nPlease change the following sentence \n> We demonstrate the effectiveness of our approach by running exact second-order optimization\n\nThe preconditioner does not use any second-order information and is instead an \"approximation\" of the global covariance matrix of the gradients.", "methods_and_evaluation_criteria": "### Evaluation of the estimator\n\nMy issue is not that the optimization results are not impressive enough. I am not denying that the proposed preconditioner (using a structured approximation given the current gradient) could be a great optimization method. But whether it does appears disconnected from the global curvature discussion at the start of the paper. \n\nThe paper claims to test the estimation error of the proposed estimator,\n> Appendix J provides additional details, along with an empirical analysis of the estimation error (see Figure 9)\n\nHowever, Appendix J and Figure 9 do not test whether the estimator is an appropriate estimator of the global curvature. Instead, they test that it produces consistent results. Those could be close, far, or completely unrelated to the global curvature.\n\nThe paper also asserts\n> We expect the error in the estimation of hyperparameters to decrease with layer size (d1 in this case).\n\nI did not find evidence or mechanism that would explain this.\n\nThe estimator claims to be able to estimate the global covariance from a single sample. This is very suspicious for any claim of globality. Taken to its extreme, the claims in the paper do not seem that far from the claim that \"in $d$ dimensions, if $x \\sim \\mathcal{N}(0, I)$, then $\\mathrm{diag}(xx^\\top)$ is a good estimator of $I$, especially as $d \\to \\infty$.\"\n\nAgain; my issue is not with the proposed optimizer or whether it works, but on attributing the reason for its performance to the global curvature. The current paper would be improved if the optimizer was instead introduced along the following lines, to make it clear that the connection is tenuous: \"We observe that the global curvature is very structured and sparse. It is possible that this structure captures the most important variations in the covariance/Hessian. Therefore we propose to use this structure as a preconditioner for our optimizer instead of the traditional Diagonal ones. It seems to work well.\"\n\nI would raise my score if the text was changed to explicitly acknowledge this disconnect, and state that\n- the \"globality\" story is an inspiration, but not a justification, for the good performance of the algorithm \n- the good performance of the algorithm, being neither exact nor second-order, does not say much about the validity of the \"global\" perspective in optimization\n\n### Evaluation of the structure\n\n> We highlight that theory predicts the overall structure, rather than the specific numerical values.\n\nAs that is the case, the visualizations could be improved as they emphasize the numerical value. This could be done by fitting the factors of the theoretical structured matrix to the observed data to show agreement on the numerical values.\n\n### Evaluation of optimizers\n\n> For SymO, we choose an exponentially decreasing learning rate\n\nThe evaluation of the optimizer claims to use a decreasing step-size \nfor the proposed optimizer. It is not specified whether the other optimizers (SGD, Adam) use a decreasing learning or other schedule. All algorithms should use the same regime; either a tuned step-size schedule or a constant step-size.", "theoretical_claims": "I have not checked the proofs in details.", "experimental_designs_or_analyses": "No strong issue in the experiment design beyond the evaluation issues listed above.", "relation_to_broader_scientific_literature": "The paper could be made stronger by discussing the relationship to the Bayesian variational inference literature in more details. This community is likely to appreciate the structure induced by the invariance, because VI maintains an estimate of the posterior distribution in the form of a  mean and covariance which are used to samples weights at every step to compute updates to the are used to sample weights at every step to compute gradients and update the mean and covariance. The standard approaches use a diagonal approximation to the covariance matrix (eg [Blundell et al., Weight Uncertainty in Neural Networks](https://arxiv.org/pdf/1505.05424)) but more expressive families are sometimes used ([e.g. Lin et al., Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations](https://arxiv.org/abs/1906.02914)", "essential_references_not_discussed": "Not that I am aware of.", "other_strengths_and_weaknesses": "The study of the impact of the invariance of the network on the gradient and curvature matrices appears novel, and the approximation obtained appear impressive. However, I fail to see why this \"global\" curvature, that enables such a simplification, helps in the context of optimization which typically relies on local quantities which need not share this structure. This disconnect between the main part of the paper and its tentative application in §3.4 onwards (as detailed above) is my main issue with the current submission.", "other_comments_or_suggestions": "### Additional details\n\nSome explanation, even if high-level, for the following phenomena would help the reader follow the paper.\n- §2.2: Why do ReLU activations satisfy fewer invariances than odd activation functions?\n- §3: Why does the addition of an extra hidden layer reduces the inter-layer correlations?\n\n### Word choices\n\n- Please change the following sentence \n  > We demonstrate the effectiveness of our approach by running exact second-order optimization\n\n  The preconditioner does not use any second-order information and is instead an \"approximation\" of the global covariance matrix of the gradients.\n  \n- The probability distribution used to sample the relevant quantities in the definition of the \"global curvature\" is not explicitly given. My understanding is that it is out of a desire to remain general as the theory only requires invariance conditions, but having a running example (for example from variational inference) would help the reader.\n\n- The use of the term \"curvature matrix\" seems heavily overloaded, as for example in §2.1 it has to encompass the Fisher, the Hessian and the covariance of the gradients.\n\n  > Various studies have employed different formulations for the curvature matrix, including the Fisher information matrix, the Gauss-Newton matrix, the Hessian matrix, and the gradient covariance matrix.\n\n  The relationship between the Hessian and the gradient covariance is unclear, and as a result \"curvature matrix\" does not seem to imply more than \"a matrix that can be used as a preconditioner\". If \"preconditioner\" was the intended meaning, I would encourage its usage instead. \n\n- The submission uses the expression \"hyperparameters of the curvature\" (and \"parameters\" in §1). I initially did not understand what it meant, as to me a \"hyperparameter\" describes a user-specified parameter controlling the behaviour of an algorithm. A more appropriate term could be the \"factors\", as theorem 1 shows that the expected matrix obeys a specific factorization.\n\n### Typos\n\nIn Lemma 2.4, I assume $f$ should be a matrix analytic function of $\\Sigma$? $f(\\Sigma) = v^\\top \\Sigma v$ appears analytic but Eq. 15 does not make sense.", "questions_for_authors": "Did I misunderstand the evaluation of the estimator?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work introduces a theoretical framework that leverages neural network symmetries for efficient and exact computation of the covariance matrix of gradients, which can be used in scaled-gradient optimization methods. By exploiting these symmetries, the computation of such matrices involves fewer degrees of freedom than naively expected. The approach reveals structural insights into curvature, influenced by activation functions, and demonstrates faster convergence on synthetic data.

The paper addresses an interesting problem and offers novel perspectives and solutions. Reviewers have provided several suggestions regarding terminology (specifically the use of "second-order"), the scope and depth of the claims, the underlying assumptions, and the experimental section. The authors have addressed many of these in their rebuttal and are encouraged to incorporate them in the next revision of the paper.
Decision: Accept (poster)