Parameter Symmetry and Noise Equilibrium of Stochastic Gradient Descent

Summary

The authors study the effect of exponential symmetries on the learning dynamics of SGD. They establish the theorem that every exponential symmetry implies the existence of a unique and attractive fixed point in SGD dynamics, which could explain phenomena in machine learning such as matrix factorization and progressive sharpening or flattening.

Strengths

The paper studies the dynamics of SGD. The authors prove a novel theorem that relates continuous symmetry to the attractive fixed point in SGD. The theorem is quite general and can have significant potential impacts on the field. The authors apply their theorem to matrix factorization, offering strong analytical and experimental support. And by my knowledge, this is a novel contribution.

Weaknesses

1. C is not guaranteed to exist. (Line 70) 2. Usually, the stability of the learning is determined by the largest eigenvalue of the Hessian instead of the trace as discussed in the paper. The trace may not be an upper bound on the largest eigenvalue as it can contain negative eigenvalues. 3. I think it’s missing dependence on $\gamma$ in proposition 5.3. 4. The authors proved that in deep linear networks, the noise will lead to a balanced norm for all intermediate layers. This is an important theorem, but it lacks empirical verification. It is important to verify that the empirics actually match the theorem, which would strengthen the paper by a lot. Such an experiment should not be hard to run. 5. There is a lack of clarity in how the analyses in Section 5.4 (Equation 24) relate to Figure 4.

Questions

1. In line 124-128, the authors state that the SGD dynamics can be decomposed into two independent parts. Why does equation (5) have no effect on the loss? I would expect the noise will result in fluctuation of the loss at the end of training, which means that the decomposition should not hold. 2. For the double rotation symmetry, why it is hard to write the matrix A explicitly? It seems that A can just be any symmetric matrix as discussed in section 5.1. Also, why do you require A to be symmetric when defining the exponential symmetry? I am wondering if the results can be generalized to a general matrix A. 3. I don’t understand why one needs to assume C* to be a constant in line 163. By definition, C is a deterministic function that only depends on $\theta$.

Rating

7

Confidence

4

Soundness

2

Presentation

3

Contribution

4

Limitations

The authors have included the limitations. For additional points, please see my comments above.

Summary

This paper considers the relationship between the behavior of SGD for solving empirical risk minimization and the exponential symmetry.

Strengths

The strength of this paper is to study the relationship (Theorem 4.3) between the behavior of SGD for solving empirical risk minimization and the exponential symmetry (Definition 4.1). Moreover, it considers a noise equilibrium (Definition 4.5) and provides explicit solutions of the noise equilibriums for some applications (Section 5). In addition, it provides some numerical results to support the theoretical analyses presented in this paper.

Weaknesses

The weaknesses are the following. Please see Questions for details. - The mathematical preliminaries are insufficient. - The assumptions seem to be strong. - The explanations of contributions are insufficient.

Questions

