Implicit bias of SGD in $L_2$-regularized linear DNNs: One-way jumps from high to low rank
Authors: Zihan Wang
Arthur Jacot
Venue: ICLR 2024 spotlight (2024)
Keywords: implicit bias
SGD
low-rank
linear networks
Primary area: learning theory
TL;DR: Low-rank bias of SGD in L2 reg. linear nets: SGD jumps from high to low rank minima, with no probability of jumping back.

Abstract: The $L_{2}$-regularized loss of Deep Linear Networks (DLNs) with
more than one hidden layers has multiple local minima, corresponding
to matrices with different ranks. In tasks such as matrix completion,
the goal is to converge to the local minimum with the smallest rank
that still fits the training data. While rank-underestimating minima
can be avoided since they do not fit the data, GD might get
stuck at rank-overestimating minima. We show that with SGD, there is always a probability to jump
from a higher rank minimum to a lower rank one, but the probability
of jumping back is zero. More precisely, we define a sequence of sets
$B_{1}\subset B_{2}\subset\cdots\subset B_{R}$ so that $B_{r}$
contains all minima of rank $r$ or less (and not more) that are absorbing
for small enough ridge parameters $\lambda$ and learning rates $\eta$:
SGD has prob. 0 of leaving $B_{r}$, and from any starting point there
is a non-zero prob. for SGD to go in $B_{r}$.

Published as a conference paper at ICLR 2024
IMPLICIT BIAS OF SGD IN LZ-REGULARIZED LINEAR DNNS: ONE-WAY WAY JUMPS FROM HIGH TO LOW RANK
Zihan Wang & Arthur Jacot
Courant Institute of Mathematical Sciences New York University New York, NY 10012, USA
1zw3508, arthur. jacot @nyu. edu
ABSTRACT
The Lz-regularized loss of Deep Linear Networks (DLNs) with more than one hidden layers has multiple local minima, corresponding to matrices with differ- ent ranks. In tasks such as matrix completion, the goal is to converge to the local minimum with the smallest rank that still fits the training data. While rank- underestimating minima can be avoided since they do not fit the data, GD might get stuck at overesttttiiinnn minima. We show that with SGD, there is always probability to jump from higher rank critical point to lower rank one, but the probability of jumping back is zero. More precisely, we define a sequence of sets B1 C B2 C... By that B, contains all critical points of rank less (and not more) that are absorbing for small enough ridge parameters and learning rates yy: SGD has prob. of leaving B,, and from any starting point there is a non-zero prob. for SGD to go in B,.
INTRODUCTION
Several types of algorithmic bias have been observed in DNNs for a range of architectures 10 34
27 28
Understanding and characterizing these types of implicit bias is crucial to understand the practical performances of Deep Neural Networks (DNNs).
We focus on Deep Linear Networks (DLNs) As WT ... W1 for 0 (W1, WL), that are known to be biased towards low-rank linear maps in a number of settings:
1. Adding Iegregularizztion to the parameters of a DLN has the effect of adding Lp -phStttteen norm (the Lp norm of the singular values of matrix) regularization to the learned matrix for 2/2 where L is the depth of the network (
2. When trained with the cross-entropy loss, Gradient Descent (GD) diverges towards infinity along direction that maximizes the margin w.r.t. the parameter norm (17), leading to form of implicit Iegueaularinatioon with the same bias towards low-rank matrices.
3. When the parameters are initialized with small variance, the network learns incrementally matrices of growing rank, thus converging low-rank solution (3 14).
This low-rank bias particularly useful in the context of matrix completion ( where the goal is to recover matrix from subset of its entries under the assumption that the full matrix is low rank. The task of finding the lowest rank matrices fitting the observed entries is NP-hard, but convex approximations
work well (5,), well as DLNs 35).
In the deep case L > 2, the Lp Schatten norm becomes non-convex (because p p= z/L < 1) and there multiple local minima in the Legreaular loss, each corresponding to matrices with different ranks (similarly with cross-entropy there could be multiple directions that locally minimize the rank). Which of these local minima/critical point will GD converge to? And what will be the final rank?
We will see how Stochastic Gradient Descent (SGD) can lead the dynamics to jump between critical points with bias towards low-rank points.

---

Published as a conference paper at ICLR 2024
1.1 CONTRIBUTIONS
In this paper, we focus on the implicit bias of SGD in Deep Linear Networks (DLNs) of depth L larger than 2 with Læ-regularization, when trained on Matrix Completion (MC) tasks.
We first describe the many critical points of the Tzzrregulriizd loss. We then split them into three groups, depending on whether they recover the 'true rank', underestimate, or overestimate it.
We show that the rank-underessimatiig critical points can easily be avoided by taking a small enough ridge X, but no such strategy exists avoid aank-oversttmmatnn points with GD.
However we show SGD has a small but non-zero probability of jumping from any parameter to a lower rank parameter, but the probability of jumping to a higher rank is zero. More precisely, define sets B, that contain all critical points of rank r or less and show that they are absorbing: the probability for SGD to leave this is zero, but the probability for SGD to move from outside of this set to inside (in sufficiently many steps) is non-zero.
This suggests that rank- overestimation can be avoided if we continue SGD training long enough (but not too long), since the rank will decrease incrementally. This illustrates the low-rank bias of SGD,
1.2 RELATED WORKS
The low-rank bias of DLNs has been observed in a number of different settings: for example as result of Iz-regularization or training with the cross-entropy loss (9), and as result of small initializations 2 23 14). These results rely on similar tools such as the balancedness condition, however the underlying training dynamics leading to sparsity are very distinct.
Motivated by the empirical observation that SGD improves generalization 20 19). there has been interest in the implicit bias of SGD. There is a line of work approximating SGD with different Stochastic Differential Equations (SDEs) (25) [3) [6 16. sometimes approximating the parameter dependent noise covariance with fixed scalar multiple of the identity thus leading to Langevin dynamics (66), and in general studying the resulting steady-state distributions 7 These SDE approximations require small learning rates (22), but approximations to capture the effect of large learning rates have been proposed too
These works however focus the bias of SGD in parameter space, showing e.g. that it can be interpreted as changing the potential/loss (7), or adding a regularization term (32). More recent work has focused the bias of SGD in diagonal linear networks (29 (36) leading sparsity effect in the vector represented by this network.
We focus the effect of SGD in the context of deep fully-connected linear networks with L2- regularization, showing that SGD strengthens the already existing low-rank bias induced by L2- regularization. To our knowledge our work also unique that it does not rely on a SDE/continuous approximation.
2 SETUP
We study Deep Linear Networks (DLNs) of depth and widths w/ din. Wi..... WL dout
Ag = WT ... W1. for the x wy-1 weight matrices Wj and parameters 0 (WI.... WL). We will always assume that the number of neurons in the hidden layers is sufficiently large we 2 min ( (dnn dout) so that any dout x din matrix A can be recovered for some parameters 0: A Ag.
2.1 MATRIX COMPLETION
We consider the Le-regularized loss
Lx = C(Ao) + + |1012
where C' is a loss on matrices such as the Matrix Completion (MC) loss
C(A) 2N
1
Σ (47 Aoss) 2
(1.))r
2

---

Published as a conference paper at ICLR 2024
where A" is the true matrix we want to recover and I {1,.... dout) } X {l..... din) is the set of observed entries of A" of size = .T1. While it is not possible in general to recover an entire matrix A from subset of its entries, it is possible if A is assumed to be low rank.
The ideal goal is to find the matrix Å with lowest rank that matches the observed entries. We define the smallest rank as the smallest integer r* such that inf A ((A)) = O. Note that one could also define r" to be the smallest integer where this infimum is attained at finite matrix A, which can be higher in MC problems where filling in infinitely large entries can allow for lower ranks fitting functions. In the main we restrict ourselves to the first definition, but we discuss the second choice and its implications in Appendix A
Since finding the minimal rank solution is NP-hard general (5), popular approximation is find the matrix Å that minimizes the MC loss with a nuclear norm regularization
min 2N
Σ (A - Au) +1
(1,))E
where the nuclear norm the sum of the singular values of A: |||||, - convex and can be efficiently minimized, and it has been shown that it recovers the true matrix A* with high probability with an almost optimal number of observations (5).)
RankA (A). This loss is
DLNs have also been used effectively in Matrix Completion, thanks their implicit low-rank bias. The importance of low-rank bias in the Matrix Completion setting, makes it ideal to study the implicit bias of SGD in DLNs.
2.2 REPRESENTATION COST
The low-rank bias of Lz-regularized DLNs can be understood in terms of the representation cost R(A: L) DLNs, which equals the minimal parameter norm required to represent a matrix with DLN of depth L:
R(A; L)
min A:A=
|10112
As observed in (9), the representation cost of DLNs equals the Lp Schatttenn norm ||Al| (the 4p norm of the singular values) of A for p 2/1.:
RankA
L) 7/2 ==
= LIA
Σ
(A)3
1=1
This implies that the L2-norm regularization in parameter space can be interpreted adding LpSShttteen norm regularization in matrix space:
min C(Ae) + > 1101 min C(A) AL AL 3/2
For shallow networks (L = 2). the representation cost equals the nuclear norm R(A;2) 2 ||| The loss has only global minima and strict saddles, thus guaranteeing convergence with probability 1 to global minimizer á of the LHS, and the represented matrix Aa then minimizes the RHS. We therefore simply recover the convex relaxation of Matrix Completion, with the advantage that the loss (0) is differentiable everywhere, so that it can be optimized with vanilla GD (35).
In the deep (L 2) case however, the representation cost R)A:
1// non-convex,
and both RHS and LHS may have distinct local minima with varying rank. We will show that the correspondence extends to local minima of the RHS and LHS, and that there always exists one local minimum with the right rank r".
But there multiple local minima and critical poinnts with different ranks, for example the zero parameters = 0 corresponding to the zero matrix Ag = 0 is always local minimum. Or there might also be critical points that overestimate the 'true rank' that want recover.
For GD with a Gaussian initialization, there is a non-zero probability to converge to any local minimum. On the other hand, will see how SGD can jump from local minima to local minima.
3

