Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descent
Anonymous Author(s)
AfﬁliationAddress email
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 theoverparametrized model can be understood as mirror descent on a different objective.
The main result here is a characterization of this phenomenon under a notiontermed commuting parametrization , which encompasses all the previous resultsin this setting.
It is shown that gradient ﬂow with any commuting parametrizationis equivalent to continuous mirror descent with a related mirror map.
Conversely,continuous mirror descent with any mirror map can be viewed as gradient ﬂowwith a related commuting parametrization.
The latter result relies upon Nash’sembedding theorem.
Introduction
Implicit bias refers to the phenomenon in machine learning that the solution obtained from lossminimization has special properties that were not implied by value of the loss function and insteadarise from the trajectory taken in parameter space by the optimization.
Quantifying implicit bias necessarily has to go beyond the traditional black-box convergence analyses of optimization algorithms.
Implicit bias can explain how choice of optimization algorithm can affect generalization [59, 41, 40].
Many existing results about implicit bias treat training (in the limit of inﬁnitesimal step size) asa differential equation or process { x ( t ) } t [-≥ 
 ⊂-] 0 ⇢ R D .
To show the implicit bias of x ( t ) , the ideais to show for another (more intuitive or better understood) process { w ( t ) } t [-≥ 
 ⊂-] 0 ⇢ R d that x ( t ) issimulating w ( t ) , in the sense that there exists a mapping G : R D !
