Criteria and Bias of Parameterized Linear Regression under Edge of Stability Regime

We thank the reviewer for the comments. We will try our best to explain the concerns raised by the reviewer.

Extension to SGD

Indeed, EoS phenomena can also occur with SGD [Cohen et al., 2022]. However, most works addressing EoS in stochastic methods are purely empirical, and very few theoretical results are known currently. This is due to the inherent difficulty in establishing convergence for GD on most practical models, and even more so for SGD, which introduces additional complexities. As for the two-layer diagonal neural networks, empirical evidences also suggest the occurance of EoS for SGD. Nevertheless, tackling the theoretical analsyssi remains a very challenging future task.

Intuition behind quadratic parameterization

As mentioned in our manuscripts, the study of two-layer diagonal neural networks (DLN) has been emerged as a very important topic in the past several years. This is because (1) DLNs are sufficiently non-linear to capture key properties of more practical deep neural networks, and (2) their simple form allows for a detailed analysis of the GD trajectory. As a result, the two-layer DLN has become one of the most well-studied toy models for implicit bias. As a result, it becomes one of the most well-studied toy model in terms of implicit bias. The two-layer DLN is equivalent to the linear regression problem where the interpolator/weight vector is forced to admit a quadratic parameterization as $\beta = w_{+}^2 - w_{-}^2$. Therefore, we believe it is a good candidate to study the condition of EoS and convergence of GD under EoS regime.

Difficulty in theoretical extension to multi-sample setting

We thank the reviewer for raising this important question. Indeed, we believe that extending the analysis to the $n$-sample case is a highly challenging task—not only for our model but also for related studies, such as [Ahn et al., 2023], [Song and Yun, 2023], among others. More generally, the inability to fully characterize the implicit bias of GD in more complex and more practical models remains a significant limitation of all the existing EoS results. The primary reason behind this is that EoS typically arises in highly non-linear loss landscapes. These objectives often exhibit intriguing geometric properties, making it difficult to track the trajctory of gradient-based methods. As a result, the majority of EoS analysis focused on toy models or the simplication of practical models.

When specializing to our model or the models considered in [Ahn et al., 2023] and [Song and Yun, 2023], the primary challenge in extending the analysis to the $n$-sample setting lies in the difficulty of effectively tracking the residual values across different sample points. As oscillation analysis plays a crucial role in these analysis, the coupling between residuals at different samples under the non-degenerate cases prevents us from obtaining meaningful results. The residuals influence each other through shared parameters, creating a highly complicated landscape that is difficult to decouple or simplify. Therefore, how to tackle the interaction between sample points will be a crucial step in future explorations.

We thank the reviewer for the feedbacks and would like to address the concern raised below.

Limitation of one-sample and $d=2$.

We would like to first discuss the first point. Indeed, our current technique does not allow to extend the theoretical analysis to the more general $n$-sample setting. As emphasized in our manuscript, the oscillation of the residual term $r_t$ is very crucial concept in our theoretical analysis. When multiple sample are allowed, the residuals evaluated at different data points become correlated and their mutual influences make the analysis intractable. This coupling of residuals introduces challenges that cannot be easily addressed with standard linear algebra tools, preventing us from fully characterizing the implicit bias of GD under the EoS regime in the multi-sample setting.

Nevertheless, this limitation does not only apply to our results, but also to many important previous works like [Ahn et al., 2023] and [Song and Yun, 2023], among others. More generally, the inability to fully characterize the implicit bias of GD in more complex and more practical models remains a significant limitation of all the existing EoS results. The primary reason behind this is that EoS typically arises in highly non-linear loss landscapes. These objectives often exhibit intriguing geometric properties, making it difficult to track the trajctory of gradient-based methods. As a result, the majority of EoS analysis focused on toy models or the simplication of practical models. Therefore, we believe our work still carries important messages to the EoS community despite the limitions shared by many other works of the same kind.

The second limitation is a restriction to the simple $d = 2$ regime. However, we emphasize that this does not diminish the significance of our current analysis, just as the analysis in [Ahn et al., 2023] remains highly insightful despite its assumptiont it assumes $d=1$. Moreover, in the one-sample case, the empirical behavior and the major characterizations of EoS under both $\eta\mu < 1$ and $\eta\mu > 1$ are consistent for any $d \geq 2$. Therefore, while linear algebra limitations prevent us from extending the analysis to higher dimensions, our results are still sufficient to capture the essential properties of EoS convergence for GD in the setting we investigate.

Extension to higher depth.

We thank the reviewer for raising this fascinating question. We believe the depth will not greatly affect the theoretical analysis in the setting of one-sample and $d=2$, given some minor changes. This is because we can always decompose $\beta_t$, $w_+^d$ or $w_-^d$ as the span of vectors $\beta_0$, $\beta_1$ and $\beta^*$, and transform the GD dynamics into simple updates of scalars, which are easy to analyze.

We thank the reviewer for their insightful feedback. We will try our best to address the concerns raised, particularly the major issue highlighted in the "Weakness" section. Since this concern is complex, we have broken our response into three parts: Q1, Q2, and Q3.

Q1. Definition of EoS

We are sorry for not clearly stating the definition of EoS in the original manuscript. In fact, throughout the whole article, we use EoS in its original sense: the phenomena that the sharpness $S_t$, i.e. the largest eigenvalue of objective function's Hessian matrix, crosses the threshold of $2/\eta$, where $\eta$ is the step-size of GD. We have demonstated empirically EoS does happen with the one-sample Diagonal Neural Network model and under the regimes of $\eta\mu<1$ and $\eta\mu>1$ through the experiments and Figure 2 in Section 3. We have added some new lines in the revised version for better clarity on this point.

We understand the reviewer's concern regarding the relationship between EoS and oscillation, and we would like to clarify why our theoretical analysis focuses on the oscillatory behavior rather than the quantity $2 - \eta S_t$. To illustrate this, it is useful to compare our approach with that of [Ahn et al., 2023], which carefully characterized the trajectory of $S_t$ in its theoretical analysis. This is because the gap between $2 - \eta S_t$ is a key quantity in its "quasi-static" analysis and plays an crucial role in establishing the oscillating convergence of GD. In contrast, our convergence analysis of EoS on the one-sample DLN model does not rely on controlling the gap $2 - \eta S_t$. Therefore, we do not put an emphasis on this in our proof. Consequently, we do not emphasize this quantity in our proof. The purpose of theory part is only to establish convergence of GD under EoS, given the fact that the occurance of EoS is already confirmed through the previous empirical section.

This distinction is subtle but important. Furthermore, in [Ahn et al., 2023], the theoretical analysis only confirmed the existence of oscillation and the fact that $\lim_{t\to\infty}\eta S_t < 2$ by a contrdiction argument. However, the justification for when and where $\eta S_t > 2$ occurs, is based on empirical evidence rather than theoretical proof. In this way, it is very similar to our approach.

Finally, from our theoretical analysis, it is easy to obtain identity $\eta S_t - 2 = 2\eta( \mu + r_t - c_x \cdot b_t) -2 = 2\eta( \mu - 1+ r_t - c_x \cdot b_t)$ using direct computation. As the result of Theorem 2, we could conclude that $\lim_{t\to\infty} \eta S_t < 2$.

Q2. $\eta\mu>1$ as the "true" EoS

We agree with the reviewer that the setting $\eta\mu>1$ represents the more "authentic" EoS regime. Consequently, the convergence analysis in this regime constitutes the primary focus and major contribution of our work. This focus is not merely because EoS and oscillations are consistently observed under $\eta\mu>1$, but also because, as discussed in Sections 1 and 4, addressing this regime poses a more significant and challenging task: running GD on our model is equivalent to a parameterized discrete dynamical system, where the parameter varies by time. While $\eta\mu<1$ corresponds to such a system where convergence can be established relatively more easily as in existing work like [Chen et al., 2023], it is extremely difficult to show the convergence of GD under regime $\eta\mu>1$ because a crucial phase transition occurs and transforms a system that initially exhibits divergent behavior into one that converges. To the best of our knowledge, this is the first result to rigorously demonstrate such a kind of convergence.

Q3. Oscillation and EoS

We agree with the reviewer that oscillation is not necessarily related to EoS for all the regimes. Nevertheless, we believe the focus of the oscillation behavior does not hurt the validity of our result. As we explained in Q1, we stick to the classical definition of EoS and showed empirically that EoS does occur under our setting, which does not depend on the sign-changing of residual $r_t$.

Moreover, we improve Theorem 2 in the latest version and provide a convergence analysis under the more important regime $\eta\mu > 1$ that no longer relies on the sign-changing assumption of $r_t$. This is at a cost of restricting $\eta \mu - 1$ to be upper bounded by a universal constant. We believe the removal of the assumption makes our result more convincing in this perpsective. The revisions are marked in red.

We sincerely thank the reviewer for the positive feedback and insightful comments. It is evident that the reviewer has a strong understanding of the relevant literature and the core contributions of our work. We appreciate this and kindly request consideration for an increased confidence score. Below, we address the reviewer's specific concerns.

Q1. Subquadratic being a sufficient condition.

We apologize for any ambiguity regarding EoS's condition. In previous works such as [Ahn et al 2023], a subquadratic loss is presented as a sufficient condition for EoS, whereas the main message of our empirical and theoretical analysis is to provide some additional necessary conditions to allow EoS to occur even with quadratic loss functions. We have modified some lines to avoid the imprecision.

Q2. How large is $\mathfrak{t}-t_0$.

We appreciate this important question, which we did not elaborate on in the paper due to multiple reasons. In brief, the length of the second phase, i.e. the gap between $\mathfrak{t}$ and $t_0$, also scales in $\log(1/\alpha)$. Specifically, in the second phase, the residual $r_t$ oscillates and its envelope remains non-shrinking because $\alpha_t < 0$ holds, leading to an almost linear increase in $b_t$ according to the update $b_{t+1} = (1 - x\eta r_t)^2 \cdot b_t.$ The linear increase continues until $b_t$ reaches the scale of $\Theta(\eta\mu - 1)$, at which $\alpha_t$ becomes positive and marks the end of the second phase. From Lemma 7. we know $b_{t_0} \propto \alpha^{\Theta(1)}$. Therefore, we can conclude that the linear increase of $b_t$ lasts for $\log(1/\alpha)$ iterations before phase transition at $t=\mathfrak{t}$.

In the revised manuscript, we have chosen to present this result in Section 3, marked in red, and provide supporting empirical evidence in Figure 4. A detailed theoretical explanation would require additional assumptions in Theorem 2. We think it is more reasonable to keep Theorem 2 with the minimal assumptions needed. Therefore, we believe it is more appropriate to focus on empirical validation in this context.

Q3. Influence of overparameterization in the EoS phenomenon.

We thank the reviewer for raising this interesting question. Actually, our empirical evidence indicate that the level of overparameterization, i.e. ratio r = d/n, does not have any significant or consistent influences over the EoS phenomena. This is also true for the cases with $r>2$ which are not covered in the current figures.