S YMMETRIES , F LAT M INIMA AND THE C ONSERVED Q UANTITIES OF G RADIENT F LOW
authors
Paper under double-blind review
A BSTRACT
Empirical studies of the loss landscape of deep networks have revealed that many local minima are connected through low-loss valleys.
Ensemble models sampling different parts of a low-loss valley have reached SOTA performance.
Yet, little is known about the theoretical origin of such valleys.
We present a general framework for finding continuous symmetries in the parameter space, which carve out low-loss valleys.
Our framework uses equivariances of the activation functions and can be applied to different layer architectures.
Importantly, we also introduce a novel set of nonlinear, data-dependent symmetries for neural networks.
These symmetries can transform a trained model such that it performs similarly on new samples.
We then show that conserved quantities associated with linear symmetries can be used to define coordinates along low-loss valleys.
The conserved quantities help reveal that using common initialization methods, gradient flow only explores a small part of the global minimum.
By relating conserved quantities to convergence rate and sharpness of the minimum, we provide insights on how initialization impacts convergence and generalizability.
We also find the nonlinear action to be viable for ensemble building to improve robustness under certain adversarial attacks.
[Figure]
Training deep neural networks (NNs) is generally a highly non-convex problem.
The loss landscape of a NN, which is shaped by the model architecture and the dataset, is generally very rugged, with the number of local minima growing rapidly with model size (Bray & Dean, 2007;
S¸ims¸ek et al., 2021).
Despite this complexity, recent work has revealed many interesting structures in the loss landscape (Kunin et al., 2021; Garipov et al., 2018).
For example, NN loss landscapes often contain approximately flat directions along which the loss does not change significantly (Freeman & Bruna, 2017; Garipov et al., 2018).
One source of flat directions is parameter transformations that keep the loss invariant (i.e. symmetries).
Specifically, moving from a local minimum in the direction of a symmetry takes us to another minimum.
In this paper, we show that continuous symmetries of the parameter space introduce many flat directions in the loss landscape, resulting in extended local minima.
It has been hypothesized that gradient descent (GD) prefers flat local minima (He et al., 2019) and that wider local minima may generalize better (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017).
Flat minima have been used to make ensemble or mixture models by sampling different parameter configurations that yield similar loss values (Garipov et al., 2018; Benton et al., 2021) which improves model performance.
However, finding such flat directions is mostly done empirically (Freeman & Bruna, 2017; Garipov et al., 2018; Draxler et al., 2018; Benton et al., 2021; Izmailov et al., 2018), with few theoretical results.
Motivated by the fact that continuous symmetries of the loss result in flat directions in local minima, we derive a general class of such symmetries.
Our key insight is to focus on equivariances of the nonlinear activation functions ; most known continuous symmetries can be derived using this framework.
However, models related by such symmetries are exactly equivalent , meaning they cannot behave differently on different inputs.
Hence, we argue that for ensembling or robustness tasks, we need to find data-dependent symmetries .
Indeed, aside from the familiar “linear symmetries” of NN, the framework of equivariance allows us to introduce a novel class of symmetries which act nonlinearly on the parameters and are data-dependent.
These nonlinear symmetries cover a much larger class of continuous symmetries than their linear counterparts, as they apply for almost any activation function and may involve larger continuous groups.
We show theoretically that such nonlinear symmetries may improve model performance on out of distribution (OOD) samples.
We provide preliminary experimental evidence that ensembles using these nonlinear symmetries are more robust to attacks.
Extended flat minima arise frequently in the loss landscape of NNs; we show that symmetry-induced flat minima can be parametrized using conserved quantities .
Furthermore, we provide a method of deriving explicit conserved quantities (CQ) for different continuous symmetries of NN parameter spaces. CQ had previously been derived from symmetries for one-parameter groups (Kunin et al., 2021; Tanaka & Kunin, 2021). Using a similar approach we both derive the CQ for general continuous symmetries. But more importantly, we show that this approach fails to find CQ for rotational symmetries.
Nevertheless, we find the conservation law resulting from the symmetry implies a cancellation of angular momenta between layers, as elaborated later.
To summarize, our contributions are:
A general framework based on equivariance for finding symmetries in NN loss landscapes.
A derivation of the dimensions of flat minima induced by symmetries.
A new class of nonlinear, data-dependent symmetries of NN parameter spaces.
Expanding prior work on deriving conserved quantities (CQ) associated with symmetries and discussing its failure for rotation symmetries. Finding a cancellation of angular momenta between layers for rotation symmetries.
parameterization of symmetry induced flat minima CQ
This paper is organized as follows.
First, we review existing literature on flat minima, continuous symmetries of parameter space, and conserved quantities.
In Section 1, we define continuous symmetries and flat minima, and show how linear symmetries lead to extended flat minima.
To illustrate our constructions, we provide examples of linear symmetries of NN parameter spaces.
We then define nonlinear, data-dependent symmetries.
In Section 2, we use infinitesimal symmetries to derive conserved quantities for parameter space symmetries, extending the results in Kunin et al. (2021) to larger groups and more activation function.
Then, we show how CQ can be used to define coordinates along flat minima.
We close with some preliminary experiments involving nonlinear symmetries, conserved quantities and a discussion of potential use cases.
R ELATED W ORK
Continuous symmetry in parameter space Overparametrization in neural networks leads to symmetries in the parameter space (Głuch & Urbanke, 2021).
Continuous symmetry has been identified in fully-connected linear networks (Tarmoun et al., 2021), homogeneous neural networks (Badrinarayanan et al., 2015; Du et al., 2018), radial neural networks (Ganev et al., 2022), and softmax and batchnorm functions (Kunin et al., 2021).
We provide a unified framework that generalizes previous findings, and identify nonlinear group actions that have not been studied before.
Conserved quantities The imbalance between layers in linear or homogeneous networks is known to be invariant during gradient flow and related to convergence rate (Saxe et al., 2014; Du et al., 2018; Arora et al., 2018a;b; Tarmoun et al., 2021; Min et al., 2021).
Huh (2020) discovered similar conservation laws in natural gradient descents.
Kunin et al.
(2021) develop a more general approach for finding conserved quantities for certain one-parameter symmetry groups, relating symmetries and conserved quantities. We use a similar approach for deriving conserved quantities related to the symmetries.
However, in contrast with prior work, we also show, for general continuous groups, when this approach will work and when it will fail. Tanaka & Kunin (2021) also studies relating continuous symmetries to dynamics of conserved quantities using the physics approach of Noether’s theorem (Noether, 1918). We briefly discuss how our approach is related to Noether’s theorem.
Topology of minimum The global minimum of overparametrized neural networks are connected spaces instead of isolated points.
We show that parameter space symmetries lead to extended flat minima.
Previously, Cooper (2018) proved that the global minima is usually a manifold with dimension equal to the number of parameters subtracted by the number of data points.
In contrast, we derive the dimensionality of the symmetry-induced flat minima and show they are related to the number of infinitesimal symmetry generators and dimension of weight matrices.
S¸ims¸ek et al.
(2021) study permutation symmetry and show that in certain overparametrized networks, the minimum related by permutations are connected.
Entezari et al.
(2021) hypothesize that SGD solutions can be permuted to points on the same connected minima.
Ainsworth et al.
(2022) develop algorithms that find such permutations.
Additional discussion on mode connectivity, ensembles, sharpness of minima and other related topics can be found in Appendix A.
C ONTINUOUS SYMMETRIES IN DEEP LEARNING
In this section, we first summarize our notation for basic neural network constructions (see Appendix C for more details).
Then we consider transformations on the parameter space that leave the loss invariant and show how they lead to extended flat minima.
1. T HE PARAMETER SPACE AND LOSS FUNCTION
The parameters of a neural network consist of weights 1 W i ∈ R n i × m i for each layer
i , where n i and m i are the layer output and input dimensions, respectively.
For feedforward networks, successive output and input dimensions match: m i = n i − 1 .
We group the widths into a tuple n = ( n L , . . . , n 1 , n 0 ) , and the parameter space becomes Param = R n L × n i − 1 × · · · × R n 1 × n 0 .
We denote an element therein as a tuple of matrices θ =
( W i ∈ R n i × n i − 1 ) Li =1 .
The activation of the i -th layer is a piecewise differentiable function σ i : R n i
→ R n i , which may or may not be pointwise.
For θ ∈ Param and input x ∈ R n 0 , the feature vector of the i th layer in feedforward network is Z
i +1 ( x ) =
W i +
σ ( Z i ( x )) , where the juxtaposition ‘ Wσ ( Z ) ’ denotes an arbitrary linear operation, including matrix product, convolution, etc.
For simplicity, we largely focus on the case of multilayer perceptrons (MLPs).
We denote the final output by F θ : R n 0 → R n L , defined as
F θ ( x ) = σ L ( Z L ( x )) .
The “loss function” L of our model is defined as:
[Equation]
where Data = R n 0 × R n L is the space of data and Cost : R n L × R n L
→ R is a differentiable cost function, such as mean square error or cross-entropy.
In the case of multiple samples, we have matrices X ∈ R n 0 × k and Y ∈ R n L × k whose columns are the k samples [-
 We-]2 , and retain the same notation with F θ : R n 0 × k
