Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descent
Authors: Zhiyuan Li
Tianhao Wang
Jason D. Lee
Sanjeev Arora
Venue: NeurIPS 2022 Accept (2022)
Keywords: implicit bias
gradient descent
mirror descent

Abstract: As part of the effort to understand implicit bias of gradient descent in overparametrized models, several results have shown how the training trajectory on the overparametrized model can be understood as mirror descent on a different objective. The main result here is a complete characterization of this phenomenon under a notion termed commuting parametrization, which encompasses all the previous results in this setting. It is shown that gradient flow with any commuting parametrization is equivalent to continuous mirror descent with a related mirror map. Conversely,  continuous mirror descent with any mirror map can be viewed as gradient flow with a related commuting parametrization. The latter result relies upon Nash's embedding theorem. 

Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descent
Zhiyuan Li' Princeton University Princeton, NJ 08540 zhi Mhiyuinni@cs princeton.edi
Tianhao Wang" Yale University New Haven, CT 06511 tianhao. wang@yale edu
Jason D. Lee Princeton University Princeton, NJ 08540
Sanjeev Arora Princeton University Princeton, NJ 08540 arora0cs www.ss..ininooooo.o edu
assolee@prinneeton edu
Abstract
As part of the effort to understand implicit bias of gradient descent in over- parametrized models, several results have shown how the training trajectory on the overparametrized model can be understood mirror descent different objective. The main result here is characterization this phenomenon under notion termed commuting parametrization, which encompasses all the previ- ous results in this setting. parametrization is equivalent continuous mirror descent with related Legendre function. Conversely, continuous mirror descent with any Legendre function can be viewed gradient flow with related commuting parametrization, The latter result relies upon Nash's 's embedding theorem.
shown that gradient flow with any commuting
Introduction
Implicit bias refers to the phenomenon in machine learning whereby the solution obtained from loss minimization has special properties that not implied by value of the loss function and instead arose from the optimization's 's trajectory through the parameter space. Quantifying implicit bias nec- essarily has go beyond the traditional black-box convergence analyses of optimization algorithms. Implicit bias can explain how choice of optimization algorithm can affect generalization [61, 42, 41].
Many existing results about implicit bias view training (in the limit of infinitesimal step size) as a differential equation or process {a(t)]ı2o C. R.. To show the implicit bias of x(t), the idea to show for another (more intuitive or better understood) process (w(t))e2 C Rd that (t) is simulating w()), in the sense that there exists a mapping G RP + Rd such that w(()) (()( Then the implicit bias of z()) can be characterized by translating the special properties of w(() back to ztt) through G. related term, implicit regularization, refers to handful of such results where particular update rules shown to lead to regularized solutions; specifically, ztt) is simulating w(t) where (t) is solution to a regularized version the original loss.
The current paper develops a general framework involving optimization in the continuous-time regime of loss Rdd >> IR that has been re-paramerrized before optimization as G(=) for some G RD Rdd Then the original loss L(w) in the w-space induces the implied loss
Equual contribution
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

