Time-Dependent Mirror Flows and Where to Find Them

12. The choice of $A_i$:

Indeed $A_i$ have to be chosen. We now have made this clear in the text, and discuss particular choices such as diagonal matrices and the broader class of matrices satisfying the alignment property. Based on the diagonal matrices, we illustrate that the type of bias for the eigenvalues of $K^TQ$ changes from $L_2$ to $L_1$. This implies for the general matrix that the type of bias changes from the Frobenius norm to the nuclear norm.

13. The shift:

The shift refers to the time-dependent Legendre function’s global minimum. This illustrates the effect of the positional bias, which we now explain in the introduction. Moreover, we have moved the discussion on the reparameterization $log(u) - log(v)$ to the appendix to improve the flow of the paper.

14. Scaling effect:

To illustrate the effect of the scaling on the type of bias we have added Figure 7, where we plot the time-dependent Legendre function.

Minor comments: We have addressed these comments in the manuscript.

2. Limitations of the framework

We have relocated the limitations of the proposed framework to Appendix B1 and expanded on these limitations by analyzing the parameterization $g(w) = \Pi_{i = 1}^k w_i$ with weight decay regularization $h(w) = \sum_{i=1}^k w_i^2$. In this analysis, we demonstrate that it lies outside the scope of the current framework, rendering Theorem 3.1 inapplicable. Nevertheless, we can transfer our insights for $m\odot w$ and $u^{2k} - v^{2k}$ to predict the implicit bias. We observe in a diagonal linear network experiment that both the type of bias and range shrinking effect are present. Specifically, when regularization is turned off, the sparse ground truth is recovered if enough regularization was applied beforehand, as observed with the $m \odot w$ parameterization. However, excessive regularization causes the range shrinking effect to dominate and prevents reaching the ground truth, as similarly observed with $u^{2k} - v^{2k}$. Note that the solution $x$ is distinct from both the ground truth and $0$, which is significant as the dynamics might get stuck at zero.

3. Novelty of insight into quadratic parameterizations

The cited papers primarily focus on the induced model capacity of regularization with reparameterization and the difference between weight decay and Frobenius regularization. However, they do not analyze the impact of explicit regularization onto the implicit bias, nor do they systematically derive the nature of the implicit bias, which as we show, changes from an $L_2$ to $L_1$ or Frobenius norm to nuclear norm during training.

In addition to our theoretical extensions, we gain novel and practical insights concerning quadratic reparameterizations:

1. The type of bias changes from an $L_2$ to $L_1$ Legendre function, as observed in our experiments. This insight can only be attained with our dynamic view. The time-dependent Legendre function changes during training, which changes the geometry of the training dynamics as described in Equation (7). This could have different implications, e.g., on the effect of early stopping, or explaining scaling laws that relate the amount of overparameterization to optimal training times (for best generalization), or guide simply dynamic tuning schedules for weight decay. Note that the type of bias is conceptually different from model capacity and is induced by a time-dependent Legendre function. We do not observe a fixed model capacity (with respect to either $L_2$ or $L_1$), as it changes from one to the other during training.
2. The effect of the regularization is determined by the time-dependent Legendre function. This allows us to turn off weight decay at the end of the training while still keeping the effect of the regularization. We have extended the LoRA experiment to illustrate this point more clearly. Now we turn off the weight decay after $200$ iterations for $2$ of the $4$ configurations. We observe that the ratio of the nuclear norm and Frobenius norm still decreases despite switched-off weight decay, effectively storing the regularization. This allows us to train longer with a desired norm ratio but unconstrained by the explicit regularization. This insight can be used to design better regularization schedules. For example, turning off the regularization at a desired ratio. In our LoRA experiments, this approach attains the models of lowest rank without drop in performance.

W2: Interpretation and relation to practical contexts:

The main practical implication of our theoretical insights is that we can change and control the implicit bias by explicit regularization, even dynamically during training. Specifically, we use our findings to derive guiding principles for sparse coding, attention, and LoRA finetuning. For instance, we gain insights into how weight decay promotes effectively low rank solutions. While attention requires sufficiently strong regularization for $L_1$ to dominate late in training, LoRA finetuning finds low rank solutions if weight decay is turned off in later training rounds. More details on the exposition follow below.

Different parameterizations: The goal of Corollary 3.2 is to explore the limitations of the proposed framework. We have moved Corollary 3.2 to the appendix and have provided an additional class of parameterizations. The additional class is described by the reparameterization $g(w) = \Pi_{i = 1}^k w_i$ and $h(w) = \sum_{i = 1}^k w_i^2$, with $k >2$. We show that this class does not satisfy the sufficient condition to apply Theorem 3.1, which characterizes the time-dependent mirror flow. Nevertheless, we can make predictions based on the developed framework, which we demonstrate in the context of diagonal linear network experiments. We observe both the type of bias and range shrinking effect for the reparameterization with $k = 3$. The type of bias prediction follows from the analyzed reparameterization $m \odot w$ and the range shrinking follows from the reparameterization $u^{2k} - v^{2k}$. We have included references to previous work where specific reparameterizations were studied to highlight the broader relevance of our framework. For instance, reparameterizations such as $m \odot w$ (Pesme et al. 2021) and $u^{2k} -v^{2k}$ (Woodworth et al. 2020) have been used to gain insights into the training dynamics of neural networks, specifically, the effect of stochasticity and overparameterization. They have also been exploited for sparsification (Jacobs & Burkholz 2024).

Pesme, S., Pillaud-Vivien, L., & Flammarion, N. (2021). Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity. Neural Information Processing Systems.

Woodworth, B.E., Gunasekar, S., Lee, J., Moroshko, E., Savarese, P.H., Golan, I., Soudry, D., & Srebro, N. (2019). Kernel and Rich Regimes in Overparametrized Models. ArXiv, abs/2002.09277.

Jacobs, T., & Burkholz, R. (2024). Mask in the Mirror: Implicit Sparsification. ArXiv, abs/2408.09966.

For the derivation of Eq. (8), we use Theorem 3.3. The main step is inverting the function $Q_a(\mu)$. We added an explanation after the equation.

Practically relevant insights: We gain two key novel insights of our intricate extension for quadratic reparameterizations:

1. The type of bias changes from an $L_2$ to $L_1$ Legendre function, as observed in our experiments. This insight can only be attained with our dynamic view. The time-dependent Legendre function changes during training, which changes the geometry of the training dynamics as described in Equation (7). This could have different implications, e.g., on the effect of early stopping, or explaining scaling laws that relate the amount of overparameterization to optimal training times (for best generalization), or guide simply dynamic tuning schedules for weight decay. Note that the type of bias is conceptually different from model capacity and is induced by a time dependent Legendre function. We do not observe a fixed model capacity (with respect to either $L_2$ or $L_1$), as it changes from one to the other during training.
2. The effect of the regularization is determined by the time-dependent Legendre function. This allows us to turn off weight decay at the end of the training while still keeping the effect of the regularization. We have extended the LoRA experiment to illustrate this point more clearly. Now we turn off the weight decay after $200$ iterations for $2$ of the $4$ configurations. We observe that the ratio of the nuclear norm and Frobenius norm still decreases despite switched-off weight decay, effectively storing the regularization. This allows us to train longer with a desired norm ratio but unconstrained by the explicit regularization. This insight can be used to design better regularization schedules. For example, turning off the regularization at a desired ratio. In our LoRA experiments, this approach attains the models of lowest rank without drop in performance.

We would like to thank you for your valuable review. We believe that we have addressed all points of criticism but would be happy to address any open issues if there should remain any. Since the discussion period is approaching its end, we would highly appreciate your feedback.

3. Motivation and significance of implicit bias framework:

The main goal of characterizing implicit bias is to understand the role of overparameterization in deep learning and how it contributes to its success. Commonly, the corresponding regularization is tied to the specifics of the parameterization, neural network architecture, and the optimizer. Yet, its strength is usually assumed to be fixed during training.

Our framework extension changes this picture, as it aims to understand the effect of explicit regularization on the implicit bias (which is almost always applied in practice). This is significant, as the explicit regularization changes the nature and strength of the implicit bias of the studied architecture and gives us a way to control and therefore exploit it. This nature and strength can even change dynamically during training. Importantly, this dynamic change is critical to address issues with range shrinking and positional bias, and thus guarantee convergence.

From a conceptual point of view, our results are of independent interest, as they extend the mirror flow framework to a time-dependent mirror flow framework.

From a practical point of view, we gain insights into modern architectural characteristics such as attention and LoRA, and how weight decay induces sparsity. To boost their performance, we propose to switch-off weight decay in the last training rounds. (Note that this switch-off enables convergence of a minimizer of the original optimization objective $f$ according to Theorem 3.2.)

We have highlighted these contributions more clearly in the introduction of our revised manuscript.

Practically relevant insights: We gain two key novel insights of our intricate extension for quadratic reparameterizations:

1. The type of bias changes from an $L_2$ to $L_1$ Legendre function, as observed in our experiments. This insight can only be attained with our dynamic view. The time-dependent Legendre function changes during training, which changes the geometry of the training dynamics as described in Equation (7). This could have different implications, e.g., on the effect of early stopping, or explaining scaling laws that relate the amount of overparameterization to optimal training times (for best generalization), or guide simply dynamic tuning schedules for weight decay. Note that the type of bias is conceptually different from model capacity and is induced by a time dependent Legendre function. We do not observe a fixed model capacity (with respect to either $L_2$ or $L_1$), as it changes from one to the other during training.
2. The effect of the regularization is determined by the time-dependent Legendre function. This allows us to turn off weight decay at the end of the training while still keeping the effect of the regularization. We have extended the LoRA experiment to illustrate this point more clearly. Now we turn off the weight decay after $200$ iterations for $2$ of the $4$ configurations. We observe that the ratio of the nuclear norm and Frobenius norm still decreases despite switched-off weight decay, effectively storing the regularization. This allows us to train longer with a desired norm ratio but unconstrained by the explicit regularization. This insight can be used to design better regularization schedules. For example, turning off the regularization at a desired ratio. In our LoRA experiments, this approach attains the models of lowest rank without drop in performance.