→ R n L × k .
We first prove properties of L that hold for any training data.
Hence, unless specified otherwise, we take a fixed batch of training data { ( x i , y i ) } ki =1 ⊆ Data , and consider the loss to be a function of the parameters only.
Example 1.1.
Two-layer network with MSE Consider a network with n = ( n, h, m ) , the identity output activation ( σ L ( x ) = ( x ) ), and no biases.
The parameter space is Param ( n ) =
R n × h × R h × m
and we denote an element as θ =
( U, V ) .
Taking the mean square error cost function, the loss function for data ( X, Y ) ∈ R n × k × R m × k takes the form L ( θ , ( X, Y ))
= 1 k ∥ Y − Uσ
( V X )
∥ 2 .
1. A CTION OF CONTINUOUS GROUPS AND FLAT MINIMA
Let G be a group.
An action of G on the parameter space Param is a function · : G × Param →
Param , written as g · θ , that satisfies the unit and multiplication axioms of the group, meaning id · θ = θ where id is the identity of G , and g 1 · ( g 2 · θ ) =
( g 1 g 2 ) · θ for all g 1 , g 2 ∈ G .
Definition 1.1 (Parameter space symmetry) .
The action G × Param → Param is a symmetry of
L if it leaves the loss function invariant, that is:
[Equation]
We describe examples of parameter space symmetries in the next section.
Before doing so, we show how parameter space symmetry leads to flat minima using the following result (see Appendix C.6):
[-Taking ddt (cid:12)(cid:12) 
-]
[-ddt (cid:12)(cid:12) 
-]
Proposition 1.2.
Suppose G × Param → Param is a symmetry of L .
If θ ∗ is a critical point (resp.
local minimum) of L , then so is g · θ ∗ for any g ∈ G .
The proof of this result relies on using the differential of the action of g to relate the gradient of L at θ ∗ with the gradient at g · θ ∗ .
We see that, if θ ∗ is a local minimum, then so is every element of the set { g · θ ∗ | g ∈ G } .
[-= 
-]
This set is known as the orbit of θ ∗ under the action of G .
The orbits of different parameter values may be of different dimensions.
However, in many cases, there is a “generic” or most common dimension, which is the orbit dimension of any randomly chosen θ .
1. E QUIVARIANCE OF THE ACTIVATION FUNCTION
In this section, we describe a large class of linear symmetries of L using an equivariance property of the activations between layers.
For accessibility, we focus on the example of two layers with output F ( x ) =
Uσ ( V x ) for ( U, V ) ∈ Param = R m × h ×
R h × n and x ∈ R n .
All results generalize to multiple layers by letting U
= W i and V =
W i − 1 be weights of two successive layers in a deep neural network (see Appendix C.5).
Let G ⊆ GL
h
( R ) be a subgroup of the general linear group, and let π : G → GL h ( R ) a representation (the simplest example is π ( g )
= g ).
We consider the following action of the group G on the parameter space Param :
[Equation]
This action becomes a symmetry of L if and only if the following identity holds:
[-= ddt (cid:12)(cid:12) 
-]
[Equation]
We now turn our attention to examples.
Example 1.2.
Linear networks A simple example of (4) is that of linear networks, where σ is the identity function: σ ( x ) = x .
One can take π ( g ) =
g and G = GL h
( R ) .
Example 1.3.
Homogeneous activations Suppose the activation σ : R h
→ R h is homogeneous , meaning that (1) σ is applied pointwise in the standard basis and (2) there exists α > 0 such that σ ( cz ) =
c
α σ ( z ) for all c ∈ R > 0 and z ∈ R h .
Such an activation is equivariant under the positive scaling group G ⊂ GL h consisting of diagonal matrices with positive diagonal entries.
Explicitly, the group G consists of diagonal matrices g = diag( c ) with c = ( c 1 , . .
. , c h ) ∈
R h> 0 .
For z = ( z 1 , . . .
, z h ) ∈
R h , we have σ ( gz ) =
(cid:80)
j σ ( c j z j )
=
(cid:80) j c αj σ ( z j ) =
g α σ ( z ) .
Hence, the equivariance equation is satisfied with π ( g ) =
g α .
Example 1.4.
LeakyReLU
This is a special case of homogeneous activation, defined as σ ( z ) =
max( z, 0) + s min( z, 0) , with s ∈ R ≥ 0 .
We have α = 1 , and π ( g ) = g .
Example 1.5.
Radial rescaling activations A less trivial example of continuous symmetries is the case of a radial rescaling activation (Ganev et al., 2022) where for z ∈ R h , we have σ ( z ) =
f ( ∥ z ∥ ) z for some function f : R → R .
Radial rescaling activations are equivariant under rotations of the input: for any orthogonal transformation g ∈ O ( h ) (that is, g T g =
I )
we have σ ( gz ) =
gσ ( z ) for all z ∈ R h .
Indeed, σ ( gz ) =
f ( ∥ gz ∥ )( gz ) =
g ( f ( ∥ z ∥ ) z ) =
gσ ( z ) , where we use the fact that ∥ gz ∥ =
z T g T gz =
z T z = ∥
z ∥ for g ∈ O ( h ) .
Hence, (4) is satisfied with π ( g )
=
g .
We arrive at our first novel result, whose proof appears in Appendix C.6.
Theorem 1.3.
The dimension of a generic orbit in Param under the appropriate symmetry group is given as follows. The cases are divided based on whether h ≤ max( n, m ) or not. Orbit Dimension Activation
Symmetry Group [Equation] Identity [Equation]
Homogeneous Positive rescaling [Equation] Radial rescaling
O ( h ) (cid:0)
[Equation]
(cid:0) [Equation] (cid:1) As an aside, we note that a familiar example where (4) is satisfied involves the permutation of neurons.
More precisely, suppose σ is pointwise and let G be the finite group of h × h permutation matrices. Then (4) holds with π ( g ) = g . However, the permutation group is finite ( 0 -dimensional), and so there are generally no flat minima.
The linear symmetries discussed above are relatively small subgroups of GL h which can potentially act on U, V . However, as we discuss next, many tasks such as domain adaptation or handling out of distribution data may involve data transformations which are not among the exact symmetries discussed above. In this section we show that, using a more general definition of equivariance, it is sometimes possible to find nonlinear symmetries which may include the entire GL h .
Need for data-dependent symmetries.
The examples in the previous section all feature symmetries of L which are independent of the input data. That is, we transform the parameters with no change in the loss: L ( g · θ , X ) = L ( θ , X ) (i.e. θ and g · θ yield models which are identical in the function space).
Now consider a task involving a domain shift. Assume we have a trained model and again consider the subnetwork F ( U, V, x ) = Uσ ( V x ) . Suppose for the linear part z
=
V x the domain shift may correspond to a linear transformation g ∈ GL h taking z to z ′ =
gz .
Now, suppose we want to know how this domain shift manifests itself in higher layers and the latent representations.
Specifically, we want to find g ′ ∈ GL h such that Ug ′ σ ( z ) =
Uσ ( gz ) .
When σ is not equivariant under g , we cannot use the results of previous section (meaning we cannot find π such that g ′ =
π ( g ) ).
While we can always solve for g ′ in terms of g using g ′ σ ( z ) = σ ( gz ) , this requires sampling z and, hence, the resulting g ′ are data-dependent .
This sampling approach will suffer from low accuracy when the hidden dimension is large.
Thus, we may prefer an analytical solution g ′ = c ( g, z ) . Indeed, such solution often exist and arise from a more general form of equivariance. General equivariance.
For common activation functions, the equivariance condition σ ( gz ) =
π ( g ) σ ( z ) in (4) holds only for g belonging to a relatively small subgroup of GL h . Specifically, for ReLU, g must be in the positive scaling group, while for the usual sigmoid activation, the equivariance only holds for trivial g = id. Yet, as in the example above with g ′ = c ( g, z ) , under certain conditions, it is possible to define a nonlinear, data-dependent action of the full GL h under which many activations, including sigmoid and ReLU, are equivariant.
The subtlety of such an action is that it is data-dependent, rather than being an action on the parameter space only.
For our general equivariance condition, we are looking for map c : GL h ×
R h
→ GL h which generalizes (4) to General Equivariance: [Equation] Setting g
= g 1 g 2 with g 1 , g 2 ∈ GL h , we can see that c satisfies a generalized form of the group multiplication axiom
[Equation]
Indeed, a map c : GL h × R h → GL h is said to be an equivariance if it satisfies (1)
c ( id h , z ) = id h for all z , and (2) c ( g 1 , g 2 z ) c
( g 2 , z )
=
c ( g 1 g 2 , z ) for all g 1 , g 2 ∈ GL h and z .
These two conditions on c translate directly into the unit and multiplication axioms of a group, generalizing π ( g 1 g 2 )
=
π ( g 1 ) π ( g 2 ) .
Every equivariance gives rise to an action of GL h on Param × R h :
Nonlinear group action:
[Equation]
This action is a symmetry preserving the loss if and only if σ ( gz )
= c ( g, z ) σ ( z ) , ∀ g ∈ GL h ∀ z ∈ R h .
Next, we show how c can be constructed. Explicit form of the equivariant map c .
For most common activations σ : R h
→ R h , it is possible to consistently select matrices c ( g, z ) that satisfy both the general equivariance axioms above and the appropriate relation with σ in (5) (see Appendix D).
A sufficient condition is that σ ( z ) is nonzero for any nonzero z ∈ R h ; this is the case for usual sigmoid, hyperbolic tangent, leaky ReLU, and many other activations. Theorem 1.4 (Algorithm to construct c ) .
An equivariance c satisfying σ ( gz )
=
c ( g, z ) σ ( z ) , ∀ g ∈ GL h , ∀ z ∈ R h , can be constructed as follows:
For the unit vector e = (1 , 0 , · · · , 0) , consider the two vectors v 1 = σ ( e ) and v 2 = σ ( ge ) . Find R be the 2D rotation in the plane Span( v 1 , v 2 ) such that Rv 1 = v 2 . Use QR decomposition to find P taking v 1 , v 2 to an orthonormal basis of the 2D plane Span( v 1 , v 2 ) .
Then, c ( g, e ) = P ( R − id 2 ) P T − id h . 5.
For z ∈ R h , find g z ∈ GL h such that z = g z e . Use equivariance to construct c ( g, z ) =
c ( g, g z e ) = c ( gg z , e ) c ( g z , e ) − 1 . See appendix D for proof.
In closing, we note that, unlike the exact symmetries of the previous section, a general equivariance may actually change the loss in the function space . However, when σ is Lipshitz continuous, we can find some bounds on how much the loss changes after the nonlinear action.
Litpshitz bounds on nonlinear group action.
Assume we use a given data point x to transform the weights using (7) to ( U ′ , V ′ , x )
=
g · ( U, V, x ) . If we use U ′ , V ′ on a different data point x ′ , it is not clear whether L ( U, V, x ′ ) and L ( U ′ , V ′ , x ′ ) are related.
This is important because we may want to know if the nonlinear respects “closeness” of data points.
Note that x or V x need not be the raw data and they may be the output of higher layers in a neural network and thus represent latent embeddings of data.
In this case, we expect distance between data points to encode important information about shared features.
Therefore, we would like to know whether the nonlinear action preserves closeness of data in some way. We will now show that when σ is Lipshitz, we can find a bound the effects of the nonlinear action on the network output.
Theorem 1.5 (Lipshitz bounds from equivariance) .
Let σ be Lipshitz continuous with ∥ σ ( a ) − σ ( b ) ∥ ≤ η ∥ a − b ∥ for a, b ∈ R h , and let c : GL h × R h
→ GL h be an equivariance satisfying (5) .
Consider the network
Uσ ( V x ) .
Let U ′ =
Uc ( g, V x ) − 1 and V ′ = gV .
For a different data point x ̸ = ˜ x ∈ R h we have
[Equation]
See Appendix 49 for proof, which simply follows from applying the Cauchy-Schwarz inequality and Lipshitz continuity a few times.
Note that on the original weights U, V , the Lipshitz condition
Uσ ( V x ) yields
[Equation]
meaning that the only change to the bound due to the nonlinear action is a factor of ∥ c ( g, V x ) − 1 ∥∥ g ∥ .
Thus, we expect the nonlinear action to not disrupt latent distribution of data significantly. Next, we turn to conserved quantities arising from symmetries and how they can be used to parametrize flat minima.
1. I NFINITESIMAL SYMMETRIES
Deriving conserved quantities from symmetries requires the infinitesimal versions of parameter space symmetries. Any smooth action of a matrix
Lie group G ⊆ GL h induces an action of the infinitesimal generators of the group, i.e., elements of its Lie algebra.
Concretely, let g
= Lie( G ) = T I G be the Lie algebra, which can be thought of as a certain subspace of matrices in gl h = R h × h .
For every matrix M ∈ g , we have an exponential map exp M : R → G defined as exp M ( t ) =
(cid:80) ∞
k =0 ( tM ) k k ! .
Given an action of G on Param , the infinitesimal action of M ∈ g is a vector field M :
Infinitesimal action of M vector field:
[Equation]
In the case of a parameter space symmetry, the invariance of L translates into the following orthogonality condition, where the inner product ⟨ , ⟩ : Param × Param → R is calculated by contracting all indices, e.g. ⟨ A, B ⟩ = (cid:80) ijk... A ijk... B ijk... .
Proposition 1.6 (Infinitesimal symmetry) .
Let G be a matrix Lie group and a symmetry of L .
Then the gradient vector field is point-wise orthogonal to the infinitesimal action of any M ∈ g :
[Equation]
We comment on case where G acts linearly on parameter space.
First, set d to be the dimension of the parameter space 3 , and make the identification Param ≃ R d by flattening matrices into column vectors.
The general linear group GL( Param ) ≃ GL d ( R ) consists of all invertible linear transformations of Param .
Suppose G is a subgroup of GL( Param ) , and hence acts on Param via matrix multiplication.
Then its Lie algebra g is a Lie subalgebra of gl d = R d × d , and M θ =
M · θ acts via matrix multiplication.
Hence, in the case of a linear action, (11) becomes:
[Equation]
Recall the linear action of G ⊆ GL h given in (3).
{+To describe the corresponding infinitesimal action, note that the representation π : G → GL h induces a representation dπ : g → gl h of the Lie algebra, defined by dπ ( M ) = ddt (cid:12)(cid:12) 
+}
π (exp M ( t )) .
Using the fact that exp − M ( t ) =
exp M ( t ) − 1 , the infinitesimal action of the Lie algebra g induced by (3) is given by:
[Equation]
Denote the Jacobian of σ : R h → R h at z ∈ R h by dσ z ∈ R h × h .
Then the infinitesimal form of the equivariance equation σ ( gz ) =
π ( g ) σ ( z ) from (3) is given by: ⟨ Mz, dσ z ⟩ = dπ ( M ) σ ( z ) .
With indices, the components are [ ⟨ Mz, dσ z ⟩ ] j = (cid:80) i
[ Mz ] i ∂ i σ ( z )
j
= [ dπ ( M ) σ ( z )] j (see Appendix C.1). For example, in the case of homogeneous σ , the equivariance equation reduces to Mz ⊙ σ
′ ( z ) =
αMσ ( z ) (where ⊙ is elementwise multiplication).
12 ∥ Y − UV X ∥ 2 , where U, V ∈ R , Y = 2 , and X = 1 . Trajectories corresponding to different values of Q intersect the minima at different points. We have shown above that continuous symmetries lead us along extended flat minima in the loss landscape.
In this section, we identify quantities that (partially) parameterize these flat minima.
We first show that certain realvalued functions on the parameter space remain constant during gradient flow (GF), a continuous version of gradient descent (GD).
We refer to such functions as conserved quantities . Applying symmetries can change the value of the conserved quantity and, therefore, conserved quantities can be used to partially parameterize flat minima.
As mentioned earlier, prior work (Kunin et al., 2021) already derived some Q for common 1-parameter symmetries in NN. We extend that result systematically to larger continuous groups. Importantly, we show that this approach fails to identify conserved quantities for orthogonal groups (rotation symmetry). Instead, the conservation law derived from symmetries yields a constraint, which can be interpreted as cancellation of angular momentum , as we elaborate below.
Gradient flow (GF) Recall that GD proceeds in discrete steps with the update rule θ t +
=
θ t − ε ∇L ( θ t ) where ε is the learning rate (which in general can be a symmetric matrix), and t = 0 , 1 , 2 , . . .
are the time steps. The idea of gradient flow (GF) is as follows: rather than update the parameter values in steps, one can instead define a smooth curve in the parameter space from a choice of initial values to the limiting local minimum. The curve is a function of a continuous time variable t ∈ R , and velocity of this curve at any point is equal to the gradient of the loss function, scaled by the negative of the learning rate.
In other words, the dynamics of the parameters under GF are given by:
[Equation]
From an initialization θ (0) at t = 0 , GF defines a trajectory θ ( t ) ∈ Param for t ∈ R > 0 , which limits to a critical point.
In this way, GF is a continuous version of GD.
Conserved quantities A conserved quantity of GF is a function Q : Param → R such that the value of Q at any two time points s, t ∈ R > 0 along a GF trajectory is the same: Q ( θ ( s ))
=
Q ( θ ( t )) .
In other words, we have dQ ( θ ( t )) /dt
= 0 .
Note that, if f : R → R is any function, and Q is a conserved quantity, then the composition f ◦ Q is also a conserved quantity.
Several conserved quantities of GF have appeared in the literature, most notably layer imbalance Q imb ≡ ∥ W i ∥ 2 −
∥ W i − 1 ∥ 2 (Du et al., 2018) for each pair of successive feedforward linear layers ( σ ( x ) = x ), and its full matrix version Q i =
W Ti W i − W
i − 1 W Ti − 1 .
It is known that imbalance and certain other conserved quantities can be related to symmetries (Huh, 2020;
Kunin et al., 2021).
Conserved quantities from symmetries We now propose a generalization of the layer imbalance results mentioned above by associating a conserved quantity to any infinitesimal symmetry. As in
Section 1.2, suppose a matrix Lie group G acts linearly on the parameter space. Then, from (12), we have the identity ⟨∇ θ L , M · θ ⟩ = 0 for any element M in the Lie algebra g . Using the gradient flow dynamics (14), this identity becomes:
[Equation]
In other words, the velocity at any point of a gradient flow curve is orthogonal to the infinitesimal action.
For simplicity, we set the learning rate to the identity: ε =
I (all results generalize to symmetric ε .)
The following proposition (whose proof is elementary and well-known) provides a way of ‘integrating’ equation (15), in the appropriate sense, in order to obtain conserved quantities:
Proposition 2.1.
Suppose the action of G on Param is linear 4 and leaves L invariant.
For any
M ∈ g , there is a conserved quantity Q M : Param → R given by Q M ( θ ) =
⟨ θ , M · θ ⟩ .
While Proposition 2.1 directly links the infinitesimal action to conserved quantities, it has the limitation that the conserved quantity corresponding to an anti-symmetric matrix M = − M T in g is constantly zero, and we do not obtain meaningful conserved quantities.
Instead, we can only conclude that flow curves satisfy the differential equation (15).
Fixing a basis ( θ 1 , · · · θ d ) for Param ≃ R d , this equation becomes (cid:80) i<
j M ij r 2 ij ˙ ϕ ij ≡ 0 where ( r ij , ϕ ij ) are the polar coordinates for the point ( θ i , θ j ) ∈ R 2 (see Appendix C.9.5).
In summary, we find:
[Equation]
Conserved quantities in gradient descent The quantities derived in Proposition 2.1 are conserved for continuous gradient flow. In gradient descent, the values of these conserved quantities may change due to the time discretization. However, the change in Q is expected to be small.
For example, in two-layer linear networks, the change of Q is bounded by the square of learning rate. Appendix H contains derivations and empirical observations of the magnitude of change in Q .
Conserved quantities parametrize symmetry flat directions.
We now observe that applying a symmetry changes the values of the conserved quantities Q M (Figure 2).
Indeed, for M ∈ g and g ∈ G , we have Q M ( g · θ ) =
Q g T Mg ( θ ) for all θ ∈ Param , so applying the group actiontransforms the conserved quantity Q M to Q g T Mg .
As discussed in Section 1, applying a symmetry g to a minimum θ ∗ of L yields another minimum g · θ ∗ ; hence applying symmetries leads to a partial parameterization of flat minima.
Note that, in general, we may lack sufficient number of Q M to fully parameterize a flat minimum.
For example, in the linear network UV x , G = GL h and flat minima generically have h 2 dimensions, whereas the number of independent nonzero Q M is h ( h + 1) / 2 , which is the dimension of the space of symmetric matrices M = M T in gl h .
Relation to Noether’s theorem
In physics, Noether’s theorem (Noether, 1918) states that continuous symmetries give rise to conserved quantities.
Recently, Tanaka & Kunin (2021) showed that
Noether’s theorem can also be applied to GD by approximating it as a second order GF.
We show that in the limit where the second order GF reduces to first order GF (14), results from Noether’s theorem reduce to our conservation law (cid:10) M θ , ∇L (cid:11) = 0 (12).
In short, using Noether’s theorem, the conserved Noether current is J M = e t/τ J 0 M with J 0 M = (cid:68) M θ , ε − 1 ˙ θ (cid:69) .
In the limit τ → 0 , using (14), J 0 M = (cid:10) M θ , ∇L (cid:11) = 0 and the conservation dJ M /dt
= 0 implies J 0 M = 0 , meaning we recover (11).
Details appear in Appendix B.2.
Examples.
We present examples of conserved quantities for two-layer neural networks.
all of which generalize in a straightforward way to the multi-layer case.
See Appendix C.9 for full derivations (which heavily rely on properties of the trace).
We adopt the notation of Section 1.3.
Example 2.1.
General equivariant activation Suppose σ is equivariant under a linear action of a subgroup G ⊆ GL h ( R ) , so that π ( g ) σ ( z ) = σ ( gz ) .
Then the two-layer network F ( z ) =
Uσ ( V z ) is invariant under G , as is the loss function.
For symmetric M ∈ g , Proposition 2.1 yields the following conserved quantity:
[Equation]
Indeed, this follows from the fact that M · ( U, V ) =
( − Udπ ( M ) , MV ) , as in (13).
Example 2.2.
Imbalance in linear layers Suppose the network is linear.
Then σ ( z ) =
z and the loss is invariant under GL h ( R ) .
For symmetric M we have the conserved quantity Q M ( U, V ) =
Tr (cid:2) ( V V T − U T U ) M (cid:3) .
Moreover, each component of the matrix V V T − U T U is conserved.
Example 2.3.
Homogeneous activation under scaling Suppose σ is a homogeneous activation of degree α . Let G =
( R > 0 ) h be the positive rescaling group, so that σ ( gz ) =
g α σ ( z ) for any g ∈ G and z ∈ R h .
Note that the Lie algebra of G consists of all diagonal matrices in gl h , so that, in particular, each M ∈ g is symmetric.
Since dπ ( M ) = αM for any M ∈ g , we obtain the conserved quantity Q M ( U, V ) =
Tr (cid:2) ( V V T − αU T U ) M (cid:3) .
Using the basis M = E kk , we see that
Q = diag (cid:2) V V T − αU T U (cid:3) is conserved (here, diag
[ A ] is the leading diagonal).
Special cases of this are LeakyReLU and ReLU with α = 1 .
Example 2.4.
Radial rescaling activations Let σ be such a radial rescaling activation.
As in
Section 1.3, the orthogonal group G
= O ( h ) commutes with σ , making it a symmetry of L .
The Lie algebra g = so h comprises anti-symmetric matrices, and so Proposition 2.1 yields no non-trivial conserved quantities.
However, using the canonical basis of g = so h given by E
[ kl ] =
E kl − E lk (so that [ kl ] indicates anti-symmetrized indices), one uses equation (15) to deduce the following novel result (see Appendix C.9): Theorem 2.2. When σ is a radial rescaling activation, we have:
[Equation]
(17) for any ( U, V ) ∈ Param , where the dots indicate derivatives with respect to gradient flow. Expanding the ( k, l ) entry of the matrix on the left-hand-side of (17), we obtain:
[Equation]
+ (cid:80) ms =1 r 2 V T ,s ; kl ˙ ϕ V T ,s ; kl
= 0 , where ( r U,s ; kl , ϕ U,s ; kl ) and ( r V T ,s ; kl , ϕ V T ,s ; kl ) are the 2D polar coordinates of the points ( U sk , U sl ) and ( V Tsk , V Tsl ) .
This is analogous to the “angular momentum” in 2D, that is: x ∧ ˙ x =
r 2 ˙ ϕ .
Intuitively, Theorem 2.2 implies that in every 2D plane ( k, l ) , the angular momenta of the rows of U and the columns of V sum to zero.
These results also apply to linear networks F ( x ) =
UV x , since rotational symmetries are linear.
A PPLICATIONS
We ran a limited set of experiments aimed at assessing the utility of the nonlinear group action and conserved quantities.
Some of the results are shown in Figure 3.
We find that Q can impact convergence rate and generalizability.
We also find the nonlinear action to be viable for ensemble building to improve robustness under some adversarial attacks.
Convergence rate and generalizability.
Conserved Q are fixed at initialization.
We find that values of Q can influence the convergence rate or the sharpness of the trained model.
In Appendix J, we [Figure] derive the relation between Q and convergence rate for two example optimization problems.
We show that initializing parameters at certain Q values accelerates convergence.
In Appendix K, we derive the relation between conserved quantities and sharpness of minima in a simple two-layer network, and show empirically that Q values affect sharpness (and possibly generalizability) in larger networks.
While Q is often unbounded, common initialization methods such as (Glorot &
Bengio, 2010) limit the values of Q to a small range (Appendix I).
As a result, only a small part of the minimum is reachable by the models.
Symmetries allow us to explore portions of flat minima which gradient descent rarely reaches.
Ensemble models.
Applying the nonlinear group action allows us to obtain an ensemble without any retraining or searching.
We show that even with stochasticity in the data, the loss is approximately unchanged under the group action.
The ensemble has the potential to improve robustness under adversarial attacks (Appendix L).  D ISCUSSION
In this paper, we have demonstrated that symmetries of the parameter space of neural networks lead to flat local minima and low-loss valleys in the loss landscape.
The linear symmetry group is determined by the choice of activation: this group consists of positive rescaling transformations in the case of homogeneous activations such as Leaky ReLU, and consists of orthogonal transformations in the case of radial rescaling activations.
Moreover, we present a non-linear, data-dependent parameter space symmetry that works for the full general linear group and a wider variety of activation functions.
Finally, we explain how parameter space symmetries also give rise to conserved quantities for gradient flow.
Such conserved quantities provide coordinates for flat local minima, and give insight into convergence rates. Our procedure for determining conserved quantities relies on infinitesimal symmetries, and is closely connected to Noether’s theorem. While our work sheds new light onto the link between symmetries and group, it contains several limitations, which merit further investigation. First, we have not been able to determine conserved quantities in the radial rescaling case, only a differential equation that gradient flow curves must satisfy.
In future work, we plan to address further the question of whether conserved quantities exist for rotation symmetries (e.g. in the radial rescaling case).
These will likely not be defined on the full parameter space, but only certain subsets.
Second, one major contribution of this paper is the nonlinear group action of Section 1.4.
However, our formulation only gurantees full GL h equivariance for batch size k = 1 .
In future work, we plan to explore more consequences and variations of this non-linear group action, with the hope of generalizing to greater batch size.
Finally, in many cases, parameter space symmetries lead to model compression: i.e., finding a lower-dimension space of parameters with the same expressivity of the original space of original space.
Such questions are beyond the scope of the current paper; this topic would be an area of future work.
C ONTENTS 1 Continuous symmetries in deep learning
The parameter space and loss function . . . . . . . . . . . . . . . . . . . . . . . .
Action of continuous groups and flat minima . . . . . . . . . . . . . . . . . . . . .
Equivariance of the activation function . . . . . . . . . . . . . . . . . . . . . . .
Nonlinear data-dependent symmetries and general equivariance . . Infinitesimal symmetries . .
2 Conserved quantities of gradient flow 3 Applications 4 Discussion
A Additional related works
B Extended derivations
B.1 Absorbing bias into weights . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.
Relation to Noether’s theorem.
C Neural networks: linear group actions
[-C.
 Running example-]C.1 Jacobians and differentials . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2 Neural network parameter spaces . . C.3 Action of continuous groups and infinitesimal symmetries . .
[-D.
-]{+C.
+}
The hidden symmetry group . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.5 Linear symmetries . .
C.6 Linear symmetries lead to extended, flat minima . . . . . . . . . . . . . . . . . . .
[-D.
-]C.6.
Flat directions in the two-layer case . . . . . . . . . . . . . . . . . . . . .
C.7 Conserved quantities . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.8 Conserved quantities from a group action . . . . . . . . . . . . . . . . . . . . . .
C.8.
Anti-symmetric case . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.9 Examples of conserved quantities for neural networks . . . . . . . . . . . . . . . .
C.9.
Conserved quantities for linear networks . . . . . . . . . . . . . . . . . .
C.9.
Conserved quantities for linear networks: two-layer case . . . . . . . . . .
C.9.
Conserved quantities for ReLU networks . . . . . . . . . . . . . . . . . .
C.9.
Conserved quantities for ReLU networks: two-layer case . . . . . . . . . .
C.9.
Conserved angular momentum for radial rescaling networks . . . . . . . .
C.9.
Conserved angular momentum for radial rescaling networks: two-layer case C.10 Jacobians: special cases . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D Neural networks: non-linear actions group actions
D.1 Set-up . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.2 Non-linear action: two-layer case . . . . . . . . . . . . . . . . . . . . . . . . . . . D.3 Non-linear action: multi-layer case . . . . . . . . . . . . . . . . . . . . . . . . . .
D.4 Discussion: increasing the batch size . . . . . . . . . . . . . . . . . . . . . . . . .
E Using nonlinearly transformed weights on different inputs
F The Lagrangian approach
F.
Lagrangian . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Second-order gradient flow conserved quantities . .
F. Addendum: more on the Lagrangian . . . . . . . . . . . . . . . . . . . . . . . . .
F.3.
Constrained uniqueness . . .
F.3.
Symmetries . .
F.3.
Limiting procedures . .
G Background: Lagrangians and Noether’s theorem
G.1 Euler-Lagrange equation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.2 First-order gradient flow cannot be derived from variational principle . . . . . . . .
G.3 Bregman Lagrangian . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.4 Relating γ to δt . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.5 Noether’s theorem . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
{+G.
+}
Relating conserved quantity to convergence rate . .
G.7 Conserved quantities in neural networks . . . . . . . . . . . . . . . . . . . . . . .
G.7.1 Two-layer linear network . . . . . . . . . . . . . . . . . . . . . . . . . .
G.7.2 Two-layer nonlinear neural networks . . . . . . . . . . . . . . . . . . . .
G.7.3 Two-layer nonlinear neural networks - alternative transformations . . . . .
H Gradient descent and drifting conserved quantities
H.1 Change in Q in gradient descent (linear layers) . . . . . . . . . . . . . . . . . . .
H.2 Empirical observations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
I Distribution of Q under Xavier initialization
J Conserved quantity and convergence rate
J. Example 1: ellipse . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.
Example 2: radial activation functions . . . . . . . . . . . . . . . . . . . . . . . . J.
Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
K Conserved quantity and generalization ability
K.1 Example: two-layer linear network with 1D parameters . . . . . . . . . . . . . . .
K.2 Experiments: two-layer networks
L Ensemble models
A A DDITIONAL RELATED WORKS
Mode connectivity / flat regions / ensembles In neural networks, the optima of the loss functions are connected by curves or volumes, on which the loss is almost constant (Freeman & Bruna, 2017;
Garipov et al., 2018; Draxler et al., 2018; Benton et al., 2021; Izmailov et al., 2018).
In these works, various algorithms have been proposed to find these low-cost curves, which provides a low-cost way to create an ensemble of models from a single trained model.
Other related works include linear mode connectivity (Frankle et al., 2020), using mode connectivity for loss landscape analysis (Gotmare et al., 2018), and studying flat region by removing symmetry (Pittorino et al., 2022).
Sharpness of minima and generalization Recent theory and empirical studies suggest that sharp minimum do not generalize well (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Petzka et al., 2021).
Explicitly searching for flat minimum has been shown to improve generalization bounds and model performance (Chaudhari et al., 2017;
Foret et al., 2020; Kim et al., 2022).
The sharpness of minimum can be defined using the largest loss value in the neighborhood of a minima (Keskar et al., 2017;
Foret et al., 2020; Kim et al., 2022), visualization of the change in loss under perturbation with various magnitudes on weights (Izmailov et al., 2018), singularity of Hessian (Sagun et al., 2017), or the volume of the basin that contains the minimum (approximated by Radon measure (Zhou et al., 2020) or product of the eigenvalues of the Hessian (Wu et al., 2017)).
Under most of these metrics, however, equivalent models can be built to have minimum with different sharpness but same generalization ability (Dinh et al., 2017).
Applications include explaining the good generalization of
SGD by examining asymmetric minimum (He et al., 2019), and new pruning algorithms that search for minimizers close to flat regions (Chao et al., 2020).
Parameter space symmetry and optimization While sets of parameter values related by the symmetries produce the same output, the gradients at these points are different, resulting in different learning dynamics (Kunin et al., 2021; Van Laarhoven, 2017).
This insight leads to a number of new advances in optimization.
Neyshabur et al.
and Meng et al.
(2019) propose optimization algorithms that are invariant to symmetry transformations on parameters.
Armenta et al.
(2020) and
Zhao et al.
(2022) apply loss-invariant transformations on parameters to improve the magnitude of gradients, and consequently the convergence speed.
The structures encoded in known symmetries have also led to new optimization methods and insights of the loss landscape.
Bamler & Mandt (2018) improves the convergence speed when optimizing in the direction of weakly broken symmetry.
Zhang et al.
(2020) discusses how symmetry helps in obtaining global minimizers for a class of nonconvex problems.
The potential relevance of continuous symmetries in optimization problems was also discussed in Leake & Vishnoi (2021), which also provides an overview of Lie groups.
B E XTENDED DERIVATIONS B. A BSORBING BIAS INTO WEIGHTS
Consider the affine function F ( X ) =
WX + b , where W ∈ R m × n , b ∈ R m and X ∈ R n × k
One can absorb b into the weight matrix by appending a row of ones 1 ∈ R k to x as its last row, defining (cid:101) X
=
( X, 1 ) ∈ R ( n +1)
× k , and appending b as the last column of the weight matrix (cid:102) W =
( W | b ) ∈ R m × ( n +1) .
We have
[Equation]
[Equation] For group actions on (cid:102) W ∈ R m × ( n +1) , we need to make sure b is not mixed with W , so that g · (cid:102) W = [Equation]
g · W | g · b ) .
The left action of GL n is defined in the natural way:
[Equation]
The right action of GL n becomes:
[Equation]
B. R ELATION TO N OETHER ’ S THEOREM .
We will now show how the approach in Tanaka & Kunin (2021) relates to our conservation law dQ/dt = (cid:68) ˙
θ , M θ (cid:69) = 0 .
Assuming a small time-step τ ≪ 1 , we can write GD as θ ( t + τ ) − θ ( t ) =
− ˜ ε ∇L ( θ ( t )) .
Expanding the l.h.s to second order in τ and discarding O ( τ 3 ) terms defines the 2nd order GF equation 2nd order GF:
[Equation]
[Equation]
Here ε = ˜ ε/τ .
To use Noether’s theorem, the dynamics (i.e. GF) must be a variational (EulerLagrange (EL)) equation derived from an “action” S ( θ ) (objective function), which for (21) is the time integral of Bregman Lagrangian Wibisono & Wilson (2015)
L
[Equation]
[Equation] where θ : R → Param is a trajectory (flow path) in Param , parametrized by t .
The variational
EL equations find the paths γ ∗ which minimize the action, meaning ∂S γ /∂γ | γ ∗ = 0 .
Noether’s theorem states that if M ∈ g is a symmetry of the action S ( θ ) (22) (not just the loss
L ( θ ) ), then the Noether current J M is conserved
Noether current:
[Equation]
Conservation:
[Equation]
Tanaka & Kunin (2021) also derived the Noether current (23), but concludes that because L ( θ , ˙ θ ) ̸ =
L ( θ ) , the symmetries are “broken” and therefore doesn’t derive conserved charges for the types of symmetries we discussed above.
However, while Tanaka & Kunin (2021) focuses on 2nd order GF, we note that our conserved Q were derived for first order GF , which is found from the τ → 0 limit of 2nd oder GF.
In this limit L → e t/τ L and thus symmetries of L also becomes symmetries of L .
When τ → 0 , 2nd order GF reduces to ε − 1 ˙ θ =
−∇L the conserved charge goes to
[Equation]
24) which means that we recover the invariance under infinitesimal action (12).
In fact, for linear symmetries and symmetric M ∈ g , J 0 M = dQ M /dt = 0 .
C N EURAL NETWORKS : LINEAR GROUP ACTIONS
In this appendix, we provide an extended discussion of the topics of Section 1, including full proofs of all results.
Specifically, after some technical background material on Jacobians and differentials, we specify our conventions for neural network parameter space symmetries.
In contrast to the discussion of the main text, we (1) assume that neural networks have biases, and (2) focus on the multi-layer case rather than just the two-layer case.
We then turn our attention to group actions of the parameter space that leave the loss invariant, and the resulting infinitesimal symmetries.
The groups we consider are all subgroups of a large group of change-of-basis transformations of the hidden feature spaces; we call this group the ‘hidden symmetry group’. We also compute the dimensions of generic extended flat minima in various relevant examples. Finally, we explore consequences of invariant group actions for conserved quantities. C. J ACOBIANS AND DIFFERENTIALS
In this section, we summarize background material on Jacobians and differentials.
We adopt notation and conventions from differential geometry.
Let U ⊂ R n be an open subset of Euclidean space R n , and let F : U → R m be a differentiable function.
Let F 1 , . .
F m : U → R be the components of F , so that F ( u ) =
( F 1 ( u ) , . . .
, F m ( u )) .
The Jacobian of F , also know as differential of F , at u ∈ U is the following matrix of partial derivatives evaluated at u :
[Equation]

