Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descent
Anonymous Author(s)
AffiliationAddress 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 flow 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 flowwith 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 [57, 39, 38].
Many existing results about implicit bias treat training (in the limit of infinitesimal step size) asa differential equation or process { x ( t ) } t {+≥ 
 ⊂+} 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 {+≥ 
 ⊂+} 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; specifically, 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 flow in the x -space is given by [Equation]
[Equation]
where ∇ x denotes the gradient with respect to x .
Using w ( t ) =
G ( x ( t )) and the fact that∇ x L ( G ( x ))
=
∂G ( x ) ⊤ ∇ L ( G ( x )) where ∂G ( x ) ∈ 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 [22, 54, 60, 4, 57, 5, 8],which give examples suggesting that gradient flow in the x -space could end up simulating a moreclassical algorithm, mirror descent (specifically, the continuous analog, mirror flow) in the wspace.
Recall that mirror flow is continuous-time limit of the classical mirror descent, writtenas d ∇ R ( w ( t )) =
−∇ L ( w ( t ))
d t where R : R d → R is a strictly convex function called mirrormap [46, 10].
Equivalently it is Riemannian gradient flow with metric tensor ∇ 2 R , an old notion ingeometry:
[Equation]
If ∂G ( x ( t )) ∂G ( x ( t ))
⊤ = ∇ 2 R ( w ( t )) − 1 for all t , then (2) becomes a simple mirror flow in thew -space as noted in [23].
Many existing results about implicit bias indeed concern reparametrizations
G that satisfy ∂G ( x ) ∂G ( x ) ⊤ = ∇ 2 R ( w ) − 1 for a stricly convex function R , and the implicitbias/regularization is demonstrated by showing that the convergence point satisfies 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 ∈ [ d ] , so here D = d .
In this case, the mirror map R must satisfy ( ∇ 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 ) =
(cid:80) di
=1 w i (ln w i − 1) .
But in general, it is hard to decide whether gradient flow for a given parametrization G can be writtenas mirror flow 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 flow, then the desired mirror map cannotexist.
However, if only one of such x can be reached by gradient flow, we must decide which x it isin order to decide the value of ∇ 2 R using ∂G∂G ⊤ . Conversely, [5] raises the following question:for what mirror map R
can the corresponding mirror flow be the result of gradient flow 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 fields we identify a notionof when a parametrization w = G ( x ) is commuting (Definition 4.1) and use it to give asufficient condition (Theorem 4.4) of when the gradient flow in the x -space governed by−∇ x L ( G ( x )) is simulating a mirror flow in the w -space with respect to some mirror mapR : R d → R , which encompasses all the previous results [22, 54, 60, 4, 57, 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 linear regression as implications of the above characterization (Corollary 4.12).
We also givenew convergence analysis in such settings (Corollary 4.10), filling the gap in previousworks [22, 57, 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 flow in the w -space with respect to some mirror map R simulates a commutative gradient flow 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 affirmative and fully generalanswer to the question of when such reparametrization functions exist, giving a full answer to 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 classification tasks with separable datawas studied in [52, 20, 43, 32, 42, 31] and for non-separable data in [29, 30], 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 [57, 20, 60, 54, 4, 5], later generalized in [8].
The sharp contrast between the so-called kernel and rich regimes [57] reflects the importance of theinitialization scale, where a large initialization often leads to the kernel regime with features barelychanging during training [26, 14, 18, 17, 2, 1, 61, 7, 58, 27], while with a small initialization, thesolution exhibits richer behavior with the resulting model having lower complexity [21, 22, 36, 49,6, 13, 38, 40, 41, 50, 53, 19].
Recently [59] 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 [37, 11, 25, 39, 16, 62] and adaptive and momentum-based methods [48, 56, 55, 33],to name a few.
Understanding mirror descent.
In the continuous-time limit as step size goes to 0, the mirror flowis equivalent to the Riemannian gradient flow.
[23] 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 flow can be reparametrized as a gradient flow.
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 ∈ R D , we denote itsi -th coordinate by u i .
For any vector u, v ∈ R D and α ∈ R , we define 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 ∪ {∞} , we denote its domainby dom R = { w ∈ R d
| R ( w ) < ∞} .
For any set S , we use int( S ) to denote its interior.
For anysubmanifold M ⊂ R D and x ∈ M , we denote by T x ( M ) the tangent space of M at x .
We assume the model has parameter vector w ∈ R d and loss function L : R d → R which is a C 1function.
Training involves a reparametrized vector x ∈ 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 flow 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.
Definition 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 ∈ R , L t ( w ) is continuously differentiable in w , and∇ L t ( w ) is piece-wise continuous in ( w, t ) .
The set of all time-dependent losses is denoted by L .
3.1 Vector fields
Vector fields are a natural way to formalize continuous-time gradient descent (a good referenceis [34]).
Given any smooth submanifold of R D , denoted by M , a vector field X on M is a continuousmap from M to R D such that for any x ∈ M , X ( x ) is in the tangent space of M at x , T x ( M ) .
Definition 3.2 (Complete vector fields, p.215, [34]) .
We say X is a complete vector field on M if andonly if for any initialization x init ∈ M , the differential equation d x ( t ) =
X ( x ( t ))
d t has a solutionon ( −∞ , ∞ ) with x (0) = x init .
For any differentiable function f : M → R , we denote by ∇ f its gradient vector field.
For example,in the special case when M = R D , ∇ f is simply the ordinary gradient of f .
For any x ∈ M , wedenote by ϕ tf ( x ) the point on M reached after time t by following the vector field −∇ f
starting at x :
(here we allow t to be negative and use the convention (cid:82) b a = − (cid:82) a b )
[Equation]
Now, let G : M → R d be a continuously differentiable parametrization function.
We denote its
Jacobian by ∂G ( x ) = ( ∇ G 1 ( x ) , ∇ G 2 ( x ) , . . .
, ∇ G d ( x )) ⊤ .
For any x ∈ M and µ ∈ R d , we define
[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.
Definition 3.3 (Canonical parametrizations) .
A canonical parametrization G : M → R d is acontinuously differentiable function that satisfies: (1) the domain M is a smooth submanifold of 133 R D , (2) the gradient vector fields ∇ G 1 , ∇ G 2 , . .
. , ∇ G d are all complete vector fields on M and (3)∂G ( x ) is of rank d for all x ∈ M .
Given a canonical parametrization G : M → R d and any initialization x init ∈ M , the set of pointsthat are reachable via gradient flow 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.
Definition 3.4 (Reachable sets) .
For a canonical parametrization G : M → R d and initializationx init ∈ M , we define the reachable set Ω x ( x init ; G ) as
[Equation]
L and
continuous x
[Equation]
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 flow
Below we introduce some basics notions for mirror descent [46, 10].
Definition 3.5 (Mirror maps, p.298, [12]) .
A function R : R d → R ∪ {∞} is called a mirrormap if and only if (1) R is strictly convex and differentiable on int(dom R ) ; (2) the gradient map∇ R : int(dom R ) → R d is surjective; and (3) for any sequence { w n } n ≥ 1 going to the boundary ofdom R , lim n →∞ ∥∇ R ( w n ) ∥ 2 = ∞ .
Then for initialization w 0 ∈ 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 flow:
[Equation]
Recall that given a differentiable function R , the corresponding Bregman divergence D R is defined as
[Equation]
We recall a well-known implicit bias result for mirror flow (which holds for mirror descent as well),which shows that for a specific type of loss, if mirror flow converges to some optimal solution, thenthe convergence point achieves some implicit regularization among all optimal solutions.
−∇ G i t i −∇
G j t j −∇ G j t j −∇ 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 ◦ G i ( x )
Theorem 3.6.
Given any data Z ∈ R n × d and the corresponding label Y ∈ R n , suppose theloss L ( w ) is in the form of L ( w )
= (cid:101) L ( Zw ) for some differentiable (cid:101) L : R n → R .
Assume thatinitialized at w 0 , the mirror flow (6) converges and the convergence point w ∞ =
lim t →∞ w ( t )satisfies Zw ∞
=
Y , then D R ( w ∞ , 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 flow is a mirror flow
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.
Definition 4.1 (Commuting parametrizations) .
A commuting parametrization G : M → R d is acanonical parametrization such that the gradient vector fields ∇ G 1 , ∇ G 2 , . .
. , ∇ G d on M satisfythat for any i, j ∈ [ d ] and any t i , t j ∈ R , ϕ t
i G i ◦ ϕ t j G j ( x ) =
ϕ t
j G j ◦ ϕ t i G i ( x ) for all x ∈ M .
When G is a commuting parametrization, we call the corresponding gradient flow (1) a commutativegradient flow .
We illustrate the concept of commuting parametrization in Figure 1.
The meaning of commuting in the above definition is built upon the flows generated by the vectorfields, which is further equivalent to that the Lie bracket of any pair of the vector fields satisfies[ ∇ G i , ∇ G j ] ≡ 0 .
In the latter sense, we say ∇ G i and ∇ G j are commuting vector fields (seep.231, [34]).
Intuitively, the Lie bracket [ ∇ G i , ∇ G j ] characterizes the change of ∇ G j along theflow generated by ∇ G i and [ ∇ G i , ∇ G j ] ≡ 0 means ∇ 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, [34]) .
Given a differentiable map G : M → R d where
M is a smooth submanifold of R D , for any i, j ∈ [ d ] , ∇ G i and ∇ G j are commuting vector fields,i.e., [ ∇ G i , ∇ G j ] ≡ 0 , if and only if any t i , t j ∈ R , ϕ t
i G i ◦ ϕ t
j G j ( x ) =
ϕ t
j G j ◦ ϕ t i G i ( x ) for all x ∈ M .
In particular, when ∇ G 1 , ∇ G 2 , . . .
, ∇ G d are ordinary gradient vectors in R D , the Lie bracketbetween any ∇ G i and ∇ G j is given by [ ∇ G i , ∇ G j ]( x )
= ∇ 2 G i ( x ) ∇ G j ( x ) − ∇ 2 G j ( x ) ∇ G
i ( x ) .
The above condition clearly holds when each G i only depends on a different subset of coordinatesof x , because we have ∇ 2 G i ( x ) ∇ G j ( x ) ≡ 0 , for all i ̸ = j ∈ [ d ] and x ∈ M as ∇ 2 G i and∇
G j live in different spaces.
We call this setting separable case and it covers all the previousexamples [22, 54, 4, 57, 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)
[57] where w ∈ R d is parametrized byw
=
u ⊙ 2 − v ⊙ 2 .
In this case, we set D = 2 d and M = { x = (cid:0) uv (cid:1) ∈ R D | u i , v i > 0 , ∀ i ∈ [ d ] } , and
G : M → R d is given by G ( x ) =
u ⊙ 2 − v ⊙ 2 .
Since each G i involves u i , v i , this is a separable caseand thus ∇ G 1 , ∇ G 2 , . .
. , ∇ G d are commuting vector fields.
Next, we present our result on characterizing any commutative gradient flow by a mirror flow.
Theorem 4.4.
Let G : M → R d be a commuting parametrization on some smooth submanifold
M ⊆ R D as defined in Definition 4.1.
For any C 1 loss function L : R d → R and any initializationx init ∈ M , let x ( t ) be given by the commutative gradient flow governed by −∇ x L ( G ( x )) :
[Equation]
Denote w ( t ) =
G ( x ( t )) , then there exists some mirror map R such that the dynamics of w ( t ) is amirror flow with respect to R , i.e.,
[Equation]
Moreover, this mirror map R only depends on the initialization x init and the commuting parametrization G , and is independent of the loss function L .
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 gradientflow 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 ∈ M , let the solution to the gradient flow governed by −∇ x L t ( G ( x )) at time t be x ( t ) , i.e., x ( t ) =
x init + (cid:82) t 0 −∇ x L t ( G ( x ( s )))
d s .
Defineµ ( t ) =
(cid:82) t {+
+}
−∇ 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 find the mirror map.
Lemma 4.6.
Let G : M → R d be a commuting parametrization.
For any initialization x init ∈ M , there exists a twice differentiable strictly convex function Q : R d → R such that ∇ Q ( µ ) =
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 ∇ 2 R ( G ( ψ ( x init ; µ )))
= ( ∂G ( ψ ( x init ; µ )) ∂G ( ψ ( x init ; µ ))
⊤ ) − 1 forall µ ∈ R d .
Now, we are ready to prove Theorem 4.4.
Proof of Theorem 4.4.
Recall that the gradient flow in the x -space governed by −∇ 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 flow satisfies x ( t ) =
ψ ( x init ; µ ( t ))where µ ( t ) =
(cid:82) t {+
+}
−∇ 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 ) satisfies d w ( t ) =
−∇ 2 R ( w ( t )) − 1 ∇ L ( w ( t ))
d t whichis exactly the mirror flow 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 finishes the proof.
Remark 4.7 (Non-commuting flow cannot be mirror flow, 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 · ∈ L , gradient flow governed by −∇ x L t ( G ( x )) can be written as a mirror flow, then G mustbe commuting parametrization.
To see why, let δ > 0 and consider a time-dependent loss given by
[Equation]
t
[Equation]
−⟨ e
[Equation]
w ⟩ 2 δ ≤ t < 3 δ
[Equation]
−⟨ e j
[Equation]
w ⟩ 3 δ ≤ t < 4 δ [Equation] t ≥ 4 δ. Suppose for any initialization x init ∈ M , d x ( t ) =
−∇ x L t ( G ( x ( t ))
d t with x (0) =
x init can be writtenas d ∇ R ( w ( t )) =
−∇ L t ( w ( t ))
d t for some mirror map R where w ( t ) =
G ( x ( t )) .
Then ∇ R (4 δ ))
=∇ R ( w (0))
+ (cid:82) t 0 −∇ L s ( w ( s ))
d s = ∇ 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 [ −∇ G j , −∇ G i ]( x init )
+ O ( δ 3 ) , whichimplies [ −∇ G j , −∇ G i ]( x init )
∈ T x init (Ω x ( x init ; G )) .
Since dim(Ω x ( x init ; G ))
=
d and G is canonical, {∇ G k ( x init ) } dk =1 form a basis of T x init (Ω x ( x init ; G )) .
Note that
= lim t [-! 
-]→ 0 G k ( x (4 t ))
− G k ( x init ) t=⟨∇ G i ( x init ) , [ −∇ G j , −∇ G i ]( x init ) ⟩ for all k [-
-]∈ [ d ] , we conclude [ −∇ G j , −∇ G i ]( x init ) must be 0 .
Thus by Theorem 4.2, we conclude G is a commuting parametrization since x init is arbitrary.
Next, we establish the convergence of w ( t ) =
G ( x ( t )) when x ( t ) is given by some commutativegradient flow.
Here we require that the mirror map R given by Theorem 4.4 is a Bregman function(see definition 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 and∇ L is locally Lipschitz.
If the strictly convex function R given by Theorem 4.4 is a Bregman functionsatisfying argmin { L ( w ) | w ∈ dom R } is non-empty, then as t → ∞ , w ( t ) converges to some w ∗such that ∇ L ( w ∗ ) ⊤ ( w − w ∗ ) ≥ 0 for all w ∈ 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 satisfiesΩ 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 ∈ [ d ]where A i ∈ R D × D are symmetric matrices and commute with each other
, i.e. , A i A j − A j A
i =0 , ∀ i, j ∈ [ 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 define Z =
( z 1 , z 2 , . . .
, z n ) ∈ R d × n .
We can apply Theorem 3.6 to obtain the implicit bias of a commutative gradient flow.
Theorem 4.11.
Suppose G : M → R d is a commuting parametrization satisfying Definition 4.1and the loss function L satisfies L ( w ) = (cid:101)
L ( Zw ) for some differentiable (cid:101) L : R n → R .
Givenany initialization x init ∈ M , consider the gradient flow d x ( t ) =
−∇ x L ( G ( x ( t ))
d t initialized atx (0) = x init .
Then there exists a convex function R (depending on x init and the parametrization G ),such that for any dataset { ( z i , y i ) } ni =1 ⊂ R d × R , if w ( t ) =
G ( x ( t )) converges and the convergencepoint w ∞ =
lim t →∞ w ( t ) satisfies Zw ∞
= Y , then R ( w ∞ ) =
min w : Zw =
Y R ( w ) , that is, gradientflow implicitly minimizes the convex regularizer R among all interpolating solutions.
Proof of Theorem 4.11.
By Theorem 4.4, w ( t ) obeys the following mirror flow:
[Equation]
Applying Theorem 3.6 yields D R ( w ∞ , G ( x init ))
=
min w : Zw =
Y D R ( w, G ( x init )) .
Therefore, for anyw ∈ 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 ( G ( x init ))
= ∇ R ( G ( ψ ( x init ; 0)))
= ∇ R ( ∇ Q (0))
=
wherethe last equality follows from the property of convex conjugate.
Combining this with the aboveinequality, we get R ( w ∞ ) ≤ R ( w ) for all w ∈ 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 / ∈ dom R , we have
R ( w ) = ∞ by definition, so R ( w ∞ ) ≤ R ( w ) holds trivially.
This finishes the proof.
Note that the identity parametrization w =
G ( x ) = x is a commuting parametrization.
Therefore, ifwe run the ordinary gradient flow 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 flow for underdetermined regression.
Furthermore, we can recover the results on the quadratically overparametrized linear model in aseries of paper [22, 57, 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 [47, 39].
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 =
(cid:0) uv (cid:1) ∈ R 2 d .
Suppose the loss function L is quasi-convex, ∇ L is locally Lipschitz and
L ( w ) = (cid:101) L ( Zw ) for some differentiable (cid:101) L : R n → R .
For any initialization x init =
(cid:0) u 0 v 0 (cid:1) such that allcoordinates of x init are non-zero, let x ( t ) be the solution to the gradient flow (1) at time t .
Then ast → ∞ , w ( t ) =
G ( x ( t )) converges to some w ∞ with Zw ∞ =
Y , and R ( w ∞ ) =
min w : Zw =
Y R ( w ) ,where the mirror map R is given by
[Equation]
Every mirror flow is a commutative gradient flow
Now we consider any C 2 mirror map R : R d → R ∪ {∞} .
Recall the corresponding mirror flow:
[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 [34]).
Then ∇ 2 R is a continuous positive-definite metric on C .
Asdiscussed previously, the above mirror flow can be further rewritten as the Riemannian gradient flowon the Riemannian manifold ( C, ∇ 2 R ) , i.e.,
[Equation]
The goal is to find a commuting parametrization G : M → R d where M is smooth submanifold of 291 R D such that the dynamics of w ( t ) =
G ( x ( t )) is given by the commutative gradient flow on x ( t )governed by −∇ x L ( G ( x )) .
To illustrate the idea, let us first suppose such a G exists and is a bijection, whose inverse is denotedby F .
Then it satisfies that ∇ 2 R ( w ) =
( ∂G ( F ( w ))
∂G ( F ( w ))
⊤ ) − 1 , which further suggests that∇ 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 , ∇ 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.
Definition 5.1 (Isometric embedding, p.15, [35]) .
Let ( M, g ) and ( (cid:102) M, ˜ g ) be Riemannian manifolds.
An isometric embedding is an embedding F : M → (cid:102) M which preserves the metric in the sense thatg is equal to the pullback of ˜ g by F .
More specifically, for any two tangent vectors v, w ∈ 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, [44, 45, 24]) .
Any d -dimensional Riemannian manifoldhas an isometric embedding in R D for some D ≥ d .
By Nash’s embedding theorem, we have a map
F : C → R D that isometrically embeds ( C, ∇ 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 satisfies the following (see Lemma D.1 and its proof):
[Equation]
for all w ∈ 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 ∪ {∞} , consider the corresponding mirror flow:
[Equation]
There exists a commuting parametrization G : M → R d such that for any initialization of the mirrorflow w (0) =
w init , it holds that w ( t ) =
G ( x ( t )) for all t where x ( t ) is given by the gradient flowgoverned by −∇ x L ( G ( x )) initialized at some 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 · ∈ L .
Given initialization w init ,let w ( t ) be the solution to
[Equation]
Since F is an isometric embedding, we have ∇ 2 R ( w ) ij = (cid:80)
[Equation]
=1 ∂F k
[Equation]
w ) ∂w i ∂F k
[Equation]
w
[Equation]
∂w j for all i, j ∈ [ d ] ,which can be rewritten in the matrix form as ∇ 2 R ( w ) = ∂F ( w ) ⊤ ∂F ( w ) . Then we have
[Equation]
Define x ( t ) =
F ( w ( t )) , and it follows that
[Equation]
where the third equality follows from (7).
Specializing L t ≡ L yields the first claim.
Next, we verify that G : M → R d is a commuting parametrization.
First, for any x ∈ M , we have∂G ( x ) ⊤ = ∂F ( G ( x ))
( ∂F ( G ( x )) ⊤ ∂F ( G ( x )))
− 1 by (7).
Since ∇ 2 R ( w ) = ∂F ( w ) ⊤ ∂F ( w ) is ofrank d for all w ∈ C , it follows that ∂F ( w ) is also of rank d for all w ∈ C , thus ∂G ( x ) is of rank dfor all x ∈ M .
Second, we show that each ∇ G j is a complete vector field on M .
For any x init ∈ M ,consider a loss L ( w ) = ⟨ e j , w ⟩ and the corresponding gradient flow is
[Equation]
x [Equation]
[Equation]
⊤
[Equation]
L ( G ( x ( t ))) d
[Equation]
= −∇ G j ( x ( t )) , so x ( t ) =
ϕ tG j ( x init ) for all t ≥ 0 .
On the other hand, w ( t ) = G ( x ( t )) satisfies 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 ∈ R and is given by ∇ R ( w ( t ))
=
e j t , so w ( t ) =
( ∇ R ) − 1 ( e j t ) .
This further implies that x ( t ) =
F ( w ( t )) is well-defined for all t ∈ R , hence ∇ G j is a completevector field.
Finally, to show the commutability of G , it suffices to use the same time-dependent lossas in Remark 4.7.
This finishes 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 infinitesimal step size, we identify a notionnamed commuting parametrization such that any gradient flow (i.e., the continuous analog of gradientdescent) with a commuting parametrization is equivalent to a mirror flow (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 flow can be characterized by a gradient flow 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 flow with a non-commuting parametrization can also bewritten as a mirror flow, 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 justification 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 are proprietary.
• 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 theChecklist 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 reflect 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] 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]
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.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[N/A] We do nothave any experiments.(b)
Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[N/A]
(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 identifiable
[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, if [Table] applicable?
[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 participants and the total amount [Table] spent on participant compensation? [N/A]
[Table]