R d such that w ( t ) =
G ( x ( t )) .
Then the implicit bias of x ( t ) can be characterized by translating the special properties of w ( t ) backto x ( t ) through G .
A related term implicit regularization refers to a handful of such results whereparticular update rules are shown to lead to regularized solutions; speciﬁcally, x ( t ) is simulating w ( t )where w ( t ) is solution to a regularized version of the original loss.
The current paper develops a general framework involving optimization in the continuous-time regimeof a loss L : R d ! R that has been re-parametrized before optimization 1 as w =
G ( x ) for some
G : R D ! R d .
Then the original loss L ( w ) in the w -space induces the implied loss L ( G ( x )) in thex -space, and the gradient ﬂow in the x -space is given by
[Equation]
where r x denotes the gradient with respect to x .
Using w ( t ) =
G ( x ( t )) and the fact thatr x L ( G ( x ))
= @
G ( x ) > r L ( G ( x )) where @ G ( x ) 2 R d ⇥ D denotes the Jacobian of G at x , thecorresponding dynamics of (1) in the w -space is
[Equation]
Our framework is developed to fully understand phenomena in recent papers [24, 56, 62, 4, 59, 5, 8],which give examples suggesting that gradient ﬂow in the x -space could end up simulating a moreclassical algorithm, mirror descent (speciﬁcally, the continuous analog, mirror ﬂow) in the wspace.
Recall that mirror ﬂow is continuous-time limit of the classical mirror descent, writtenas d r R ( w ( t )) =
r L ( w ( t ))
d t where R : R d ! R is a strictly convex function called mirrormap [48, 10].
Equivalently it is Riemannian gradient ﬂow with metric tensor r 2 R , an old notion ingeometry:
[Equation]
If @ G ( x ( t )) @ G ( x ( t )) >
= r 2 R ( w ( t ))  1 for all t , then (2) becomes a simple mirror ﬂow in thew -space as noted in [25].
Many existing results about implicit bias indeed concern reparametrizations
G that satisfy @ G ( x ) @ G ( x ) > = r 2 R ( w )  1 for a stricly convex function R , and the implicitbias/regularization is demonstrated by showing that the convergence point satisﬁes the KKT conditionsneeded for minimizing R among all minimizers of L .
A concrete example is w i ( t ) =
G i ( x ( t )) =( x
i ( t )) 2 for all
i 2 [ d ] , so here D = d .
In this case, the mirror map R must satisfy ( r 2 R ( w ( t )))
1 =@ G ( x ( t ))
@ G ( x ( t ))
> = 4diag
(( x 1 ( t )) 2 , . . .
, ( x d ( t )) 2 ) =
4diag( w 1 ( t ) , . . . , w d ( t )) which suggests
R is the classical negative entropy function, i.e. , R ( w ) =
P di =1 w i (ln w i  1) .
But in general, it is hard to decide whether gradient ﬂow for a given parametrization G can be writtenas mirror ﬂow for some mirror map R , especially when D > d and G is not an injective map.
In suchcases, there could be multiple x ’s mapping to the same G ( x ) yet having different @ G ( x ) @ G ( x ) > .
If more than one of such x can be reached by gradient ﬂow, then the desired mirror map cannotexist.
However, if only one of such x can be reached by gradient ﬂow, we must decide which x it isin order to decide the value of r 2 R using @ G @ G > .
Conversely, [5] raises the following question:for what mirror map R can the corresponding mirror ﬂow be the result of gradient ﬂow after somereparametrization G ?
Answering the questions in both directions requires a deeper understanding ofthe impact of parametrizations.
The following are the main contributions of the current paper:
• In Section 4, building on classic study of commuting vector ﬁelds we identify a notionof when a parametrization w = G ( x ) is commuting (Deﬁnition 4.1) and use it to give asufﬁcient condition (Theorem 4.4) of when the gradient ﬂow in the x -space governed byr x L ( G ( x )) is simulating a mirror ﬂow in the w -space with respect to some mirror mapR : R d ! R , which encompasses all the previous results [24, 56, 62, 4, 59, 5, 8].
Moreover,the mirror map is independent of the loss L and depends only on the initialization x init andthe parametrization G .
• We recover and generalize existing implicit regularization results for underdetermined linearregression as implications of the above characterization (Corollary 4.12).
We also givenew convergence analysis in such settings (Corollary 4.10), ﬁlling the gap in previousworks [24, 59, 8] where parameter convergence is only assumed but not proved.
• In the reverse direction, we use the famous Nash embedding theorem to show that every
mirror ﬂow in the w -space with respect to some mirror map R simulates a commutativegradient ﬂow under some embedding x = F ( w ) where F : R d ! R D and the parametrization G is the inverse of F .
(Theorem 5.3)
This provides an afﬁrmative and fully general answer to the question of when such reparametrization functions exist, giving a full answerto questions raised in a more restricted setting in [5].
Related work
Implicit bias.
With high overparametrization as used in modern machine learning, there usuallyexist multiple optima, and it is crucial to understand which particular solutions are found by the optimization algorithm.
Implicit bias of gradient descent for classiﬁcation tasks with separable datawas studied in [54, 22, 45, 34, 44, 33] and for non-separable data in [31, 32], where the implicit biasappears in the form of margin maximization.
The implicit bias for regression problems has also beenanalyzed by leveraging tools like mirror descent [59, 22, 62, 56, 4, 5], later generalized in [8].
The sharp contrast between the so-called kernel and rich regimes [59] reﬂects the importance of theinitialization scale, where a large initialization often leads to the kernel regime with features barelychanging during training [28, 16, 20, 19, 2, 1, 63, 7, 60, 29], while with a small initialization, thesolution exhibits richer behavior with the resulting model having lower complexity [23, 24, 38, 51,6, 15, 40, 42, 43, 52, 55, 21].
Recently [61] give a complete characterization on the relationshipbetween 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., stochasticgradient descent [39, 11, 27, 41, 18, 64] and adaptive and momentum-based methods [50, 58, 57, 35],to name a few.
Understanding mirror descent.
In the continuous-time limit as step size goes to 0, the mirror ﬂowis equivalent to the Riemannian gradient ﬂow.
[25] showed that a partial discretization of the lattergives rise to the classical mirror descent.
Assuming the existence of some reparametrization function,[5] showed that a particular mirror ﬂow can be reparametrized as a gradient ﬂow.
Our paper showsthat such reparametrization always exists by using Nash’s embedding theorem.
Preliminaries and notations
For any positive integer n , we denote { 1 , 2 , . . . , n } by [ n ] .
For any vector u 2 R D , we denote itsi -th coordinate by u i .
For any vector u, v 2 R D and ↵ 2 R , we deﬁne u  v = ( u 1 v 1 , . .
, u D v D ) >and u  ↵ =
(( u 1 ) ↵ , . .
, ( u D ) ↵ ) > .
For any integer k , we say a function f is C k if it is k timescontinuously differentiable.
For any convex function R : R d ! R [ { 1 } , we denote its domainby dom R = { w 2 R d
| R ( w ) < 1 } .
For any set S , we use int( S ) to denote its interior.
For anysubmanifold M ⇢ R D and x 2 M , we denote by T x ( M ) the tangent space of M at x .
We assume the model has parameter vector w 2 R d and loss function L : R d ! R which is a C 1function.
Training involves a reparametrized vector x 2 R D , which is a reparametrization of w suchthat w =
G ( x ) for some differentiable parametrization function G .
More precisely, we assume thatthe reparametrized vector x takes value in a smooth submanifold of R D denoted by M , so wheneverwe say G : M ! R d is a parametrization function, by default M is a smooth submanifold of R D .
Given an initialization x init , we consider running gradient ﬂow on x :
[Equation]
From now on, we follow the convention that d is the dimension of the original parameter w and D isthe dimension of the reparametrized x .
We also refer to R d as the w -space and R D as the x -space.
Some of the proofs and counterexamples also use the following notion of time-dependent loss.
Deﬁnition 3.1 (Time-dependent loss) .
A time-dependent loss is a function L t ( w ) that maps ( w, t ) {+
+}
[-∈ 
-]
R d ⇥ R to a real number such that for any t 2 R , L t ( w ) is continuously differentiable in w , andr L t ( w ) is piece-wise continuous in ( w, t ) .
The set of all time-dependent losses is denoted by L .
3.1 Vector ﬁelds
Vector ﬁelds are a natural way to formalize continuous-time gradient descent (a good referenceis [36]).
Given any smooth submanifold of R D , denoted by M , a vector ﬁeld X on M is a continuousmap from M to R D such that for any x 2 M , X ( x ) is in the tangent space of M at x , T x ( M ) .
Deﬁnition 3.2 (Complete vector ﬁelds, p.215 in [36]) .
We say X is a complete vector ﬁeld on Mif and only if for any initialization x init 2 M , the differential equation d x ( t ) =
X ( x ( t ))
d t has asolution on ( 1 , 1 ) with x (0) = x init .
For any differentiable function f : M ! R , we denote by r f its gradient vector ﬁeld.
For example,in the special case when M = R D , r f is simply the ordinary gradient of f .
For any x 2 M , wedenote by  tf ( x ) the point on M reached after time t by following the vector ﬁeld r
f starting at x :
(here we allow t to be negative and use the convention [Equation]
[Equation]
Now, let G : M ! R d be a continuously differentiable parametrization function.
We denote its
Jacobian by @ G ( x ) = ( r G 1 ( x ) , r G 2 ( x ) , . . .
, r G d ( x )) > .
For any x 2 M and µ 2 R d , we deﬁne
[Equation]
where  denotes the composition of functions.
3.2 Parametrizations
We focus on a class of parametrization functions that do not change the expressiveness of the originalmodel.
So when it comes to implicit bias, we require by default the original model is expressiveenough to have multiple optimal solutions.
Deﬁnition 3.3 (Canonical parametrizations) .
A canonical parametrization G : M ! R d is acontinuously differentiable function that satisﬁes: (1) the domain M is a smooth submanifold of R D , (2) the gradient vector ﬁelds r G 1 , r G 2 , . . .
, r G d are all complete vector ﬁelds on M and (3)@ G ( x ) is of rank d for all x 2 M .
Given a canonical parametrization G : M ! R d and any initialization x init 2 M , the set of pointsthat are reachable via gradient ﬂow with respect to some time-dependent loss is a subset of M thatdepends on G and x init .
Since "reachable" is a transitive relationship, the notion of reachable setsdivide the domain M into equivalent classes.
Deﬁnition 3.4 (Reachable sets) .
For a canonical parametrization G : M ! R d and initializationx init 2 M , we deﬁne the reachable set ⌦ x ( x init ; G ) as⌦ x ( x init ; G ) = ⇢ x ( t )  8 t 2 R , L · 2 L and continuous x ( · ) : R ! R d s.t.
d x
[Equation]
s ) /
[Equation]
( G ( x ( s )))
[Equation]
R , x (0) = x init  . The above reachable set in the x -space further induces the corresponding reachable set in the w -spacegiven by ⌦ w ( x init ; G )
=
G ( ⌦ x ( x init ; G )) .
3.3 Mirror descent and mirror ﬂow
Below we introduce some basics notions for mirror descent [48, 10].
Deﬁnition 3.5 (Mirror maps, p.298 in [13]) .
A function R : R d ! R [ { 1 } is called a mirror 146 map if and only if (1) R is strictly convex and differentiable on int(dom R ) ; (2) the gradient mapr R : int(dom R ) ! R d is surjective; and (3) for any sequence { w i } 1 i =1 going to the boundary of 148 dom R , lim [Equation]
Then for initialization w 0 2 int(dom R ) , mirror descent with step size ⌘ updates as follows:
[Equation]
In the limit of ⌘ ! 0 , the above discrete update becomes the continuous mirror ﬂow:
[Equation]
Recall that given a differentiable function R , the corresponding Bregman divergence D R is deﬁned as
[Equation]
We recall a well-known implicit bias result for mirror ﬂow (which holds for mirror descent as well),which shows that for a speciﬁc type of loss, if mirror ﬂow converges to some optimal solution, thenthe convergence point achieves some implicit regularization among all optimal solutions.
r G i t i r
G j t j r G j t j r G
i t i x  t i G i ( x )  t j G j ( x )  t
i G i   t j G j ( x ) =
t
j G j   t i G i ( x )
Theorem 3.6.
Given any data Z 2 R n ⇥ d and the corresponding label Y 2 R n , suppose theloss L ( w ) is in the form of L ( w ) =
e L ( Zw ) for some differentiable e L : R n ! R .
Assume thatinitialized at w 0 , the mirror ﬂow (6) converges and the convergence point w 1 = lim t !
1 w ( t )satisﬁes
Zw 1 =
Y , then D R ( w 1 , w 0 ) =
min w : Zw =
Y D R ( w, w 0 ) .
See Appendix B for a proof.
The above theorem is the building block for proving the implicit biasinduced by any commuting parametrization in overparametrized linear models (see Theorem 4.11).
Every commutative gradient ﬂow is a mirror ﬂow
We now formalize the notion of commuting parametrization.
We remark that M is a general smoothsubmanifold of R D that the reparametrized vector x takes value in.
Deﬁnition 4.1 (Commuting parametrizations) .
A commuting parametrization G : M ! R d is acanonical parametrization such that the gradient vector ﬁelds r G 1 , r G 2 , . .
, r G d on M satisfythat for any i, j 2 [ d ] and any t i , t j 2 R ,  t
i G i   t j G j ( x ) =
t
j G j   t i G i ( x ) for all x 2 M .
When G is a commuting parametrization, we call the corresponding gradient ﬂow (1) a commutativegradient ﬂow .
We illustrate the concept of commuting parametrization in Figure 1.
The meaning of commuting in the above deﬁnition is built upon the ﬂows generated by the vectorﬁelds, which is further equivalent to that the Lie bracket of any pair of the vector ﬁelds satisﬁes[ r G i , r G j ] ⌘ 0 .
In the latter sense, we say r G i and r G j are commuting vector ﬁelds (see 172 p.231, [36]).
Intuitively, the Lie bracket [ r G i , r G j ] characterizes the change of r G j along the 173 ﬂow generated by r G i and [ r G i , r G j ] ⌘ 0 means r G j is invariant along  tG i ( x ) , thus justifyingthe commutability between  t i G i and  t j G j .
The equivalence is summarized in the following theorem,which can be applied to check whether a parametrization is a commuting parametrization or not.
Theorem 4.2 (Adapted from Theorem 9.44 in [36]) .
Given a differentiable map G : M ! R d where
M is a smooth submanifold of R D , for any i, j 2 [ d ] , r G i and r G j are commuting vector ﬁelds,i.e., [ r G i , r G j ] ⌘ 0 , if and only if any t i , t j 2 R ,  t i G i
t
j G j ( x ) =
t
j G j   t i G i ( x ) for all x 2 M .
In particular, when r G 1 , r G 2 , . . .
, r G d are ordinary gradient vectors in R D , the Lie bracketbetween any r G i and r G j is given by [ r G
i [Equation]
The above condition clearly holds when each G i only depends on a different subset of coordinatesof x , because we have r 2 G i ( x ) r G j ( x ) ⌘ 0 , for all
i 6 = j 2 [ d ] and x 2 M as r 2 G i andr G j live in different spaces.
We call this setting separable case and it covers all the previousexamples [24, 56, 4, 59, 5].
For concreteness, we give the example below.
Example 4.3 (Quadratically overparametrized model) .
Consider the case of the overparametrizedlinear model (also known as diagonal linear network)
[59] where w 2 R d is parametrized byw
=
u  2  v  2 .
In this case, we set D = 2 d and M = { x =  uv  2 R D | u i , v i > 0 , 8 i 2 [ d ] } ,and G : M !
R d is given by G ( x ) =
u  2  v  2 .
Since each G i involves only u i and v i , this is aseparable case and thus r G 1 , r G 2 , . . .
, r G d are commuting vector ﬁelds.
Next, we present our result on characterizing any commutative gradient ﬂow by a mirror ﬂow.
Theorem 4.4.
Let M be a smooth submanifold of R D and G : M !
R d be a commuting parametrization as deﬁned in Deﬁnition 4.1.
For any initialization x init 2 M , consider the gradient ﬂow for anytime-dependent loss function L t : R d ! R :
[Equation]
Denote w ( t ) =
G ( x ( t )) for all t  0 , then the dynamics of w ( t ) is a mirror ﬂow with respect to themirror map R given by Lemma 4.6, i.e.,
[Equation]
Moreover, this R only depends on the initialization x init and the parametrization G , and is independent of the loss function L t .
The proof of the above theorem relies on two lemmas that highlight the special properties ofcommuting parametrizations.
Lemma 4.5 tells us that the point reached by a commutative gradientﬂow is determined by the integral of the negative gradient of the loss along the trajectory.
Lemma 4.5.
Let G : M ! R d be a commuting parametrization and consider any timedependent loss L t .
For any initialization x init 2 M , let the solution to the gradient ﬂow governed by r x L t ( G ( x )) at time t be x ( t ) , i.e., x ( t ) =
x init + R t 0 r x L t ( G ( x ( s )))
d s .
Deﬁneµ ( t ) =
R t [-
-]0 r
L t ( G ( x ( s )))
d s , then it holds that x ( t ) =
( x init ; µ ( t )) .
Based on Lemma 4.5, the following key lemma reveals the essential approach to ﬁnd the mirror map.
Lemma 4.6.
Let G : M ! R d be a commuting parametrization.
For any initialization x init 2 207 M , there exists a twice differentiable strictly convex function Q : R d ! R such that r Q ( µ ) =
208 G ( ( x init ; µ )) .
Moreover, the convex conjugate of Q , denoted by R , is a mirror map such thatint(dom R ) =
⌦ w ( x init ; G ) and r 2 R ( G ( ( x init ; µ )))
= ( @ G ( ( x init ; µ )) @ G ( ( x init ; µ ))
> )  1for all µ 2 R d .
Now, we are ready to prove Theorem 4.4.
Proof of Theorem 4.4.
Recall that the gradient ﬂow in the x -space governed by r x L ( G ( x )) is
[Equation]
Using w ( t ) =
G ( x ( t )) , the corresponing dynamics in the w -space is
[Equation]
Next, by Lemma 4.5, we know that the solution to the gradient ﬂow satisﬁes x ( t ) =
( x init ; µ [Equation] t ))where µ ( t ) =
R t [-
-]0 r
L t ( G ( x ( s )))
d s .
Therefore, applying Lemma 4.6, we get a mirror map R withdomain ⌦ w ( x init ; G ) such that
[Equation]
for all t  0 .
Then the dynamics of w ( t ) satisﬁes d w ( t ) = r
2 R ( w ( t ))  1 r L ( w ( t ))
d t which isexactly the mirror ﬂow with respect to R initialized at w (0) =
G ( x init ) .
Note that the result of
Lemma 4.6 is completely independent of the loss function L , thus the mirror map R only depends onthe initialization x init and the commuting parametrization G .
This ﬁnishes the proof.
Remark 4.7 (Non-commuting ﬂow cannot be mirror ﬂow, when dimension matches.) .
Suppose
G : M ! R d is a canonical parametrization with dim( ⌦ x ( x 0 ; G ))
=
d and for any time-dependentloss L · 2 L , gradient ﬂow governed by r x L t ( G ( x )) can be written as a mirror ﬂow, then G must 224 be commuting parametrization.
To see why, let  > 0 and consider a time-dependent loss given by
[Equation]
Suppose for any initialization x init 2 M , d x ( t ) =
r x L t ( G ( x ( t ))
d t with x (0) =
x init can bewritten as d r R ( w ( t )) =
r L t ( w ( t ))
d t for some mirror map R where w ( t ) =
G ( x ( t )) .
Thenr R (4  ))
= r R ( w (0))
+ R t 0 r L s ( w ( s ))
d s = r R ( w (0)) , which further implies that w (4  ) =
w (0) , and thus G ( x (4  ))
=
G ( x init ) .
One can show that x (4  ) =
x init +  2 [ r G j , r G i ]( x init ) +
O (  3 ) , which implies [ r G j , r G i ]( x init )
2 T x init ( ⌦ x ( x init ; G )) .
Since dim( ⌦ x ( x init ; G ))
=
dand G is canonical, { r G k ( x init ) } dk =1 form a basis of T x init ( ⌦ x ( x init ; G )) .
Note that
=lim t {+! 
+} G k ( x (4 t ))
G k ( x init ) t= hr G i ( x init ) , [ r G j , r G i ]( x init )
{+i for all k 
 [ d ] , we conclude 232 [ r G j , r G i ]( x init ) must be 0 .+} Thus by Theorem 4.2, we conclude that G is a commutingparametrization since x init is arbitrary.
Next, we establish the convergence of w ( t ) =
G ( x ( t )) when x ( t ) is given by some commutativegradient ﬂow.
Here we require that the mirror map R given by Theorem 4.4 is a Bregman function(see deﬁnition in Appendix A).
The proofs of Theorem 4.8 and Corollary 4.9 are in Appendix C.
Theorem 4.8.
Under the setting of Theorem 4.4, further suppose that the loss L is quasi-convex andr L is locally Lipschitz.
If the strictly convex function R given by Theorem 4.4 is a Bregman functionsatisfying argmin { L ( w ) | w 2 dom R } is non-empty, then as t ! 1 , w ( t ) converges to some w ⇤such that r L ( w ⇤ ) > ( w  w ⇤ )  0 for all w 2 dom R .
Moreover, if the loss function L is convex,then w ( t ) converges to a global minimizer.
Corollary 4.9.
Under the setting of Theorem 4.8, if the reachable set in the w -space satisﬁes⌦ w ( x init ; G )
=
R d , then R is a Bregman function and all the statements in Theorem 4.8 hold.
Corollary 4.10.
Under the setting of Theorem 4.8, if M = R D + and G i ( x ) = x > A i x for all i 2 [ d ]where A
i 2 R D ⇥ D are symmetric matrices and commute with each other, i.e. , A i A j  A j A
i
=0 , 8 i, j 2 [ d ] , then R is a Bregman function and all the statements in Theorem 4.8 hold.
4.1 Case study: underdetermined linear regression
Next, we specialize to underdetermined linear regression problems to showcase our framework.
Setting: underdetermined linear regression.
Let { ( z i , y i ) } ni =1 ⇢ R d ⇥ R be a dataset of size n .
Given any parametrization w = G ( x ) , the output of the linear model on the i -th data is z > i G ( x ) .
Thegoal is to solve the regression for the label vector Y =
( y 1 , y 2 , . . .
, y n ) > .
For notational convenience,we deﬁne Z =
( z 1 , z 2 , . . .
, z n ) 2 R d ⇥ n .
We can apply Theorem 3.6 to obtain the implicit bias of a commutative gradient ﬂow.
Theorem 4.11.
Let M be a smooth manifold of R d and G : M ! R d be a commuting parametrization as deﬁned in Deﬁnition 4.1. Suppose the loss L satisﬁes L ( w ) = e
L ( Zw ) for some differentiable e L : R n ! R .
For any x init 2 M , consider the gradient ﬂow d x ( t ) =
r x L ( G ( x ( t )))
d twith x (0) = x init .
Then there exists a convex function R (given by Lemma 4.6, dependingonly on x init and the parametrization G ), such that for any dataset { ( z i , y i ) } ni =1 ⇢ R d ⇥ R , if 259 w ( t ) =
G ( x ( t )) converges and the convergence point w 1 = lim t !
1 w ( t ) satisﬁes Zw 1 =
Y , then 260 R ( w 1 ) =
min w : Zw =
Y R ( w ) , that is, gradient ﬂow implicitly minimizes the convex regularizer Ramong all interpolating solutions.
Proof of Theorem 4.11.
By Theorem 4.4, w ( t ) obeys the following mirror ﬂow:
[Equation]
Applying Theorem 3.6 yields D R ( w 1 , G ( x init ))
=
min w : Zw =
Y D R ( w, G ( x init )) .
Therefore, forany w 2 int(dom R ) such that
Zw =
Y , we have
[Equation]
which can be reorganized as
[Equation]
Note that by Lemma 4.6, we have r R ( G ( x init ))
= r R ( G ( ( x init ; 0)))
= r R ( r Q (0))
=
wherethe last equality follows from the property of convex conjugate.
Combining this with the aboveinequality, we get R ( w 1 )  R ( w ) for all w 2 int(dom R ) such that
Zw = Y .
By the continuityof R , this property can be further extended to the whole dom R , and for any w / 2 dom R , we have
R ( w ) = 1 by deﬁnition, so R ( w 1 )  R ( w ) holds trivially.
This ﬁnishes the proof.
Note that the identity parametrization w =
G ( x ) = x is a commuting parametrization.
Therefore, ifwe run the ordinary gradient ﬂow on w itself and it converges to some interpolating solution, thenthe convergence point is closest to the initialization in Euclidean distance among all interpolatingsolutions.
This recovers the well-known implicit bias of gradient ﬂow for underdetermined regression.
Furthermore, we can recover the results on the quadratically overparametrized linear model in aseries of paper [24, 59, 8] and strenghthen their results by removing the convergence assumption intheir theorems by Corollary 4.10.
Corollary 4.10 is also more general than convergence results on
Example 4.3 in [49, 41].
Corollary 4.12 (Quadratically overparametrized linear model) .
Consider the quadratically overparametrized linear model where the parametrization is given by w = G ( x )
= u  2  v  2 withx =
uv  2 R 2 d .
Suppose the loss function L is quasi-convex, r L is locally Lipschitz and
L ( w ) = e L ( Zw ) for some differentiable e L : R n !
R whose only global minimizer is Y .
Forany initialization x init =
u 0 v such that all coordinates of x init are non-zero, let x ( t ) be the solutionto the gradient ﬂow (1) at time t .
Then as t ! 1 , w ( t ) =
G ( x ( t )) converges to some w 1 with
Zw 1 = Y , and R ( w 1 ) =
min w : Zw =
Y R ( w ) , where the mirror map R is given by
[Equation]
Every
mirror ﬂow is a commutative gradient ﬂow
Now we consider any C 2 mirror map R : R d ! R [ { 1 } .
Recall the corresponding mirror ﬂow:
[Equation]
We denote the interior of the domain of R by C , then C is a convex open set of R d , hence a smoothmanifold (see Example 1.26 in [36]).
Then r 2 R is a continuous positive-deﬁnite metric on C .
Asdiscussed previously, the above mirror ﬂow can be further rewritten as the Riemannian gradient ﬂowon the Riemannian manifold ( C, r 2 R ) , i.e.,
[Equation]
The goal is to ﬁnd a commuting parametrization G : M ! R d where M is smooth submanifold of R D such that the dynamics of w ( t ) =
G ( x ( t )) is given by the commutative gradient ﬂow on x ( t )governed by r x L ( G ( x )) .
To illustrate the idea, let us ﬁrst suppose such a G exists and is a bijection, whose inverse is denotedby F .
Then it satisﬁes that r 2 R ( w ) =
( @ G ( F ( w ))
G ( F ( w ))
> )  1 , which further suggests thatr 2 R ( w ) = @ F ( w ) > @ F ( w ) .
Note that this corresponds to expressing the metric tensor using anexplicit map F , which is further equivalent to embed the Riemannian manifold ( R d , r 2 R ) into a
Euclidean space R D in a way that preserves its geometry.
This refers to a notion called isometricembedding in differential geometry.
Deﬁnition 5.1 (Isometric embedding, p.15 in [37]) .
Let ( M, g ) and ( f M, ˜ g ) be Riemannian manifolds.
An isometric embedding is an embedding F : M ! f M which preserves the metric in the sense thatg is equal to the pullback of ˜ g by F .
More speciﬁcally, for any two tangent vectors v, w 2 T x ( M )we have g ( v, w ) =
˜ g
(d F ( v ) , d F ( w )) .
Nash’s embedding theorem is a classic result in differential geometry that guarantees the existence ofisometric embedding of any Riemannian manifold into a Euclidean space with a plain geometry.
Theorem 5.2 (Nash’s embedding theorem, [46, 47, 26]) .
Any d -dimensional Riemannian manifoldhas an isometric embedding in R D for D = max { d ( d + 5) / 2 , d ( d + 3) / 2 + 5 } .
By Nash’s embedding theorem, we have a map F : C !
R D that isometrically embeds ( C, r 2 R )into ( R D , I D ) .
Denote M = F ( C ) , i.e., the embedding of C in R D .
We denote the inverse of F by
G : M ! R d , which satisﬁes the following (see Lemma D.1 and its proof):
[Equation]
for all w 2 C .
This G is exactly the commuting parametrization that we seek for.
Theorem 5.3.
For any C 2 mirror map R : R d ! R [ { 1 } , consider the corresponding mirror ﬂow:
[Equation]
There exists a commuting parametrization G : M ! R d such that for any initialization of the mirrorﬂow w (0)
= w init , it holds that w ( t ) =
G ( x ( t )) for all t where x ( t ) is given by the gradient ﬂowgoverned by r x L ( G ( x )) initialized at the same x init , i.e.,
[Equation]
Proof of Theorem 5.3.
We follow the notation used in the previous discussion in this section andconsider a more general case where we have a time-dependent loss L · 2 L .
Given initialization w init , 319 let w ( t ) be the solution to
[Equation]
Since F is an isometric embedding, we have r 2 R ( w ) ij =
[Equation]
for all i, j 2 [ d ] ,which can be rewritten in the matrix form as r 2 R ( w ) = @
F ( w ) > @
F ( w ) .
Then we haved w ( t ) =  ( @ F ( w ( t )) > @
F ( w ( t )))  [Equation]
w
[Equation]
t [Equation] d t. Deﬁne x ( t ) =
F ( w ( t )) , and it follows that
[Equation]
where the third equality follows from (7).
Specializing L t ⌘ L yields the ﬁrst claim.
Next, we verify that G : M ! R d is a commuting parametrization.
First, for any x 2 M , we have@ G ( x ) > = @
F ( G ( x )) ( @ F ( G ( x )) > @
F ( G ( x )))  1 by (7). Since r 2 R ( w ) = @ F ( w ) > @
F ( w ) is ofrank d for all w 2 C , it follows that @ F ( w ) is also of rank d for all w 2 C , thus @ G ( x ) is of rank dfor all x 2 M .
Second, we show that each r G j is a complete vector ﬁeld on M .
For any x init 2 M ,consider a loss L ( w ) = h e j , w i and the corresponding gradient ﬂow is
[Equation]
>
[Equation]
L
[Equation]
G ( x ( t ))
[Equation]
= r G j ( x ( t )) , so x ( t ) =
tG j ( x init ) for all t  0 .
On the other hand, w ( t ) = G ( x ( t )) satisﬁes that
[Equation]
where the last equality follows from (7).
Therefore, rewriting the above as a Mirror Flow yields
[Equation]
to which the solution exists for all t 2 R and is given by r R ( w ( t ))
=
e j t , so w ( t ) =
( r R )  1 ( e j t ) .
This further implies that x ( t ) =
F ( w ( t )) is well-deﬁned for all t 2 R , hence r G j is a completevector ﬁeld.
Finally, to show the commutability of G , it sufﬁces to use the same time-dependent lossas in Remark 4.7.
This ﬁnishes the proof.
Conclusion and future work
We presented a framework that characterizes when gradient descent with proper paramterizationbecomes equivalent to mirror descent.
In the limit of inﬁnitesimal step size, we identify a notionnamed commuting parametrization such that any gradient ﬂow (i.e., the continuous analog of gradientdescent) with a commuting parametrization is equivalent to a mirror ﬂow (i.e., the continuous analogof mirror descent) in the original parameter space with respect to a mirror map that depends only onthe initialization and the parametrization.
Conversely, we use Nash’s embedding theorem to show thatany mirror ﬂow can be characterized by a gradient ﬂow in the reparametrized space with a commutingparametrization.
Using our framework, we recover and generalize results on the implicit bias ofgradient descent in a series of existing works, including a rigorous and general proof of convergence.
However, it remains open whether gradient ﬂow with a non-commuting parametrization can also bewritten as a mirror ﬂow, which, in other words, requires a precise necessary condition.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section X.•
Did you include the license to the code and datasets?
[No] The code and the data areproprietary.
• Did you include the license to 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 instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions 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] Thispaper only studies the theoretical properties of optimization algorithms.(d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] We explicitlyclarify the assumptions for each result.(b)
Did you include complete proofs of all theoretical results?
[Yes] Part of the proofsappears in the main context and others are deferred to appendix.
If you ran experiments...
[Table] (a) Did you include the code, data, and instructions needed to reproduce the main exper- [Table] imental results (either in the supplemental material or as a URL)?
[N/A] We do not [Table] have any experiments.
[Table]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they [Table] were chosen)?
[N/A]
[Table]
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
[Table]
ments multiple times)?
[N/A]
[Table]
(d) Did you include the total amount of compute and the type of resources used (e.g., type
[Table]
of GPUs, internal cluster, or cloud provider)?
[N/A]
[Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[N/A]
We do not use any
[Table]
existing assets.
[Table]
(b) Did you mention the license of the assets?
[N/A]
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A]
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
We do not use crowdsourcing nor conduct research with human
[Table]
subjects.
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to