The differential dF u defines a linear map from R n to R m , that is, an element of R m × n .
Observe that if F itself is linear, then, as matrices, dF u = F for all points u ∈ U .
If G : R m → R p is another differentiable map, then the chain rule implies that, for all u ∈ R n , we have:
[Equation]
In the special case the m = 1 , the differential is a 1 × n row vector, and the gradient ∇ u F of F at u ∈ R n is defined as the transpose of the Jacobian dF u :
[Equation]
(cid:12)(cid:12)
[Equation]
∂F∂x n (cid:12)(cid:12) u (cid:3) T = (cid:18)
∂F ∂x i (cid:12)(cid:12)(cid:12)(cid:12) u (cid:19) n i = ∈ R n
C. N EURAL NETWORK PARAMETER SPACES
Consider a neural network with L layers, input dimension n 0 = n , output dimension n L = m , and hidden dimensions given by n 1 , . .
. , n L − 1 .
For convenience, we group the dimensions into a tuple n = ( n 0 , n 1 , . . . , n L ) .
The parameter space is given by:
[Equation]
We write an element therein as a pair θ =
( W , b ) of tuples W =
( W L , . . .
, W 1 ) and b = ( b L , . . .
, b 1 ) , so that W i is a n i × n i − 1 matrix and b i is a vector in R n i for i = 1 , . . .
, L .
When n is clear from context, we write simply Param for the parameter space.
Fix a piecewise differentiable function σ i : R n i
→ R n i for each i = 1 , . . .
, L .
The activations can be pointwise (as is conventionally the case), but are not necessarily so.
The feedforward function
[Equation]
corresponding to parameters θ =
( W , b ) ∈ Param with activations σ i is defined in the usual recursive way.
To be explicit, we define the partial feedforward function F θ ,i =
F i : R n → R n i to be the map taking x ∈ R n to σ i ( W i F i − 1 ( x ) + b i ) , for i = 1 , . . .
, L , with F 0 = id R n 0 .
Then the feedforward function is F = F L .
The “loss function” L of our model is defined as:
[Equation]
where Data = R n 0 × R n L is the space of data (i.e., possible training data pairs), and Cost : R n L × R n L → R is a differentiable cost function, such as mean square error or cross-entropy.
Many, if not most, of our results involve properties of the loss function L that hold for any training data.
Hence, unless specified otherwise, we take a fixed batch of training data { ( x i , y i ) } ki =1 ⊆ Data , and consider the loss to be a function of the parameters only.
The above constructions generalize to multiple samples.
Specifically, instead of x ∈ R n 0 and y ∈ R n L , one has matrices X ∈ R n o × k and Y ∈ R n L × k whose columns are the k samples.
Additionally, one uses the Frobenius norm of n L × k matrices to compute the loss function.
The i -th partial feedforward function is F θ ,i : R n 0 × k
→ R n i × k , and we have the feedforward function
F θ : R n 0 × k
→ R n L × k .
(We use the same notation as in the case k = 1 ; the number of samples will be understood from context.)
Example C.1.
Consider the case L = 2 with no biases and no output activation.
The dimension vector is ( n 0 , n 1 , n 2 ) =
( n, h, m ) , so the parameter space is Param ( n, h, m )
=
R h × n × R m × h .
Taking the mean square error cost function, the loss function for data ( X, Y ) ∈ R n × k × R m × k takes the form L ( θ , ( X, Y ))
= 1 k ∥ Y − Uσ
( V X )
∥ 2 , where θ =
( W 2 , W 1 )
=
( U, V ) ∈ Param .
C. A CTION OF CONTINUOUS GROUPS AND INFINITESIMAL SYMMETRIES
Let G be a group.
An action of G on the parameter space Param is a function · : G × Param →
Param , written as g · θ , that satisfies the unit and multiplication axioms of the group, meaning
I · θ = θ where I is the identity of G , and g 1 · ( g 2 · θ ) =
( g 1 g 2 ) · θ for all g 1 , g 2 ∈ G .
Recall that we say an action G × Param → Param is a symmetry of Param with respect to L if it leaves the loss function invariant, that is:
[Equation]
The groups within the scope of this paper are all matrix Lie groups, which are topologically closed subgroups G ⊆ GL n ( R ) of the general linear group of invertible n × n real matrices.
Any smooth action of such a group induces an action of the infinitesimal generators of the group, i.e., elements of its Lie algebra.
Concretely, let g
=
Lie( G ) =
T I G be the Lie algebra, which can be thought of as a certain subspace of matrices in gl n = R n × n , or (equivalently) as the tangent space 6 .
at the identity I of G .
For every matrix M ∈ g , we have an exponential map exp M : R → G defined as exp M ( t ) =
(cid:80) ∞
k =0 ( tM ) k k ! .
Given an action of G on Param , the infinitesimal action of M ∈ g is a vector field M :
Infinitesimal action of M vector
[Equation]
[Equation]
Hence, the value of the vector field M at the parameter value θ is given by the derivative at zero of the function t (cid:55)→ (exp M ( t ) · θ ) .
In the case of a parameter space symmetry, the invariance of L translates into the following orthogonality condition, where the inner product ⟨ , ⟩ : Param × Param → R is calculated by contracting all indices, e.g. ⟨ A, B ⟩ = (cid:80) ijk... A ijk... B ijk... .
Proposition C.1 (Infinitesimal symmetry) .
Let G be a matrix Lie group and a symmetry of L .
Then the gradient vector field is point-wise orthogonal to the infinitesimal action of any M ∈ g :
[Equation]
Proof.
{+The gradient is the transpose of the Jacobian (see Section C.1), so the left-hand-side becomes d L θ (cid:0) ddt (cid:12)(cid:12) 
+}
(exp M ( t ) · θ ) (cid:1) .
We compute:
[Equation]
where the first equality follows by the chain rule, the second equality uses the invariance of L , and the third equality follows since L ( θ ) does not depend on t .
Next, we comment on the case of a linear action.
Observe that the parameter space is a vector space of dimension d = dim( Param )
= (cid:80) Li
=1 n i (1 + n i − 1 ) .
Hence, there is an isomorphism
Param ≃ R d which flattens any tuple θ =
( W , b ) into a vector in R d .
We can identify the group
GL( Param ) of all invertible linear transformations of the parameter space with the group GL d ( R ) of invertible d × d matrices, and the Lie algebra of GL( Param ) with gl d = R d × d .
Suppose G acts linearly on the parameter space.
Then we can identify 7 G with a subgroup of GL( Param ) , acting on Param with matrix multiplication.
Similarly, we can identify the Lie algebra g of G with a Lie subalgebra of gl d = R d × d .
In this case, the infinitesimal action is given by matrix multiplication:
M θ = M · θ .
We recover Equation (12).
C. T HE HIDDEN SYMMETRY GROUP
Consider a neural network with L layers and dimensions n .
The hidden symmetry group corresponding to dimensions n is defined as the following product of general linear groups:
[Equation]
An element is a tuple of invertible matrices
g =
( g 1 , . .
, g L − 1 ) , where g i ∈ GL n i .
Consider the action of the hidden symmetry group on the parameter space given by
[Equation]
where g 0 = id n 0 and g L = id n L .
This action amounts to changing the basis at each hidden feature space.
The Lie algebra of GL n hidden is
[Equation]
and the infinitesimal action of the tuple M ∈ gl n hidden is given by:
[Equation]
where we set M 0 and M L to be the zero matrices.
Example C.2.
In the case L = 2 , with dimension vector n = ( n, h, m ) and no biases.
The hidden symmetry group is GL h with Lie algebra gl h .
The action of the group and the infinitesimal action of the Lie algebra are given by:
[Equation]
C. L INEAR SYMMETRIES
Consider a feedforward fully-connected neural network with widths n = ( n 0 , . .
. , n L ) , so that the parameters space consists of tuples of weights and biases θ =
( W i ∈ R n i × n i − 1 , b i ∈ R n i ) Li
=1 .
For each hidden layer 0 < i < L , let G i be a subgroup of GL n i , and let π i : G i → GL n i ( R ) be a representation (in many cases, we take π i ( g ) = g ).
Hence the product G = G 1 × G L − 1 is a subgroup of the hidden symmetry group GL n hidden .
Define an action of G on Param via ∀ g
[Equation]
g 1 , . . .
[Equation]
g L
[Equation]
{+where g 0 and g L are the identity matrices id n 
+}
and id n L , respectively.
This is a version of the action defined in (28), with the addition of the twists resulting from the representations π i .
We now consider the resulting infinitesimal action.
For each i , the representation π i induces a
Lie algebra representation d ( π i ) : g i → gl n i .
The infinitesimal action of the Lie algebra g = g 1 × · · · × g L induced by 29 is given by:
[Equation]
The proof of the first part of the following Proposition proceeds by induction, where the key computation is that of from (4).
The second part relies on (12).
Proposition C.2.
Suppose that, for each i = 1 , . . .
, L , the activation σ i intertwines the two actions of G i , that is, σ i ( g i z i ) =
π i ( g i ) σ ( z i ) for all
g i ∈
G i , z i ∈ R n i .
Then:
(Combined equivariance of activations)
The action of G =
G 1 × · · · × G L defined in 29 is a symmetry of the parameter space.
(Infinitesimal equivariant action)
The action of g = g 1 × · · · × g L defined in 30 satisfies ⟨∇ θ L , M · θ ⟩ for all θ ∈ Param
and all M ∈ g .
Proof.
Let g
= ( g i )
L − 1 i =1 ∈ G , so that g i ∈ G i ⊆ GL n i .
{+As usual, set g 
+}
= id n 0 and g L = id n L .
Also set π 0 and π L to be the identity maps on GL n 0 and GL n L , respectively.
Fix parameters θ and an input value x ∈ R n .
We show by induction that the following relation between the partial feedforward functions holds:
[Equation]
for i = 0 , . . .
, L .
The base step is trivial.
For the induction step, we use the recursive definition of the partial feedforward functions:
[Equation]
Hence θ and g · θ define the same feedforward function.
Since the loss function depends on the parameters only through the feedforward function, the first claim follows.
The second claim is a consequence of Proposition 1.6.
Note that two-layer case of the above result amounts to (4) (the argument can be simplified in that case).
While Proposition C.2 is stated for feedforward networks, it can easily be adopted to more general settings, such as networks with skip connections and quiver neural networks.
Denoting the
Jacobian of σ i : R n i
→ R n i at z ∈ R n i by d ( σ i ) z ∈ R n i × n i , the infinitesimal form of version of σ i ( gz ) =
π i ( g ) σ i ( z ) is
[Equation]
When the activation is pointwise, we have Mz ⊙ σ
′ i ( z ) =
d ( π ) i
( M ) σ i ( z ) , where ⊙ denotes elementwise multiplication.
We now illustrate Proposition C.2 through examples in the two-layer case with input dimension n , hidden dimension h , hidden activation σ , and output dimension m .
Example C.3.
Linear networks For linear networks, we have σ ( x ) = x .
One can take π ( g ) =
g and G = GL h
( R ) .
Example C.4.
Homogeneous activations Suppose the activation σ : R h
→ R h is homogeneous , so that (1) σ is applied pointwise in the standard basis, and (2) t there exists α > 0 such that σ ( cz ) =
c
α σ ( z ) for all c ∈ R > 0 and z ∈ R h .
These σ are equivariant under the positive scaling group G ⊂ GL h consisting of diagonal matrices with positive diagonal entries.
For g ∈ G , we have g = diag( c ) is a diagonal matrix with c
=
( c 1 , . .
. , c h ) ∈
R h> 0 .
For z = ( z 1 , . . .
, z h ) ∈
R h , we have σ ( gz ) =
(cid:80)
j σ ( c j z j )
=
(cid:80) j c αj σ ( z j ) =
g α σ ( z ) .
Hence, the equivariance condition holds with π ( g ) =
g α .
Since dπ ( M ) = αM for any element M of the Lie algebra g of G , the infinitesimal version of rescaling invariance of homogeneous σ becomes Mz ⊙ σ
z ) =
αMσ ( z ) .
Example C.5.
LeakyReLU
This is a special case of homoegeneous activation, defined as σ ( z ) =
max( z, 0) − s min( z, 0) , with s ∈ R > 0 .
We have α = 1 , and π ( g ) = g .
Since σ ( z ) = zσ ′ ( z ) , infinitesimal equivariance becomes Mz ⊙ σ ′ ( z )
=
Mσ ( z ) .
Example C.6.
Radial rescaling activations A less trivial example of continuous symmetries is the case of a radial rescaling activation (Ganev et al., 2022) where for z ∈ R h \ { 0 } , σ ( z ) =
f
( ∥ z ∥ ) z for some function f : R → R .
Radial rescaling activations are equivariant under rotations of the input: for any orthogonal transformation g ∈ O ( h ) (that is, g T g =
I )
we have σ ( gz ) =
gσ ( z ) for all z ∈ R h .
Indeed, σ ( gz ) =
f ( ∥ gz ∥ )( gz ) =
g ( f ( ∥ z ∥ ) z ) =
gσ ( z ) , where we use the fact that ∥ gz ∥ =
z T g T gz =
z T z = ∥
z ∥ for g ∈ O ( h ) .
Hence, (4) is satisfied with π ( g )
=
g . C. L INEAR SYMMETRIES LEAD TO EXTENDED , FLAT MINIMA In this section, we show that, in the case of a linear group action, applying the action of any element of the group to a local minimum yields another local minimum.
This fact is a corollary of a more general result; in order to describe it and remove ambiguity, we include the following clarifications.
Let G be a matrix Lie group acting as a linear symmetry.
Fix a basis ( θ 1 , . .
, θ d ) of the parameter space.
The gradient ∇ θ L of the loss L at a point θ ∈ Param in the parameter space as another vector in Param ≃ R d , whose i -th coordinate is the partial derivative ∂
L ∂θ i (cid:12)(cid:12) θ .
Hence, it makes sense to apply the group action to the gradient: g · ∇ θ L .
We regard vectors in Param ≃ R d as column vectors with d rows.
Thus, the transpose of any vector is a row vector with d columns.
In the case of the gradient, its transpose at θ matches the Jacobian d L θ ∈ R 1 × d of L (see Appendix C.1), that is: d L θ =
( ∇ θ L ) T .
Alternative notation for the Jacobian is d L θ 0 =
∂
L ∂ θ (cid:12)(cid:12) θ 0 , where we now use θ as a dummy variable and θ 0 ∈ Param as a specific value.
As noted above, we are interested in matrix Lie groups G ⊆ GL
d ( R ) =
GL( Param ) , and assume that the matrix transpose g T belongs to G for any g ∈ G .
These assumptions hold in all examples of interest.
We have the following reformulation of Proposition 1.2:
Proposition C.3.
Suppose the action of G on the parameter space is linear and leaves the loss invariant.
Then the gradients of L at any θ 0 and g · θ 0 are related as follows:
[Equation]
If θ ∗ is a critical point (resp.
local minimum) of L , then so is g · θ ∗ .
Sketch of proof.
Let T g : Param → Param be the transformation corresponding to g ∈ G .
The the Jacobian d L θ 0 is given by:
[Equation]
where we use the definition of the Jacobian, the invariance of the loss ( L ◦ T g = L ), the chain rule, and the linearity of the action.
The result follows from applying T g − 1 on the right to both sides, and taking transposes (see Appendix C.1).
The last statement follows from the invariance of L under the action of G , and the fact that ∇ θ ∗ L = 0 at a critical point θ ∗ of L .
We conclude that, if θ ∗ is a critical point, then the set { g · θ | g ∈ G } belongs to the critical locus.
This set is known as the orbit of θ under the action of G , and is isomorphic to the quotient
G/ Stab G ( θ ) , where Stab G ( θ ) = { g ∈ G | g · θ = θ } is the stabilizer subgroup of θ in G .
In the case of a linear action, the orbit is a smooth manifold.
While the results above imply that the critical locus is a union of G -orbits, they do not imply, in general, that the critical locus is a single G -orbit.
They also do not rule out the case that the stabilizer is a somewhat ‘large’ subgroup of G , in which case the orbit would have low dimension.
However, in many cases, there is a topologically dense subset of parameter values θ ∈ Param whose orbits all have the same dimension. We call such an orbit a ‘generic’ orbit. We now turn our attention to examples of two-layer networks where such a generic orbit exists. C.6. F LAT DIRECTIONS IN THE TWO - LAYER CASE
Recall that the parameter space of a two-layer network is Param
= R m × h ×
R h × n , where the dimension vector is ( m, h, n ) , and we write elements as ( U, V ) .
The action of GL h is g · ( U, V ) =
( Ug − 1 , gV ) .
Let Param ◦
⊆
Param be the subset of pairs ( U, V ) where each of U and V have full rank.
This is an open dense subset of Param , and is preserved by the GL h -action.
Proposition C.4.
The GL h -orbit of each element of Param ◦ has dimension
[Equation]
Proof.
Fix ( U, V ) ∈ Param ◦ .
Suppose h ≤ n , so that h 2 − max(0 , h − n ) max(0 , h − m )
=
h 2 .
Then V ∈ R h × n defines a surjective linear map, so has a right inverse V † ∈ R n × h .
If g ∈ GL h belongs to the stabilizer of ( U, V ) , then we have gV =
V .
Applying V † on the right to both sides, we obtain g = id h .
Thus, the stabilizer of ( U, V ) is trivial, and the orbit has dimension equal to the dimension of the group, namely h 2 .
The case h ≤ m is similar.
Now suppose h > max( n, m ) .
In this case, h 2 − max(0 , h − n ) max(0 , h − m )
=
h ( n + m ) − nm .
Set U
=
[ id m 0] ∈ R m × h and V
= (cid:20) i
d n 0 (cid:21) ∈
R h × n , so that the last h − m rows of U 0 are zero, and the last h − n columns of V 0 are zero.
Then, by the rank assumption, there exists g 1 ∈ GL h such that Ug 1 = U 0 , and there exists g 2 ∈ GL h such that g − 1 2 V
=
V 0 .
Without loss of generality, we can take g 1 and g 2 such that det( g 1 ) > 0 and det( g 2 ) > 0 .
Thus, both g 1 and g 2 belong to the component of the identity in GL h .
Next, consider the action of G = GL h on full rank matrices in R m × h and R h × n individually.
We have that Stab G ( U ) =
g 1 Stab G ( U 0 )
g − 1and Stab G ( V ) =
g 2 Stab G
( V 0 )
g − 1 2 .
The stabilizer in
GL h of the pair ( U, V ) ∈ Param ◦ can be written as:
[Equation]
Since g 1 and g 2 belong to the connected component of the identity, the dimension of ( g 1 Stab G ( U 0 ) g − 1 1 ) ∩ ( g 2 Stab G
( V 0 )
g − 1 2 ) is equal to the dimension 8 of Stab G ( U 0 )
∩ Stab G ( V 0 ) .
Hence we reduce the problem to computing the dimension of Stab G ( U 0 ) ∩ Stab G ( V 0 ) .
The dimension γ
To this end, observe that a matrix g belongs to Stab G ( U 0 ) (resp. Stab G ( V 0 ) )
if and only if is of the form:
[Equation]
where the lower left ∗ ∈ R ( h − m ) × m and the lower right ∗ ∈ GL h − m (resp.
upper right ∗ ∈
R n × ( h − n ) and lower right ∗ ∈ GL h − n ) are arbitrary.
If m ≥ n , taking the intersection amounts to considering matrices of the form:
[Equation]
where the rows and columns are divided according to the partition h = n + ( m − n ) + ( h − m ) .
If m ≥ n , taking the intersection amounts to considering matrices of the form:
[Equation]
where the rows and columns are divided according to h
= m +( n − m )+( h − n ) .
In both cases, the dimension of the intersection is ( h − n )( h − m ) =
h 2 − hn − hm + nm .
We obtain the dimension of the orbit as: h 2 − ( h − n )( h − m ) = h ( n + m ) − nm.
Recall that the symmetry group for homogenous activations is the coordinate-wise positive rescaling subgroup of GL h , consisting of diagonal matrices with positive entries along the diagonal. We denote this subgroup as T + ( h ) . Similarly, the symmetry group for radial rescaling activation is the orthogonal group O ( h ) . For linear networks, the activation is the identity function, so the symmetry group is all of GL h .
Corollary C.5.
The orbit of a point in Param ◦ under the appropriate symmetry group is given by:
Type of activation
[Table]
Linear
GL h h 2 − max(0 , h − n ) max(0 , h − m )
Homogeneous
T + ( h ) min( h, max( n, m ))
Radial rescaling
O ( h ) (cid:40)(cid:0) h 2 (cid:1) if h ≤ max( n, m ) (cid:0) h 2 (cid:1) − (cid:0) h − max( m,n )(cid:1) otherwise
Proof.
Adopt the notation of the proof of the above Proposition.
The stabilizer in T + ( h ) of ( U 0 , V 0 ) is the intersection of the stabilizer in GL h of ( U 0 , V 0 ) and T + ( h ) .
This intersection is easily seen to have dimension max(0 , h − max( n, m )) .
Subtracting this from dim( T + ( h ))
=
h , we obtain the result for the homogeneous case.
For the orthogonal case, the stabilizer in O ( h ) of ( U 0 , V 0 ) is the intersection of the stabilizer in GL h of ( U 0 , V 0 ) and O ( h ) .
This intersection has dimensionif h ≤ max( n, m ) and (cid:0) h − max n,m(cid:1) otherwise.
Subtracting from dim( O ( h ))
= (cid:0) h 2 (cid:1) , we obtain the result for the radial rescaling case.
C. C ONSERVED QUANTITIES
[-where, for convenience, we set g 
-]
We now turn our attention to gradient flow and conserved quantities.
[-and b ( h ) i − 
-]
In this section, we give a formal definition of a conserved quantity.
[-47) for some constant λ > 
-]
[-k P kj = 
-]
Let V
[-Case of m = 
-]
= R d be the standard vector space of dimension d .
Suppose L : V → R is a differentiable function.
Let Flow t : V → V be the flow for time t along the reverse gradient vector field, so that:
[Equation] Note that Flow 0 is the identity on V , and, for any s, t , the composition of Flow s and Flow t is
Flow s + t .
We will write v ( t ) for Flow t ( v ) , so that ˙ v ( s ) =
−∇ v ( s ) L . A conserved quantity is a function Q : V → R that satisfies either of the following equivalent conditions:
For any t , we have Q ( v ( t ))
=[-
-]
Q ( v ) . {+Let ˙ Q ( v ) = ddt (cid:12)(cid:12) 
 [ t (cid:55)→ Q ( v ( t ))] be the derivative of Q along the flow.+}
