On the Role of Momentum in the Implicit Bias of Gradient Descent for Diagonal Linear Networks

Reference

[1] Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang. Implicit regularization in heavyball momentum accelerated stochastic gradient descent. ICLR, 2023.

[2] Barrett & Dherin. Implicit Gradient Regularization.

[3] Kovachki & Stuart. Continuous Time Analysis of Momentum Methods, 2021.

[4] Wilson et al. A Lyapunov analysis of momentum methods in optimization, 2016.

[5] Shi et al. Understanding the acceleration phe- nomenon via high-resolution differential equations, 2018.

[6] Rumelhart et al. Learning internal representations by error propagation, 1986.

(3). On "Major technical concern"

We did not check the proof techniques of Theorem 4 in Kovachki and Stuart (2021) when proving Proposition 1, and we thank the reviewer for pointing this similarity out. We fix the corresponding statement in the revision (page 4 of the revision).

We thank the reviewer for the suggestion regarding the importance of stating the deviation of continuous counterpart from the true discrete trajectory. In the revision, we clearly state the $\mathcal{O}(\eta)$ approximate order for the continuous version and fix the corresponding description from "implicit bias of HB and NAG" to "implicit bias of the $\mathcal{O}(\eta)$ approximate continuous version of HB and NAG", or more concretely HB and NAG flow.

• Regarding Question (A): We fix the corresponding description following the suggestion of the reviewer in the revision.

• Regarding Question (B):

  1. $\mathcal{O}(\eta^2)$ continuous approximation.

     The reason of adopting the $\mathcal{O}(\eta)$ approximate continuous version of momentum-based methods is because we aim to analyze the role of momentum in the implicit bias of widely-studied gradient flow, which is also an order $\mathcal{O}(\eta)$ approximate continuous version of GD, for diagonal linear networks, and the analysis should be conducted in the same order of approximation.

     For the $\mathcal{O}(\eta^2)$ approximate continuous version, based on the results of [1], we believe that our results could be generalized to the $\mathcal{O}(\eta^2)$ case following exactly the same approach as in the current work, by replacing the current $\mathcal{O}(\eta)$ approximate second-order ODE with the $\mathcal{O}(\eta^2)$ approximate version of HB in [1]. We thank the reviewer for the suggestion of $\mathcal{O}(\eta^2)$ continuous approximation and we discuss this in the revision and leave such generalization to future works.

  2. The use of second-order ODE.

     As pointed out by Kovachki and Stuart (2021), the second-order ODE, though being the $\mathcal{O}(\eta)$ approximate continuous version of HB and NAG, contains considerably more qualitative information about the dynamics. In addition, the second-order ODE is also a common choice in the literature for studying momentum-based methods, e.g., Wilson et al. (2016). Therefore, we choose the second-order ODE.

(4). On observations from section 3.3.

As discussed in our paper, this phenomenon has been observed by, for example, Kovachki and Stuart (2021) and Rumelhart et al. (1986). Section 3.3 mentioned the observations to show that our theoretical results and numerical results match previous ones and further reveal how the implicit bias of momentum-based methods depend on the hyper-parameters. We also state clearly in the revision that [1] also made such kind of observations.

(5). On insights in terms of convergence and saddle to saddle jump issue.

On one hand, whether momentum-based methods indeed accelerate convergence for diagonal linear networks (which is a non-convex problem) is still not clear and the re-scaled gradient perspective suggested by the reviewer is not sufficient since one can simply choose an effective learning rate for GD. To the best of our knowledge, acceleration effects of momentum-based methods need careful and delicate analysis and many works are for convex problems, e.g., Shi et al. (2018) and Wilson et al. (2016).

On the other hand, as shown in our previous response, momentum-based methods do not always return solutions with better generalization performance than GD.

Therefore, we leave such discussion for future works, since also our focus is the precise characterization of the role of momentum in the implicit bias of gradient flow.

Questions Part

1. Please see our response to point Weakness 3 above.

2. As discussed in our response to 1 and 2 of the Weakness part above, our work precisely characterizes the implicit bias of momentum-based methods and the corresponding dependence on the initialization of parameters and gradients, and further reveals the conclusion that whether momentum-based methods find solutions with better generalization performance than GD depends on the two-fold effect of momentum. These can not be covered only by the analysis in [1] since it is model-agnostic and does not consider other important sources of implicit bias. Thus it is incomplete, or inconsistent with the numerical experiments, to study the role of momentum on the implicit bias of GD solely based on [1], e.g., it is unclear how the initialization of parameters and gradients can affect the implicit bias of momentum-based methods. We believe that, though being a general analysis, only knowing the preference of flatter trajectory of momentum-based methods than GD ([1]) does not solve the implicit bias of momentum-based methods once and for all.