(L) = L(G())) in the z-space, and the gradient flow in the z-space is given by²
dr(()= = -
(1)
Using (() = G(x(l)) and the fact that V(L (x) = aG()) T V VL(G(=)) where aG(x) E RdxD denotes the Jacobian of G at v; the corresponding dynamies of (1) in the w-Space is
dw(() - (G(x(()dzz() - -0((((()) (L(((())d
Our framework is developed to fully understand phenomena in recent papers [26, 58, 64. 4. 61, 5, 8], which give examples suggesting that gradient flow in the r-space could end up simulating a more classical algorithm, mirror descent (specifically, the continuous analog, mirror flow) in the w-- space. Recall that mirror flow is continuous-time dV ((t(t)) - [[(w(t))dt where R. IE > R U {oo} is strictly convex function [49, 10], which is called mirror map or Legendre function in literature. Equivalently it is Riemannian gradient flow with metric tensor 2 ft, old notion in geometry:
limit of the classical mirror descent, written as
dw(t) = - V2
(3)
If there exists a Legendre function R such that (G(x()))C(z(() - (² for all t, then (2) becomes a simple mirror flow in the w-space. Many existing results about implicit bias indeed concern reparametrizatioos G that satisfy (G(z)aG(z) T = V V²R(w)) for strictly convex function R. and the implicit bias/regularrization demonstrated by showing that the convergence point satisfies the KKT conditions needed for minimizing R among all minimizers of the loss d. ((=(t))(((z(()
L.. A concrete example is that wx(()) G((z(()) = (xx(1))2 for all E [d]. so here In this case, the Legendre function R must satisfy (V (V²((((t))-
4diag( (t))))))))))),,))))))))))))))))) ..... (z²) 4diag(w) (w1 (f),...,,,,))) wa(()) which suggests R is the classical negative entropy function, Le., R(w) Ei= w; wi -1).
However, in general, is hard to decide whether gradiens flow for a given parometrization G can be written as mirror flow for some Legendre function R. especially when D > d and G is not an injective map. In such cases, there could be multiple z's mapping to the same G(x) yet having different OG (z)))((zTTT If more than one of such x can be reached by gradient flow, then the desired Legendre function cannot exist. If only one of such x can be reached by gradient flow, we must decide which it is in order to decide the value of V R using oGaG Conversely, [5] raises the following question: for what Legendre function R can the corresponding mirror flow be the result of gradient flow after some reparametrization G? Answering the questions in both directions requires a deeper understanding of the impact of parametrizations.
The following are the main contributions of the current paper:
In Section 4, building on classic study of commuting vector fields we identify a notion of when a parametrization w = G(x) is commuting (Definition 4.1) and use it give a sufficient condition (Theorem 4.8) and a slightly weaker necessary condition (Theorem 4.9) of when the gradient flow in the x-space governed by (L respect to some Legendre function R I R. This condition encompasses all the previous results [26, 58, 64, 4, 61, 5, 8]. Moreover, the Legendre function is independent of the loss L and depends only on the initialization Tinit and the parametrization G.
(L o G) is simulating a mirror flow in the w-space with
We recover and generalize existing implicit bias results for underdetermined linear regression as implications of the above characterization (Corollary 4. 4.1.). We also give new convergence analysis in such settings (Theorem 4. 15), filling the gap in previous works [26, 61, 8] where parameter convergence is only assumed but not proved.
the reverse direction, we use the famous Nash's embedding theorem to show that every mirror flow in the w-space with respect to some Legendre function R simulates a gradient flow with commuting parametrization under some embedding x F(w) where F Rdd - R and the
"wwo examples from recent years, where G does not change expressiveness of the model, involve (a) overparametrized linear regression where the parameter vector w is reparametrized (for example as 02 y0 [61]) and (b) deep linear nets [6] where a matrix W is factorized W Wi ..... Wr. where each We is the weight matrix for the ê-th layer.
"To avoid such an issue, [5] has to assume all the preimages of G at w have the same 0G(DG) and recent
paper [23] assumes that G is injective.
2

---

parametrization G' is the inverse of F (Theorem 5.1). This provides an affirmative and fully general answer to the question of when such reparametrization functions exist, giving full answer to questions raised in more restricted setting in [5].
2 Related work
Implicit bias. With high overparametrization as used in modern machine learning, there usually exist multiple optima, and it is crucial understand which particular solutions are found by the optimization algorithm. Implicit bias of gradient descent for classification tasks with separable data was studied in [55, 24, 46, 35, 45, 34] and for non-separable data [32, 331, where the implicit bias appears in the form of margin maximization. The implicit bias for regression problems has also been analyzed by leveraging tools like mirror descent [61. 24, 64, 58, 4, 5]. later generalized in [8].
The sharp contrast between the so-called kernel and rich regimes [61] reflects the importance of the initialization scale, where large initialization often leads the kernel regime with features barely changing during training [30, 16, 20, 19, 2, 1, 65, 7, 62, 31], while with a small initialization, the solution exhibits richer behavior with the resulting model having lower complexity [25, 26, 39, 52, 6, 15, 41, 43, 44, 53, 56, 22]. Recently [63] gave a complete characterization on the relationship between initialization scale, parametrization and learning rate in order to avoid kernel regime.
There are also papers on the implicit bias of other types of optimization algorithms, e.g., stochastic gradient descent [40, 11, 29, 42, 18, 66] and adaptive and momentum methods [51, 60, 59, 36], to name a few.
Understanding mirror descent. In the continuous-time regime, the mirror flow is equivalent to Riemannian gradient flow with the metric tensor induced by the Legendre function. [27] showed that partial discretization of the latter gives rise to the classical mirror descent. Assuming the existence of some reparametrization function, [5] showed that a particular mirror flow can be reparametrized as a gradient flow. Our paper shows that such reparametrization always exists by using Nash's s embedding theorem. [23] generalized the equivalence result of [5] to discrete updates.
3 Preliminaries and notations
Notations. We denote N as the set of natural numbers. For any n N, we denote {1, 2, ......... .n. by [n]. For any vector E RD,, we denote its i-th coordinate by uj For any vector For any
R
and a E I., we define u @ v (ujV),... uDUD)" T and u uaa = ((u))".... E NU {oo}, we say a function f is C* if it is l times continuously differentiable, and use C* (M) to denote the set of all C* functions from M to R. We use a to denote the composition of functions, egg. f g(x) = [((x))) For any convex function R: TRD - IR U {oo}, we denote its domain by dom = (w E R (P((()) < oo). For any set S, we denote its interior by (S) and its closure by S.
We assume that the model has parameter vector w E IRd and C² loss function L : Rd + R. Training involves a reparametrized vector I E R.,, which is reparametriizaion of w such that w G(z) for some differentiable parametrization function G, and the objective is LGG())). From now on, we follow the convention that d is the dimension of the original parameter w and D is the dimension of the reparametrized x. We also refer to Rdd as the w-space and RRP as the r-space.
In particular, we interested in understanding the dynamics of gradient flow under the objective L o G on some submanifold M C R.P. Most of our results also generalize to the following notion of time- deme-dependent loss.
Definition 3.1 in time t and continuously differentiable in w E Rd, that is, there exist & E N, = t < t/2 < .... < t--1 = oo and C loss functions [1)), (2) (A) such that for each E [k]] and all E [ta, t++1),
loss). A time dependent loss L, (w) is a function piecewise constant
L, (w) = (w)),,
Vw E Rd
We denote the set of such time-dependent loss functions by L.
3

---