Then ˙ Q ≡ 0 .
The gradients of Q and L are point-wise orthogonal, that is, ⟨∇ v Q, ∇ v L⟩ = 0 for all v ∈ V .
The equivalence of (1) and (2) are immediate.
To show the equivalence of the third and second statements, let v ∈ V and compute:
[Equation]
where we use the definition of the flow in the second equality, and the chain rule in the third.
We note that, if f : R → R is any function, and Q is a conserved quantity, the f ◦ Q is also a conserved quantity.
Additionally, any linear combination of conserved quantities is again a conserved quantity.
Let Conserv( V , L ) denote the vector space of conserved quantities for the gradient flow of L : V → R .
For any v ∈ V , there a map:
[Equation]
taking a conserved quantity to the value of its gradient at v .
By the above discussion, the map ∇ v is valued in the kernel of the differential dL v .
C. C ONSERVED QUANTITIES FROM A GROUP ACTION
Let G be a subgroup of the general linear group GL d ( R ) .
Thus, there is a linear action of G on
V = R d .
Suppose the function L is invariant for the action of G , that is,
[Equation]
Let g =
Lie( G ) be the Lie algebra of G , which is a Lie subalgebra of gl d = R d × d .
The infinitesimal action of g on V is given by g × V → V , taking ( M, v ) to Mv .
Proposition C.6.
Let L : V → R be a G -invariant function, and let M ∈ g .
For any v ∈ V , the gradient of L and the infinitesimal action of M are orthogonal:
[Equation]
Suppose γ : ( a, b ) → V is a gradient flow curve for L .
Then:
[Equation]
Suppose M T belongs to g .
Then the function
[Equation]
is a
conserved quantity for the gradient flow of L .
Proof.
For the first claim, observe that the invariance of L implies that the left diagram commutes: G [Equation]
[Equation]
where inc : G → R d × d is the inclusion (which passes through the inclusion of G in to GL d ( R ) ), ev v : R d × d → R d be the evaluation map at v , and the left vertical map is the constant map at { v } .
Indeed, the clockwise composition is g (cid:55)→ L ( gv ) , which is equal to the constant map at g (cid:55)→ L ( v ) .
The chain rule implies that taking Jacobians at the identity 1 of G results in the commutative diagram on the right.
where g is the Lie algebra of G , identified with the tangent space of G at the identity; the tangent space of the vector space R d at v is canonically identified with R d ; and the the zero appears because the tangent space of a single point is zero.
The derivative of the inclusion map is the inclusion g (cid:44) → gl d = R d × d , while the derivative the evaluation map is itself as it is a linear map.
Hence, for M ∈ g , we have:
[Equation]
The first claim follows.
The second claim is consequence of the first claim, together with the definition of a gradient flow curve.
For the third claim, we take the derivative of the composition of Q M with a gradient flow curve γ :
[Equation]
Both terms in the last expression are constantly zero by the second claim.
Hence Q M is constant on any gradient flow curve, and so it is a conserved quantity.
We summarize some of the results and constructions of this section diagrammatically.
Let g sym denote the vector space of symmetric matrices in g (this is not a Lie subalgebra in general).
Observe that g ∩ g T is the set of all M ∈ g such that M T ∈ g .
Let Infin G ( V , L ) denote the vector space of infinitesimal-action conserved quantities for the gradient flow of the G -invariant function L : V → R .
We have:
[Equation]
where the map g ∩ g T → g sym takes M to its symmetric part 12 (cid:0) M + M T (cid:1) , while the map g sym (cid:44) → g ∩ g T is the natural inclusion.
We note that g ∩ g T is the Lie algebra of the group G ∩ G T , while g sym is in general not a Lie algebra.
By definition, the vector space Infin G ( V , L ) is the image of the map M → Q M defined on g ∩ g T .
It is straightforward to verify the following result:
Corollary C.7.
The map M → Q M establishes an isomorphism of vector spaces: g sym ∼ −→
Infin G ( V , L ) .
As discussed in Section ??
, applying a symmetry g to a minimum θ ∗ of L yields another minimum g · θ ∗ . Using flattened θ ∈ R d it is easy to show that acting with g changes some Q M ( θ ) = θ T M θ .
Let g = exp M ′
( t )
≈
I + tM ′ , with M ′ ∈ g and 0 <
η ≪ 1 be a g ∈ G close to identity.
We have
Q M ( g · θ ) =
Q M + t θ T (cid:16) M ′ T M + MM ′ (cid:17) θ
+ O ( η 2 ) .
Thus, whenever M ′ T M + MM ′ ̸ = 0 , applying g changes the value of Q M .
Therefore, Q M can be used to parameterize the flat minima.
However, for anti-symmetric M , we could not find nonzero Q explicitly.
C.8. A NTI - SYMMETRIC CASE
Suppose M ∈ g is anti-symmetric, so M = − M T .
Let γ : ( a, b ) → R d be a gradient flow curve.
Write γ in coordinates as γ =
( γ 1 , . . . , γ d ) .
Proposition C.6 implies that (˙ γ ( t ))
T Mγ = 0 .
Hence we have:
[Equation]
where r ij = r ij ( t ) is equal to (cid:112) γ
i ( t ) 2 − γ j ( t ) 2 and θ ij = θ ij ( t ) is the angle between the i -th coordinate axis and the ray from the origin to the projection of γ ( t ) to the ( i, j ) -plane.
One verifies the last equality using the definition of θ ij in terms of the arctangent of the quotient γ j /γ i .
We see that ( r ij ( t ) , θ ij ( t )) are the polar coordinates for the point ( γ i ( t ) , γ j ( t ))
∈ R 2 .
Case d = 2 .
Then M = (cid:20) [Equation] for some nonzero a ∈ R , and so:
[Equation]
where r ( t ) and θ ( t ) are polar coordinates.
Setting the final expression equal to zero, we obtain that θ ( t ) is constant along any flow line γ ( t ) that begins away from the origin.
Case d = 3 .
Then M =
[Equation]
for some a, b, c ∈ R , and so:
[Equation]
C. E XAMPLES OF CONSERVED QUANTITIES FOR NEURAL NETWORKS
We now compute conserved quantities for gradient flow on neural network parameter spaces in the case of linear, homogeneous, and radial networks.
In each case, we state results first in the for a general multi-layer network, and then for the running example of a two-layer network.
Throughout, ⊙ denotes the Hadamard product of matrices, defined by entrywise multiplication.
We also set τ ( M ) to be the sum of all entries in a matrix M .
We note that, for square matrices M and N of the same size, τ ( M ⊙ N ) =
Tr( M T N ) , which is the same as the inner product of the flattened versions of M and N .
The notation for the running example of a two-layer network is as follows.
We set the input and output dimensions both equal to one, hidden dimension equal to two, and no bias vectors.
The hidden layer activation is σ : R 2 → R 2 .
The parameter space is Param
= R 2 × 1 × R 1 × 2 , with elements written as a pair of matrices: ( U, V ) =
(cid:18) [ uu 2 ] , (cid:20)
[Equation]
v
[Equation]
(cid:21)(cid:19) .
The hidden symmetry group is GL 2 , with action given by:
[Equation]
The Lie algebra gl 2 of GL 2 ( R ) consists of all two-by-two matrices.
C.9. C ONSERVED QUANTITIES FOR LINEAR NETWORKS
Suppose a neural network with L layers has the identity activation σ i
= id n i in each layer, so that the resulting network is linear.
Then it is straightforward to verify that the networks with parameters g · ( W , b ) and ( W , b ) have the same feedforward function.
Consequently, the loss is invariant for the group action: its value the original and transformed parameters is the same for any choice of training data.
(As we will see below, for more sophisticated activations, one needs to restrict to a subgroup of the hidden symmetry group to achieve such invariance.)
Suppose M ∈ gl n hidden is such that M i ∈
gl n i is symmetric for each i .
The conserved quantity implied by Proposition C.6 is:
[Equation]
We examine these conserved quantities in the following convenient basis for the space of symmetric matrices in gl n hidden .
For j = 1 , . . .
, L and { k, ℓ } ⊆ { 1 , . .
. , n j } , set:
[Equation]
ℓ where E ( n j ) kℓ is the elementary n j × n j matrix with the entry in the k -th row and ℓ -th column equal to one, and all other entries equal to zero.
Then one computes:
[Equation]
In other words, we take the sum of the following three terms: the product of the k -th and ℓ -th entries of the bias vector b j , the dot product of the k -th and ℓ -th rows of W j , and the dot product of the k -th and ℓ -th columns of W j +1 . In particular, we see that every entry of the matrix [Equation]
is a conserved quantity valued in gl n i rather than in R .
Additionally, we have a moment map :
[Equation] C.9. C ONSERVED QUANTITIES FOR LINEAR NETWORKS : TWO - LAYER CASE
In the two layer case of a linear network, we have that that the single hidden activation is the identity: σ = id 2 : R 2 → R 2 .
The hidden symmetry group is GL 2 with Lie algebra all of gl 2 .
The space of symmetric matrices in gl 2 is spanned by the matrices:
E 11 = (cid:20)
[Equation]
0(cid:21) ,
[Equation]
(cid:20) 00(cid:21) , and E (1 , 2) = (cid:20) 01(cid:21) . The corresponding conserved quantities are:
[Equation]
Thus, we obtain a three-dimensional space of conserved quantities. (Since GL 2 also contains the orthogonal group O (2) , Equation 36 below holds along any gradient flow curve.) C.9. C ONSERVED QUANTITIES FOR R E LU NETWORKS The pointwise ReLU activation commutes with positive rescaling, so we consider the subgroup of the hidden symmetry group consisting of tuples of diagonal matrices with positive diagonal entries, that is:
[Equation]
This subgroup, also known as the positive coordinate-wise rescaling subgroup, is isomorphic to the product ( R > 0 ) (cid:80) L − 1 i =1 n i .
Its Lie algebra is spanned by the elements E ( j ) kk defined above, for j = 1 , . . .
, L − 1 and k = 1 , . . . , n j .
The conserved quantity implied by Proposition C.6 is:
[Equation]
In other words, we take the sum of the following three terms: the square of the k -th entry of the bias vector b j , the norm of the k -th row of W j , and the norm of the k -th column of W j +1 .
C.9. C ONSERVED QUANTITIES FOR R E LU NETWORKS : TWO - LAYER CASE
In the two-layer case, the positive rescaling group is:
[Equation] positive [Equation] The Lie algebra of G is the two-dimensional space of diagonal matrices in gl 2 (with not necessarily positive diagonal entries).
In other words, g is spanned by the matrices E 11 = (cid:20) 10(cid:21) and E 22 = (cid:20) 00(cid:21) .
One computes the conserved quantities corresponding to these elements as:
[Equation]
C.9. C ONSERVED ANGULAR MOMENTUM FOR RADIAL RESCALING NETWORKS
Suppose each σ i is a radial rescaling activation σ i ( z )
=
λ i ( | z | )
z , where λ i : R → R is the rescaling factor.
Each such activation commutes with orthogonal transformations, so we consider the subgroup of the hidden symmetry group consisting of tuples of orthogonal matrices:
[Equation]
The Lie algebra of this subgroup consists only of anti-symmetric matrices, and so there are no infinitesimal-action conserved quantities.
However, given an anti-symmetric matrix M i ∈ gl n i for each i , any gradient flow curve satisfies the following differential equation (encoding conservation of angular momentum): [Equation] (cf. Section C.8.1).
An equivalent way to write this equation is:
[Equation]
Indeed, one uses the facts that τ ( A ⊙ B ) =
Tr( A T B ) ,
Tr( A T ) =
Tr( A ) , and Tr( AB ) =
Tr( BA ) , for any two matrices A, B of the appropriate size in each case.
Using a basis of anti-symmetric matrices, one can show that the matrix
[Equation]
is equal to zero:
ν i ( W , b ) = 0 .
Note that ν i depends on taking derivatives with respect to the flow.
In fact, ν i is more properly formulated as a function on the tangent bundle T ( Param ) of
Param , which is then evaluated on the gradient flow vector field.
Similarly, we have a moment map T ( Param ) → gl ∗ n hidden , and the gradient flow vector field is contained in the preimage of zero.
We omit the details.
A basis for the space of anti-symmetric matrices in gl n hidden is given by:
[Equation]
where j = 1 , . . .
, L − 1 , and k, ℓ ∈ { 1 , . .
. , n j } satisfy k < ℓ .
The differential equation corresponding to E ( j ) k ≤ ℓ is given by:
[Equation]
where ( r b j ; kℓ , θ b j ; k,ℓ ) are the polar coordinates of the image of b j under the projection R n j → Rwhich selects only the k -th and ℓ -th coordinates.
Similarly, for any pair matrix entries we have a projection R n j × n j − 1 → R 2 and can take the polar coordinates of the image of W j under this projection.
C.9. C ONSERVED ANGULAR MOMENTUM FOR RADIAL RESCALING NETWORKS : TWO - LAYER CASE
In the two-layer radial rescaling case, suppose the dimension vector is ( n, h, m ) , and that there are no bias vectors.
For U ∈ R m × h , V ∈ R h × n and M ∈ so ( h ) , we have: (cid:68) ˙
θ , M · θ (cid:69) =
(cid:68) ( ˙ U, ˙ V ) , ( − UM, MV ) (cid:69) = − Tr( ˙ U T UM ) +
Tr( ˙ V T MV ) =
[-g  g(cid:21) ∈ GL 2 ( R ) | g 
-]Tr( V
[Equation]
M
[Equation]
Hence we obtain the differential equation:
[Equation]
In the case where ( n, h, m ) = (1 , 2 , 1) , we have the two by two orthogonal group:
[Equation]
The Lie algebra of G consists of anti-symmetric matrices in gl 2 , and contains no non-zero symmetric matrices.
Hence, we do not obtain any conserved quantities from the infinitesimal action in this case.
However, using the element [Equation] ∈ g , we obtain that the following differential equation holds along any gradient flow curve:
[Equation]
where ( r U , θ U ) are the polar coordinates of ( u 1 , u 2 ) ∈ R 2 , and similarly for ( r V , θ V ) .
Note that the left-hand side of Equation 36 is a function of t ; so if γ : ( a, b ) → V is a gradient flow curve, then a more precise version of the equation is (cid:0) r [Equation] U ◦
[Equation]
for all t .
C. J ACOBIANS : SPECIAL CASES
We conclude this appendix with a side remark on special cases of the Jacobian formalism. Manifolds. Suppose M and N are smooth manifolds, and suppose F : M → N is a smooth map. The differential of F at m ∈ M is a linear map between the tangent spaces:
[Equation]
The map dF m is computed in local coordinate charts as the Jacobian of partial derivatives.
If G : N → L is another smooth map, then the chain rule becomes d ( G ◦ F ) m = dG F ( m ) ◦ dF m , for any m ∈ M . Matrix case. Suppose L : R m × n → R is a differentiable function.
In this case, we regard the
Jacobian at W ∈ R m × n as an n × m matrix: dL W =   [Equation] 11 (cid:12)(cid:12)(cid:12)(cid:12) W [Equation]
21 (cid:12)(cid:12)(cid:12)(cid:12) W · · · [Equation] 1 (cid:12)(cid:12)(cid:12)(cid:12) W [Equation]
12 (cid:12)(cid:12)(cid:12)(cid:12) W
[Equation] 22 (cid:12)(cid:12)(cid:12)(cid:12) W · · · ∂L∂w m 2 (cid:12)(cid:12)(cid:12)(cid:12) W ... [Equation] n (cid:12)(cid:12)(cid:12)(cid:12) W ∂L∂w
[Table]
n (cid:12)(cid:12)(cid:12)(cid:12) W · · ·
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Table]