We thank the reviewer for the detailed comments and valuable suggestions. Below we address your concerns in detail.

Weaknesses part

(1). On "Major novelty issues".

We were not aware of the nice related work [1] and we thank the reviewer for the suggestion of it. We include the discussion of this work in the revision and we fix the inaccurate sentence "which is still missing in current literature".

After carefully reading [1], we would like to clarify that we disagree with the reviewer on the aspect of novelty issues. The main reason is that the analysis in [1] is model-agnostic and does not consider other important sources that affect the implicit bias such as model architectures (at least incompletely by only utilizing the calculation of gradients) and initialization, while our work focuses on precisely characterizing the implicit bias of momentum-based methods and its explicit dependence on the architecture and the initialization of both parameters and gradients, which can not be captured solely by the analysis in [1] and support our novelty.

Furthermore, indeed [1] studied HB in a general setting, however, as discussed in the IGR paper [2], only relying on such general analysis is not complete for analyzing the implicit bias. For an instance, the solution of the ODE $\dot{\theta} = -\nabla L(\theta)$ (according to [1, 2], both momentum-based methods and GD can have this form) certainly depends on the initial conditions ––– the initialization of $\theta$ and perhaps the initialization of gradients when considering momentum-based methods. We would like to highlight that overlooking these important sources of implicit bias and directly extending [1] to a more specific setting would lead to an incomplete conclusion that may even contradict the results of numerical experiments (Fig.1), as we will explain in the next point for the case of diagonal linear networks.

(2). On "extension of [1] to diagonal linear networks".

We appreciate the reviewer for providing this interesting extension.

In our opinion, however, the extension of the results in [1] to diagonal linear networks is incomplete for the following reasons:

• It can not capture the effects of the initialization of both parameters and gradients on the implicit bias of momentum-based methods, which are important since the transition from kernel to rich regime of diagonal linear networks depends on the initialization of parameters and initialization of gradients are crucial for momentum-based methods, as discussed in our response to the point 1.
• To the best of our knowledge, the derived form of weight-decay in the review (which is an $\ell_2$ norm) does not provably promote sparseness of model parameters $u$ and $v$ (which should be described by $\ell_1$ norms). In addition, we would also like to mention that both GD and momentum-based methods can return $\ell_2$-norm solution, which is inconsistent with the argument of the reviewer that the implicit regularization effects $||\nabla L (w)||^2$ of momentum-based methods (the argument can also be applied to GD which also has similar regularization effects) promote sparseness.
• It only partly considers the model architecture and can not characterize the exact differences between the properties of solutions for momentum-based methods and that of GD.

As a consequence, the conclusion that "sparser solutions correspond to ... to flatter minima[s]" is not always true. In fact, based on our results in Theorem 1, compared to GD, whether momentum-based methods induce solutions that have better generalization performances for sparse regression depends on the two-fold role of momentum in the implicit bias of GD: the initialization mitigation effect and the explicit dependence on the initialization of gradients.

For example, when the initialization is unbiased (Definition 1 in page 5) and the initialization of gradient is not negligible, momentum-based methods may have worse generalization performance than GD, as shown in Fig.1 where $D ( \theta(\infty), \theta^* ) = || \theta(\infty) - \theta^* ||^2_2$ is larger when the diagonal linear network is trained by momentum-based methods than by GD for certain range of initialization scale $\|\xi\|_1$. Another example is Fig.2(b) where advantages of momentum-based methods over GD on the generalization disappear as the initialization becomes less biased due to the fact that the dependence of solutions of momentum-based methods on the initialization of gradients outperforms the initialization mitigation effects in such case.

These numerical experiments are not consistent with the conclusion of the simple extension of [1], indicating that it is incomplete to characterize the implicit bias of momentum-based methods solely based on analysis of [1]. Therefore, we believe that such inconsistency between numerical experiments and the incomplete extension of [1] actually support the novelty and significance of our work.

We thank the reviewer so much for the further comment and insightful questions. Below we address your concerns.

On point 1 of the comment

We appreciate the reviewer for the argument and we would like to clarify that we were aware of that higher order approximation of GD depends on $\eta$.