- SGD for minimizing $L(\theta) = E_z [\ell (\theta,z)]$ is defined by (1): $d \theta_t = - \nabla L(\theta_t) dt + \sqrt{2 \sigma^2 \Sigma (\theta_t)} d W_t$ using $\sigma^2 = \eta/(2S)$, where $\eta > 0$ is a step-size and $S > 0$ is a batch size. My first question is "Which is a problem considered in this paper: the expected risk minimization or the empirical risk minimization?" $L$ is defined by using the expectation $E$ (Lines 56-57), however, Line 57 says "the empirical risk function." - Line 19 says that SGD with $\sigma = 0$ implies GD. I do not understand it, since $\sigma^2 = \eta/(2S) = 0$ iff $\eta = 0$, which contradicts $\eta > 0$. - Line 21: I do not understand what $\sigma^2 \ll 1$ is. - Line 58: Please define $\Sigma_v (\theta)$ explicitly. - Line 93: a --> an - Eq (7): Is $-4 \gamma$ in (7) replaced with $-2 \gamma$? - Line 142: I do not understand the definition of "fixed point." Please define "fixed point" explicitly. A fixed point $x$ of a mapping $T \colon R^d \to R^d$ is defined by $x = T(x)$. The definition of "fixed point" considered in this paper seems to be different from a fixed point of a mapping (see Theorem 4.3). - Eq (8): I do not understand the relationship between $\ell (\theta, z)$ and $\ell$ defined by (8). Why do we consider $\ell$ in (8) not $\ell (\theta, z)$? I think that the main objective of this paper is to minimize $L(\theta) = E_z [\ell (\theta, z)]$. - Eq (9): We cannot define (9) when $\gamma = 0$. Please define $C(\theta^\star)$ explicitly. - Eq (11): $\theta'$ is not defined. - Section 5.1: This section considers a matrix factorization: Given $A$, find $U$ and $W$ such that $A = UW$. The assumption of $A$ seems to be strong such that $A$ is close to the identity matrix $I$ (Lines 204 and 205). - Proposition 5.1: The assumption "the symmetry (11) holds" (based on $A \approx I$) seems to be strong. Please provide practical examples of empirical risk minimization and SGD satisfying (11). - Figure 2: There are not definitions of the matrix size, a batch size $S$, $\eta$, and $\ell (\theta,x)$. - Line 222: $\epsilon$ is not defined. - Theorem 5.2: Is the theorem based on (11)? If so, please provide some examples of $\dot{C}_B = 0$ under (11). - Figure 3: There are not definitions of $d_x$, $d_y$, $\epsilon$, $S$, and so on. Moreover, I do not understand why a warm-up with $\eta_\min = 0.001$, $\eta_\max = 0.008$, and the switching steps $5000$ is used. Please justify the setting from the theoretical results. - ... - Additional comments: I would like to know new insights that come from the theoretical results in this paper. This is because I am not surprised at the results. For example, can we apply the results to training more practical deep neural networks on practical datasets? Then, can we say new insights for training practical deep neural networks? Moreover, can we estimate appropriate step-size $\eta$, batch size $S$, and hyperparameter $\gamma$ using the theoretical results before implementing SGD? Anyway, I would like to know contributions of this paper.

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

N/A

Summary

The paper takes important steps toward showing that the noise in SGD pushed the dynamics along symmetry directions toward a unique fixed point. This differs in an important way from GD and gradient flow (GF) where one can show that in the presence of continuous symmetries (called exponential symmetries in the paper) there will be conserved quantities $C$, such as layer imbalance. These conserved quantities ensure that runs starting with different values of $C$ end up on different points along a minimum manifold (with some caveats). This paper shows that SGD noise causes drift toward a fixed point on a minimum manifold.

Strengths

This is potentially a very significant work. This is one of the few works I know which focuses on the effect of SGD noise on the symmetry directions and derives a concrete dynamics for this. I think the approach is at least partly original. Although there are many related works (see weaknesses), this paper is the only one which derives concrete solutions for the noisy dynamics along the symmetry directions. There are some corners that need to be ironed out. But if the results survive, this work could answer many questions about the implicit bias and generalization in SGD.

Weaknesses

There exist multiple other works addressing the same questions. Two particular ones with very similar claims that the authors may have missed are: [1] Chen, Feng, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. "Stochastic collapse: How gradient noise attracts sgd dynamics towards simpler subnetworks." Advances in Neural Information Processing Systems 37 (2023). [2] Yang, Ning, Chao Tang, and Yuhai Tu. "Stochastic gradient descent introduces an effective landscape-dependent regularization favoring flat solutions." Physical Review Letters 130, no. 23 (2023): 237101. I ask the authors to clarify the distinction between their work and the above. [1] is also about the effect of SGD noise on the symmetry directions, afaik. [2] also derives the effective dynamics of SGD as a Langevin equation and shows a result similar to yours, namely that along symmetric minima, the noise drives the system toward the least sharp minimum.

Questions