[Equation]
n where w ij are the matrix coordinates. If F : R → R m × n is a differentiable function, we regard its Jacobian at s ∈ R as a m × n matrix: dF t =   dF 11 dt (cid:12)(cid:12)(cid:12)(cid:12) s dF 12 dt (cid:12)(cid:12)(cid:12)(cid:12) s · · · dF 1 n dt (cid:12)(cid:12)(cid:12)(cid:12) s dF 21 dt (cid:12)(cid:12)(cid:12)(cid:12) s dF 22 dt (cid:12)(cid:12)(cid:12)(cid:12)
s · · · dF 2 n dt (cid:12)(cid:12)(cid:12)(cid:12) s ...
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation] (cid:12)(cid:12)(cid:12)(cid:12) [Equation] where F ij : R → R are the coordinates of F .
Then the chain rule becomes:
[Equation]
In other words, the derivative of the composition L ◦ F at s ∈ R is the trace of the product of the matrices dL F ( s ) ∈ R n × m and dF s ∈ R m × n .
D N EURAL NETWORKS : NON - LINEAR ACTIONS GROUP ACTIONS
In this section, we consider a non-linear action of the hidden symmetry group on the parameter space.
This action has the advantage that exists for a wider variety of activation functions (such as the usual sigmoid, which has no linear equivariance properties), and that it is defined for the full general linear group.
However, in constrast to the linear action, the non-linear action is datadependent: the transformation of the weights and biases depends on the input data.
We adopt the notation of Section C.2. In particular, consider a neural network with L layers and widths n = ( n 0 , n 1 , . . . , n L ) .
The parameter space is given by: [Equation] We write an element therein as a pair θ = ( W , b ) of tuples W =
( W L , . . .
, W 1 ) and b = ( b L , . . .
, b 1 ) .
Fix activations σ i : R n i
→ R n i for each i = 1 , . . .
, L .
Let
[Equation]
be the feedforward function corresponding to parameters θ =
( W , b ) ∈ Param with activations σ i .
Taking the parameters into account, we form the following extension of the feedforward function
F
[Equation]
One can also define the extension of the partial feedforward function (cid:101) F i : Param × R n → R n i as ( θ , x ) (cid:55)→ F θ ,i ( x ) .
D. N ON - LINEAR ACTION : TWO - LAYER CASE
Consider a two-layer network with dimension vector ( m, h, n ) , no bias vectors, and no output activation.
The parameter space is hence Param = R m × h ×
R h × n .
We are interested in the case when the preimage of the zero vector under the hidden layer activation σ : R h
→ R h contains at most the zero vector, that is, σ − 1 (0) ⊆ { 0 } .
This condition holds, for example, in the cases of the usual sigmoid activation, hyperbolic tangent, and leaky ReLU (but not the usual ReLU).
Lemma D.1.
Suppose that σ − 1 (0) ⊆ { 0 } .
Then there exists a map c : GL h × (cid:0) R h \ { 0 } (cid:1) → GL h satisfying the following identities for any nonzero z ∈ R h and g, g 1 , g 2 ∈ GL h :
[Equation]
Before giving a proof of this result, we state a consequence. The non-degenerate locus is defined as:
[Equation]
Corollary D.2.
Suppose that σ − 1 (0) ⊆ { 0 } .
Then there is an action of GL h on the non-degenerate locus given by:
[Equation]
Moreover, this action leaves the extended feedforward function unchanged, that is:
[Equation]
[-where, for convenience, we set g 
 = id n and g L = id n L . In words, for i = 1 , . . . , L − 1 , we first take the product-]Proof of Lemma D.1.
Any z ∈ R h \ { 0 } can be written as z = ge , for some g ∈ GL h , where e = (1 , 0 , . . . , 0) is the first basis vector of R h . Define a map c ( − , e ) : GL h
→ GL h by setting c ( g, e ) to be any matrix that maps σ ( e ) to σ ( ge ) , with the condition that c ( g 1 , e ) =
c ( g 2 , e ) whenever g 1 e = g 2 e . We discuss a possible construction of c ( − , e ) below.
Next, we define a map: [Equation]
where γ ∈ GL h is any element taking e to z . We note that the definition of c ( g, z ) does not depend on the choice of γ taking e to z since we have required c ( g 1 , e )
= c ( g 2 , e ) whenever g 1 e = g 2 e .
To verify the identities, fix γ ∈ GL h such that γe =
z . Then:
[Equation]
Similarly,
[Equation] where we use the fact that g 2 z = g 2 γe .
And finally:
[Equation]
where we use the defining property of c ( − , e ) twice.
Proof of Corollary D.2.
The identity relating c and σ implies that c ( g, z ) − 1 σ ( gz ) =
σ ( z ) for all non-zero z ∈ R h ; the proof of the corollary follows from the case
z = V x . We remark on a possible way to define the map c ( − , e ) : GL h → GL h . Let g ∈ GL h .
If σ ( ge ) is a multiple of σ ( e ) , take g to be the corresponding scalar matrix. Otherwise, set u e = σ ( e ) | σ ( e ) | and u ge = σ ( ge ) | σ ( ge ) | to be the unit vectors in the directions of σ ( e ) and σ ( ge ) , and let t ∈ GL h be the unique positive rescaling that takes σ ( e ) and σ ( ge ) to u e and u ge , respectively. Let q ∈ GL h be the (unique?) rotation that takes u e to u ge and fixes the codimension 2 hyperplane orthogonal to σ ( e ) and σ ( ge ) .
Then g = qt takes σ ( e ) to σ ( ge ) . If ˜ σ is continuous, then c ( − , e ) can be chosen to be continuous as well.
Finally, we remark on a differential-geometric interpretation of the construction of this section.
One can regard σ as a section of the trivial bundle on R h \ { 0 } with fiber R h .
The map c defines a
GL h -equivariant structure on this bundle such that σ is an equivariant section. Indeed, the action of
GL h on the total space (cid:0) R h \ { 0 } (cid:1) × R h is given by g · ( z, a ) =
( gz, c
( g, z ) a ) , and the equivariance of σ is precisely the condition c ( g, z ) σ ( z )
= σ ( gz ) .
D. N ON - LINEAR ACTION : MULTI - LAYER CASE
Let n = ( n 0 , n 1 , . .
. , n L ) be a dimension vector for a feedforward network with L layers. The parameter space is:
[Equation]
So for each layer i , we have a matrix W i ∈ R n i × n i − 1 and vector b i ∈ R n i . We write θ =
( W i , b i )
Li = for a choice of parameters. Fix activations σ i : R n i
→ R n i .
Set:
[Equation]
to be the function defined recursively as:
[Equation]
The (extended) feedforward function is then F ( θ , X ) = σ L ( Z L ( θ , X )) . Define the non-degenerate locus as:
[Equation] × R n ) [Equation]
= { [Equation] θ ,
[Equation]
[-where F i − 1 ( X ) † denotes a left inverse of F i =
-]| Z [Equation] [Equation]
, x [Equation] = 0 for i = 1 , . . . , L − 1 } .
Proposition D.3.
Suppose that, for i = 1 , . . . , L − 1 , the activation σ i : R n i
→ R n i satisfies σ − 1 i (0) ⊆ { 0 } . Then there is an action of the hidden symmetry group GL n hidden on the non-degenerate locus given by:
[Equation]
where c i : GL n i × ( Param × R n i ) ◦
→ GL n i is as in Lemma D.1 for i = 1 , . . . , L − 1 (replacing h with n i and σ with σ i ; for i {+= 0 , we have g 
 = id n 0 and c 0+} ( g 0 , x ) is the identity matrix). Moreover, this action preserves the extended feedforward function.