We would like to first clarify that, in our response, we did not entirely oppose the reviewer that the $\mathcal{O}(\eta^2)$ continuous approximation for HB is a justifiable choice. In our view, if it is compared to the $\mathcal{O}(\eta^2)$ continuous for GD, e.g., the modified flow in the IGR paper [2], we agree with the reviewer that it is a justifiable choice to characterize the role of momentum on the implicit bias for diagonal linear networks, and it is indeed a promising future direction to discuss this choice.

In this sense, we would like to still emphasize that the same order of approximations for both momentum-based methods and GD should be used to make a "fair" comparison on their implicit bias for diagonal linear networks. For an instance, the generalization of the current work (following almost exactly the same technique as in the current one) to $\mathcal{O}(\eta^2)$ approximate version of HB using results of [1] should compare with the generalization of previous gradient flow results for diagonal linear networks [3, 4] to $\mathcal{O}(\eta^2)$ approximate version of GD, not the usual $\mathcal{O}(\eta)$-close version gradient flow, using the modified gradient flow presented in the IGR paper [2]. This is similar to the spirit in [1] where the authors of [1] compared $\mathcal{O}(\eta^2)$ continuous approximation of HB with $\mathcal{O}(\eta^2)$ version of GD (the modified flow), rather than the usual $\mathcal{O}(\eta)$-close version gradient flow.

Therefore, we argue that the main reason why we use the $\mathcal{O}(\eta)$ continuous approximation for HB (and NAG) is that, in this paper, we aim to precisely characterize the the role of momentum in the implicit bias of the widely-studied gradient flow, which is the $\mathcal{O}(\eta)$ continuous approximate version of GD, for diagonal linear networks. And one can naturally generalize the current work to the $\mathcal{O}(\eta^2)$ case following similar techniques, which is a promising future direction.

On point 2 of the comment

We thank the reviewer for the clarification "I am mentioning sparse ... $u$ and $v$" of the argument in the original comment. And we agree with the reviewer that the derived form of the weight-decay might promote sparseness of $\theta$ ($s$ in the comment of the reviewer).

We would like to kindly clarify that the reason why we mentioned "GD and momentum-based ... $\ell_2$-norm solution" is that only using the model-agnostic analysis for diagonal linear networks, e.g., IGR in [2] and IGR-M in [1], might not capture the transition from rich regime (sparse solution) to kernel regime by tuning the initialization of parameters (perhaps also gradients when considering momentum-based methods), given the argument that the re-scaled weight-decay promotes sparseness of $\theta$. This indicates the importance of studying other sources of implicit bias simultaneously and novelty and significance of our results.

Reference

[1] Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang. Implicit regularization in heavyball momentum accelerated stochastic gradient descent. ICLR, 2023.

[2] Barrett & Dherin. Implicit Gradient Regularization.

[3] Azulay et al. On the implicit bias of initialization shape: beyond infinitesimal mirror descent.

[4] Woodworth et al. Kernel and rich regimes in overparametrized models.

4. On Fig.1

   We would like to clarify that whether HB/NAG induce solutions with better generalization performance than GD for sparse regression depends on the two-fold role of momentum on the implicit bias: the initialization mitigation effect (beneficial for generalization), which is due to the integral in Eq.(9) on page 6, and the extra explicit dependence on the initialization of gradients (possibly harmful for generalization), which is due to the term of the order $\eta$ in $\mathcal{R}$ in Eq.(8) on page 6.

   When the dependence on the initialization of gradients outperforms the initialization mitigation effect, it is possible that solutions of HB/NAG would generalize worse than that of GD, which is the case of Fig.1(a) when $\theta(0) = 0$ and $\mathcal{R}^{GF} = 0$. On the other hand, as we decrease the scale of the initialization of gradients, solutions of HB/NAG would generalize better than that of GD, as shown in Fig.1(c). Therefore Fig.1 does not contradict the theory.

   Furthermore, when the initialization is biased, both $\mathcal{R}$ of HB/NAG and $\mathcal{R}^{GF}$ are nonzero. In this case, as we increase the relative significance of the initialization of gradients as shown in Fig. 2(b) on page 8, the advantages of HB/NAG over GD about the generalization performance gradually disappear, which supports the theory.

Reference

Gunasekar et al. Characterizing Implicit Bias in Terms of Optimization Geometry.

Azulay et al. On the implicit bias of initialization shape: beyond infinitesimal mirror descent.

Woodworth et al. Kernel and rich regimes in overparametrized models.

Shi et al. Understanding the acceleration phe- nomenon via high-resolution differential equations.

Wang et al. Does Momentum Change the Implicit Regularization on Separable Data?

Jelassi et al. Towards understanding how momentum improves generalization in deep learning.

We thank the reviewer so much for the valuable comments and interesting questions. Blow we answer your questions.

Weaknesses Part

1. On missing related works.

   We greatly appreciate the reviewer for suggesting these related works. Below we discuss these related works.

   • Comparison to Gunasekar et al. (2018) and Wang et al. (2022). We would like to clarify that we discussed Gunasekar et al. (2018) in page 2 (paragraph below the question Q) where we mentioned that Gunasekar et al. (2018) revealed that there is no difference between the implicit bias of momentum-based methods and that of GD for linear regression problem. In addition, Wang et al. (2022) studied the linear classification problem and showed that momentum-based methods converge to the same max-margin solution as GD for single-layer linear networks, i.e., they share the same implicit bias. These works confirmed that momentum-based methods does not enjoy possible better generalization performance than GD for single-layer models. Compared to these works, our results reveal that momentum-based methods will have different implicit bias when compared to GD for diagonal linear networks, a deep learning models, indicating the importance of the over-parameterization on the implicit bias of momentum-based methods.

   • Comparison to Jelassi et al. (2022). Jelassi et al. (2022) studied classification problem and also showed that momentum-based methods improve generalization of a linear CNN model partly due to the historical gradients. The setting of our work is different from that of Jelassi et al. (2022): our work focuses on regression problems and diagonal linear networks. In addition, there are also differences between the conclusion of our work and that of Jelassi et al. (2022), in the sense that we conclude that momentum-based methods does not always lead to solutions with better generalization performance than GD, which depends on whether the initialization mitigation effect of momentum-based methods (interestingly this effect can also be regarded as coming from the historical gradients as Jelassi et al. (2022)) outperforms their extra dependence on initialization of gradients. Therefore, the momentum-based method is not always a better choice than GD.

   We add the discussion of these related works in the revision (Page 2, 3, and 12 in the revision) as suggested by the reviewer.

2. On difference between GD and HB/NAG.

   Indeed, GD contains the same range of implicit bias as HB/NAG. We would like to first clarify that HB/NAG have better implicit bias than GD when their initialization mitigation effects outperform their extra dependence on the initialization of gradients for diagonal linear networks.

   And if this condition is satisfied, when stating that HB/NAG have better implicit bias than GD, we mean that HB/NAG converge to solutions with better generalization performance for sparse regression (since the solutions of momentum-based methods are closer to the $\ell_1$-norm solutions due to the effective initialization mitigation effect) than GD under exactly the same initialization and architecture, i.e., if all the other conditions are the same, then solutions returned by HB/NAG generalize better than GD for sparse regression. As an example of indicating such difference between HB/NAG and GD, Wang et al. (2022) showed that momentum-based methods converge to the same max-margin solution as GD for linear classification problem (thus they have same implicit bias), while this is not the case for regression problem with diagonal linear networks --- momentum-based methods converge to different solutions compared to GD, leading to the conclusion that they have different implicit bias.

3. On setting $\mu = 0$ in page 7.

   We would like to clarify that $\mathcal{R}$ in page 7, which is derived under the unbiased initialization (Definition 1 on page 5), is only for momentum-based methods since it is of the order of $\eta$ which does not exist for GF where no term of the order $\eta$ appears, i.e., $\mathcal{R}^{GF} = 0$ under unbiased initialization $\theta(0) = 0$ (please see Eq.(7) in page 5, or Azulay et al. (2021); Woodworth et al. (2020)).

   The question raised in the review when setting $\mu = 0$ is because only doing so is not sufficient and one should set $\mu = 0$ and neglect all terms of the order $\mathcal{O}(\eta)$ simultaneously to back to the GF equation from the second-order momentum ODE, which leads to $\mathcal{R}$ = 0 in page 7. We would also like to kindly mention that only setting $\mu = 0$ commonly could not lead various second-order momentum ODEs to be the GF equation, e.g., various high-resolution momentum ODEs in Shi et al. (2018).

We thank the reviewer a lot for the reply. Below we address your concern.

Indeed using GD with a smaller initialization might achieve similar regularization effect as HB/NAG. The main harm of using GD with a smaller initialization is the saddle point escape issue. In particular, very small initialization scales (especially those required for leaving the kernel regime in wide models) lead to the issue that these scales correspond to the initialization close to a saddle point that might be difficult to escape. This has been discussed in, for example, Nacson et al. (2022).

In our case, to achieve good generalization, a smaller initialization is required for GD compared to HB/NAG, however, this is problematic from an optimization perspective since $u = v = 0$ is a saddle point and, as a result, it is highly likely that using smaller initialization scale leads to a larger saddle point (the vicinity of zero) escape time of GD, which may even not converge when the initialization is too small. As a comparison, HB/NAG could start from a relatively larger initialization than GD to alleviate or avoid the saddle point escape problem to achieve good generalization. This indicates the difference between GD and HB/NAG

Reference

Mor Shpigel Nacson et al. Implicit Bias of the Step Size in Linear Diagonal Neural Networks, ICML 2022.

We thank the reviewer for the valuable comments and the appreciation of our work. Below we answer your questions.

On Weaknesses Part

1. On constraints of the architecture.

   Diagonal linear networks exhibit many important and interesting properties of complex neural networks such as the transition from kernel regime to rich regime, therefore we choose this insightful model to study the momentum-based methods. Typically the analysis for standard deep linear networks is more complicated. And, based on our results, we believe the generalization to standard deep linear networks, though being harder, are also possible and it is also interesting to investigate whether the initialization mitigation effects of momentum-based methods are general across different architectures.

2. On equivalence to standard diagonal linear networks.

   We thank the reviewer for pointing this out. We add the related discussion of the equivalence between the diagonal structure used in our paper and the standard diagonal linear networks (blue fonts on page 25 in the revision).

I thank the authors for the response. My rating remains.

We thank the reviewer for the valuable comments and questions. Below we address your concerns.

Weaknesses Part

(1). On neglecting second-order terms.

We would like to clarify that neglecting the second-order terms for the ODE is eligible. This is because the momentum ODE is established to the order of $\eta$, i.e., the momentum ODE in Eq.(4) is in fact $\alpha \ddot{\beta} + \dot{\beta} + \nabla L / (1 - \mu) = \mathcal{O}(\eta^2)$. Thus all terms of order higher than $\eta$ (including the order $\eta^2$) appearing in anywhere else should also be dropped. To be consistent with this convention, the $\eta^2$ terms are ignored in our proof. For example, from Eq.(4) we have $\eta \dot{\beta} = - \eta\frac{\nabla L}{1 - \mu} + \mathcal{O}(\eta^2),$ which is eligible (it states that momentum adds a perturbation of the order $\eta$ to re-scaled gradient flow) and can also be derived from the discrete update rules by following exactly the same approach of the proof for Proposition 1 and ignoring $\eta^2$ terms in the discrete relation Eq.(14).

We mentioned this point in the proof and we thank the reviewer for pointing out that we missed this in the statement of theorems. We fix this in the revision by stating clearly in Proposition 1 that we ignore all terms of the order higher than $\eta$.

One can try to generalize the analysis to include terms of the order higher than $\eta$ by using the momentum ODE with higher order terms, e.g., Eq.(13) on page 11 of Kovachki and Stuart (2021) that includes arbitrarily higher order terms, and following our current approach where one should now preserve these higher order terms. Since most momentum-based modellings of momentum-based methods are truncated to the order of $\eta$, we also focus on this choice.

(2). On the order of the continuous approximation of HB and NAG and inaccurate statement of Theorem 1.

The continuous counterpart of HB and NAG used in our paper (Eq.(4)) is $\mathcal{O}(\eta)$ approximate continuous version of discrete HB and NAG, as shown in Kovachki and Stuart (2021). We thank the reviewer for the suggestion of inaccurate statement of Theorem 1, and we fix this in the revision by carefully and clearly stating the current implicit bias is for the $\mathcal{O}(\eta)$ approximate continuous version of HB and NAG, or more concretely HB and NAG flow, rather than directly for the discrete HB and NAG.

We adopt the $\mathcal{O}(\eta)$ approximate continuous version of momentum-based methods since we aim to analyze the role of momentum in the implicit bias of widely-studied gradient flow, which is also an order $\mathcal{O}(\eta)$ approximate continuous version of GD, for diagonal linear networks, and the analysis should be conducted in the same order of approximation. The analysis can be naturally generalized to the order $\mathcal{O}(\eta^2)$ approximate continuous version of HB by replacing the $\mathcal{O}(\eta)$ approximate continuous version of HB (Proposition 1) with the $\mathcal{O}(\eta^2)$ ones in Kovachki and Stuart (2021) (Eq.(20)) and following similar approach as in the current work.

Reference

Kovachki and Stuart. Continuous Time Analysis of Momentum Methods. arXiv: 1906.04285, 2019.