There are many details int he computations and proofs which I think need to be clarified. One aspect which I think is an error is related to the Lie algebra of the symmetries discussed in the paper. Most derivations in the paper hold for symmetric Lie algebra elements (i.e. generators of scaling and hyperbolic directions in the group) but fail for rotations. A relevant work which studies the parameter space more thoroughly is: [3] Zhao, Bo, Iordan Ganev, Robin Walters, Rose Yu, and Nima Dehmamy. "Symmetries, flat minima, and the conserved quantities of gradient flow." ICLR (2023). My questions are as follows: line 49: ref [40] may be wrong (it doesn't discuss conservation laws) maybe another paper by Zhao is meant: Zhao et al ICLR 2023 "Symmetries, flat minima, and the conserved quantities of gradient flow". sec 4: important caveat about conserved quantities is that the antiderivative C may not always exist. An important case is when the symmetry is a rotation. (see Zhao et al "symmetries, flat minima ..." ICLR 2023) In this case, the Noether charge is zero. 1. line 78: eq 5: to clarify, $\nabla^T C \nabla \ell(\theta, z) = 0$ is always zero due to eq 2 and 3, right? Is that why both terms in eq 30 (App B) vanish? However, there are some statements in Appendix B which seem inaccurate to me: a. when $\Sigma = GG^T$, the matrix $G$ need not be square and $G = \sqrt{\Sigma}$ is just one of the possible $G$. What is your argument for this $G$ being the correct $G$? b. I don't agree with your statement that "Because $\Sigma$ and $\sqrt{\Sigma}$ share eigenvector..." eq 31 would vanish. Trivial example: the 2x2 matrices $A=diag(1,-1)$ and the identity matrix share eigenvectors, but the vector $v=(1,1)$, $Av = 0$, but $Iv \ne 0$. I think the real reason is that $\Sigma = GG^T$. Doing SVD, $G = USV^T$ and $\Sigma = US^2U^T$. Then eq 30 yields $\nabla^T C U S= 0$. Assuming the covariance is nondegenerate (i.e. $S_{ii}\ne 0$ for all $i$), we must have $\nabla^T C U = 0$ which also implies $\nabla^T C G = 0$. 2. line 83: at a minimum, some symmetry orbits (i.e. invariant manifold) may collapse to a point, e.g. rotation symmetry in a quadratic loss as in linear regression $\|\theta x -y \|^2$. 3. line 104: Instead of "double rotation" this should be called double linear transformation, becuase the group is not just $O(h)$, but rather $GL(h)$ 4. eq 6: As noted above, this Noether charge vanishes for rotations $SO(h)$ because $A$ is anti symmetric and hence $\theta^T A \theta=0$ (see Zhao ICLR 2023). 5. line 114: why? Why the explicit weight decay? Why in defining the symmetry do you ignore $\|\theta\|^2$ which doesn't have rescaling or $GL(n)$ symmetry, only rotation symmetry $O(n)$? 6. Eq 7: what happens to the cross terms of $\theta^t \nabla \ell$ in $\Sigma_\gamma$? 7. line 118: eigenvalues and vectors of A can be complex (e.g. for Lie algebra of rotation), so $n^T$ needs to be $n^\dagger$. 8. line 214-215: I don't understand how you arrive at $\dot{C} = \mathrm{Var}[r]C$. I think you mean when the hidden dimension is 1, not the output, right? I hidden is $d>1$, let $\sigma=\Sigma_{u_i}$, then for $B= E_{kl}+E_{lk}$ we have $\mathrm{Tr}[\Sigma_{u_i}B] = \sigma_{lk}+\sigma_{kl}$, which are off-diagonal terms in the covariance matrix, unless $k=l$. 9. eq 13: When $\gamma\ne 0$, eq 13 does not have the symmetry (11) for symmetric Lie algebra element $B$. The reason is $\|W\|_F$ is invariant under rotations where $A^{-1} = A^T$. The Lie algebra of rotations (infinitesimal $A = I+\rho B$) has anti-symmetric $B$. 10. line 235: under what condition will the covariances $\Sigma_x$ and $\Sigma_\epsilon$ be isotropic? I don't think this happens often in SGD. I think you need adaptive gradient optimizers or Newton's method to have isotropic $\Sigma_\epsilon$. 11. eq 39,40: $\nabla^T_{\theta_\lambda} \ell = [\nabla_{\theta_\lambda} \ell]^T e^{-\lambda A^T}$, so we don't have (40) and in (41) there should $e^{-\lambda A^T}A e^{-\lambda A}$. 12. line 531: $Z_++Z_-$ will have complex eigenvalues when $A$ is not symmetric. Diagonalizing $A$, your arguments about $Z_\pm$ only hold for symmetric $A$.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

There are some gaps in the derivations (see questions) which I hope the authors can fill. Additionally, the experiments seem to only look at two aspects: alignment of weights in matrix factorization (which is already known under weight decay, afaik), and bias toward less sharp minima (again, known). I would have liked to see a more direct test of the mechanism the paper claims is driving this movement toward flatter minima. At least, the balancing of the weights or some other direct result from the paper would have been desirable.

Reply to the authors' comments

I also thank you for replies on my comments. I misread the manuscript. I am sorry for my first review. I still have some minor concerns for your replies. - I do not understand why $S \to \infty$ is done. Since $S$ is a fixed batch size, we should replace $S$ with $S_t$? - I do not know why $\sigma$ is less than $1$. Do you mean that $\sigma_t = \eta/(2S_t)$ is less than $1$ for a sufficiently large $t$?