Proof.
The fact that the action is well-defined follows from the identities satisfied by each c i (cf. Lemma D.1). One verifies by induction that Z i ( g · ( θ , x ))
= g i Z i ( θ , x ) for i = 0 , 1 , . .
, L . Hence, F ( g · θ , x ) = σ L ( Z L ( g · ( θ , x )))
=
σ L ( g L Z L ( θ , x )) = σ L ( Z L ( θ , x )) = F ( θ , x ) and so the extended feedforward function is preserved under this action.
D. D ISCUSSION : INCREASING THE BATCH SIZE
In this section, we discuss difficulties in adopting the construction of the previous sections to cases where the batch size is greater than one.
Fix a batch size k , so that the feature space of the hidden layer is R h × k . By abuse notation, we write σ : R h × k → R h × k for the map applying σ column-wise. We say that σ preserves full rank matrices if σ ( Z ) is full rank for any full-rank matrix in Z ∈ R
h × k . As a final piece of notation, let (cid:0) R h × k (cid:1) ◦ ⊆ R h × k be the subset of full rank matrices.
Lemma D.4. Suppose that k ≤ h , and that σ preserves full-rank matrices. Then there exists a map c : GL h × (cid:0) R h × k \ { 0 } (cid:1) → GL h satisfying the following identities for any nonzero Z ∈ R h × k and g, g 1 , g 2 ∈ GL h : c ( [Equation]
, Z [Equation] = [Equation]
[Equation]
[Equation]
[Equation]
We omit a proof of this lemma, as it proceeds along the same lines as that of Lemma D.1.
A key tool is the fact that, for k ≤ h , any two matrices in (cid:0) R h × k (cid:1) ◦ are related by an element of GL h .
This lemma implies that, for a multi-layer network, if σ i preserves full rank matrices in R n i × k for each i , then there is a non-linear group action as in Proposition D.3, where the appropriate version of the non-degenerate locus is:
[Equation]
However, as the following examples show, the condition that σ preserves full rank matrices is not satisfied in the case of common activation functions.
Example D.1. For k > 1 , the column-wise application of the usual sigmoid activation does not preserve full rank matrices. For example, for k = 2 , take: [Equation]
σ
[Equation]
(cid:1)(cid:21) ≃ (cid:20) − 1 .
− 0 .− 0 .1 .(cid:21)
Then det( σ ( Z ))
{+= 
 while det( Z ) = − 2 . 0862 . For k > 1 , the column-wise application of hyperbolic tangent does not preserve full rank matrices. To see this, set k = 1 and consider:+}
[Equation]
(cid:0) 45 (cid:1) (cid:21) ≃ (cid:20) 0 . 2027 0 .0 .
4236 1 .(cid:21)
Then det(tanh( Z ))
{+= 
+}
while det( Z ) = 0 . 0432 .
Let s be a real number with 0 < s < 1 .
The corresponding leaky ReLU activation function is given by σ ( z ) =
sz min(0 , z ) + z max(0 , z ) .
For k > 1 , the column-wise application of leaky ReLU tangent does not preserve full rank matrices.
Indeed, for k = 2 , set: Z = (cid:20) s
[Equation]
−s (cid:21)
Then det( σ ( Z ))
= det (cid:18)(cid:20) s − s − s
s (cid:21)(cid:19) =
while det( Z ) = s 2 − 1 ̸ = 0 .
Finally, in the case k > h , the action of GL h on full rank h × k matrices is not transitive. Hence, there will generally be no matrix in GL h taking σ ( Z ) to σ ( gZ ) , even if both are full rank. E U SING NONLINEARLY TRANSFORMED WEIGHTS ON DIFFERENT INPUTS The transformation U ′
=
c ( g − 1 , gV X ′ ) depends on the input data batch X via Z = V X Ideally, we would like to find transformed weights U ′ , V ′ which could be used in place of the original U, V for any input data. Therefore, we want to understand under what circumstances a U ′
( Z ′ ) computed using a given batch of input data Z ′ =
V X ′ would also work on a different batch Z
= V X , meaning
[Equation]
In general, (44) is not satisfied by arbitrary Z, Z ′ and σ .
But batches of input can be related to each other, like being samples from an underlying distribution.
We will now explore whether, under reasonable assumptions, (44) can be approximately satisfied.
We will use three ingredients for this.
First, we will assume σ is Lipschitz continuous, meaning | σ ( a ) − σ ( b ) | ≤ η | a − b | , allowing us to relate the differences in the output of σ to differences of inputs. Second ingredient is an assumption about the statistics of Z ′ being representative of the underlying distribution. Lastly, we will use the Cauchy-Schwartz inequality (CSI) in multiple occasions.
Let us first review the matrix for of the
CSI.
Cauchy-Schwartz inequality (CSI) for matrices.
For real vectors u, v ∈ R h with the Euclidean inner product ⟨ u, v ⟩ = u T v ∈ R , the CSI states that
[Equation]
u and v can also be rows and column of two matrices.
Let U ji denote row i and column j of U .
We have
[Equation]
where U i means row i and V k means column k .
Summing (46) over i and k we get
[Equation]
[Equation]
which is the matrix form of CSI.
Lipschitz continuous nonlinearity.
Assume σ is Lipschitz continuous meaning for a, b ∈ R we have
[Equation]
Define the action of σ on matrices to be elementwise, meaning σ ( A ) ji = σ ( A ji ) .
Then the Lipschitz condition can be extended to the Frobenius norms as follows.
For A, B ∈ R m × n we have
[Equation]
For parts of the derivation we also need the second assumption that σ (0) = 0 .
This results in
[Equation]
Representative batches We will now show that if if the samples (columns) in the new input Z are close to samples in Z ′ used to compute U ′
( Z ′ ) , then the equivariance condition is satisfied and
U ′
( Z ′ ) can be used as weights for batch Z as well.
Assume that the input batch Z ′ =
V X ′ ∈
R h × b ′ has samples which are close to the samples in the new input Z ∈ R h × b .
This means that for each column Z k ∈ R h , there exists a column Z ′
µ ( k ) ∈
R h close to it, meaning ( Z b = { 1 , . . .
, b } )
[Equation]
[Equation]
{+for some constant λ > 
+}
Here µ : Z b → Z b ′ is a map.
We can capture the mapping µ into a matrix
P ∈ { 0 , 1 } b × b ′ as
[Equation]
[Equation]
Note that because σ is {+element-wise and (cid:80) k P kj = 
 (i.e. each Z k is mapped to a single Z ′ µ ( k ) ) we have+}
[Equation]
Using (52) in the Lipschitz bound (49) we get
[Equation]
Using CSI (47) and (52) yields
[Equation]
[Equation]
Finally, we can use this result to address (44).
We want to check if U ′ ( Z ′ )
σ ( gZ ) is approximately the same as Uσ ( Z ) .
We will now show that it is approximately Uσ ( Z ′ ) P =
Uσ ( Z )
+ O ( ηλ ) (from (55)).
To do so, we look at the difference between these two term
[Equation]
) where we used σ ( gZ ′ ) − 1 σ ( gZ ′ ) = I b (assuming b ≤ h ). (55) and CSI we have
[Equation]
Therefore we have
[Equation]
where ∥ σ ( Z ′ )
P − σ ( Z ) ∥ ≤ ηλ √ b from (55)
[Equation]
F T HE L AGRANGIAN APPROACH F. L AGRANGIAN
We identify the tangent bundle TM with R d × R d , writing elements as pairs ( x, y ) .
Given V as a above and λ ∈ R × , define a Lagrangian 9 :
[Equation]
Observe that at t = 0 and y = 0 , we recover V , that is, L λ (0 , x, 0)
=
V ( x ) .
It is straightforward to show that a curve f : ( a, b ) → M satisfies the Euler–Lagrange equation for L λ if and only if:
[Equation]
While Equation 60 reduces to the gradient flow equation 14 when λ = 0 , note that it does not make sense to set λ = 0 directly in the expression for the Lagrangian L λ .
We discuss some constrained uniqueness properties of this Lagrangian in an addendum to these notes.
Orthogonal symmetries of V are also symmetries of the Lagrangian L λ .
More precisely, let ρ − 1 ( O ( d )) be the subgroup of G consisting of all g ∈ G such that the image ρ ( g ) defines an orthogonal transformation of R d .
Lemma F.1.
Suppose g ∈ ρ − 1 ( O ( d )) .
Then, for all λ ̸ = 0 , and for all ( x, y ) ∈ TM ≃ R d × R d , we have:
[Equation]
F. S ECOND - ORDER GRADIENT FLOW CONSERVED QUANTITIES
For U ∈ g and λ ∈ R × , set:
[Equation]
Given f : ( a, b ) → R d , we write Q U,λ ◦ ˙ f for the map t (cid:55)→ Q U,λ ( t, f ( t ) , ˙ f ( t )) .
The following lemma shows two facts.
First, the function Q U,λ vanishes on the path in TM defined by differentiating an integral curve for the gradient flow.
Second, when U is anti-symmetric, the function Q U,λ is a conserved charge for the Lagrangian L λ .
Proposition F.2.
Let U ∈ g and λ ∈ R × .
Let f : ( a, b ) → M be a smooth curve.
If f satisfies the gradient flow equation for V (Equation 14), then:
[Equation]
If f satisfies the Euler–Lagrange equation for L λ (Equation 60), then:
[Equation]
where U sym is the symmetric part of the matrix U ∈ gl d .
Proof.
The first claim is a direct consequence of the definition of Q U,λ .
For the second claim, we have:
[Equation]
An standard (elementary) calculation shows that ˙ f ( t ) T U ˙ f ( t ) = ˙
f ( t ) U sym ˙
f ( t ) .
If f satisfies Equation 60, then we can write ¨ f in terms of ˙ f and ∇ f ( t ) V to obtain:
[Equation]
The term f ( t ) T U ˙ f ( t ) cancels, while f ( t ) T U (cid:0) ∇ f ( t ) V (cid:1)
= 0 .
The second claim follows.
Corollary F.3.
Suppose U ∈ gl
d is anti-symmetric.
If f satisfies the Euler–Lagrange equation for
L λ (Equation 60), then:
[Equation]
Corollary F.3 is a special case of Noether’s theorem, since, as noted above, orthogonal symmetries of V are also symmetries of L λ .
We also invoke the fact that the Lie algebra of the orthogonal group consists of anti-symmetric matrices, so that the Lie algebra of the subgroup ρ − 1 ( O ( d )) of G can be identified the Lie subalgebra of g consisting of those U ∈ g such that U is skew-symmetric.
For clarity, we perform the computation of the proof of Proposition F.2 in alternative notation:
[Equation]
F. A DDENDUM : MORE ON THE L AGRANGIAN F.3. C ONSTRAINED UNIQUENESS
Conjecture F.4.
Consider functions of the form
L : R × TM
[Equation]
R
[Equation]
for non-vanishing coefficient functions g, h : R → R × .
The Lagrangian L λ is the unique such function that satisfies the following two conditions: (1) For all x ∈ M , we have L (0 , x, 0)
=
V ( x ) .
(2) The Euler–Lagrange equation for L is equivalent to Equation 60.
Proof.
The two conditions imply the following identities:
[Equation]
From these, one concludes that g ( t ) =
λ 2 e t/λ and h ( t ) =
e t/λ .
Remark Note that L λ is not the unique Lagrangian that satisfies the following two conditions: (1)
For all x ∈ M , we have L λ (0 , x, 0) =
V ( x ) .
The Euler–Lagrange equation for L λ limits to the gradient flow equation for V as λ → 0 .
F.3. S YMMETRIES
Any symmetry of V (not necessarily orthogonal) can be used to obtain an asymptotic symmetry of
L λ in the sense of first setting t = 0 and then taking λ → 0 .
Specifically, for g ∈ G , we have:
[Equation]
which limits to V ( x ) as λ → 0 , just as L λ (0 , x, y )
=
V ( x
[Equation]
+ λ 2 | y | 2 does.
F.3. L IMITING PROCEDURES
Let F : R × × ( a, b ) → R d be a smooth map.
Writing f λ ( t ) for F ( λ, t ) , we see that F encodes a smooth family of curves f λ indexed by λ ∈ R × .
We say that a curve f : ( a, b ) → R d is a limit of the family f
λ if lim λ
→ 0 f λ ( t ) = f ( t ) for every t ∈ ( a, b ) .
For a given family, a limit may or may not exist; if it exists, it is unique.
Lemma F.5.
Suppose f λ satisfies the Euler–Lagrange equation for L λ (Equation 60), and suppose the limit f = lim λ
→ 0 f λ exists.
Then f satisfies the gradient flow equation for V (Equation 14).
Proof.
If the limit exists, then ˙ f
= lim λ
→ 0 ˙ f λ , and similarly for the second derivative.
The result now follows from the fact that the Euler–Lagrange equation for L λ limits to the gradient flow equation for V as λ → 0 .
Example F.1.
Suppose V : R d → R is given by V ( x ) = [-
-]12 | x | 2 , so that ∇ x V = x for every x ∈ R d .
Given x ∈ R d , set:
[Equation]
( x
[Equation]
x Then f ( x ) λ (0)
= x and f ( x ) λ satisfies the Euler–Lagrange equation for L λ .
The limit of f ( x ) λ as λ →is given by: f
[Equation]
d [Equation] (cid:55)→ e − t x i.e., the unique flow line for the reverse gradient vector field of V starting at x .
Remark Suppose each f λ satisfies the Euler–Lagrange equation for L λ , and that the limit f = lim λ
→ 0 f λ exists.
By Corollary F.3, we have that
[Equation]
for U antisymmetric.
This implies that Q U,λ ◦ ˙ f
= K for some constant K .
Using the definition of
Q U,λ , we have f λ ¯ U ˙ f λ =
e − t/λ K .
Taking the limit as λ → 0 , we obtain f ¯ U ˙ f = 0 , and so
Q U,λ ◦ ˙ f
=for any λ .
Thus we recover Part 1 of Proposition F.2 in the anti-symmetric case.
We emphasize that we have proven this result by more elementary means already.
Anyway, this proof only works for anti-symmetric U , and for flows that are limits of solutions to the Euler–Lagrange equation.
G B ACKGROUND : L AGRANGIANS AND N OETHER ’ S THEOREM G. E ULER -L AGRANGE EQUATION
The action L is defined as the integral of a Lagrangian L :
[Equation]
From the stationary action principle, L is minimized.
In other words, we want ∂ L ∂x = 0 .
Consider a small change ( δx ) in x .
The first-order Taylor expansion of L is
[Equation]
We can remove dδx µ dt by integration by parts:
[Equation]
Then,
[Equation]
Recall that we want δ L = 0 .
If we assume δx = 0 on the boundary, the second integral evaluates to
and we get the Euler-Lagrange equation:
[Equation]
G. F IRST - ORDER GRADIENT FLOW CANNOT BE DERIVED FROM VARIATIONAL PRINCIPLE
Consider the Lagrangian
[Equation]
[Equation]
To minimize this Lagrangian, we use the Euler-Lagrange equation,
[Equation]
This equation does not contain first-order terms of x .
Therefore, we cannot get the original gradient flow dynamics from the variational principle.
G. B REGMAN L AGRANGIAN
The Bregman Lagrangian (Wibisono & Wilson, 2015) provides a framework for modeling dynamics in neural networks with different learning rules.
The form of this Lagrangian that we will be using is
[Equation]
where η : R d ⊗ R d
→ R ≥ 0 is a degenerate Riemannian metric (positive semi-definite).
In this paper, we consider the following version of the Bregman Lagrangian:
[Equation]
71) where 1 γ is the friction constant, x is the trainable parameters in our optimization problem, and V ( x ) is the loss function.
For simplicity, we consider ϵ
=
I in the main text.
From the Euler-Lagrange equation (69), the equation of motion of this Lagrangian is
[Equation]
Simplifying the above equation, we have
[Equation]
G. R ELATING γ TO δt
We show that the coefficient γ in the Lagrangian defined in (22) is proportional to the time step δt in gradient descent.
The second order expansion of x is
[Equation]
From this expansion, δx = δt ˙ x + δt 2 2 ¨ x .
On the other hand,
[Equation]
75) so δx = − δtε
dVdx .
Equating the δx obtained in the two methods, we have
[Equation]
Comparing this with the equation of motion derived from (22), we have
[Equation]
G. N OETHER ’ S THEOREM
If we consider x, ˙ x that is a solution of the system, the first integral in (66) becomes 0, and the remaining term in δ L is
[Equation]
From Stoke’s theorem,
[Equation]
Setting δ L = 0 gives us the Noether’s theorem:
[Equation]
The conserved charge is J = δx µ ∂L ∂ ˙ x µ .
The term ∂L ∂ ˙ x µ = p µ is the conjugated momentum.
G. R ELATING CONSERVED QUANTITY TO CONVERGENCE RATE
The Euler-Lagrange equation for our Lagrangian (22) gives the following dynamics:
[Equation]
When γ −→ ∞ ,
[Equation]
Then, we write the change in x in terms of the change in the coefficient of the basis:
[Equation]
Assume that ε
= I .
Substitute ˙ x back, we have
[Equation]
Recall that the conserved quantity for the second-order gradient flow is
[Equation]
[Equation]
where ˙ Q GFi = x T S
i ˙ x = x T S i ˙ θ j T j x is the time derivative of the conserved quantity for the firstorder gradient flow, and J i =
x T A i ˙ x
= x T A i ˙ θ j T
j x is the angular momentum.
Then, the components of Q i appears in dVdt in the following way:
[Equation]
G. C ONSERVED QUANTITIES IN NEURAL NETWORKS G.7. T WO - LAYER LINEAR NETWORK
In this section, we find the conserved quantity corresponding to the following Lagrangian:
[Equation]
88) where the parameters are U ∈
R m × h and V ∈ R h × n , and the loss L is a function of UV .
For simplicity, we assumed that the learning rate matrix ε = I .
Let T i be an element of the Lie algebra of GL h .
For g ∈ GL h , the transformation U −→ Ug − 1 , V −→ gV leaves L unchanged.
When g is close to identity, we can use the first-order expansion of g and write the transformation as U ( I − ηT i ) and ( I + ηT i ) V , where η ∈ R is a small number.
Putting
U and V together in a matrix, we have
[Equation]
Let S i ,
A
i be the symmetric and antisymmetric parts of T i , that is T
i =
S i
+ A
i .
Then, δX =
[Equation]
where σ z = (cid:20) [Equation] − [Equation]
is the Pauli matrix. From Noether’s theorem, the conserved quantity is:
[Equation]
Let ∇ L = ∂L∂W , where W ba =
U µa V bµ .
Assume that γ −→ ∞ so that we can approximate the firstorder gradient flow.
Then ˙ U = − ∂L∂U
and ˙ V = − ∂L∂V .
Below is the explicit form of ˙ X in this limit.
Note that ∂L∂W
[Equation]
Substituting in ˙ X , the first term of J becomes
[Equation]
The second term of J is
[Equation]
Therefore, in the limit of first-order gradient flow, tr [ σ z XS i ˙ X T ] = ˙
Q GF = 0 .
In summary, we derived the explicit form of angular momentum and conserved quantity of the two-layer linear reparametrization, verified that angular momentum equals 0, and rediscovered the conserved quantities we had before.
G.7. T WO - LAYER NONLINEAR NEURAL NETWORKS
In this section, we find the conserved quantity in the dynamics described the following Lagrangian: [Equation]
[Equation]
where U ∈ R m × h and V ∈ R h × n are parameters, X is data, σ is an element-wise function, and L is a function of Uσ ( V X ) .
For simplicity, we again assumed that ε = I .
Let T i be an element of the Lie algebra of GL h .
One action that leaves f unchanged is:
[Equation]
[Equation]
We first derive δU and δV .
The action of infinitesimal g ε = I + ε · T yields
[Equation]
[Equation]
For small ϵ ,
[Equation]
⊙ here denotes element-wise product.
Then, the action of g ε yields
[Equation]
For a set of parameters θ =
( U, V, X ) ∈ Θ ,
[Equation]
Combining U, V into a single matrix,
[Equation]
Let ∇ L = ∂L∂W , where W ba =
U µa σ ( V νµ X bν ) .
Assume that γ −→ ∞ .
Then ˙ U = − ∂L∂U
and ˙ V = − ∂L∂V .
Below is the explicit form of ˙ X in this limit.
[Equation]
From Noether’s theorem, the conserved quantity is
[Equation]
Substitute in δX and ˙ X , we verify that J = 0 :
[Equation]
Note that the second term in the 4th line can be evaluated by first taking the element-wise product of 4 matrices with same dimensions, then summing over all indices.
The 5th line uses that fact that matrix element-wise product is associative and commutative.
Again, we obtained that in the limit of first-order gradient flow ( γ −→ ∞ ) , the conserved quantity for the second-order gradient flow J =for any T i .
Conserved quantities For conserved quantity Q and infinitesimal generators T , we have
[Equation]
[Equation]
and
[Equation]
where w represent the list of all parameters ( U, V ).
By matching terms, one way to find a conserved quantity is to find a Q that satisfies ∂Q∂U = δU and ∂Q∂V = δV .
Therefore, one set of conserved quantities is
[Equation]
where the second term Q V does not depend on U and is the solution to ∂Q V ∂V ij =
[Equation]
T WO - LAYER NONLINEAR NEURAL NETWORKS - ALTERNATIVE TRANSFORMATIONS
Another transformation that leaves f unchanged is:
[Equation]
We first derive δU and δV .
The action of infinitesimal g ε = I + ε · T yields
[Equation]
To simplify notation, define f : R h × k −→ R k × h , x
(cid:55)→ [ σ ( x )] − 1 .
Then
[Equation]
For small ε ,
[Equation]
The last term can be evaluated to
[Equation]
The fourth line uses the derivative of matrix inverse.
Substitute back and let H =
V X , we have
[Equation]
⊙ here denotes element-wise product.
Then, the action of g ε yields
[Equation]
Combining U, V into a single matrix,
[Equation]
Let ∇ L = ∂L∂W , where W ba =
U µa σ ( V νµ X bν ) .
Assume that γ −→ ∞ .
Then ˙ U = − ∂L∂U
and ˙ V = − ∂L∂V .
Below is the explicit form of ˙ X in this limit.
[Equation]
The conserved quantity is
[Equation]
Substitute in δX and ˙ X ,
[Equation]
Conserved quantities The conserved quantities corresponding to the new group action can be found in a similar way.
Note that δU now depends on both U and V .
We seek to simplify the differential equations by removing δU ’s dependence on V .
Let
[Equation]
where F = (cid:0) [ TH ⊙ σ ′ ( H )] σ ( H ) − 1 (cid:1) − 1 .
Then dQ/dt is still 0.
Recall that H =
V X , so F is a function whose value only depends on V, X , and T .
One set of conserved quantities is
[Equation]
(121) where the second term Q V does not depend on U and is the solution to ∂Q V ∂V ij = [Equation] F [Equation]
X
[Equation]
T TV
[Equation]
H G RADIENT DESCENT AND DRIFTING CONSERVED QUANTITIES
While gradient flows are well approximated by gradient descent Elkabetz & Cohen (2021), the conserved quantities of gradient descent are no longer conserved in gradient flow due to noninfinitesimal time steps.
However, with small learning rate, we expect the change in the conserved quantities to be small.
In this section, we first prove that the change of Q is bounded by the square of learning rate for two layer linear networks, then show empirically that the change Q is small for nonlinear networks.
H. C HANGE IN Q IN GRADIENT DESCENT ( LINEAR LAYERS )
Proposition H.1.
Consider the two layer linear network, where U ∈ R m × h , V ∈ R h × n are the only parameters, and the loss function L is a function of UV .
In gradient descent with learning rate η , the change in the conserved quantity Q =
Tr (cid:2) U T U − [Equation] V T
[Equation]
at step t is bounded by
[Equation]
Proof.
Let U t and V t be the value of U and V at time t in a gradient descent.
The update rule is
[Equation]
Consider the two layer linear reparametrization W = UV .
[Equation]
Substituting in U t +1 and V t +1 , expanding Q t +1 , and subtracting by Q t , we have
[Equation] [Equation]
Note that
[Equation]
and similarly
[Equation]
Therefore, (125) simplifies to
[Equation]
and the variation of Q in each step is bounded by the convergence rate:
[Equation]
[Equation]
H. E MPIRICAL OBSERVATIONS
In gradient flow, the conserved quantity Q is constant by definition.
In gradient descent, Q varies with time.
In order to see how applicable our theoretical results are in gradient descent, we investigate the amount of variation in Q in gradient descent using two-layer neural networks.
Since Q is the difference between the two terms f 1 ( U ) =
Tr[ U T U ] and f 2 ( V ) =
(cid:80) a,j (cid:82) V aj x 0 dx σ ( x ) σ ′ ( x ) , we normalize Q by the initial value of f 1 ( U ) and f 2 ( V ) , i.e.,
[Equation]
and denote the amount of change in ˜ Q as
[Equation]
[Equation]
We run gradient descent on two-layer networks with whitened input with the following objective
[Equation]
[Equation]
where σ is the identity function, ReLU, sigmoid, or tanh.
Y ∈ R 5 × 10 , U ∈ R 5 × 50 and V ∈ R 10 ×have random Gaussian initialization with zero mean.
We repeat the gradient descent with learning rate 0.1, 0.01, and 0.001.
The variation ∆ ˜ Q ( t ) and loss is shown in Fig.4.
The amount of change in Q is small relative to the magnitude of f 1 ( U ) and f 2 ( V ) , indicating that conserved quantities in gradient flow are approximately conserved in gradient descent.
The error in Q grows with step size, as ∆ ˜ Q ( t ) is larger with the largest learning rate we used, although it has the same magnitude as those of smaller learning rates.
We also observe that Q stays constant after loss converges.
[Figure]
I D ISTRIBUTION OF Q UNDER X AVIER INITIALIZATION
We first consider a linear two-layer neural network UV X , where U ∈ R m × h ,
V ∈ R h × n , and
X ∈ R n × k .
We choose the following form of the conserved quantity:
[Equation]
Xavier initialization keeps the variance of each layer’s output the same as the variance of the input.
Under Xavier initialization (Glorot & Bengio, 2010), each element in a given layer is initialized independently, with mean 0 and variance equal to the inverse of the layer’s input dimension:
[Equation]
The expected value of Q is
[Equation]
Figure 5 shows the distribution of Q for 2-layer linear NN with different layer dimensions.
For each dimension tuples ( m, h, n ) , we constructed 1000 sets of parameters using Xavier initialization.
The centers of the distributions of Q match Eq.
[Figure]
Next, we consider the nonlinear two-layer neural network Uσ
( V X ) , where σ : R −→ R is an element-wise activation function.
For simplicity, we assume whitened input ( X = I ).
We choose the following form of the conserved quantity:
[Equation]
Figure 6 shows the distribution of Q for 2-layer linear NN with different nonlinearities, each with 1000 sets of parameters created under Xavier initialization.
The shapes of the distributions are similar to that of linear networks.
The value of Q is usually concentrated around a small range of values.
Since the range of Q is unbounded, the Xavier initialization limits the model to a small part of the global minimum.
[Figure]
J C ONSERVED QUANTITY AND CONVERGENCE RATE
The values of conserved quantities are unchanged throughout the gradient flow.
Since the conserved quantities parameterize trajectories, initializing parameters with certain conserved quantity values accelerates convergence.
For two-layer linear reparametrization, Tarmoun et al.
(2021) derived the explicit relation between layer imbalance and convergence rate.
We derive the relation between conserved quantities and convergence rate for two example optimization problems and provide numerical evidence that initializing parameters with optimal conserved quantity values accelerates convergence.
J. E XAMPLE 1: ELLIPSE
We first show that the convergence rate is related to the conserved quantity in a toy optimization problem.
Consider the following loss function with a ∈ R :
[Equation]
Assuming gradient flow,
[Equation]
Then w 1 , w 2 are governed by the following differential equations:
[Equation]
where w 1 0 , w 2 0 are initial values of w 1 and w 2 .
We can find conserved quantities by using an ansatz
Q = f
( w i 1 w k 2 ) and solving ∇ Q · ∇ L = 0 for i, k .
Below we use the following form of conserved quantity:
[Equation]
To show the effect of Q on the convergence rate, we fix L (0) and derive how Q affects L ( t ) .
Let
L (0) =
w 21
+ aw 22
= L 0 .
Let w 2 0 continue to be an independent variable.
Then w 21
=
L 0 − aw 22 0 .
Substitute in w 21 0 , the loss at time t is
[Equation]
and Q becomes
[Equation]
The derivative of L in the direction of Q is
[Equation]
In general
[Equation]
∂ Q L ( t ) ̸ = 0 , meaning that the loss at time t depends on Q .
Since we have fixed the initial loss, the convergence rate L ( t ) − L (0) also depends on Q .
Special cases where ∂ Q L ( t ) =include a = 1 (circle), a = 0 (collapsed dimension), and certain initializations such as w 2 0 =(local maximum of gradient magnitude).
J. E XAMPLE 2: RADIAL ACTIVATION FUNCTIONS
In this example, we find the conserved quantities and their relation with convergence rate for twolayer reparametrization with radial activation functions under spectral initialization.
Define radial function g : R m × n −→ R m × n as
[Equation]
where | W i | = (cid:0)(cid:80)
k W 2 ik (cid:1) 12 is the norm of the i th row of W , and h : R −→ R outputs a scalar.
Consider the following objective:
[Equation]
with spectral initializations
[Equation]
where Φ , Ψ come from the singular value decomposition Y =
ΦΣ Y Ψ T , and U
, V 0 are random diagonal matrices.
Proposition J.1.
Under the gradient flow U = −∇ U L and V = −∇ V L , the following quantity is an invariant:
[Equation]
Proof.
Since g is a radial function on rows and Ψ T is an orthogonal matrix, g ( V T Ψ T )
=
g ( V T )Ψ T .
With spectral initialization, the loss function can be reduced to only involving diagonal matrices:
[Equation]
Since V is a diagonal matrix, g is now an element wise function on V .
Let W =
Ug ( V T ) .
The gradients for U and V are
[Equation]
where g ′ ( x ) = dg ( x ) /dx is the derivative of the nonlinearity.
Additionally, since L does not depend on Φ and Ψ ,
[Equation]
Since the rows of Φ [Equation] Ψ are orthogonal,
[Equation]
Φ and Ψ are not changed in gradient flow, so ∂Q∂U
= ∂Q ∂U Φ T and ∂Q∂V =
∂Q ∂V Ψ T .
Define inner product on matrices as ⟨ X, Y ⟩ =
Tr[ X T Y ] .
For Q to be a conserved quantity, we need ⟨∇L , ∇ Q ⟩ = 0 : [Equation]
[Equation]
Following the same procedure as for elementwise functions, to have a Q which satisfies (150) it is sufficient to have
[Equation]
For simplicity, let f ( U, V ) =
Then, (151) is satisfied by
[Equation]
Tarmoun et al.
(2021) shows that the conserved quantity Q appears as a term in the convergence rate of the matrix factorization gradient flow.
We observe a similar relationship between Q and convergence rate when the loss function is augmented with a radial activation function, as shown in the following proposition.
Proposition J.2.
Consider the objective function and spectral initialization defined in Proposition
J.1.
Let h ( | W i | ) =
| W i | − 2 , and X =
Ug ( V T ) =
ΦΣ X Ψ T .
Then, the eigencomponent of X approaches the corresponding eigencomponent of Y at a rate of
[Equation]
(153) where σ X i = diag (Σ X ) i , σ Y i
=
diag (Σ Y ) i , and λ i
= ¯ U 2 ii + ¯ V 2 ii are conserved quantities.
Proof.
Similar to Tarmoun et al.
(2021), components can be decoupled, and we have a set of differential equations on scalars:
[Equation]
We also have
[Equation]
Let σ Xi
=
u
i g ( v i ) .
Then
[Equation]
Since V is a diagonal matrix, g is now an element wise function on V .
Specifically, g ( v i ) =
1 v i .
According to Proposition J.1, the following quantity is invariant:
[Equation]
Since any function of the invariant is also invariant, we will use the following form:
[Equation]
and define
[Equation]
Using the g that we defined,
[Equation]
[Equation] In order to relate σ X and Q , we first write u i and v i as functions of σ Xi ad Q using (159) and (160): u 2 i
= λ
i σ Xiσ Xi 2 + 1 , v 2 i
= λ
i σ Xi 2 + 1 .
Then, substitute u i , v i , g ( v i ) , and dg ( v i ) dv i into (156), and we have ˙
σ Xi =
[ σ i − u
i g ( v i )] (cid:34) g ( v i )
+ u 2 i (cid:18) dg ( v i ) dv
i (cid:19) 2 (cid:35) = (cid:2) σ Yi − u i g ( v i ) (cid:3) (cid:34)(cid:18)v i (cid:19)+ u 2 i (cid:0) − v − 2 i (cid:1)(cid:35) = (cid:2) σ Yi − u i g ( v i ) (cid:3) (cid:2) ( v 2 i ) − 1 + u 2 i ( v 2 i ) − 2 (cid:3) = (cid:2) σ Yi − σ Xi (cid:3)   (cid:32) λ i σ Xi 2 +(cid:33) −+ λ
i {+σ Xiσ Xi 
+}
+(cid:32) λ i σ
Xi[-
-]
{++(cid:33) − 2   = (cid:2) σ Yi − σ Xi (cid:3) (cid:34) σ Xi 
+} +λ i + σ Xi 2 ( σ Xi 2 + 1) λ i (cid:35) = (cid:2) σ Yi − σ Xi (cid:3) (cid:34) σ X i 4 + 2 σ Xi
+λ i (cid:35) =
λ i ( σ Yi − σ Xi )
( σ Xi 2 + 1)(162)
Proposition J.2 relates the rate of change in parameters ˙
σ Xi and the conserved quantity λ i .
To get a more explicit expression of how λ i affects convergence rate, we will derive a bound for | σ Yi − σ Xi | , which describes the distance between trainable parameters to their desired value.
Proposition J.3.
The difference between the singular values of Ug ( V T ) and Y is bounded by
[Equation]
Proof.
Note that
[Equation]
Consider the following two differential equations, with same initialization a (
[Equation]
=
b
[Equation]
[Equation]
[Equation]
In these equations, both a and b moves from a (0) =
b (0) to σ monotonically.
Since ˙ a ≥ ˙ b at every a = b , a will always be closer to σ than b does.
We can explicitly solve for b , which yields b ( t ) =
σ + ( b (0) − σ ) e − tλ .
Then the distance between b and σ is | b − σ | = | b (0) − σ | e − tλ .
Using | b − σ | , we can bound | a − σ | :
[Equation]
Therefore,
[Equation]
Since λ is a conserved quantity, its value set at initialization remains unchanged throughout the gradient flow.
Therefore, we are able to optimize the convergence rate by choosing a favorable value for λ at initialization.
In this example, smaller λ i ’s lead to faster convergence.
J. E XPERIMENTS
We compare the convergence rate of two-layer networks initialized with different Q values.
We run gradient descent on two-layer networks with whitened input with the following objective
[Equation]
where σ is the identity function, ReLU, sigmoid, or tanh.
Matrices Y ∈ R 5 × 10 , U ∈ R 5 × 50 and
V ∈ R 10 × 50 have random Gaussian initialization with zero mean.
We repeat the gradient descent with learning rate 0.1, 0.01, and 0.001.
The learning rate is set to 10 − 3 , as we do not observe significant changes in the shape of learning curves at smaller learning rates.
U and V are initialized with different variance, which leads to different initial values of Q .
As shown in Fig.7, the number of steps required for the loss curves to drop to near convergence level is correlated with Q in both linear and element-wise nonlinear networks.
This result provides empirical evidence that initializing parameters with optimal values for Q accelerates convergence.
We then demonstrate the effect of conserved quantity values on the convergence rate of radial neural networks.
Fig.8 shows the training curve for loss function defined in Proposition J.2.
We initialize parameters U ∈ R 5 × 5 and V ∈ R 10 × 5 with 4 different values of Q and the learning rate is set to 10 − 5 .
As predicted in Eq.
163, convergence is faster when Q =
Tr[ U T U + V T V ] is small.
K C ONSERVED QUANTITY AND GENERALIZATION ABILITY
Conserved quantities parameterize the minimum of neural networks and are related to the eigenvalues of the Hessian at minimum.
Recent theory and empirical studies suggest that sharp minimum
[Figure] [Figure]
do not generalize well (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Petzka et al., 2021).
Explicitly searching for flat minimum has been shown to improve generalization bounds and model performance (Chaudhari et al., 2017;
Foret et al., 2020; Kim et al., 2022).
We derive their relationship for the simplest two-layer network, and show empircally that conserved quantity values affect sharpness.
Like convergence rate, a systematic study of the relationship between conserved quantity and generalization ability of the solution is an interesting future direction
[Figure]
K. E XAMPLE : TWO - LAYER LINEAR NETWORK WITH 1D PARAMETERS
We again consider the two-layer linear network with loss L = 12 ∥ Y − UV X ∥ 2 .
For simplicity, we work with one dimensional parameters U, V ∈ R and assume X =
Y = 1 in this example.
We show that at the point to which the gradient flow converges, the eigenvalues of the Hessian are related to the value of the conserved quantity.
The gradients and Hessian of L are
[Equation]
At the minima, U, V are related by UV X = Y .
Recall that Q =
U 2 − V 2 is a conserved quantity.
From the above two equations, we can write U, V as functions of Q .
Taking the solution U = (cid:113)
[Equation]
Q
[Equation]
(cid:112)
[Equation]
2 + 4) and substitute in X
=
Y = 1 , we have
[Equation]
and the eigenvalues of H are
[Equation]
We have shown that Q is related the eigenvalues of the Hessian at the minimum.
Since the eigenvalues determines the curvature, Q also determines the sharpness of the minimum, which is believed to related to model’s generalization ability.
The result in this example can also be observed in Figure 1, where the minimum of the Q = 0 trajectory lies at the least sharp point of the loss valley.
K. E XPERIMENTS : TWO - LAYER NETWORKS
The goal of this section is to explore the relation between Q and the sharpness of the trained model.
We measure sharpness by the magnitude of the eigenvalues of the Hessian, which are related to the curvature at the minima.
We use the same loss function (168) in Section J.3.
The parameters are
U ∈ R 10 × 50 and V ∈ R 5 × 50 , each initialized with zero mean and various standard deviations that lead to different Q ’s. We first train the models using gradient descent.
We then use the vectorized parameters in the trained model to compute the eigenvalues of the Hessian.
The linear model extends the example in Section K.1 to higher dimension parameter spaces.
700 out of the 750 eigenvalues are around
(with magnitude ≤ 10 − 3 ) , which verifies the dimension of the minima in Proposition C.4.
After removing the small eigenvalues, the center of the eigenvalue distribution correlates positively with the value of Q (Figure 9(a)).
In models with nonlinear activations,
Q is still related to eigenvalue distributions, although the relations seem to be more complicated.
[Figure]
L E NSEMBLE MODELS
In neural networks, the optima of the loss functions are connected by curves or volumes, on which the loss is almost constant (Freeman & Bruna, 2017; Garipov et al., 2018;
Draxler et al., 2018;
Benton et al., 2021; Izmailov et al., 2018).
Various algorithms have been proposed to find these lowcost curves, which provides a low-cost way to create an ensemble of models from a single trained model.
Using our group actions, we propose a new way of constructing models with similar loss values.
We show that even with stochasticity in the data, the loss is approximately unchanged under the group action (Appendix L).
This provides an efficient alternative to build ensemble models, since the transformation only requires random elements in the symmetry group, without any searching or additional optimization.
We implement our group actions by modifying the activation function between two consecutive layers.
Let H
=
V X be the output of the previous layer.
The group action on the weights U, V is g
[Equation]
[Equation]
(172) where π ( g, H )
=
σ ( H ) σ ( gH ) † .
The new activation implements the symmetry group action
Uσ
[Equation]
H )
→ Uπ ( g, H ) σ ( gH ) (173) by wrapping the transformations around an activation function σ ′ ( x ) =
π ( g, x ) σ ( gx ) , so that
Uσ ′ ( H ) =
Uπ ( g, H ) σ ( gH ) .
We test the group action on CIFAR-10.
The model contains a convolution layer with kernel size 3, followed by a max pooling, a fully connected layer, a leaky ReLU activation, and another fully connected layer.
The group action is on the last two fully connected layers.
After training a single model, we create transformed models using g = I + εM , where M ∈ R 32 × 32 is a random matrix and ε controls the magnitude of movement in the parameter space.
We then use the mode of the transformed models’ prediction as the final output.
We compare the ensemble formed by group actions to four ensembles formed by various random transformation.
Let g
= I + εM .
The random baselines are:
• ‘group’:
( U, V ) (cid:55)→ ( Uπ ( g, H ) , gV ) .
This is the model created by group actions.
• ‘ g − 1 ’:
( U, V ) (cid:55)→ ( Ug − 1 , gV ) .
• ‘random’:
( U, V ) (cid:55)→ ( Ug ′ , gV ) , where g ′
=
I + εD and D is a random diagonal matrix.
• ‘shuffle’: ( U, V ) (cid:55)→ ( Uπ ′ ( g, H ) , gV ) , where π ′ ( g, H ) is constructed by randomly shuffling π ( g, H ) .
• ‘interpolated permute’ or ’perm interp’:
( U, V ) (cid:55)→ ( U (cid:16)
[Table]
(cid:17)
[Table]
V ) ,
where
S ∈ R 32 × 32 is a random permutation matrix.
[Figure]
[Table]
[Figure] [Figure]
Figure 10 shows the accuracy of the ensembles compared to single models.
The ensemble formed by group actions preserves the model accuracy for small ε and has smaller accuracy drop at larger ε .
The ensemble model also improves robustness against Fast Gradient
Signed Method (FGSM) attacks (Figure 11).
Under FGSM attacks with various strength, the ensemble model created using group actions consistently performs better than the baselines with random transformations.
However, the same improvement is not observed under Projected Gradient Descent (PGD) attacks.