---

Published as a conference paper at ICLR 2024
2.3 STOCHASTIC GRADIENT DESCENT
We consider SGD with replacement, that is at each time step / an index (ft; is sampled uniformly from the index set I, independently from the previous iterations. The parameters are then updated according to the learning rate 1)
1.+1 = (1 (1-111200 2mp (1-2n))6,
2 Vo (Ai - -Ao,ūije)
Note that due to the Iz-regularization there remains noise even at the local minimizers, in contrast without Iz-regularization there is neither noise nor drift at the global minima of the loss. Thus with L2 -n-regularization the dynamics never completely stop, making it possible for SGD to jump from one local minimum to another.
Remark 2.1. A number of previous works have approximated SGD by GD with Gaussian noise, the simplest of which is to approximate SGD by Langevin dynamics. Under this approximation, there is always likelihood of jumping from local minimum any other local minimum, with a higher likelihood of going to (and staying at) local minima with lower loss. Our theoretical results show completely different behavior, where SGD may have non-zero probability jumping from one local minimum to another, but zero likelihood of jumping back. Furthermore the likelihood of SGD visiting certain local minimum will not scale with the loss of that local minimum, but rather its rank. This further shows that the Langevin approximation of SGD is inadequate.
appears that what leads to this difference in dynamics, is that the covariance of the noise is ciw-dimenioonl and highly anisotropic, with a lot of noise along directions that keep or lower the rank, and almost no noise along directions that increase the rank.
MAIN RESULTS
We will first give description of the loss landscape of I-zrregllarized DLNS and then state our main result, which says that SGD has non-zero probability of jumping from parameters of certain rank to lower rank, and that once in the neighborhood of low rank critical point, the probability of reaching higher rank is zero.
3.1
L2 LEGULALLRDD Loss LANDSCAPE
The correspondence of the minimizers of (x) and Cx(A) = C(A) + + | extends their local minima:
Theorem 3.1. If I is a local minimum of (x),), then Ag is local mininum of CA)). Conversely, if Å is a local minimum of Cx(A) then there is a local minimum of Lx(() such that À = Ag-
The critical points of the Legregularized loss balanced, i.c. WTW, W WT_ for all 1,
(see Appendix [A.. This implies that all weight matrices have the same singular values and the same rank r. We may therefore define the rank of a critical point Ô as the rank r of any weight matrix W. which also matches the rank of the represented matrix Ag-
In general, there are several distinct local minima/critical points with different ranks. The origin 0 0 is always local minimum (since the unregularized loss has vanishing first L derivatives at 0, it becomes a local min even for any X > 0), furthermore for small enough ridge X, there always is a local minimum that finds the minimal rank required fit the observed entries:
Proposition 3.1. Consider matrix completion problem with true matrix A and observed entries T. For all X, there is rank r" local minima 0 (() of (0) such that lim yo C(Ao())) = 0.
Note that finding a fitting matrix of minimal rank is known to a NP-hard problem in general (5). which means that it should in general be hard to find this local minimum. There are three types of problematic local minima/critical points:
Sann-unddrestimating critical points: these critical points such the origin 0 0 with rank lower than the minimal rank r", so that the represented matrix Ag cannot fit the observed entries. GD with small enough ridge and learning rate will avoid such rank underestimating minima:
4

---

Published as a conference paper at ICLR 2024
Proposition 3.2. Given initialization du such that unregularized (A 0) gradient descent converges to a global minimum, there a constant с >> 0 such that for all small enough ridge X and learning rate 1), regularized GD 0, for all t 2 the -*-th largest singular value of the matrix is non-zero: 8,- (Ao.) 2 c.
There exists multiple papers proving convergence a global minimum under different assumptions 13 26),, but empirical evidence seem to suggest that convergence to global minimum is much more general, and we expect more convergence results to come.
Rank-overestimating critical points: these have a larger rank than r" and the represented matrix Ag can fit the observed entries (with small O()) error). These are harder to avoid, suggesting that the NP-hardness of finding optimal rank r* fitting matrix can be related to avoiding these points. might happen that there are no ovressttrstimmtinn minima/critical points, in which case GD can recover the minimal rank solution easily. but from ovnesseerestimating points appear and how SGD manages to avoid them.
will focus on settings where these
'Bad' Rank-estimating local minima: Finally there might exists local minima (or non-strict saddles where GD/SGD might also get stuck) with the right rank r" that do not the true matrix. This paper will not focus how to avoid these points, but we believe that it could be proven that such minima are rare for large input and output dimensions, using similar tools as (4 18). The noise of SGD might also help avoid these minima/saddles. In numerical experiments, never encountered such minima, but we did encounter both other types.
3.2 ONE-WAY JUMPS FROM HIGH TO Low RANK
We now show how SGD helps avoiding aank-verrstimattnn critical points. More precisely we show under conditions on the learning rate 1) and ridge À that there is always a (small) likelihood of jumping to a lower rank, but the probability of jumping back zero. This suggests a strategy: train the network with small ridge to guarantee convergence to a critical point of at least the right rank, and then take advantage of the SGD noise to lower the rank until finding the right rank.
For our analysis, define family of regions B, IR. of parameters 0 that are:
1. 61 aaappooximately balanced: for all layers e, ||WTW: - We- We-
2. 62, a-aaproxxmmaely rank (or less): for all e, E RankW, So (s,(W, We)) S + e2 where 8; (A) the i-th singular value of A and (a is twice differentiable function such that (0) 0, (r) = 1 for x > a, 0 (x) < for some constant K and (x) 0.
3. C-bounded: ||Will? c.
Note that we chose the function fa to be differentiable, and to satisfy fa) = 0 and fa (x) = Yz 2 a. This yields a notion of approximate rank that converges the true rank Y O. One
#22a x), a
example of such approximate rank is fa())
a
Since all critical points i of the Lz-regularized loss are balanced, the set B, contains all critical points of rank r or less for C' large enough and all 61, 62 2 0, and it contains no critical point of higher rank for e2 and a small enough. These sets allows us separate critical points by rank, with small neighborhood.
Proposition 3.3. For any critical point ô Br, we have Eiz1 (a(s((A)))1)) sr+
RankAs
Proof. Since is balanced, Ag S where S E Rdouexdin is the diagonal matrix of RankW: fa(s;(W, We)) $ r + E2. Ag satisfies singular values of all We's. Since for any L. E RankA, Sa (s (A2/T)) Sr+ter 62.
We can now state our main result, which says that the set B, is absorbing for all r. i.. SGD starting from anywhere will always end up at some time inside B, and then never leave it:
5

---