3.1 Manifold and vector field
Vector fields are a natural way formalize the continuous-tinn -ime gradient descent (a good reference is [38D, Let M be any smooth submanifold of RD. A vector field X on M is a continuous map from M to IR" such that for any
I E M, (z) is in the tangent space of M at x, which is denoted by dr
M.M. Formally, T.M- = ar t=
V smooth curves Y IR + M, (0) = =).
Definition 3.2 (Complete vector field; p.215, [38]). Let M be smooth submanifold of R and X be a vector field on M. We say X a complete vector field on M if for any initialization Tinit M, the differential equation (x((t) X (z(t))dt has a solution on (-o0, 20) with z0)) = Ziii..
Equipping the smooth submanifold M R with a metric tensor g. we then have a Riemannian manifold (M,g) where for each x E M, 9r : T2M x Ta M + R is a positive definite bilinear form. In particular, the standard Euclidean metric g corresponds to I. (u, v) = u² for each æ E M and u T,M,, under which the length of any are on M is given by its length as a curve in R
For any differentiable function f M + R, we denote by ygf its gradient vector field with respect to metric tensor g. More specifically, V (z) is defined as the unique vector in RD such that yrf(() E T.M and df(x(e)) l-- = Ix ( (zx),, dz(0) le-o) for any smooth curve R - - with (0)) J. Throughout the paper, we assume by default that the metric on the submanifold M C RD is inherited from (RD, 9.), and we will use V as a shorthand for Vgʃ- If M is an open set of IR f is then simply the ordinary gradient of
For any x E M and C'l function f M + R, we denote by of (x) the point on M reached after time t by following the vector field f starting at x, i.e. the solution at time (when exists) of
dØ' = - V (o)))
(x)==
We say (x) is well-defined at time when the above differential equation has solution time t. Moreover, for any differentiable function M + Rd, we define its Jacobian by
(x(x) (VXi (x), VX2 V
Definition 3.3 (Lie bracket). Let M be smooth submanifold of RR Given two Ca vector fields Y,Y on M, we define the Lie bracket of X and Y as [X, (z) == OY
3.2 Parametrizations
We use the term parametrization to refer to differentiable maps from a smooth submanifold of Rº (x-space) to Rd (w-space). We reserve G' to denote parametrizations, and omit the dependence on G for notations of objects related to G when it is clear from the context.
The following notion of regular parametrization plays an important role in our analysis. and it is necessary for our main equivalence result between mirror flow and gradient flow with commuting parametrization. This is because if the null space of (G is non-trivial, i.e., it contains some vector u / O, then the gradient flow with parametrization obviously cannot simulate any mirror flow with nonzero velocity in the direction of u.
Definition 3.4 (Regular parametrization). Let M be a smooth submanifold of IR., regular parametrization G M + Rd is a C¹ parametrization such that @G(x) is of rank d for all M.
Note that regular parametrization can become irregular when its domain is changed. For example, G(xx) x² is regular on R,, but it is not regular on IR as (G = 0.
Given a C2 parametrization G M + Rd, for any I E M and μ E Rdd, we define
(b;;; (4) = d' o ei o... &eg (x)
(4)
when it is well-defined, i.e., the corresponding integral equation has solution. For any M, we define the domain of vxx;;) as
U(x) { E Iedd (x;;;; 11) is well-defined].
(5)
When every VG, a complete vector field on M as in Definition 3.2, we have ((()) = Rd. However, such completeness assumption is relatively strong, and most polynomials would violate it. For

---

example, consider G(z) = 28 for z E R", then the solution to (t) = (t) explodes in finite time for each E (d. To relax this, we consider parametrizations such that the domain of the flows induced by its gradient vector fields is pairwise symmetric. More specifically, we define
(x)()) { (s, () E R2 oG. & (r) is well-defined]
for any x E M and i,j E [d], and we make the following assumption.
Assumption 3.5. Let M be smooth submanifold of RD and G M - Rd be a parametrization We assume that for any и E M and E [d], d (x) is well-defined for E (T) such that either limæ-r, |105 (z:)ll2 = oc or T+ oc and similarly for T_. Also, we assume that for any E M and i [d]. it holds that (x)( [(t, s) E IRa (s, (z)())) i.. dG. oec, (2) is well-defined if and only d'c, o oe. (x) is.
Indeed, under Assumption 3.5, show that for any x E M, is hyperrectangle in R i., (6)
u(x) (x) (2)
[I( where each [1) R is an open interval.
See Lemma C.I and its proof in Appendix C. Next, for any initialization Vinit. E M, the set of points that are reachable via gradient flow under some time-dependent loss (see Definition 3.1) with parametrization G is subset of M that depends on G and Finit.
Definition 3.6 (Reachable set). Let M be smooth submanifold of TR For any C² parametrization G: M + gd and any initialization Finit E M, the reachable set S2 (Tiit: G) is defined as
Q2 (Ziniti G) ={ {0% Lyog
LLoc Wk ENE [M, E C² (²
It is clear that the above definition induces transitive "reachable" relationship between points on M, and it is also reflexive since for all L E C² (Rd) and > 0, oLoG o o(-L)oG is the identity map on the domain of o² LeGG. In this sense, the reachable sets are orbits of the family of gradient vector fields (V G) E E C¹l (RR)), iee. the reachable sets divide the domain M into equivalent classes. The above reachable set in the r-space further induces the corresponding reachable set in the w-space given by S2 (Finiti G) (Fmiii G)).
In most natural examples, the parametrization G is smooth (though this is not necessary for our results), and by Sussman's 's Orbit Theorem [57], each reachable set Ω2 (Zimii G) is an immersed submanifold of M. Moreover, it follows that Ω2 (Xiiit: G) can be generated by (VG.) ie., Q₂ (Finit; G) = { &G
& (Fimit) Vk: E N, Vi E [k), [d), 14 2 0}.
3.3 Mirror descent and mirror flow
Next, we introduce some basic notions for mirror descent [49, 10). We refer the readers to Appendix B for more preliminaries on convex analysis.
Definition 3.7 (Legendre function and mirror map). Let Rt I2 + [Rf } be a differentiable convex function. We say R is a Legendre function when it satisfies that (1) R is strietly convex int (dom R). and (2) for any sequence (w;) going to the boundary of dom R, lim;-+20 [V = sc. In particular, we call R mirror map if R further satisfies that the gradient map V R int R) - Rd is surjective (see p.298 in [13]).
Given a Legendre function R Rdd - R U (oo), for any initialization we = Wint E int R), mirror descent with step size n updates as follows:
vR(u++1) V((ww) 17 L(x)).
Usually R required to be surjective so that after a discrete descent step in the dual space, it can be projected back to the primal space via ( (VR)- -l Nonetheless, as long as V R(wx)) 7/ (VL())) is in the range of R.R, the above discrete update is well-defined. In the limit of n + 0, (7) becomes the continuous mirror flow:
dVR(w(t)) - VVLLLL(WEE (L(((t))dt
(8)
5

---

08 (x)
9c: о 06 (x) = 08, o (1)
-VG
i
4
-VG
VG,
96, (x)
Figure Illustration of commuting parametrizations. Suppose G M Rd is commuting parametrization satisfying Assumption 3.5, then starting from any E M. first moving along - VV for time then moving along VG, for time t, yields the same result as first moving along -VG for time l then moving along - VG; for time i; does, i.e., 66. o0g (x) od, & (r).
Given a differentiable function R. the corresponding Bregman divergence DR is defined as
Da(w, w') R(w) R(w') (VR(w'), w-)>
We recall a well-known implicit bias result for mirror flow [24] (which holds for mirror descent as well), which shows that for specific type of loss, if mirror flow converges to some optimal solution, then the convergence point minimizes some convex regularizer among all optimal solutions.
Theorem 3.8. Given any data Z E R and corresponding label Y E RR", suppose the loss L(w)
in the form of L(w) L(Zw) for some differentiable L R" + R. Assume that initialized at (()) = Winit, the mirror flow (8) converges and the convergence point Woo = lim w(t) satisfies Z wx = Y, then D DR(Woo wo) min Du Dn(w, wo).
See Appendix C for a proof. The above theorem is the building block for proving the implicit bias induced by any commuting parametrization in overparametrizee linear models (see Theorem 16).
4 Every gradient flow with commuting parametrization is a mirror flow
4.1 Commuting parametrization
We now formalize the notion of commuting parametrization. We remark that M is smooth submanifold of R.P, and is the domain of the parametrization G.
Definition 4.1 (Commuting parametrization). Let M be a smooth submanifold of R.D A Ca parametrization G' :M- + Rd is commuting in subset S C M if and only if for any i.j [d]. the Lie bracket [VG, VG,](x) 0 for all E S. Moreover, we say Gs commuting parametrization if it is commuting in the entire M.
In particular, when M is open subset of Rd, (VG,) are ordinary gradients in RD_ and the Lie bracket between any pair of VG; and VG, is given by [V,, VG,[(z) = (T)
(=G,())VG,(1t). This provides an easy way to check whether G is commuting or not.
The above definition of commuting parametrizations builds upon the differential properties of the gradient vector fields (VG,) ) == where each Lie bracket VV,, VG,] quantifies the change of VG; along the flow generated by VG;- Indeed, the above characterization of "commuting" is further equivalent to another characterization in the integral form (Theorem 4.2), as illustrated in Figure
Theorem 4.2. Let M be smooth submanifold of R and G M Rd be C² parametrization satisfying Assumption 3.5. For any i. E [d]. [VG, (G,](z) = 0 for all x E M if and only if for any x. M, it holds that é'. D',, (x) oc, &G. (a) for all (s, E Ixxx) x [x),,, where Ix) and I, (a)) are the time domains of &G, (x) and d'a, (x) as defined in (6).
The commuting condition clearly holds when each G; only depends on a different subset coordi- nates of Jr, because we then have V-
= 0 for any distinct i, E [d] as V²G, and VG; live in different subspaces of TR We call such G separable parametrizations and this case covers all the previous examples [26, 58, 4, 61, 5]. Another interesting example is the quadratic parametriza- tion: We parametrize E Rd by G' RD + Rd where for each E [d]. there is a symmetric matrix
We further discuss the existence of non-separable commuting parametrizations in Appendix A.2.

---

A; E RDD such that G,(z)= 2 Then each ,G,, = - - A,A))a and thus G' is commuting parametrization if and only if matrices {A, commute.
For conereteness, we analyze two examples below. The first one is both a separable parametrization and commuting quadratic parametrization, while the second one is quadratic but non-commuting.
Example 4.3 (1 o02 parametrization, [61]). Parametrize
Rdd by D = 2d, and the parametrization G is given by G(x) = u u u for x = (*) E RD. Since each (x) involves only u, and V/ G is a separable parametrization and hence commuting parametrization. Meanwhile, each (z)z)) is quadratic form in æ, and it can be directly verified that the matrices underlying these quadratic forms commute with each other.
n 002 Here
Example 4.4 (Matrix factorization). As counter-example consider two parametrizations for matrix factorization: G(U) = UUT and G(U,V) V) = UV where U, V E Rdxr and d 2 2,r 2 1. These are both non-commuting quadratic parametrizations. Here we only demonstrate for the parametrization G(U) = UU and V(U.V) = UV follows a similar argument. For each i,j E [d], we define Faj E Rdd as the one-hot matrix with the (i, entry being and the rest being O, and denote E. = 1 EEi + Ejj). For r = 1, we have Gi)(U) = - for any i.j [d]. so G is quadratic parametrization. Note that En = E # ZE, = Eig for all distinct i E [d], which implies that [VG,,, VG,] # O, so G is non-commuting. More generally, we can reshape U as vector T- = U U E BRdd where each U.j is the j-th column of U, and the resulting quadratic form for the -entryy of G(U) corresponds to a olock-diagonal matrix:
(U) (T) diag EE, E.y)U
Therefore, Gig does not commute with Ga due the same reason as in the rank-1 1 case.
Remark 4.5. This non-commuting issue for general matrix factorization does not conflict with the theoretical analysis in [26] where the measurements are commuting, or equivalently, only involve diagonal elements, as {G are indeed commuting parametrizations. [26] is the first to identify the above non-commuting issue and conjectured that the implicit bias result for diagonal measurements can be extended to the general case.
4.2 Main equivalence result
Next, we proceed to present our analysis for gradient flow with commuting parametrization. The following two lemmas highlight the special properties commuting parametrizations. Lemma 4.6 shows that the point reached by gradient flow with any commuting parametrization is determined by the integral of the negative gradient of the loss along the trajectory.
Lemma 4.6. Let M be smooth submanifold of RDD and G M - IRd be a commuting parametriza- tion. For any initialization Tinit E M, consider the gradient flow for any time-dependent loss L, E L as in Definition 3.1: dr(t) -V ()ax(t)dt, x(0) Finit- Further define ((l) LG,
Suppose μ(l) E U(rma) for all L E (O, where T IR U (oc), then it holds that at())= (Fimiti u(t)) for all E (0,T).
Based on Lemma 4.6, the next key lemma reveals the essential approach to find the Legendre function. Lemma 4.7. Let M be a smooth submanifold of R and G : M - IRd be a commuting and regular parametrization satisfying Assumption 3.5. Then for any Finit E M, there exists a Legendre function Q QR .2 + R I }oc} such that VQ(n) = G(g((imit n)) for all E U ((initt). Moreover, let R be the convex conjugate of Q, then R is also a Legendre function and int R) = Qi (Jiniti G) and 2² V2I (xiniti))))) p))) = (aG(v(Zinit) 1 ()) n)) for all μ E ((Iinit).
Next, we present our main result on characterization of gradient flow with commuting parametrization. Theorem 4.8. Let M be a smooth submanifold of RR and G M + Rd be commuting and regular parametrization satisfying Assumption 3.5. For any initialization Zinit E M, consider the gradient flow for any time-dependent loss function L, Ed R:
dx((l) = (L (z)())d
x<()= Finit-
Define w(() = G(z(()) for all t 2 O, then the dynamics w(() is a mirror flow with respect to the Legendre function R given by Lemma 4.7, i.e.,
dV R(w(t)) = -L, (L,((t))
w( = G(
7

---

Moreover. this R only depends on the initialization Finit and the parametrization G. and is indepen- dent of the loss function L..
Theorem 4.8 provides a sufficient condition for when a gradient flow with certain parametrization G is simulating a mirror flow. The next question is then: What are the necessary conditions the parametrization G so that it enables the gradient flow to simulate a mirror flow? We provide a (partial) characterization of such G in the following theorem.
Theorem 4.9 (Necessary condition on smooth parametrization to be commuting). Let M be a smooth submanifold TR and G M - IR be smooth parametrization. for any Tmi M, there is Legendre function R such that for all time-depen dependent loss L E L. the gradient flow under L o G initialized at Finit can be written as the mirror flow under L with respect R, then G must be regular parametrization, and it also holds that for each x E M,
Lie (8G) a C ker(@G(z)),
(9)
where Lie2 == span (II[v,,, VG VG V.,. 2 [6],5i E [d]} is the subset of the Lie algebra generated by (VG.)== only containing elements of order higher than K, and (z)) is the orthogonal complement of span(( (ean((G;())){
in R..
Note the necessary condition in (9) is weaker than assuming that G is a commuting parametrization, and we conjecture that it is indeed sufficient.
Conjecture 4.10. The claim in Theorem 4.8 still holds, if we relax the commuting assumption to that 2ie (aG)) S (z))) for all r E M.
With the above necessary condition (9), we can formally refute the possibility that one can use mirror flow to characterize the implicit bias of gradient flow for matrix factorization in general settings, as summarized in Corollary 4.11. It is also worth mentioning that [40] constructed concrete counter example showing that the implicit bias for commuting measurements, that gradient flow finds the solution with minimal nuclear norm, does not hold for the general case, where gradient flow could prefer the solution with minimal rank instead.
Corollary 4.11 (Gradient flow for matrix factorization cannot be written as mirror flow). For any d E N. let M be an open set in Rdxr and G M + Raxd be a smooth parametrization given by C(U)=UU) Then there exists initial point Umi E M and a time- dependent loss Lt such that the gradient flow under L, G starting from Una cannot be written as mirror flow with respect to any Legendre function R under the loss Lt.
The following corollary shows that gradient flow with non-commuting parametrization cannot be mirror flow, when the dimension of the reachable set matches that of the w-space.
Corollary 4.12. Let M be smooth submanifold of D whose dimension is at least d. Let G M - Rdd be a regular parumetrization such that for any Timit E M, (1) Og(aiisss G) is a submanifold of dimension d, and (2) there is Legendre function R such that for any time- dependent loss T E L. the gradient flow governed by- (L o G) with initialization Vint can be written as a mirror flow with respect to R. Then G must be commuting parametrization.
Next, we establish the convergence of w(()= G()(()) when at()) is given by some gradient flow with the commuting parametrization G. Here we require that the convex function R given by Lemma 4.7 is Bregman function (see definition in Appendix B). The proofs of Theorem 4. 13, Corollary 4.14 and Theorem 4.15 are in Appendix D.
Theorem 4.13. Under the setting of Theorem 4.8, further assume that the loss L is quasi-convex,
V is locally Lipschitz and argmin (T [[(w) w E dom RJ is non-empty where R: Rd - TR U (ox) is the convex function given by Lemma 4.7. Suppose R is Bregman function, then converges to some w" such that V VL(w*) (w w²) Z 0 for all w E dom R. Moreover, if the loss function L is convex, then w(t) converges to a minimizer in dom R.
oc, w((l)
Corollary 4.14. Under the setting of Theorem 4. 4.13, if the reachable set in the w-space satisfies Dep ( G) Rdd, then R is Bregman function and all the statements in Theorem 4.13 hold. Theorem 4.15, Under the setting of Theorem 4.13, consider the commuting quadratic parametrization
G. R + I2d where each Gz(z) 1xx² AqZ. for symmetric matrices A₁, A2...., Ad E DxD that commute with euch other, i.e., A:A; A,A; = O for all i E [d]. For any Tinis E RR if {AA )d are linearly independent, then the following holds:
8

---

(a) For all E R,, (D(mii p1) exp <E i= (;A4) where exp(-) is the matrix exponential defined as exp(A) = Ek-o A (b) For each j [d] and all μ Rd, (c) ((P) = (Tinti p a Legendre function with domain Rd..
R..
Gy(bb(amit; pi) Sinit exp Ei A; Aiiin
(d) R is a Bregman function with dom R range VQ where range VQ is the range of VQ, and thus all the statements in Theorem 4. hold.
4.3 Solving underdetermined linear regression with commuting parametrization
Next, we specialize to underdetermined linear regression problems to showcase our framework. Setting:
underdetermined linear regression. Let {2; ) C Rd x I be a dataset of size Given any parametrization G. the output of the linear model on the i-th data is =; G(z). The goal is to solve the regression for the label vector = (yh, ...,....8 yn) For notational convenience, we define = (21, 22, . 2n) E _dx".
We can apply Theorem 3.8 to show the implicit bias of gradient flow with commuting parametrization. Theorem 4.16. Let M be smooth submanifold R2 and G M + Rdd be commuting and regular parametrization satisfying Assumption 3.5. Suppose the loss function L satisfies L(w) L(Zw) ) for some differentiable T. RRn - TR. For any initialization Tinit E M, consider the gradient flow
dx((() VL ()(x(())d
=x()) Zinit.
There exists a convex function (given by Lemma 4.7, depending only on Tinit and G), such that for functiinn
any dataset {22 y:) C Rd x R. if (t) = G(x((t) converges as t - ca and the convergence point w = lim w(t) satisfies Zw = Y, then R( = minw:2==Y R(w), that is, gradient flow implicitly minimizes the convex regularizer R. among all interpolating solutions.
Note that the identity parametrization we run the ordinary gradient flow on w itself and it converges to some interpolating solution, then the convergence point is closest the initialization in Euclidean distance among all interpolating solutions. This recovers the well-known implicit bias of gradient flow for underdetermined regression.
G()))= x is commuting parametrization. Therefore, if
recover the results the quadratically overparametrrized linear model studied in a series of papers [26, 61, 8]. as summarized in the following Corollary 4.17. Note that their results assumed convergence in order to characterize the implicit bias, whereas framework enables to directly prove the convergence as Theorem 4. 15. The convergence guarantee here also more general than existing convergence results for Example 4.3 in [50, 42].
Furthermore,
Corollary 4.17. Consider the underdetermined linear regression problem with data Z E Rddxa and ER.. Let L IR" + RR be a differentiable loss function such that is quasi-convex, VL is locally Lipschitz, and Y E Ra is its unique global minimizer. Consider solving ming, L(Zw) by running gradient flow L(w) L(Zw) with the quadratic parametrization
G(x) u 02 02 where = =())) E Rdd, for any initialization Finit E RRd: dz(t) = -V(L (ztt))dt x00)= = Finit- Then as + oc, w(()= = G(a(tt) converges to some Woo such that Zw = Y and R(w.o) = min:::2== R(w) where R is given by
R(w)= = 1" Wi arcsinh
Wi
u + 4ug 20g - w, In uo
220,
Vo.
5 Every mirror flow is gradient flow with commuting parametrization
For any smooth Legendre function R: Fdd R U (oo), recall the corresponding mirror flow:
dVR(w((t) - (tw()))
Note that int (dom R) is convex open set of E hence smooth manifold (see Example 1.26 in [38]), and V is a continuous dositive-definite metric on int int(dom R). As discussed previously in (3), the above mirror flow is the Riemannian gradient flow the Riemannian manifold (int(dom R). V R). The goal is to find parametrization G U - Rd, where U is an open set of TR D such that the dynamics of (t) = Gar(t) can be induced by the gradient flow on s(() governed by - (z). Formally, have the following result:

---

Theorem 5.1. Let R. R + RRU foo be a smooth Legendre function. There exist a smooth submanifold of RD denoted by M, an open neighborhood U of M and a smooth and regular purametrization U + Rd such that for the mirror flow under any time- tome-depndeet loss L with any initialization Winit iint(dom R)
dV RV)()())) = - VLo
m0() = Winit.
(10)
it holds that (t) G(x(()) for all t 2 where x(t) is given by the gradient flow under Le G: dx(()= = ( (z)))))
=²0) Tinit
(11)
where Finit satisfies (z Winit. Moreover. ler G M be the restriction of G on M, then Glas is a commutiin and regular parametrization and åG = dG| on M, which implies x(t) E M for all t 2 0. If R. is further a mirror map. then {VGr | are complete vector fields on M.
The proof of Theorem 5.1 can be found in Appendix E. To illustrate the idea, let us first suppose such smooth and regular parametrization G exists and is a bijection between the reachable set Og (Fiiit G) C IR' and int (dom R), and denote its inverse by F. It turns out that we can show OF(w) (FF() = (0G (Fo))() where the second equality follows from the relationship between R and G as discussed in the introduction on (2). Note that this corresponds to expressing the metric tensor V using an explicit map F, which is further equivalent to embedding the Riemannian manifold (int(dom It), V Rt into a Euclidean space (R.,, .7) a way that preserves its metrie. This refers to a notion called isometric embedding in differential geometry.
(w=)) (F))))"((())
V R(w)
Definition 5.2 (Isometric embedding). Let (M, g) be a Riemannian submanifold of I2 An isometric embedding from (M, g) to (R?.g) g is differentiable injective map F: : M + RD' that preserves the metric, i... for any two tangent vectors v, w M it holds that gu (v. w) IT (
Nash's embedding theorem is classic result in differential geometry that guarantees the existence of isometric embedding of any Riemannian manifold into Euclidean space with plain geometry. See Appendix A.1 for additional discussion on construction of G given a Legendre function R.
Theorem 5.3 (Nash's s embedding theorem, [47, 48, 28]). Any d-dimensional Riemannian manifold has isometric embedding to (IR- 77 for D max (d(d 5)/2, d(d 32 + 5).
As another way understand Theorem 4.8, note that (² is the Riemannian gradient of L on the Riemannian manifold (int (dom R), V²R). It is well-known that gradient flow is invariant under isometric embedding, and thus we can use Nash's 's embedding theorem rewrite the Riemannian gradient flow on (int (dom R), gd) as that on P.)).
6 Conclusion
We presented a framework that characterizes when gradient descent with proper paramterization becomes equivalent to mirror descent. In the limit of infinitesimal step size, we identify a notion named commuting parametrization such that any gradient flow (i.c., the continuous analog of gradient descent) with a commuting parametrization is equivalent to a mirror flow (i.c., the continuous analog of mirror descent) in the original parameter space with respect to a Legendre function that depends only on the initialization and the parametrization. Conversely, we use Nash's 's embedding theorem to show that any mirror flow be characterized by a gradient flow in the reparametrized space with a commuting parametrization. Using our framework, we recover and generalize results on the implicit bias of gradient descent series of existing works, including a rigorous and general proof of convergence. We also provide a necessary condition for the parametrization such that gradient flow in the reparametrized space is equivalent mirror flow in the original space. However, the necessary condition is slightly weaker than the commuting condition and it is left for future work to close the gap.
Acknowledgement
This work was supported by NSF, DARPA/SRC, Simons Foundation, and ONR. ZL acknowledges support of Microsoft Research PhD Fellowship and JDL acknowledges support of the ARO under MURI Award 1911NF-11-110300 the Sloan Research Fellowship, NSF CCF 2002272, NSF IIS 2107304, ONR Young Investigator Award, and NSF CAREER Award 2144994.
10

---

References
[1] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances in neural information processing systems, 2019.
[2] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization.. In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[3] Felipe Alvarez, Jérôme Bolte, and Olivier Brahic. Hessian riemannian gradient flows in convex programming. SIAM journal on control and optimization, 43(2):477-501, 2004.
[4] Ehsan Amid and Manfred Warmuth. Winnowing with gradient descent. In Conference on Learning Theory, pages 163-182. PMLR, 2020.
[5] Ehsan Amid and Manfred KK Warmuth. Reparameterizing mirror descent as gradient descent. Advances in Neural Information Processing Systems, 33:8430-8433 2020.
[6] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[7] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang- Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019.
[8] Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E. Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference on Machine Learning, pages 468-477. PMLR, 2021.
[9] Heinz H Bauschke, Jonathan M Borwein, et al. Legendre functions and the method random bregman projections. Journal of convex analysis, 4(1):27-67, 1997.
[10] Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-175. 2003
[11] Guy Blane, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by omstein-uhlenbeck like process. In Conference on learning theory, pages 433-513. PMLR, 2020.
[12] Lev M Bregman. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-217, 1967.
[13] Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231-357,201
[14] Yair Censor and Arnold Lent. An iterative row-action method for interval convex programming. Journal of Optimization theory and Applications, 34(3):321-353, 1981.
[15] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, pages 1305-1338. PMLR, 2020.
[16] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. arXiv preprint arXiv: 18X2::/8/2.07956. 2018.
[17] Jean-Pierre Crouzeix. A relationship between the second derivatives of a convex function and of its conjugate. Mathematical Programminng 13(1):364-365, 1977.
[18] Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimiz- ers. arXiv preprint arXiv:2106.06533 2021.

---

[19] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pages 1675-1685. PMIR, 2019.
[20] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054, 2018
[21] Robert L Foote. Regularity of the distance function. Proceedings of the American Mathematical Society, 9255)1153-555, 1984.
[22] Rong Ge, Yunwei Ren. Xiang Wang, and Mo Zhou. Understanding deflation process in over- parametrized tensor decomposition. Advances Neural Information Processing Systems, 34, 2021.
[23] Udaya Ghai, Zhou Lu, and Elad Hazan. Non-convex online learning via algorithmic equivalence. arXiv preprini urXiv:2205. 15235, 2022.
[24] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pages 1832-1841, PMLR, 2018.
[25] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems,
31, 2018.
[26] Suriya Gunasekar, Blake Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nathan Srebro. Implicit regularization in matrix factorization. In 2018 Information Theory and Applications Workshop (ITA), pages 1-10. IEEF, 2018.
[27] Suriya Gunasekar, Blake Woodworth, and Nathan Srebro. Mirrorless mirror descent: natural derivation of mirror descent. In International Conference Artificial Intelligence and Statistics, pages 2305-2313. PMLR, 2021.
[28] Matthias Gunther. Isometric embeddings of riemannian manifolds, kyoto, 1990. In Proc. Intern. Congr. Math., pages 1337-1143. Math. Soc. Japan, 1991.
[29] Jeff Z. HaoChen, Colin Wei, Jason D Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.0m 086800 2020.
[30] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv: 2018 8.
[31] Arthur Jacot, François Ged, Franck Gabriel, Berfin Şimşek, and Clément Hongler. Deep linear networks dynamics: Low-rank biases induced by initialization scale and 12 regularization. arXiv preprint arXiv:2106. 15933, 2021.
[32] Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint arXiv: 1rXiv:1803.07300 2018
[33] Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory, pages 1772-1798. PMLR, 2019.
[34] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 17 17176-17186 Curran Associates, Inc., 2020.
[35] Ziwei Ji and Matus Telgarsky, Characterizing the implicit bias via primal-dual analysis. In Algorithmic Learning Theory, pages 772-804. PMLR, 2021.
[36] Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. In International Conference on Machine Learning, pages 4860-4869. PMLR, 2021.
[37] Serge Lang- Introduction to differentiable manifolds. Springer Science & Business Media, 2006.
12

---

[38] John M Lee. Introduction to Smooth Manifolds. Springer, 2013.
[39] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in wer-paramettrized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pages 47. PMLR, 2018.
[40] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint (rXiv:1907.0459) 2019.
[41] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2020.
[42] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022. URL http:///pppneeeeeee..teee htte://forrumn..sss11142222
[43] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv: 1906.0580, 2019.
[44] Kaifeng Lyu, Zhiyuan Li, Runzhe Wang. and Sanjeev Arora. Gradient descent two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34, 2021.
[45] Edward Moroshko. Blake Woodworth, Suriya Gunasekar, Jason Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. Advances neural information processing systems, 11:2218--2299, 2020.
[46] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR, 2019.
[47] John Nash. C1 isometric imbeddings. Annals of mathematics, pages 383-396, 1954.
[48] John Nash. The imbedding problem for riemannian manifolds. Annals of mathematics, pages 20 60-63, 1956.
[49] Arkadij Semenovië Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
[50] Scott Pesme, Loucas Pillaud-Vivin and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34, 2021.
[51] Qian Qian and Xiaoyuan Qian. The implicit bias of adagrad on separable data. Advances in Neural Information Processing Systems, 32, 2019.
[52] Noam Razin and Naday Cohen. Implicit regularization in deep learning may not be explainable by norms. Advances in neural information processing systems, 33:21 174-211777, 2020.
[53] Noam Razin, Asaf Maman, and Nadav Cohen. Implicit regularization in hierarchical tensor factorization and deep convolutional neural networks. arXiv preprint 2xXi::2001.11229 2022.
[54] Ralph Tyrell Rockafellar. Convex analysis. In Convex analysis. Princeton university press, 2015.
[55] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 1981:222222277, 2018
[56] Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 34, 2021.
13

---

[57] Héctor Sussmann. Orbits of families of vector fields and integrability of distributions. Trans- actions of the American Mathematical Society, 180:171-188, 1973.
2019. [58] Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. Advances in Neural Information Processing Systems,
[59] Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu. The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning. pages 118499-18553 PMLR, 2021.
[60] Bohan Wang, Qi Meng. Huishuai Zhang, Ruoyu Sun, Wei Chen, and Zhi-Ming Ma. Momentum doesn't t change the implicit bias. arXiv preprint arXiv:21 10.03891, 2021.
[61] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-367. PMLR, 2020.
[62] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian pro- cess behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint 192.v4490,. 2019.
[63] Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pages 11727-11737, PMLR, 2021.
[64] Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view implicit bias in training linear neural networks. arXiv preprint arXiv:2010.02501, 2020.
[65] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine Learning, 109(3):467-492, 2020.
[66] Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, Dean Foster, and Sham Kakade. The benefits of implicit regularization from sgd in least squares problems. Advances in Neural 2021.
Information Processing Systems,
Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO]] to [Yes], [No], or [N/A] You are strongly encouraged to include justification to your answer, either by referencing the appropriate section of your paper or providing a brief inline description. For example:
Did you include the license the code and datasets? [Yes] See Section X.
Did you include the license to the code and datasets? [No] The code and the data are proprictary.
Did you include the license the code and datasets? [N/A]
Please do not modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A] This paper only studies the theoretical properties of optimization algorithms.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
14

---

2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] We explicitly clarify the assumptions for each result.
(b) Did you include complete proofs of all theoretical results? [Yes] Part of the proofs appears in the main context and others are deferred to appendix.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or URL)? [N/A] We do not have any experiments.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [N/A]
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [N/A]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [N/A] We do not use any existing assets.
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you's 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A] We do not crowdsourcing conduct research with human subjects.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
15


Reviews:
[{"review_id": "sZBd5hCHhRg", "reviewer": "Reviewer_eah9", "summary": "For a reparametrisation $w = G(x)$ where $x_t$ follows a gradient flow on some loss $L \\odot G$, the authors give a sufficient condition on the parametrisation (i.e. commuting parametrisation) so that $w_t$ to corresponds a mirror flow on $L$ and with a related mirror map.  They then show that any mirror flow can be seen as the commutative reparametrisation of a gradient flow.", "questions": "- for a fixed commutating parametrisation $G$, is it possible to understand the effect of $x_{init}$ on the associated mirror map $R$\n- you provide sufficent conditions, do you think they are necessary ? or not ? \n- is Nash's embedding constructive ? can we know F, G, or the dimension D ?\n- do you think the claim in Remark 4.7 still holds if we restrict ourselves to losses which are non time-dependent ? \n\n", "limitations": "N/A", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I think this is a very good paper.  It is very well written: \n- the motivations are clear and relevant as it sheds light on when we can hope to obtain implicit regularisation results using mirror flow type of analysis\n- the setup and preliminaries are very clear\n- All the results are introduced and commented, examples are also provided\n- though I am not familiar with all the mentioned notions, they all seem well defined and the math seems rigorous\n\nThe paper is very well written and I believe the general framework they introduce as well as the lemmas / tools they provide are relevant: as they help understand the link between gradient flow, reparametrisation and Mirror flow.\n\n\nMinor comments: \n- in the main contributions in the intro: \"we also five new convergence analysis in such settings, filling the gap in previous works [22, 57, 8] where parameter convergence is only assumed but not proved.\" is an overclaim as the convergence of mirror descent / flow is well understood and has been investigated for years (though I agree that the mentioned works seem to have overlooked them. \n- I would have appreciated an example illustrating Lemma 4.5 and Lemma 4.6 , for the $u^2 - v^2$ parametrisation (or even simply $G(x) = x$)  for example: what is $\\Psi(x, \\mu)$ in this case ? and $\\mu_t$? and $G(\\Psi(x_{init}, \\mu))$ ? this would help the reader gain more intuition / understanding", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "U17bThc-TQ-", "reviewer": "Reviewer_94Rv", "summary": "This work establishes a link between minimizing a loss function using continuous-time mirror descent (i.e. mirror flow) and minimizing a reparameterized version of the same loss function using continuous-time gradient descent (i.e. gradient flow) to explain the implicit bias of gradient flow in reparameterized models. The authors introduce a class of reparameterization mappings named commuting parameterizations which satisfy certain commutativity conditions of the Jacobian. The authors show that for any commuting parameterization, there exists a mirror function in the original domain such that the trajectory of gradient flow of a loss function with that parameterization coincides with the trajectory of mirror flow of the same loss function in the original domain with the associated mirror function. Furthermore,  they also give a converse result by showing the existence of a commuting parameterization for any mirror function such that gradient flow in the reparameterized space is equivalent to mirror flow in the original space. The authors also show that several past works on implicit bias/regularization can be considered as specific cases of different commuting parameterizations. \n", "questions": "I would appreciate it if the authors clarify the following points:\n\n1. Is it possible to provide an order-wise relationship between $D$ and $d$? More explicitly, given mirror function $R$, can you tell how large the dimension of the reparameterization space $D$ must be with respect to $d$?\n2. Given a mirror function $R$, how easy and feasible is it to construct a commuting parameterization $G$?\n3. Does it make a difference in terms of finite-time (i.e. early stopping) convergence rate if one performs the training in x-space with gradient flow or w-space with mirror flow?\n", "limitations": "The authors could make this paper more impactful if they can provide a discrete-time analysis or a finite-time (i.e. early stopping) convergence rate analysis. It would also be interesting to know the order-wise relationship between dimensions $d$ and $D$. The authors could also support their results by simulations performed with practically relevant models (deep nets, kernel models, etc.) to show efficiency of using (possibly sparse) parametrizations in optimizing these models.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The authors give a sufficient condition to check whether gradient flow for a given reparametrization can be written as a mirror flow for some mirror map in a possibly smaller dimensional domain. The condition simply relies on checking commutativity of the row vectors of the Jacobian of reparametrization map $G:M\\subset \\mathbb{R}^{D} \\to \\mathbb {R}^d$, i.e., $[\\nabla G_i, \\nabla G_j] (x)=0$ for any $i,j\\in[d]$ for all $x \\in M$. They also show that several examples of reparameterization satisfy this condition.\n\nWeaknesses:\n1. The nature of the proof of Theorem 4.4 seems to be existential rather than constructive. It doesn't seem to be straightforward to construct a commuting reparametrization $G$ given a mirror map $R$ and vice versa.\n2. Although the analysis is valid for continuous-time flows, the authors do not address the fidelity of this equivalence in discrete-time setting which is more practically relevant. \n3. The theoretical analysis in this paper could have been supported with extensive simulations on certain benchmarks. \n\n\nAlthough the paper is well written and easy to understand, it mostly builds on Amid and Warmuth (2020) by generalizing to non-injective parametrizations. This makes it an incremental contribution as there is neither much technical novelty nor extensive simulations to show practical relevance and to draw more insights. ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EW2HoALRXZU", "reviewer": "Reviewer_puj1", "summary": "This paper focuses on implicit bias of gradient descent for reparametrized models.\nIn this setting, a loss L over parameters w (eg weights of a neural network) is minimized, where these parameters w write w=G(x).\nHence, they consider the gradient flow of L in the x space. In the weight space, the gradient flow writes with a preconditioner (partial G(x) partial G(x))^T.\nIt is known that if the latter preconditioner is equal to a Hessian map nabla^2 R(w)^-1 where R is strictly convex, then the latter flow is a Riemannian gradient flow (RGF) with respect to this metric tensor nabla^2 R.\n\nIn this paper, the authors tackle the question of characterizing the gradient flow , i.e. when given a parametrization G, we can consider the assocciated \nRGF with mirror map R; and conversely, for what R does one have a corresponding parametrization G.\n\nThe main class at study are commutative parametrizations, that holds for instance when each coordinate of G depends on different subsets of coordinates (separable case). Th 4.4 is the first main result of the paper and states that if G is commutative, then there exists a mirror map R (depending on xinit and G) such that the dynamics of w is a mirror flow. The proof relies at some point on some sufficient condition (given in Lemma C.1, already highlighted in [42]) under which a vector field is the gradient of a convex function. Convergence of the flow is proved if G is a linear reparametrization (G(x)=xAx) in Corollary 4.10. To characterize implicit bias of gradient descent, the authors use Th 3.6 which characterizes implicit bias (ie properties of the solution) of a mirror flow.\nFinally the authors study the reciprocal question, i.e. for a mirror map R, does there exist a corresponding commutative reparametrization G, and provide a positive answer in Theorem 5.3 that holds generally.\n", "questions": "- to what extent the models at study (ie commutative parametrizations) are restrictive. Could the authors clarify this point and the limitations of this hypothesis ? how large is the class of models considered?  is it limited to diagonal neural nets? are more complicated architectures ruled out because of the negative result of Remark 4.7?\n- th 4.4: intuitively, why the mirror map nabla R should depend on xinit?   \n- th 4.11: while the result is interesting, since it says that the solution is a minimizer of the convex regularizer R, this function R remains implicit/is not known if I understood well ? in corollary 12 it is explicit.\n- l295: the proposed formula for nabla^2 R (with the inverse of G = F) look quite different from the one in Lemma 4.6 involving psi (hence gradient flows of G ), could the authors clarify the link between the mirror map in section 4 vs section 5? \n\n\n", "limitations": "Limitations are too briefly evoked (e.g. in the conclusion) but they should be discussed more extensively in my opinion.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nThe paper brings interesting contributions to the growing literature on implicit bias of optimization algorithms, an important problem in the machine learning literature nowadays. I think the characterization of implicit bias of gradient descent through the mirror map R in the reparametrization setting at study, explicitly when it is more explicit in Th 4.11 and Cor 4.12, is novel and interesting. \nThe paper is quite well written and the results appear technically sound (I read the proofs in the main, but I did not read the proofs in detail in the appendix).\n\nWeaknesses: I think that the paper lacks clarity in its current state. It would be better with more interpretation and discussion of the results, see questions below. Moreover, due to the limitations of the reparametrization model at study and the lack of empirical support, the impact of the paper may be limited. ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept