Implicit Bias of Mirror Flow for Shallow Neural Networks in Univariate Regression

Authors: Shuang Liang, Guido Montufar
Primary area: optimization
Keywords: implicit bias, overparametrized neural network, mirror descent, univariate regression, lazy training

Abstract:
We examine the implicit bias of mirror flow in least squares error regression with wide and shallow neural networks. For a broad class of potential functions, we show that mirror flow exhibits lazy training and has the same implicit bias as ordinary gradient flow when the network width tends to infinity. For univariate ReLU networks, we characterize this bias through a variational problem in function space. Our analysis includes prior results for ordinary gradient flow as a special case and lifts limitations which required either an intractable adjustment of the training data or networks with skip connections. We further introduce \emph{scaled potentials} and show that for these, mirror flow still exhibits lazy training but is not in the kernel regime. For univariate networks with absolute value activations, we show that mirror flow with scaled potentials induces a rich class of biases, which generally cannot be captured by an RKHS norm. A takeaway is that whereas the parameter initialization determines how strongly the curvature of the learned function is penalized at different locations of the input space, the scaled potential determines how the different magnitudes of the curvature are penalized.

Decision: accept (spotlight)
Reviewer scores: 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1729602511923

thread_type:
official_review

<thread_object>
creator: reviewer_8kz2

<content>
summary:
The paper studies the implicit bias of mirror flow in univariate two-layer networks in the lazy regime. For a broad class of potential functions, they show that the implicit bias at the limit of infinite width is similar to the implicit bias of mirror flow. They also provide a characterization of this implicit bias in function space which generalized prior results. Moreover, they introduce scaled potential functions, where training is in the lazy regime but not in the kernel regime, and characterize its implicit bias.

strengths:
The implicit bias of mirror descent/flow has been studied mostly for linear models. Here, the authors extend our understanding of mirror flow to the case of univariate two-layer networks in the lazy regime, and also give a more precise characterization of the implicit bias in function space for gradient flow in this setting. I think that the contribution is of interest to the implicit bias community, and continues two lines of work: the works on the implicit bias of mirror descent, and the works on the implicit bias of univariate shallow networks.

I believe that the contribution is significant and hence recommend acceptance.

weaknesses:
It would be nice to extend Theorem 2 beyond Assumption 1, which restricts the class of potential functions covered here, and also to relax the many assumptions in Theorem 8. But I understand that the authors needed these restrictions for technical reasons.

questions:
In both Theorem 2 and Theorem 8:
“with high probability over the random parameter initialization,
there exist constants C_1, C_2 such that …”. It sounds like the constants C_1,C_2 may depend on the initialization (which also depends on n). Did you mean to write “there exist constants C_1, C_2 such that w.h.p. over the initialization …”?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730709988899

thread_type:
official_review

<thread_object>
creator: reviewer_qtju

<content>
summary:
The paper establishes various implicit biases for wide shallow networks trained via mirror flow in the context of univariate regression which roughly fall into two categories: (i) parameters do not move far from initialization while asymptotically minimizing risk (ii) the asymptotic solution (asymptotic with respect to time) can be characterized via a variational problem.

strengths:
Characterizing the asymptotic solution of mirror flow via a variational problem is interesting.

weaknesses:
The work only considers univariate models. 


minor comment:
Considering the work only deals with mirror flow, I believe the title should be remedied to "Implicit bias of Mirror Flow ...".

questions:
1) All the convex potentials bake in the initial parameter $\hat \theta := \theta_0$. I wonder how much of the implicit bias results are affected by the choice of $\hat \theta$. If $\hat \theta$ was arbitrary (or any point other than $\theta_0$), how do the lazy training results change (i.e. is it still true that the parameters do not move far from initialization)? Of course when the coordinate function $\phi(x) = \frac 1 2 x^2$ (i.e. GF), the choice of $\hat \theta$ does not matter but for $\phi(x) = x^2 + x^4$, the mirror flow update explicitly penalizes movement away from $\hat \theta$ and hence it less clear whether lazy training should occur. 

2) Why is the mirror flow update scaled by $\eta$? Do similar results hold if $\eta$ is dropped?

3) Is the restriction to the univariate case only needed for the variational characterization of the asymptotic solution? 
(i.e. can the univariate requirement be dropped for establishing lazy training results and risk minimization?)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730956064024

thread_type:
official_review

<thread_object>
creator: reviewer_ex1t

<content>
summary:
The paper studies the implicit bias of mirror flow on shallow univariate neural networks in the standard parameterization. The authors show that the networks are in the lazy regime (only the second layer weights move non-negligibly), and characterize the representation norm in two different regimes. In the "unscaled" regime, the mirror map plays a minor role, only through its Taylor expansion around 0, leading to an NTK-like implicit bias similar to GD. In the "scaled" regime, where the potentials are rescaled by width, a richer implicit bias arises involving the corresponding Bregman divergence, which is no longer a RKHS norm.

strengths:
The paper is well written, and the obtained implicit bias results for scaled potential is quite interesting as it involves new norms that have not been used before in the context of neural networks, to my knowledge.

weaknesses:
The significance is unclear: while the obtained norms are interesting, it is unclear if mirror flow is a relevant method in the study of neural networks. The lazy training regime suggests that it may be good to reframe the paper more generally outside of neural networks (which in practice are often in non-lazy regimes), for instance mirror descent on random features models, which could be of interest to the statistical learning and kernels community more broadly. The analysis is also somewhat incremental over existing implicit bias results.

questions:
* Given that the networks are in lazy regimes and similar to random feature models, is the setting related to this paper on random feature models with Lp penalties on the weights? https://arxiv.org/abs/2103.15996

* Do you have a sense of whether such mirror flow analyses extend to feature learning regimes, e.g. by switching to a mean-field parameterization? What would be the resulting implicit bias?
</content>
</thread_object>

</thread>
<thread>
creation date:
1739038213876

thread_type:
public_comment

<thread_object>
creator: ~shuang_liang6

<content>
title:
Camera-Ready Revisions

comment:
We sincerely thank the reviewers and the Area Chair for their valuable suggestions and feedback. Following the acceptance of the paper, we have made the following revisions in the camera-ready version addressing their final comments:

* We rephrased the interpretation of the experiments on scaled potentials (Figure 2), to better illustrate that “the parameter initialization distribution determines the strength of the curvature penalty over different locations of the input space and the potential determines the strength of the penalty over different magnitudes of the curvature”. 

* We revised the discussion on representation cost in the related work section to improve accuracy and completeness.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies the implicit bias of mirror flow in shallow (one-hidden-layer) and wide neural networks, focusing on ReLU and absolute value activations. For a class of “unscaled” potentials, the authors show that the ReLU network under mirror flow stays in the lazy and kernel regime and (perhaps surprisingly) has an implicit bias identical to gradient flow in the infinite width limit. The paper further provides a variational characterization of the trained function in the univariate input case. The authors also consider “scaled” potentials, and show that training an absolute-value-activated NN under mirror flow now stays in the lazy regime but not necessarily in the kernel regime. For the univariate input case, the paper also offers a variational characterization of the implicit bias.

The reviewers found the results novel and interesting. The paper offers a concrete characterization on the implicit bias of mirror flow on shallow and wide networks. To me, it interesting to know that for standard unscaled potential functions the implicit bias of mirror descent ends up being the same as gradient flow in the infinite width limit, which is in contrast to other results on finite number of parameters. Also, for scaled potentials, what I find the most interesting is that the model is in the lazy regime BUT can be in a non-kernel regime—personally, this is something that I haven’t seen in the literature. The variational characterization is nice, and the experiments look very much aligned with theoretical analysis. Overall, it is deemed that the paper provides strong contributions to the implicit bias community.

Of course, there are some limitations in the analysis, as acknowledged by the authors in the conclusion section: results holding for 1) shallow and (infinitely) wide networks, 2) univariate inputs, 3) specific first layer weight initialization, and 4) under some restrictions on the potential function. 

However, my overall evaluation is that the merits of this paper significantly outweigh the shortcomings, and I recommend acceptance of this paper.