Published as a conference paper at ICLR 2024
Theorem 3.2. For any 2 0, À 0, any large enough C and small enough 61, C2, a, 77 the set B, is closed
0, B, 0t+1 E B,
and r 2 and any parameters 0e there is a time T = 12(A - 17 (i.e up to log terms) such that
P (O++T E B,) Z mmin(din dout)
thus for any starting point SGD will eventually reach By:
P(ET: Ourr E B,( =1
Proof. (sketch) (1) The closedness of the set of er-approximately balanced parameters follows from the fact that in the gradient flow limit > 0, the balancedness errors W We Wi-1W2 decay exponentially
a (w? We We-1 - (w? We- To guarantee a similar decay with SGD, we simply need to the control the O terms.
Given e--approximatel balancedness, the closedness of the ap-ppoxximaeely rank or less parame- ters follows from the fact that the dynamics resulting from the minimization of the cost C(Ag) are very slow along the smallest singular vectors of Ao (I but the Lz-regularization term pushes these small singular values towards zero. For small enough singular values, this second force dominates, thus leading to a decay towards zero.
(2) Under the event Ar that in the steps from & to t + T - all the random entries (i., j.) are sampled from the same r of the dout columns, one can show that the dout r other columns of WA decay exponentially to approximately 0, implying an approximate rank of r or less. The probability
of that event is least
dow!
This shows the implicit bias of SGD towards low-rank matrices in matrix completion: SGD can avoid any oanestttresnnitting minima/critical point given enough training steps.
Explicit bounds C, at é1, E2, 7) can be found in Appendix B The bounds are rather complex, but
we give here an example of acceptable rates in terms of A: C~>-1,a~~ x ~ 2 21++1
62 X and X X __2, These rate suggest that extremely small learning rate 7j is necessary, especially for large depths L. thus making the likelihood of a jump appear very small. This seems in contradiction with empirical observations that larger depths tend to make these jumps more likely. We believe our bounds could be made tighter, in particular when comes to the dependence on the depth T. to better reflect our empirical observations.
We expect this result to generalize to other tasks. The first part of the theorem (the closedness of B,) should generalize to costs such the MSE loss and others, and the second part too, under the event that one samples from the same r training points over T time steps for the MSE loss, or sample from the same classes for classification tasks.
A limitation however is that the second part of the result relies on the fact that we sample the observed entries independently at each time with possible replacement. In practice, the dataset is randomly shuffled and taken in this random order, so that every observed entry is chosen exactly once during each epoch. This would force the jumps to happen within an epoch, which may not be possible depending on the problem.
Another limitation is the average time required to observe one of our predicted jumps can easily be absurdly large. To observe a jump in reasonable time, one also needs rather large learning rates n, leading to very noisy dynamics. This makes periodic learning rate choices attractive, with large periods allowing for jumps to lower-rank region, and low 7) periods allowing for SGD to settle around a local minimum.
Nevertheless, our result also shows that the common approach of approximating SGD with SDE such as Langevin dynamics and studying the stationary distribution (usually with full support over the parameter space) is misleading. In contrast, result implies that any stationary distribution must have support inside B, thus under-estimating the true rank general. It is thus crucial to understand the distribution of SGD intermediate times, when the rank has yet collapsed to 1.

---

Published as a conference paper at ICLR 2024
30
20
1004
-0.4
1000
2000
3000
4000
5000
5000
0
1000 2000
1003
400D
5200
6002
7000
time
time
Figure Annealing Schedule: DLN with L 3. W1 = 100 on the 2 x 2 MC problem with == 0.25 [Light blue] SGD with == 0.1 and 7 0.2 (1) 0.03 for the first 500 steps to avoid explosion) [dark blue] different times, we create offssoots with À 0.001 and 0.02 to fit the data. [red] The ratio of the second to first singular value of Ag the large X, 17 path. We see a jump around time 2000, where the output matrix becomes rank 1. The offshoots created hefore this jump fail to fit the missing entry while those created after succeed.
Figure 2: Effect of depth: We study the effect
of on the MC problem with e = 0.1. We train 5
networks of each depths == 3 [blue] and T. 4 [red] with the same schedule: = 0.03, X = 0.1 until & 500, then r 0.25, X 0.1 until f
5000 and finally = 0.05, À 0.001 until the end. We see that the five depth L networks are unable to jump in this time, while all five depth L networks jump different times during the first 5000 SGD steps.
3.2. NONLINEAR NETWORKS
Since linear networks are simplification of nonlinear networks, it is natural to wonder whether the results presented here could be generalized the nonlinear case. We identify two possible strategy to generalize our results:
First along the lines of ( which observes similar phenomenon where SGD is naturally attracted to symmetric regions of the loss (where for example two neurons are identical or one neuron is dead) in nonlinear networks. The L2 regullaariiztion is known to make these region more attractive (I5). which could have compound effect with SGD. In DLNs, the regions of low rank that we prove are attractive can also be interpreted as neighborhoods of symmetric/ invariant regions.
Second, recent work has shown that L regularized ReLU DNNs with large depths are biased towards minimizing a notion of rank over nonlinear functions, the Bottleneck rank (12). We have hope that results could be extended prove similar low-rank bias with this new notion of rank. This is further motivated by the observation that such large depth networks exhibit a Bottleneck structure (II) where the middle layers of the network behave approximately like linear layers.
NUMERICAL EXPERIMENTS
For our numerical experiments, we want to find a Matrix Completion problem that GD cannot solve but SGD can. In particular, we want to find a setup where GD converges with high probability to rank-overestimating minimum, and where SGD can jump from this minimum to a lower rank minimum in reasonable amount of time.
is rather difficult to find setup that lies between the regimes where both GD and SGD work and where neither work. This is in line with previous work in the bias of SGD (9): in diagonal networks a value (determined by the initialization) determines a transition between a sparse and non-sparse regimes, and SGD has the effect of pushing this value towards the sparse regime; this have significant sparsity effect if the original value was at the transition between regimes, but little effect it was far into either regimes.
7

---

Published as a conference paper at ICLR 2024
10-
0.6
6-0
250
1000
0.2
0.0
10-1
10.
1.000
2000
30000
4000
0000
7000
Figure 3: Annealing accross For a range of G, we train networks (I. = 4,w 100) with an annealing schedule in Figure and plot the test loss divided by the test loss when putting zeros in the missing entries. The four networks are trained for to steps with high noise, followed by 4000 steps in the low noise regime. Without noise phase (tg = 0) the network fails to recover the rank I solution. Larger in allow the network to recover for even smaller c.
Figure 4: Periodic Schedule: DLN with L= 3, W1 100 the 2 x MC problem with 0.2. We plot runs of SGD with periodic learning rates alternating between S 0.1, À 0.001 and 1/ = 0.4, X = 0.1. We see that the different trials make jumps during the high 7), À period. After the jump. SGD will settle at a low test етrоr in the slow periods, allowing us to identify when the jump happened.
We choose a MC problem, inspired by (30), that allows tune the difficulty of finding sparse solution. We observe 3 out of 4 entries of a x 2 matrix:
( 1 r)
Filling the missing entry with c-1 leads to rank matrix. The smaller c is, the larger the missing entry that needs to be filled in needs to be.
In La-regularized DLNs with L > there are three local minima: the rank 0 minimum at the origin which can easily be avoided, a set of minima that learn the rank 1 solution, and a set of rank-overestmmating minima that learn rank 2 solution by filling the missing entry with small value.
For small values, GD almost always converges a rank-overestimating minimum (see Figure In such setup, SGD can outperform GD by jumping to rank 1 solution. To achieve a jump in reasonable amount of time, we need the ridge parameter À and the learning rate r to be large. But such a choice of large X, prevent SGD from minimizing the train error.
We investigate two strategies to take advantage of both the jumping properties of large X, n and fitting properties of small X, n7:
In Figure I we run SGD with large X, r for some time tg. waiting for jump and then switch to small values of X, 17 for convergence. Another specificity is that we take a small learning rate for the first few steps, because SGD diverges if we start with a too large learning rate directly at initialization, whereas large learning rates are possible after few steps (we do not have theoretical explanation for that).
'Annealing' Schedule:
We test different switching times to small X, 17 values, and we see clearly that if we switch after the jump at time 2000, we obtain rank 1 solution, but we switch before the jump then training fails and recovers rank 2 solution.
By changing we can tune the difficulty of finding the rank solution. We see in Figure 3 that the smaller e, the longer one needs wait for jump, and thus the longer one needs to stay in the high noise setting. We also see that without a high noise period (i.e. when we are close to GD) the network fails to recover the rank solution even for =
8

---

Published as a conference paper at ICLR 2024
Periodic Schedule: Another strategy it alternate between large and small X, 17. We see in Figure 4 how the jumps all happen during the large X. 77 periods. is also interesting to see that even after SGD has settled in the vicinity of local minimum in one of the small X, 1 periods, SGD can still jump to another minimum in a subsequent large X, 1) period.
Finally we also study the effect of depth in Figure 2 and observe that depth increases the probability of jumps. We train networks of depths L = 3 and Z. = -4 on the x 2 MC task with c = 0.1. While for the choice 0.25, depth L = 3 network was able jump in a reasonable amount of time, for this smaller choice of e we do not observe a jump (even with the same hyper-parameters, In contrast, the deeper networks L the likelihood of a jump.
jump reasonable amount of time, suggesting that depth increases
5 CONCLUSION
We have given description of the loss landscape of I-z-regularized DLNS, giving a classification of its critical points by their rank. We have then shown that SGI) has a non-zero probability of jumping from any higher rank critical point to a lower rank one, but it has a zero probability of jumping in the other direction. We observe these jumps empirically, To our knowledge, this is the first description of the low-rank bias of SGD in the context of fully-connected linear networks with two more hidden layers.
Our analysis is also significamtly different from previous approaches that rely approximating SGD with a continuous stochastic process, and/or studying of the limiting distribution this continuous process. It appears that the phenomenon of absorbing sets of different ranks cannot be recovered with a continuous approximation, and the jumps we describe happen before SGD has reached its limiting distribution. This puts into question the adequacy of the continuous approximation and limiting distribution assumption.
REFERENCES
[1] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. International Conference on Learning Representations, 2019.
[2] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 244-253. PMLR, 10-15 Jul 2018.
[3] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[4] Emmanuel Candes and Benjamin Recht. Exact matrix completion via convex optimization. Communications of the ACM, 55(6):111-11, 2012.
[5] Emmanuel J Candès and Benjamin Recht. Exact matrix completion via convex optimization. Foundations of Computational mathematics, 76):777777 2009.
[6] Emmanuel J Candès and Terence Tao. The power of convex relaxation: Near-optimal matrix completion. IEEE Transactions on Information Theory, 56(5):2053- 2010.
[7] Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference. converges to limit cycles for deep networks. In 2018 Information Theory and Applications Workshop (ITA), pages 1-10. IEEE, 2018
[8] Feng Chen, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. Stochastic collapse: How gradient noise attracts sgd dynamics towards simpler subnetworks, 2023.
[9] Zhen Dai, Mina Karzand, and Nathan Srebro. Representation costs of linear neural networks: Analysis and design. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Newral Information Processing Systems, 2021.

---

Published as a conference paper at ICLR 2024
[10] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 1832-1841. PMLR, 10-15 Jul 2018.
[11] Arthur Jacot. Bottleneck structure in learned features: Low-dimension vs regularity tradeoff, 2023.
[12] Arthur Jacot. Implicit bias of large depth networks: notion of rank for nonlinear functions. In The Eleventh International Conference on Learning Representations, 2023.
[13] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kemel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems 31, pages 8580-8589. Curran Associates, Inc., 2018.
[14] Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to- saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity, 2022.
[15] Arthur Jacot, Fugene Golikov, Clément Hongler, and Franck Gabriel. Feature learning in Iz-regularized dnns: Atracctonnrepullion and sparsity. In Advances in Neural Information Processing Systems, volume 36, 2022.
[16] Stanistaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima sgd. arXiv preprint arXiv: 1111111144440 1.04623, 2017.
[17] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 17176-17186. Curran Associates, Inc., 2020.
[18] Raghunandan H Keshavan, Andrea Montanari, and Sewoong Oh. Matrix completion from few entries. IEEE transactions on information theory, 56(6):2980-299) 2010.
[19] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ArXiv, abs/16099.eele 048866, 2016.
[20] Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klbert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pages 9-48. Springer, 2012.
[21] Adrian S Lewis and Hristo S Sendov. Twice differentiable spectral functions. SIAM Journal on Matrix Analysis and Applications, 23(2):368-386, 2001.
[22] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pages 2101-2110. PMLR, 2017.
[23] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. International Conference on Learning Representations, 2020.
[24] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Toward a theory of optimization for over- parameterized systems of non-linear equations: the lessons of deep learning. arXiv preprint arXiv:2003.00300 2020.
[25) Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. arXiv preprint arXiv: 04289, 2017.
[26] Hancheng Min, Rene Vidal, and Enrique Mallada. On the convergence of gradient flow on multi-layer linear models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 Proceedings of Machine Learning Research, pages 24850 24887. PMLR, 23-29 Jul 2023.
10

---

Published as a conference paper at ICLR 2024
[27] Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 22182-22193, Curran Associates, Inc., 2020.
[28] Greg Ongie, Rebecca Willett, Daniel Soudry, and Nathan Srebro. A function space view of bounded norm infinite width relu nets: The multivariate case. In International Conference on Learning Representations, 2020.
[29] Scott Pesme, Loucas Pillaud- Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 1921009223.. Curran Associates, Inc., 2021.
[30] Noam Razin and Nadav Cohen. Implicit regularization in deep learning may be explainable by norms. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 21174-21 Curran Associates, Inc., 2020.
[31] Horn Roger and R Johnson Charles. Topics in matrix analysis, 1994.
[32] Samuel L Smith, Benoit Dherin, David GT Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. arXiv preprint arXiv:2101.12176, 2021.
[33] Samuel L. Smith and Quoc V. Le. A bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018.
[34] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1):2822-2878, 2018
[35] Ruoyu Sun and Zhi- thi-Quan Luo. Guaranteed matrix completion via non-convex factorization. IEEE Transactions on Information Theory, 62(16) (2(11):65356677) 2016.
[36] Loucas Pillaud Vivien, Julien Reygner, and Nicolas Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Conference on Learning Theory, pages 2127-2159. PMLR, 2022.
[37] G.A. Watson. Characterization the subdifferential of some matrix norms. Linear Algebra and its Applications, 170:33-45, 1992.

---

Published as a conference paper at ICLR 2024
The Appendix is organized follows:
Section
A
contains the proofs of Theorem 3.1 and Propositions 3.1 and 3.2 of the main.
Section
B
then describes how Theorem 3.2 of the main can be split into two statements:
Theorems
B.
and B.2
Section
state some preliminary result for the proofs.
Section
proves the first part of Theorem 3.2 from the main, that Theorem B.1
Section
proves the second part of Theorem 3.2 that is Theorem B.2
Section
states and proves more general version of Proposition 3.3 of the main.
A Loss LANDSCAPE
Proposition A.1. Let be a critical point of the loss Lx. then is balanced, i.e. WoW? =
W W241.
Proof. At critical point, we have
W WIVC(AA) .. ....N....WWW............. 2AW, 0.
Thus
1 WT W. W WT. We+1.
WWW/-
2A
Proposition A.2. Let 0 be critical point of the loss Lx, then:
i is local minimum if and only Ag is local minimum of Ox.
0 is strict saddle/maximum if and only Ag is strict saddle/maximum..
Proof. We know that any critical point of the Le-regularized loss is balanced. The parameters = (W1....,Wa Wz) are therefore of the form
= U.U
for some d diagonal S (where d min(din, dout) } and we x matrices Ue with orthonormal columns (U? Ug Ta).
(0) For any sequence of matrices A2. converging to A with SVD decompositions A; Ü.S. ỤT (chosen so that Ü. S; and V. converge to the SVD decomposition Ag U,SU we can construct parameters 0, with weight matrices
Wi = U.S. V.T
W. = U.StU
Wz = ú,st UT-1-
We have (1) Ao, A, (2) - - (3) |0,11 = T. 1|4:11//2 and therefore Cx(A)) = Lx(O.).
(1a) If 0 is not a local minimum, there is sequence 0 + 0 with (x,)) < Lx),,, thus the sequence A; Ag, converges Ag and (A.))) (0,) (0) Cx(Ag), implying that A is not local minimum.
(1b) If Ag is a local minimum, there is sequence A; Ag with Cx(A.) < <Cx(Ag), by point
(0), we construct a sequence 6, - 9 such that L, (0,) = CA,, < C2(Ag) = Lx(0), proving that not a local minimum.
12

---

Published as a conference paper at ICLR 2024
(2a) 0 is strict saddle/maxiimum if there is sequence 8, - 0 such that
lim Lx(0:) (x 0.
||0 0112
Our goal now to show that the sequence Ag; that converges to Ag satisfies the same property, First
note that
Cx(Ae,) - C2(Ag)
Cx(0:) Lx
||0 012
lim 7-000
lim
|A.. AAlli
1000
110 -01
||Ab. AgllT
since (A(A)) S (0.) and C2(Ag) = (0).
Since Ag is twice differentiable w.r.t. 8, we know that ||Ag. - Agll p - cl1; -81 -0 00 - -812) and thus
10:- 0|
|00 1112
1 2
As. Agll?
(cll0, बll O (100, - -(1p))
(c+ O (100: - -01)
4
for large enough This implies that
C2(Ao)) C2(Ag)
Cx(0.) (x) 1
lim 6-000
lim
|0: 01 c² 0,
Agg -
10000
where we used the fact that for large enough (x(0.)-2x(0) 18,-0113
is negative.
This implies that Ag is a strict saddle/maximum.
(2b) Ag is a strict saddle if there sequence A; Ag with
lim Cx(A,) - (Ag)
< 0.
[A A
For T. 2 we may assume that A, has the same rank as A₂ for large enough à: a matrix cannot be approached with matrices of strictly lower rank (the best rank approximation matches the l largest singular values leading to an error larger than the k + largest singular value), and it is approached with strictly larger rank the regularization term || 3/2 would be strictly larger.
We now construct a sequence 0; - as in (0). Consider the map ø that maps matrices A = ÜS²T the neighborhood of Ag with the same rank to the parameters
Wi = U.
W₂ U,
Wz = ūstuT
We have ((A,)) = 0, and ((Ag) = And since á is differentiable at Ag along directions that do not change the rank (37). we have
lim (0,) (x) lim Cx(A.) - C2(Ag)
||A, - -Agle < <. 1A,-Aolla. [(0(A.) - 6 (Ap)ll
0 8
Finally we prove the existence of a minimum with the minimal rank required to fit the data: Proposition A.3, Consider a matrix completion problem with true matrix A and observed entries For all X, there is rank y* local minima ((2) of (0) such that limx C(Ae())) 0.
13

---

Published as a conference paper at ICLR 2024
Proof. For all X, take the matrix A()) minimize the cost (A) amongst the set of matrices of rank r" or less. The regularization ensures that the infinimum inf A:RankASr Cx(A) is attained at a finite matrix A.
The matrix A(A) is also a (possibly non-global) minimum of the non-restricted loss: is locally optimal amongst matrices of the same rank, and for any matrix A' of rank r", define A'gr- to be the closest rank rd matrix (which is obtained by setting all but the -"ttt largest singular values of A' to zero), we have
CA(A) = -2(AE))) + ALIIA' A'2--11/2 -(A')))((A)))))) )].
2/
We have that C^ (A's..)) 2 CA(A()) and in small enough neighborhood of A(A), one can guarantee that the second term dominates the third one since C is differentiable (if 8,++1 > 0 the r" 1-th largest singular value of A', then the second term is at least ALS while the second one is at most el @(Rank A' 8,++ where < oc bounds the I Lipschitzness of C' in the neighborhood).
By Proposition A.2 there are corresponding local minima (() of (0( such that A(A) Ae(a)-
By the definition of the minimal rank r", we know that inf A:RankAgr Cx(A) and thus limx C(Aa()) O.
A.1 AVOIDING RANK-UNDERESTIMATTIN CRITICAL POINTS
With small enough ridge À and learning rate one can guarantee that GD will avoid all rank- underestimating local minima:
Proposition A.4 (Proposition 3.2 in the main). Given an initialization Oo such that unregularized (A = 0) gradient descent converges to a global minimum, there is constant >> O such that for all small enough ridge X and learning rate n, regularized GD for all I to the -thth largest singular value of the matrix is non-zero: Sy- (Aa.) 2 c.
Proof. Since unregularized GD converges to global minimum with zero loss, there is time such that ((0)).2=0) inf A.RankA C(A).
We know that I.grregularized GD 0x (f) converges to unregularized GD 0(t) À Y 0 for any fixed time t. Then for all sufficiently small À have L (0t,.)) S 1/ inf A:RankA C(A).
Lemma A.5 also tells that GD is non-increasing for small enough so that for all 2 lo we also have
(O(,,,)))) S 1 inf C(A).
ArRankAcr
Let us now show that there is constant e > 0 such that the r" largest singular value 8, (A of Ag..., is lower bounded by c.
Assume by contradiction that this is not the case, then for all e, there matrix A. with s, s, Se for all 2 r" and with (x(A)) < 2 inf ARRakkAer C(A).
We know that the matrices A. are bounded for all e since ||A.12/1. 2/L I 1CA(A.) S 2X inf RankArr C(A) co. There is thus convergent sequence Ae Aeg.... with En + as + 0o, whose limit À = lim Ae, is finite, has rank strietly less than r" and satis- fies C(A) S Cx(A) < 2 A.RnkAAC C(A) which yields contradiction with the fact that inf A:RankACr C(A) > 0.
Here are a few Lemmas required for the previous proof: Lemma A.1. The loss gradient (.) in the ball B(0, D) is Lipschitz with constant 2L N (2D2L-2 cD4-2 2X).
Proof. For all parameters with ||0|| < D, we have [Wllr & D and thus ||All & D Further- more for any two parameters [W..... Wr) and Φ (V1,..., Vi.) with [|0]|, |o] D,
14

---

Published as a conference paper at ICLR 2024
have
||Ag A0llF S I... W.(( Vi)ll ... 1(W,, v.).vc-I...Vill Vill
- D Σ
| - Vell
8=1
VVDD"=111 110 +01-
Now since VC(As)) M o (Ao- -A*), this implies that
|[CC(A))]F - N |A + N
2 ||
S
Dz. D +c
N
and
2 VLDD-1 ||0 611
- - N
| - Apll 1
N
Now using these facts, we get that the difference
||V[x(0) (+) is upper bounded by
W WIVC(AA)T .. WZ, ViT1 VEVCCA ... V2 + 2AW; -2AV/
DDL-2 W Vill |VC(AA)] DD-1 (VC(A)) vC(A)) +22 |e: Vella)*
kett
(D2L-2 + |W - Valle + 2 N VT. 110 21 + 22 We -Vell
t=1
(T((2221 2DD1-2 22) 110 - -b1)
N
t=1 Kif
4T. (T. 1I
(2D2L-2 + cDL-2 + 22)²11 110 211
N2
and thus
2L 2L 11-1)
- VZ, <
2T + cD" 222))++2) 110 Φ11
N
Proposition A.5. Given initial parameters Go. we have that GD with ridge X and learning rate 17 H for D (22l and H 2I.v(T.-1) 22122-22 + cD--2 2X) remains inside the ball
N
B(0, D) and has a non-increasing regularized loss.
Proof. Given the parameters at any time step 0, in the ball B(0, D) the loss along the gradient 0, (0,;) (for any small enough so that this next step remains inside the ball B(,, D)) satisfies
(0, (V2.))V2(00)) - (x(0))= : [ - -
S <-> +1 IVCA(0000II 7 S <=1\\V2((0011 (1)
- -2. PVL) -V.00-VV dp
for H
2L
21 + -. + 2X). Thus for any < n the loss is non- in--reeerasn::
N
Li ())) (0. < 0.
15

---

Published as a conference paper at ICLR 2024
Choosing D D= v (a( we know that 8g E [00, D) since 110,11² < Ex(0u)
A= and all subsequent steps O, with a learning rate 7 < H will remain inside the ball, since (00) is non-increasing along the
continuous path made up of the concatenation of the segments [0c, 0z+1] it can never leave the ball since | < (x(0x) < (*(a))) (where / is a real value that can lie on any segment).
A.1.1 DIFFERENT NOTION OF OPTIMAL RANK
In settings where the infimum inf A ArRankAgr- C(A) is not attained, we can define another notion of smallest rank 8" > r" to be the smallest integer where this infimum is attained. One could wonder under which conditions one can avoid minima with rank < F*. similar result can be proven, though we require an additional assumption (and for simplicity, it is proven for GF instead of GD):
Proposition A.6. Given initialization Go such that unregularized (A 0) gradient flow (GF) converges global minimum Ozx such that the loss is 3-PI. neighborhood of foo, then for small enough, regularized GF converges to minimum rank no smaller than f*.
Proof. If we let the ridge À go to zero we have that GF trained with >-weight decay or,2 converges to GF without weight decay: 0,,2 + 0g as À ' O. We can therefore choose a time to large enough
/ /Ex lies in the neighborhood dc where the loss is 3-PL. We can apply Lemma A.2 to obtain that at time Tx the loss will be below Ako and one easily check that kn is bounded X ' (since bt converges to the unregularized GF 0to.>=0 as A > 0). Since the loss will only decrease after that, we know that GF will converge to local minimum with (A) loss.
and small enough An such that for all À S Ao. the ball ((Oto.): R) for R
3
Let us now assume by contradiction that there is a sequence a > A >... with X, + 0as n + oo such that the minimum O that GF with ridge X, converges to is rank-underestimating for all n, i.c RankA < f Since Xn (,.
(0.....)..) = O(A,), we know that the 0c, are bounded, which implies the existence of convergent subsequence that converges parameters Ō which by continuity of + Ag and + (0) satisfies RankA, F* and [(0) which is in contradiction with the assumption that y* is the smallest fitting rank.
Lemma A.2. Let the loss L satisfy the 3-PL inequality ( [[V2(0))1 2 BC(0)) in a ball of radius
R (Cx( around initialization 8, for some X. then there is a time Tx S log b ko for vB
3
kg- (2 (3C(00)) (1161l + R) + 3 (116ol + R3²)). where GF OT,,A on the La-regularized loss
LA satisfies Lx (x( Ako-
Proof. Let TR. be the first time gradient flow 0G,, leaves the ball of radius R, we will describe the dynamics before TR, and then show that TR, is larger than the time Tx we are interested in.
Inside the ball, we have
|V, (IVCC 2A 1101)2
2 (V23C(0) 2A [101)
2 23C(O) 2 23C(0)2X 11011
2 23Cx(0) 2A (2V28C(80) (116ol + R) + 3 + -B))
2 (2Cx(0)) Aho),
for ko (2/23C(00) (1161) + R) + 3 (110ol + 12)2)
Let Ty be the first time that L. (0x,)) = Ako. then for all 1. S min (Tx, TRA)
a, (01.)) (0
16

---

Published as a conference paper at ICLR 2024
Ako S 2 (Ex(0>))
2
S (BLA((
which implies that (0(,,,) - (x(D) 8t and thus that T3 5 log
log fa/2
under the condition
that this is smaller than TR,
Let us now show that TR, 2 Tx, by showing that |10 0oll < R.
TI [V
110TAA- - -bl1 -
dt
0
(00)
VLL(O=>)))I1 dr
where we did a change of variable in time to ((T) which chosen so that Lx (0(()))) -2(00) -
which implies that (T) VLx (0<()) (so that 0, (x((()))))))) = - [[ -1 as needed). We can now further bound
[[X(O=(TaN)1 dr S
(Ota) dt
Cx(00)
1
VB Jo
V VCx(0)) dr
V Ko
VCx(00)
VB
(0,(
vB
R.
While the PL inequality condition might be unexpected, it is actually satisfied almost all global minima:
Proposition A.7. Given global minimum 0 of network with widths we din + dout for all É == 1,..., L 1, then.for all > 0 there is closeby global minimum 0', i.e. 110 - 0'll S such that the loss satisfies the PL inequality in a neighborhood of 0'
Proof. W.l.ogg,, let us assume that din S dout, then it is possible to change the parameters infinitesi- mally to make WL-1 ... W1 full rank while keeping the outputs As Wr ... W, unchanged (by only changing WL-1- W1 orthogonally to ImW7 J. which is possible since w1-1 2 din + dout).
We now choose a neighborhood of 0' such that the smallest singular value of WT-1 Wi is lower bounded by some X > O. For any parameters P in this neighborhood, the loss satisfies the B = inequality:
2x N PL.
|
1
Σ
WT
WH [M O (A Ao)] WT W
N
1
|| O (A* Ao)] WF
WF_1!
N2
2
2
|
@ (A* (A)1
N2
21²
C(0).
N
17

---

Published as a conference paper at ICLR 2024
The PL-inequality is typically satisfied in the NTK regime (13 24), but in the Saddle-to-Saddle regime (4) it seems that GF converges to the vicinity of a minima that does not satisfy the PL inequality (minima that balanced and low-rank typically do not satisfy it), so that the PL-inequality might only be satisfied in a small neighborhood and with a small constant B. This suggests that in settings where the notions of minimal rank r" and F do not agree, the question of which minima GF converges to might be dependent on the regime of training we are in, with the NTK regime leading rank no less than f* and the Saddle-to-Saddle regime leading rank no less than r* least for reasonable values of X.
B Low RANK BIAS
In Theorem 5, there are two statements: (f) 02 & Brei.e. then lc+1 E Brel.e..C and (2) with positive probability such that there exists time T such that OT Breisz,.C. The following theorems give the formal expression of the two statements.
Theorem B.1. For weight W,, = I,..., .....,,, L, L 2 3, let
Ba = 10: 1 < C, [WWW, - Wi+1ll+
1= 1-.,..,
for any x E Rd, where fo twice differentiable function such that (0) = 0. fa(x) = for x > a I fx(x) S for some constant K and ["( O. Denote
where 222 C1/2X. Define Fx(x) = Σx
B (0: Fa a (W, wT) S r + E2, 1= I...., L}.
where a maps a matrix to its singular values and a -
((0++04)
Then for any E E2 >
such that < 1/2 and V S
Angg 6inKL(r+1)C T= /2(C++0))
if ([)(1) B = Bc. O Brae then
stochastic gradient descent iteration with learning rate
C2 (01C-+CC4CC-1 + x-C) 4(01 + CL) XZE1
2λE1
5 min
XAEZ
2(r + +11
64nK (1)(221)) + (
satisfies O(t (1)++)e E B, where n is the maximal widths and heights of weight matrices.
Theorem B.2. For any initialization Gg denote Co == maxisisL || if
Ae 4(2(C) CF)CK + X²Co) 4(C1 + CL) 2X²C
C
n S min
and C. 2 C. then for any time T = To + Ti satisfying To 2 bog((4(n-r(C)/terl)) we have 2nd
log(2Co/E1) and Ti 2 nd
T. min miindinn dont)
P(Or Br.seac)) N
C PRELIMINARIES OF PROOFS
C.1 FACTS IN LINEAR ALGEBRA
Fact C.1. ||A||+ < |A||~|BI| where || ||+ represents Frobenius norm norm.
Fact C.2. For matrices A, B satisfy AB is square, have [Tr(AB)] < ||
Let 01 2 02 .... be the singular values of a matrix A E Rxx" where min(m,n) n). We have following facts.
Fact C.3. o (AB) < (B) and or (AB) a. (AAo) for any
18

---

Published as a conference paper at ICLR 2024
Fact C.4 (Theorem 3.3.13 in (3)). For square matrix A E R let A 2 2 An be the eigenvalues. Then we have
k
(2((A)) - Σ (A)
for k = ..2.....,n and p > O.
C.2 SPECTRAL FUNCTION
For function f RD + R that preserves permutation, we consider the function f X, where X(A) represents all eigenvalues of symmetric matrix A RRxx. We define Diagu be the diagonal matrix with its entries equal to and diag A = (A11,..., Ann). The following lemmas gives the first and second order derivatives of f o X.
Lemma C.1 (Lemma 3.1 from (1)). f differentiable at point MA) if and only if X differentiable at A. Moreover, have
VI () V
where U a orthogonal matrix satisfying (ADing
For a decreasing sequence μ E RR", where
111 Mkl Pk++1 Mlz > Mky+1 1k,,
denote 1, = -K--1 + I...., ,. for 1 = I,...,r. For a twice differentiable function f, we define vector 6(1) as
b²(n) = ((H(()
if Lt = 1, for any # q E
(1)
(ppp()) (M)
and matrix A(w)
if = J if j E Ii otherwise.
Ag(())
bi(m) Fim0)-f(( Pe-
Lemma C.2 (Theorem 3.3 from (1). f twice differentiable at point M(A) if and only if f o is twice differentiable at A. Moreover, we have
V²( A)(A)() = 4 (H, where A WDingM(A)WT and H W HW.
(A(A(A), f ff),
D PROOF OF THEOREM B.1
For stochastic gradient descent, the parameter updates as 7
! ! Ao..osi?
o (1 1)
(2)
Then for each I, 1-th layer's weight Wi updates as
wi(++ wi(t) - (t)+ ... ...W((())) Go Wi(t) ... Wt) + AWI()), (3)
where Gg.sy matrix where the (i.j)-th entry is Ag,i5 - At and other entries are 0. In the proofs below, we will omit the iteration t for convenience (for example W; = Wt)t)). We denote T WT,, W. G, ... WT
19

---

Published as a conference paper at ICLR 2024
D.1 APPROXIMATE BALANCE
First we give a lemma that bounds |[oo,ill
Lemma D.1. For any C, if Wi(t) I for any = I...., L, then GGo, - 2(11 + 2 Ch).
Proof. By Fact C.1 we have ||Aal? - II [|il] C Then
||Go,sllE = (Ao,is - 2 < (Ag + (41))² 2 (l|Aeli- 11: Ci) s 2( -1).
(4)
Proposition D.2. For any C2 C.
2X if (l( - C for any = 1.... L, then stochastic gradient satisfies [ + 1>ll. I C for
descent iteration with learning rate 1j S +2 1 1-1,..., L.
C2
Proof. If [Wil] I C, we have
[[ 1)1 ||Wil|? 2n (AIIWill? Tr(W, + 112 AWal (5)
We estimate each part in the equation separately. We have
Tr( Tr(W, ... W. Ge.s)) (((A)Go.t) = Ao, - Aij)
...
2 - - (A3) 1 Cr.
4
By Lemma D.1
we have
[ + AW1 2 |Goo.s 11
II
[ + 1
< (2(C) -0
+
Then we have
| + 1)11 & (1 - 2 | + 5 + (2C
When
C 42111+6406L and C²> G 2x.
[[ + 1)11 < (1- 2n + S.
Proposition D.3. For any e, C. O. (lWW())) Wi+1(1)
( W(+1(1))1 S e and (1) S
C' for all I, then stochastic gradient descent iteration with learning rate 17 S 4(C++6200--4x6 satisfies (t 1) 1) Wi+1(t 1)T Wi+r(t+ 1)ll2 S e L.....L L.
248
Proof. We first compute the update of W.W,T:
Wi(t + 1) = ((1-((1)))) - - 17TH)
= (1 (11-nn))W.w,"" ((1 (1-nA)n(WT))) + +TWWT) + TT....
Similarly, we have
W++1(t + 1) 7 (t++)
(1 - W (+1 W41 10:-1--111111111 (1-n)n(Wii117h + TII, W1+1 + i
Since
TWW WTT, ... WI .., W.I,T T4
20

---

Published as a conference paper at ICLR 2024
and
W (I:WT)* (WiTTT++1)* Ti-: W1+1,
we have
Wi(L + () 11
= =11- - Wi41 + (f(TT T--1)
Then
Wi(++ 1) - -W+1 + 1 11ll2
W
((1-1A) |1]-))>\ww... W4 Wi+1 112 +1 |TT T1+2
(6)
< 7 ( n
II kg/
[ +
II IW&l
ke1+1
S 211e x2>2E+ (C
When 77
2 4[C++04)c
1ス2
1)Wit+ 11 W1+1(1++ 1) W/+1 (t++ 1)l12 5 22jee + 2n/s E.
With Proposition D.2 and Proposition D.3 we have the following statement that 0,41 approximate balance and the weight of each layer bounded.:
Theorem D.4. For any ε > 0 and C < 1/2x, (t)) E Bc,.. then the stochastic gradient descent
iteration with learning rate 1) S min Tr2(1+0000 (o 1xOD -01+0"00-110 X251
2
satisfies O(
1) E Bc.
D.2 APPROXIMATE RANK-T
In this section, we prove the following theorem that the weight Wi(++1) of each layer is approximately rank-r.
Theorem D.5. For any E1,52 > O such that E2 < 1/2 and v S 32nL(r+1)C /XC++OL) if
Anez
the number of layers L 2 3 and 0(t) E B, then stochastic gradient descent iteration with learning rate S min maximal widths and heights of weight matrices.
Xaca
2(r+1) satisfies 0(t + 1) E Bree where n. is the
Proof. We denote by rı the minima of height and width of W, and the singular value decomposition
Wi
U, SiV, where U, and Vi are orthogonal matrices. Let Sa and F, be as defined Theorem By Taylor's expansion, for any 1 we have
B.
Fo o (Wi(L + W)](+1) + 1))
=Fa o (W, Wi W7 7i - nHं Wi - 2nAW, Wi+ + + +WG) (T) W nWW (TTW))
( (Wi))
=F, (W Wi) (V(F'
(V(Fa oW, Wi), 21),2,,,TTw)
+ (V(F. (W, Wi), AW3) (
(T) + +WW))
+ v2 (W, W; + *m nA),
where y (0, 1) and A WWT -T, W: 2AW, Wi (TT AW) ((T))T(TTT))) + AW)). By Lemma C.I we have
(FF. o OW, Wi) VT, diag (f((s)) ... fa(,)
where $1, ..... Sr, are the entries of diagonal matrix S.. Denote diag (f) ..... ((()))))).....())) =2 Then
W Fa 0a (Wi(t Wi(t + 1)) s Σ fa( 2x) + Tr (v.
21

---

Published as a conference paper at ICLR 2024
where B is the O(1²²) term. Now we can estimate the trace term
Tr (((T((5))(ii)T..)
Ir (WIT... ... .... ... W.GG..) $ WT. ... ..,WWI Vi5o ... W. 1Goo,slll
Wi .....WII W.. ...WW $ V2( +04
Vica v Tr (IN.
... WI, WSW WW .. WIL ...
= +Ch) Tr (Wi-1.. WiTT WIL W. ... ..W.. ... (M)
(8)
where the first inequality is from Fact
Let Eg = E (W
W k+1 WR+ WW )r- for h < I and E, E (WW Wi - WiK+IWEH1D (W,, W for Then we have
k-1
(WaWI)) = -(WTT+ Wa+1)* + Es
for k < and
(W, (2) = (Wa- W. + Ex
for k > I. Thus,
Tr W1 W.W. WI. M ... W.W. (.(5))
< Ir wi) FA Wi) -T Σ Tr(Ex))
(9)
Tx (FT + Σ (r(Ex))
where
@:=W-11 WA ... W., I. W.W F.
for k < and
Ex = (W, (W,,Tw)'--y
... (W
WW
for k > I.
Second term in
Denote & = C.V. J. (S2)
We define an operator S(A) equals to the sum [Tr(E,] < S(Ex). Since | < |[Walle < vo, and ER 2 S (L k + CL-k for > by D.4 Since
of all singular values of A. Then by Fact C.4
|E&lla S ke, C for
for < and [Colle s [VT (5 11 < Tr(f (S²)), we have ||Cx|12 I (
(
k 1)81 for I. Thus, by Fact C.3
Σ Tr(Ea)) & (SL)
Σ
||Ce|12
< (SEL(S3))(((((((9) ()L-)
Σ
k + ++=+(=
(10)
kel
I c
Σ
fl (s?)
2 E1
22

---

Published as a conference paper at ICLR 2024
First term in
By Fact C.4 and Fact
C.3
we have
Ir (1( < (.
- E(aa6))) (1
- Σ f2
Then we have that
Tr (.l.1...WWW.. WW,, (i
WI W2 .. W
F1
L
fa (3))ss))) CT.
f 2s)))²
2 51
Σ ( -² L LVE1 fa(
==1
We denote
= E {o))s
V (5)ss and E
L LEEE D 1 (57). Then
oo (Wi(t wiitt+ + 1)) Σ fa())) -2n8 + 2n + -.
i=1
When En fo (5 > r + 82/2, we have 21j0 + 2n + 3 5 0 by Lemma D.7 and when
E So (of) I r + 82/2, we have 2n16 + 21 + s 82/2 2 by Lemma D.8 Then Fo
(ii( + 1) Wi(( + 1)) 6
D.3 BOUNDS ON ERROR TERMS
Note that we have (x) E [0, 1). f.(()) (x) E (0. and (x < 0. Lemma D.6. With same conditions and notations in Theorem
D.5 the 0 term
B s (2()
+120)
fa (s)).
1=1
Proof. As defined in Theorem D.5
B = (V (W,
(TI + AWI))
(11)
+ (W, Wi Wi+ Wn+ mAI.
We bound the two terms in separately.
First term
II..
By Fact C.3 and the proof of Proposition D.2 we have
(V(F) (W, Wt), 1²( + AW1) T (Zi + AWi)
(T || + AW/l
(12)
<2/ - ch + X²) Σ (2(
1=1
Second term in П.]. By Lemma C.2
(2 (W, Wi + mn)) A
= v2 ('F,(o(W, Wi + Ad)dddaa diagA] + (Alo(W, Wi + mm))), & A
where A 1.A.TTT T and
(3 (.) 6 8
if but 3? 33.
Ay Wt + ))))
if डे # otherwise,
23

---

Published as a conference paper at ICLR 2024
where $1....5, Sr, are the eigenvalues of W Wi yn. Since (W Wi + 0n
diag(fa (3)( Sa (3(
($22) (3) with all entries non-positive, we have
Fo Wi + mn diagA] - 0. Moreover, since (a(x) is concave, all entries of (W, Wi)) are non-positive. Thus,
(A(o(W, Wt + )))),, A Á) 0.
Overall, 3 5 21C C f
Lemma D.7. With same conditions and notations in Theorem D.5 when + E S Y r+ E2, we have -21 + 24j + 3 S O.
Proof. Define g f/ (x) (xx V2(Ci + Ch)abi2) on x 2 0. Then since L 2 3, for
- fa (z) = when x > a and d v2 + Cz)aa
a S (2(0++04)
1/2 2 when x S a.
Thus, g(x) 2 for any x 2 0. Since
Em= g(8)), we have 2 0. Note that there are at most r sy such that fa() N 2(r+1)
grt1 =: My. Otherwise,
2r + 2 = +
Σ fa( 2
E2-
Specifically, when E [a(
+ 2r/82/2. we have
Σ
fa(s))
(a - r 82
(13)
2
<.,
1=1
Since (a) - x) 21 1 (s(()) by concavity, we have
fa 2 1 1- (a) 1 sa(x)
(14)
We also have fa(x) = S (s)dd v Kz according so (x) K. Then we have
fa(x)
(15)
K
Moreover, by L 1>3 3 and the concavity of fI(,,, we have
1
(a)
2 fz(
2
(16)
2
By equation Eq.
[14]
15
and (6). we have
g(x) \()
1 (1-
(c))))
Afa(x)
(17)
2K
8K(r + 1)2
Thus, we have
g(5))
g(sT - -1)>
Σ fa(o?)
8K(r + 1)2 16K (r 1) (18)
2 Aaa
(.f)
<
where the last inequality is by 13
Note that Σn f2(
< 2n: <
2n Then
since V <
Anea
we have E S
32k(r+1y² and
Aeg
since " <
64n MAnK(r++1)CC
V2(C1+C5))
+X-CJ we have 3 S (TK((TT)" Thus,
Bink
A 8K + 1)2 16K + 1)2 16K(r + 1)2
Ae2
A
2276 + 27 + 3 3
+
+
= =0. (19)
24

---

Published as a conference paper at ICLR 2024
Lemma D.8. With same conditions and notations in Theorem D.5 we have 226 + 2n + 3 S 2
when Ei= fa(s?) + E2/2,
Proof. Note that 5 2 0, ES
and B S 16K(++1] 32K(++1)=
A
Since 1 Σ
4k(r+1)²
we have
A
A
Ae2
82
-2n0 2n₁ 3 5 16K(r+
S
(20)
16K(r + 1)3
2
E PROOF OF THEOREM B.2
In Theorem B.2 T = To + Ti The following statements explain the change of 0, during first To
iterations and last 7i iterations respectively.
Theorem E.I. For any initialization Co. denote Co maxisIsL Will if n S C1
min| +C(C)+ +x66) 10++0000--2200 } and C z x
de
C2 then for any time T 2
log(2Ca/E1) we have
Or E Bc
Proof.
Similar to the proof of D.2 if 4(2(C++C 4x²c.) and C 2 Gt we have
G
Wi((++))11 S (1 21) n1..
If (1/ C, then [[){ < (1 - - ( Otherwise, || 1)11 S C. Thus,
there exists t To such that |Wi]] v C for any when To
log(Co/C) log Co
C
log (1 17A).
7
After all weights satisfy || < C. C.W.W WWT, Wirlll s 2C. Similar to the proof of D.3 we have
1) 1) - W/+1 (1 1) 1 Wi+ 1)ll2 |[Wi(L 1)W +
(1 -11))) WT,, Wi+ill + +41/(C1 + 04)04-1
(21)
(1 - (111)))\WWW,
+ 4n + c4)CL-1
dC When 7) - TCC,+04y0 12x=C we have [[ + 11 + 1)T - W411 + 1) WT
1)lle < (1 nA) max{ log ( (6)) log (1 - na), OIotI, E Bc....
Wit W 1111 for any I. Then for Ti > Ing(2(/e1)
Theorem E.2. For any parameter 0, E Be, satisfying E S 2((r-))(--)) then for any T
as
log (o((2K(nn)))/((ee)) we have 24p
( P(1TT E Breves, E BB,,)) mintfims dout)
T
Proof. For the true matrix A", the number of columns is din and the number of rows is dout- Let min(d,,, d. Without loss of generality can assume that din i.e. there are columns. We consider the r columns with most observed entries and denote the the set of these entries by J. J) 2 E. Then the event that all steps s from / to + T - 1, random entries (f.,j.) ) are all sampled from J has probability least (x)² Under this event, we consider the weight of first layer Wi. We have
Then [J] 2 ///l and for each step s, the probability of sampling from J is P((is,)))
Wi + = (1 - (m))(s)) ... WL((s) Ge
Then
T
- (1-))((--\)))))
Σ (1 - n W2(++ ... ... T Geusstender"
25

---

Published as a conference paper at ICLR 2024
Since (t++) E J, the non-zero entry of Sed is located on the columns supporting J, for any 8 ,....; T. Thus, W2(t + ... (L(t 8) Go only has non-zero entries those columns. Then the r+11's singular value of W (t+T) satisfies o, (Wi(t+)) & (1-1A) VC for any i r.
For > 1, we have || W - W Will2 E1. Then lo, (W, Wi) o v $1 for
any i, i.c (w; 0, (W, 1- Then for any I, we have o, (W)())) < (1 -1-))1)2 C+ (1 - 1)81.
When E - (1 and 2 1or(22(n-))0)/te)))
052
2 log
2k(n-r)C /2log - 77), we
ae2
2nd
have o, (Wi())) & K(n-r)
a²2
for any i >> r. Then for any > r,
(w( (W)((()))) v E2
since Sa S K/a. Thus,
Fa Fo s Σ "
(a(o.(()(TT)) sr+
If n dout, the proof is the same by selecting rows with most observed entries.
F Low RANK PROPERTY OF As
In Proposition 4, we show that for any minimizer i By Cs,C, it is approximate rank-r. In fact, any general parameter 0 E Br.e en.0 is approximate rank-r or less: Proposition F.1. For any parameter E Br, we have
Rank Ag
KnL² ch-1s1-
(s( Aa))
2a
Moreover, i S Knl.GL-T, we have
2003
Rank
Σ
fa(s:(A) Ab)) S 252.
Proof. Since (WWWT) = (W W*+1)* Eg and Ex Ei (WI+ W ) 1-W,)"-(WWWWW
W. 44 )k-i, we have
AgA WL WWW. WI (WLWZ)² Σ Ex.
where & -WL- W&+1 F. k+1 WT. Then by Taylor's expansion, have
(a Ao)) = Fa oo
- PPooo((W +))
(22)
-P =Paoo(WW...)) + (V o (WW. + +V . o) ((W, ,E)[E
where y (0, 1) and E = i >> Rank Ag.
Ex. Note that the Rank Ag I rı, let s,AA3 Ao) 0 for k I
First term in (22.. F. ((W fa( where (sI...., 3ry} are the singular values of Wi. Then since fa is non-decreasing, (s fa(s)) for 8; I and (a())
(x for 8, > 1> Thus, F. a (WW(WW7)) < F. (W S 52.
26

---

Published as a conference paper at ICLR 2024
Second term in 22 Then by Fact C,3 and C.4
By Lemma C.1
V(Fa =
= (f (s.") ,.....
)...)....... (s
we have
(V () S o (VVV) 11812
s <|18||2
fa (s,
Krl
Kn
1 a
|111 S
|11|z
a
Since || S kC* EI, we have ||E I kCL-1eg 51. Then
L
||]|2 - ||E|F - Σ [Elll Ch-
k S
Ch-
k=
Thus, (V(F. (VV S Knl' CL-181-
Third term in 22
By Lemma
(FFaoo) + E
P (a((W + 78)) IdiagE, diagE] + (AVo
where E = ÜLEū Since the entries of Fr (o((WWWZ) + 7E)) and xlo((W2WZ(+++88 are all non-positive (follows the proof in D.6,, we have
V2 o ow) + $
Therefore, we can add the three terms up and have
Rank Ag
KnI.
[a(s:(A] (o)) + +
ch-Ls1- E1.
2a
27


Reviews:
[{"review_id": "vGeIJn81sO", "reviewer": "Reviewer_M8aA", "summary": "This work shows that when applied to matrix completion with deep linear networks, SGD can transition from local minima with higher ranks to solutions with lower ranks, while transitions in the opposite direction are zero. Crucially, this results depends on the gradient distribution of SGD which leads to drastically different outcomes than what common SDE-based models for SGD exhibit. The authors further provide numerical experiments that exhibit the predicted transitions in practice.", "questions": "As you point out, the common Langevin-based models predict a transition among each pair of points and thus fall short to capture the phenomenon shown by your result. Do you know if they would nevertheless exhibit a systematic low-rank bias (i.e. the transition toward lower-rank solutions being much more likely than toward higher-rank solutions)?", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This work provides an interesting theoretical insight on the distinction between stochastic and deterministic gradient descent. I find it especially exciting that it provides a concrete example how the gradient distribution of SGD enables phenomena that are not apparent from SDE-based models.", "weaknesses": "I can not think of a weakness. It's a very nice paper in my opinion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HC6Wv0wKOo", "reviewer": "Reviewer_NbUi", "summary": "This paper analyzes the matrix completion task with deep linear neural networks. It shows that the critical points that are not local minima can be avoided with a small ridge. And it shows that GD cannot avoid overestimating minima but  SGD can jump from any minimum to a lower rank minimum.", "questions": "1. In the remark before section 3.2, it's said that it's possible that GD can recover the minimal rank solution easily. Can you say something about this case? \n\n2. I have a concern about the constant in Theorem 3.2. It is said $\\lambda$ and $C$ are large enough, but an example of acceptable rates in terms of $\\lambda$ is $C \\sim \\lambda^{-1}$. Is it contradictory?\n\n3. In the proof of Theorem 3.2, r columns with the most observed entries are taken.  What if all the columns have the same number of observed entries? Will the $d_{out}-r$ other columns of $W_L$ decay exponentially? I don't see the relation between rank $r$ and the number of observed entries here.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "I think analyzing the training dynamics of the deep linear neural networks on the matrix completion task is a very interesting problem. This paper provides insights on the advantages of using SGD to get a low-rank minima and provides experimental results.", "weaknesses": "1. I feel the statement of theorems is not very clear. It uses a lot of ''small enough'', \"large enough\".  I think the statement should be more rigorous. \n\n2. This paper claims that it shows GD can avoid rank-underestimating minima by taking a small enough ridge $\\lambda$. But Proposition 3.2 is for Gradient Flow with a very strong assumption. I believe there is a gap. \n\n3. For the function $f_\\alpha$, it takes a very specific form. The authors claim that changing $f_\\alpha$ with similar properties should not affect the results. I don't see the reason why the condition cannot be extended to more general functions. I believe it could improve the results.\n\n4. A small suggestion is that the proof in the appendix is not easy to read. I think it can be more organized and add more explanation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eVTpeBHrwM", "reviewer": "Reviewer_fYKu", "summary": "This paper studies the optimization landscape of regularized stochastic gradient descent applied to matrix completion with linear networks problems (which is equivalent to matrix completion with a $2/L$ Shatten norm regularizer). Several properties of the optimization landscape are proved, including the fact that the only critical points of the optimization problem over the factor matrices (minimizing $\\mathcal{L}_{\\lambda}(\\theta)$ must be local minima of the optimization problem over the full matrix $A$, unless they are strict saddle points in the original problem.\n\n In addition, it is shown that if gradient flow converges to a global minimum, then a version of gradient flow with a sufficiently small regularization parameter will converge to a minimum with a larger rank than the ground truth. Arguably the most significant final result is that stochastic gradient descent jumps from high rank local minima to lower rank local minima, with the jumps being one directional: one cannot return to a higher rank region after entering a lower rank region. Here, the lower rank regions should be understood as defined on page 5, in an approximate sense. Throughout the proofs, the fact that the local minima of the optimization problem over $\\mathcal{L}_\\lambda(\\theta)$ must be balanced (cf. Proposition A.1). An approximate version of this condition is also present in the definition of the low absorbing low rank spaces in the main results of the paper.", "questions": "1. In the third line of page 13, ou mention that the quantity in the limit is strictly positive but possibly infinite. Apologies if I  lack some background knowledge but could you explain your reasoning there? It is not at all obvious to me. \n\n2 At the beginning of the proof of proposition A5 in the first equation, should the infimum run over  $Rank A>r$ instead of $Rank A<r $ as written?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The main paper is well-written and the results appear generally sound. The results are of great importance to the field and interest to the community. **This is highly non trivial and important work**.", "weaknesses": "Although the main paper is well written, the **proofs are not reader-friendly** at all. \n\nThe writing of the proofs is very terse and laconic, omitting many details. Although this is reminiscent of some great pure mathematics papers that were ahead of their time and I enjoyed the challenge some of the time, I strongly believe this style should only be considered acceptable if there is absolutely zero tolerance for any errors or inaccuracies whatsoever. I don't think the proofs actually stand up to this amount of scrutiny: there are **at least a few typos, minor errors and imprecisions** in the subset of the proofs I was able to look at, and since a lot of information is left out for the reader to figure out, the additional presence of even a small number of actual errors dramatically expands the \"search space\" from the point of view of the reader.  I would really like to see a substantial revision of the paper with more detailed and careful proofs (and maintaining my score is conditional on that). \n\nFor instance, in page 12, point \"(0)\", the definition of the $U_i, V_i$ is not really consistent: the index under the $U,V,S$ is used both to mean the iteration step in the sequence and the position in the product $W^L...W^1$. \n\nIn addition, in page 13, consider the following statement the authors make \" as $\\lambda\\rightarrow 0$, the critical points of the loss move continuously. Consider a continuous path of critical points, as $\\lambda\\rightarrow 0$, it converges to...\"\nAlthough the argument makes sense intuitively, filling in the gaps with rigorous proofs is definitely beyond the scope of what can be expected of the reader to do. At least some citations are a minimum. I doubt that simple continuity is enough to guarantee convergence (even if a subsequence converges, the path could oscillate widely), probably the only way to rigorously prove the statement is to use a quantitative version of the statement relying on calculus of variations. \n\nThis is not the only example. In point (1) in page 14, the authors say \"the singular value ..... must converge to a non zero eigenvalue\".  It is not clear **why this is the case**, or why the the *singular value* turns into an *eigenvalue* after convergence. Far more details are required. \n\nIn the middle of page 14, it is hard to imagine that the equation $U_{\\ell,i}(\\lambda)U_{\\ell-1,i}(\\lambda)$ can be correct without **at least a transpose missing**. Of course, the lack of a rigorous and consistent definition of $U_{\\ell,i}$ does not help here. \n\nAt the bottom of page 14: the line starting with \"other directions\" ends with \" $L-1,)$\" and a few lines below we have the equation $U_\\ell^\\top dU_{\\ell}+ = -dU_\\ell^\\top U_\\ell$. What does \"+=\" mean here? The same issue is present in many other parts of the paper, including in the third line of text on page 15.\n\nTowards the end of Appendix A in page 17, the term \"saddle to saddle\" is mentioned with absolutely no explanation or citation. \n\nIn the middle of page 13, the authors use the fact that \"a matrix cannot be approached with matrices of strictly lower rank\", which is true but should probably warrant a citation since the equivalent statement is not true for tensors. \n\n\nThe proof of Proposition A4 is very hard to make sense of without further information: the first sentence is \"\"let A(\\lambda) be path of global minima restricted to the set of matrices of rank $r^*$ or less.\" how do you construct the path? Even for $L=2$, there can be a continuous set of global minima of local intrinsic dimension higher than 2, how do you use the axiom of choice to construct a \"path\"? \nSentences such as \"going along directions that increase the rank of $A(\\lambda)$, the regularization term increases at a rate of $d^{2/L}$ for $d$ the distance\" definitely need more mathematical details. \n\nSimilarly, the statement about $\\phi$ being differentiable in the directions which do not change the rank should be made more precise (although I agree with it, probably at least a citation to [1] is a minimum) \n\nFor proposition A.5, the proof starts with the following sentence \"We know that L2 regularized GF $\\theta_\\lambda(t)$ converges to unregularized GF $\\theta(t)$ as $\\lambda\\rightarrow 0$\". There are two parameters here, $\\lambda$ and $t$, is the convergence uniform over all $t$?\n\n\n\n\n\n\n\n\n\n\n\n\n\n========more minor points:=====\n\nMany apologies if I am being picky but as a relative outsider to optimization literature, even the statement that the point 0 is a critical point was not immediately  obvious to me (perhaps either a calculation of the gradient or a mention of the fact that $L>1$ would help). \n\nIn the bottom of page 13, the equation before equation (1) is presumably the end of a sentence, thus the next line should be rewritten. Below, that \"no such thing happen\" should be \"no such thing happens\"\n\nSome citation for Fact C.4 (Ky Fan?) would be nice. \n\n\nIn page 19, just before the beginning of Section D.1. Do the authors mean $G_{\\theta,ij}$ instead of $G_{\\theta,j}$?\n\n\nJust above equation (6), $\\|W_\\ell|^2$ should be $\\|W_\\ell\\|^2$ and the sentence is missing a period. \n\n\n\n\n\n\n\n\n[1] Characterization of the subdifferential of some matrix norms, G.A. Watson. 1992, linear algebra and its applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JuuESbN1Ic", "reviewer": "Reviewer_Gs3p", "summary": "This papers considers the problem of minimizing parameters of a deep linear network with the matrix completion loss. Specifically the authors considered an $\\ell_2$ regularized version of this problem and show that with certain conditions on the learning rate and the regularization parameter, SGD can jump from a high-rank local min to a low-rank local min, while it cannot jump from a low-rank local min to a high-rank local min.", "questions": "Please see previous section. Also, I wonder if the proof in this paper for Theorem 3.2 also works for GD? In other words, is $B_r$ also absorbing for GD?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I think the theoretical results of this paper are interesting because it gives a nice characterization of the implicit bias of GD/SGD for deep linear networks. The proof that SGD can \"jump\" from high-rank to low-rank local minima is new to me and I think its a good step towards understand the training dynamics for deep neural networks.", "weaknesses": "I think the main weakness of this paper is that both the theoretical results and the experiments require specific conditions on the $\\ell_2$ regularization parameter $\\lambda$ and the learning rate $\\eta$, which seem to be a bit artificial. For example, the requirement that $\\lambda$ is large in Theorem 3.2 seem to artificially cause $\\|\\theta\\|$ to decay more quickly, thus biasing it towards low numerical rank. In the numerical experiments, a similar annealing technique is used run SGD with a large $\\lambda$ and $\\eta$, before switching to smaller parameters. \n\nMy main point is that the implicit bias observed in this paper could be a result of a deliberate choice of parameters, instead of a natural property of SGD and GD. I hope the authors can clarify this point.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)