Implicit Bias of Mirror Flow for Shallow Neural Networks in Univariate Regression
Authors: Shuang Liang
Guido Montufar
Venue: ICLR 2025 Spotlight (2025)
Keywords: implicit bias
overparametrized neural network
mirror descent
univariate regression
lazy training
Primary area: optimization
TL;DR: Univariate wide shallow networks trained with mirror flow are biased towards smooth interpolants and scaled mirror potentials induce biases affecting the different magnitudes of the curvature.

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.

Published as a conference paper at ICLR 2025
IMPLICIT BIAS OF MIRROR FLOW FOR SHALLOW NEU- RAL NETWORKS IN UNIVARIATE REGRESSION
Shuang Liang UCLA 1 angghuungggg ucla edu
Guido Montúfar
UCLA & MPI MIS
montufar@math ucla edu
ABSTRACT
We examine the implicit bias of mirror flow in least squares error regression with wide and shallow neural networks. For 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 variational problem in function space. Our analysis includes prior results for ordinary gradient flow as a special case and lifts limitations which required either intractable adjustment of the training data or networks with skip connections. We further introduce 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.
INTRODUCTION
The implicit bias of a parameter optimization procedure refers the phenomenon where, among the many candidate parameter values that might minimize the training loss, the optimization procedure biased towards selecting one with certain particular properties. This plays an important role in explaining how overparameterized neural networks, even when trained without explicit regularization, can still learn suitable functions that perform well on new data (Neyshabur et al. 2015) Zhang et al., 2017). As such, the concrete characterization of the implicit biases of parameter optimization in neural networks and how they affect the solution functions is one of the key concerns in deep learning theory and has been subject of intense study in recent years (see, e.g., Ji & Telgarsky
2019 Lyu & Li 2020 Williams et al. 2019 Chizat & Bach 20200] Sahs et al.) 2022 Frei et al. [023 Jin & Montúfar
2023 In this work consider an important class of parameter optimization procedures that have remained relatively underexplored, namely mirror descent as applied to solving regression problems
with overparametrized neural networks.
Mirror descent is broad class of first-order optimization algorithms that generalize ordinary gradient descent (Nemirovskij & Yudin, 1983). It can be viewed as using general distance-like function, defined by the choice of a convex potential function, instead of the usual Euclidean squared distance to determine the update direction in the search space (Beck Teboulle, 2003). The choice of the geometry of the parameter space of learning system is crucial, as it can affect the speed of convergence and the implicit bias of parameter optimization (Neyshabur et al. 2017). The implicit bias of mirror descent has been studied theoretically for linear models (Gunasekar et al.) 201 Sa Sun et al. 2023) Pesme et al. 2024), matrix sensing (Wu & Rebeschini] (002), and also for nonlinear models at least in local sense (Azizan et al. 2022). However, are not aware of any works characterizing the implicit bias of mirror decent for neural networks describing more than the possible properties of the parameters also the nature of the returned solution functions. This is the gap that target in this article. We focus on the particular setting of mirror flow applied to least squares error regression with shallow and infinitely wide neural networks.
We show for networks with standard parametrization that, under a broad class of potential functions, mirror flow displays lazy training, meaning that even though the loss converges to the parameters

---

Published as a conference paper at ICLR 2025
remain nearly unchanged. We show in turn that mirror flow has the same implicit bias as gradient flow when the network width tends to infinity. This is surprising because, in contrast, for dinte-dimensioall parameter models mirror flow has been observed to exhibit a different implicit bias than gradient descent (see, e.g.. Gunasekar et al. 20188] [Sun et al. 2023). We then proceed to characterize the particular solution function that is returned by mirror flow for wide ReLU networks, by expressing the difference between the network fiunctions before and after training, g(z) = furained (x) - finitial (z)). as the solution to variational problem. By examining the function representation cost, we show that mirror flow is biased towards minimizing weighted L norm of the second-orde derivative of the difference function, / p (xx))))) where the weight p- is determined by the network parameter initialization, and two additional terms that control the first-order derivative. This includes as a special case a previous result for gradient flow and ReLU activations by Jin & Montúfar (2023). Our analysis also removes a key limitation of that result, which required either an intractable linear adjustment of the training data or otherwise adding skip connections to the network. We observe that other previous works on the implicit bias of gradient descent suggested that shallow ReLU networks in the kernel regime could be well-approximated by a cubic spline interpolator of the training data (Williams et al., 2019 Heiss et al.] 2019 (ahs et al., 2022) but overlooked the need of linearly adjusting the training data or adding skip connections. We illustrate that even a slight modification of the training data by applying a constant shift to all training labels can lead significantly different solutions upon training which are biased not only in terms of curvature but also slopes.
To explore training algorithms capable of implementing different forms of regularization depending on the parameter geometry, we introduce mirror descent with scaled potentials. These are separable potentials (0)) En o(0k) whose input is scaled by the network width as (nn)). We show that for networks with standard parametrization, mirror flow with scaled potentials also exhibits lazy training but is not in the kernel regime. In this case the training dynamics can become fundamentally different from that of gradient flow, for which the lazy and kernel regimes fully overlap (see, e.g., Lee et al. 2019). For networks with absolute value activations, we show that mirror flow with scaled potentials induces a rich class of implicit regularizations. The difference function minimizes a functional of the form f Do (x)p (z),, O)p(z (r), where D& denotes the Bregman divergence induced by G. Notably, for non-homogeneous potential functions ø, the resulting regularization is not homogeneous and therefore it cannot be expressed by an RKHS norm. Overall our results show that, whereas the choice of parameter initialization distribution determines how strongly the curvature is penalized at different locations in the input domain via the choice of the potential function ø determines how the different magnitudes of the curvature are penalized.
1.1 CONTRIBUTIONS
The goal of this article is to provide insights into the implicit bias of mirror flow for shallow and wide neural networks with standard parametrization. Our contributions can be summarized as follows.
We show for a broad class of potential functions that, mirror flow, when applied to networks with general input dimension, displays lazy training. We show that in this case the implicit bias of mirror flow is equivalent to that of gradient flow when the number of hidden units tends to infinity.
We characterize the implicit bias of mirror flow for univariate networks in function space. Con- cretely, we express the difference of the trained and initialized network function as the solution to a variational minimization problem in function space (Theorem 2).
We introduce class of scaled potentials for which we show that mirror flow, applied to networks with general input dimension, falls in the lazy regime but not in the kernel regime. We characterize the corresponding implicit bias for infinitely wide networks and observe it is strongly dependent on the potential and thus different from ordinary gradient flow.
Finally, we characterize the implicit bias of mirror flow with scaled potentials in function space for univariate networks with absolute value activation function (Theorem()) Our results show that the choice of the potential biases the magnitude of the curvature of the solution function.
1.2 RELATED WORKS
Implicit bias of mirror descent The implicit bias of mirror descent has been studied for under- determined linear models. In classification problems with linearly separable data, [Sun et al. (2022) show that mirror descent with homogeneous potentials converges in direction max-margin
2

---

Published as a conference paper at ICLR 2025
classifier with respect to a norm defined by the potential. Pesme al. (2024) extend this work to non-homoge honooooggenuuus potentials. For regression, [Gunasekar et al. (201 Sa): Azizan & Hassibi (2019)
show that mirror descent is biased towards the minimizer of the loss that closest to the initialization with respect to the Bregman divergence defined by the potential. Azizan et al. (2022) extend this result nonlinear models, alheit placing a strong assumption that the initial parameter sufficiently close to a solution. For matrix sensing, Wu & Rebeschini (2021) show mirror descent with spectral hypentropy potential tends to recover low-rank matrices. [i al. (2022) connect gradient flow with commuting reparameterization to mirror flow, and show that both have identical trajectories and implicit biases. This connection further studied byChou al. (2023) specifically for linear models.
Implicit bias of gradient descent for overparametrized networks A seminal line of works in theoretical deep learning investigates the training dynamics of overparameterized neural networks in the so-called lazy regime (Du et al. 2019 Lee al., 2019; Lai al., 2023). Gradient descent training in this regime minimizes the training loss but, remarkably, the network parameters remain nearly unchanged. In turn, the network becomes equivalent to a kernel model (Jacot al. 2018 A particular consequence is that in the lazy regime the implicit bias of gradient descent can be described as the minimization of an RKHS norm (Zhang et al., 2020). Williams et al. (2019); Sahs et al. (0022) derived the explicit form of this RKHS norm for univariate shallow ReLU networks. However, Williams et al. (2019) implicitly constrain the computation to a subspace in the RKHS consisting of functions with specific boundary conditions, and Sahs et al. (2022) placed an incorrect assumption on the behavior of the initial breakpoints of the network in the infinite width limit. Jim & Montúfar (2023) obtained function description of the implicit bias of gradient descent for univariate and multivariate shallow ReLU networks where they explicitly characterize the effect of parameter initialization. However, their result requires either networks with skip connections or an intractable linear data adjustment. The latter, noted by the authors, is also equivalent to computing the RKHS norm in function subspace. Beside the kernel regime, we observe that a complementary line of works studies the dynamics and implicit bias in the active regime, also known as the rich or adaptive regime, where the hidden features change during training and the networks do not behave like kernel models (Mei et al. |019| Williams et al. 2019; Chizat & Bach 2020; Jacot et al. 2021 Mulayoff
et al. 2021 Shevchenko et al.) 2022] Boursier et al. 2022] Chistikov et al.) 2023 Qiao et al.) 2024).
Representation cost For parameterized model and a given cost function, e.g., a norm, on the parameter space, the representation cost of a target function refers to the minimal cost required to implement that function. Existing work has revealed a close connection between representation cost and the implicit bias of gradient-based optimization methods, the latter can often be described by specific parameter norm minimization problems. A number of works have studied the representation cost for neural networks with varying architectures (Gunasekar al. 2018b Savarese al.) 2019 Ongie et al., 2020 Boursier & Flammarion 2003 Jacot] | Parkinson et al. 2024). Among these, Savarese al. (2019); Ongie et al. (2020); Boursier & Flammarion (2023) have focused on infinitely wide and shallow ReL.U networks, similar to our work. While these works focused on the 7.1 norm on the output weights as the cost function, we considered the 7.2 norm and encompasses broad family of cost functions induced by Bregman divergence. Moreover, instead of studying representation cost in static and optimization-free manner, we provide detailed analyses of the optimization dynamics of mirror flow and characterize the implicit bias with concrete convergence guarantees.
PROBLEM SETUP
We focus on fully-connected two-layer neural networks with d input units, hidden units and one output unit, defined by:
Σ ( (wf x -b,) + d,
f(x, 0)
(I)
A=
where a (2) is a nonlinear activation function. We write vec (W b, a, d) E R2, with p nld 2) + 1. for the vector of all parameters, comprising the imput weights W = (WI...., wx) E Rnxd input biases b- (b,, ...,. b,,) E R.", output weights a = (aI..... an) E R", and output bias d E TR. We assume the parameters of the network are initialized by independent samples of mutually independent random variables W, B, A and D as follows:
1
Wk ~ W = 1n by ~ B, ag
A, k k=
n. d d~ D.
(2)
3

---

Published as a conference paper at ICLR 2025
We assume W is uniform the unit sphere in R and B, A, D are sub naubGaussian random variables. Further, assume that B has a continuous density function pg and that the random vector (W. B) is symmetric, i.e., (W, (3) and (-W, -B) have the same distribution. We use μ for the distribution of (V, and vec( (W, b, a, d) for the initial parameter vector.
In Appendix g we further discuss our settings on the network parametrization and initialization. Specifically, we compare the parametrization in (T). which is known the standard parametrization. to the NTK parumetrization (Jacot et al.) 2018), and discuss the effect of the Anti-Symmetrical Initialization (Zhang et al., 2020), which is common technique to ensure zero function output at initialization, on the implicit bias of mirror descent training.
Consider a finite training data set { D:) C R IR. Consider the mean squared error L(0)= = 1 En (zz 8) y1) Given a strictly convex and twice differentiable potential functtiol & on the parameter space, the Bregman divergence (Bregman 1967) with respect to d is defined as Do (0. 0¹) ())) - (00)) - (V 0.0 0*) The mirror descent parameter update (Nemirovskij
2m
& Yudin 1983 Beck & Teboulle 2003) for minimizing the loss T.0)) is defined by:
A(t + 1) = arg min n/6"", VL(((())) Do (0', (()),
6'
where the constant 17 > is the learning rate. When the learning rate takes infinitesimal values (we consider 1) To for some constant 1/o > 0 and n + oo), (3) can be characterized by the mirror flow:
0()) = - 1 (V20((( (VL((tt).
The mirror flow 4 can be interpreted as a Riemannian gradient flow with metric tensor V For convenience, provide in Appendix F brief introduction to the Riemannian geometry of the parameter space endowed with metric V
The dynamics for the predictions can be given follows. Let y [y1..... ... Pm R denote the vector of the training labels and f(t) denote the vector of the network prediction on the training input data points time i, i.e., f(t) (ft)),((t),,. J 0(t)) E R By the chain rule,
di. (()) = -000=-))000000---000000000)00001 - y): where (1) 2n
(5)
Here Jof(1) (a) O()). [V)]==/((,,,)))....... Vol(zm, ((1))=² E RR denotes the Jacobian matrix of the network output ft) with respect to the network parameter 0. We see in (5) that different choices of & induce different modifications of the Gram matrix (t)
To fix a few notations, we write [n] = ...,,,...777, n) for n E N. We use (*2 for the space of times continuously differentiable real-valued functions on Ω and C(22) for the space of continuous real-valued functions on f2. We use supp(-) to denote the support set of a function or a distribution. We say that sequence of events {&,, JeEN holds with high probability if P(Ex)
-> co.
3
MAIN RESULTS
We present our main theoretical results, followed by an overview of the proof techniques.
3.1 IMPLICIT BIAS OF MIRROR FLOW WITH UNSCALED POTENTIALS
We consider the following assumptions on the potential function.
Assumption The potential function tb R2 + satisfies: (i) can be written the form (0) = Σh= (0xx - 0x), where & is a real-valued function on I; (ii) & is twice continuously differentiable: (iii) (2 positive definite for all 9 E R
Potentials that satisfy (i) are referred to separable in the literature, they operate independently on each coordinate of the input. We consider separable potentials since they have a natural definition as
grows and can be conveniently extended to infinitely dimensional spaces (see details in Section (3.3).
Items (ii) and (iii) are common in the mirror descent literature (Sun et al.) 2023 Pesme et al. 2024). Assumption i is quite mild and covers broad class of potentials including: (1) )(x)
4

---

Published as a conference paper at ICLR 2025
(which recovers ordinary gradient descent); (2) d(x) = cosh(2);; and (3) the hypentropy function ((²) p3(x)) x arcsinh(z/9)) 2 for 3 (Ghai et al. 2020). In particular and significantly, we do not require bomogeneity. Assumption I does not cover the unnormalized negative entropy (xx) log I and the corresponding natural gradient descent (NGD). However, since the second derivative of the hypentropy function, 1/v2² + 8 converges to that of the negative entropy, 1/2,, + O can still obtain insights into NGD by applying our results to hypentropy potentials. We refer to potentials that satisfy Assumption T as unscaled potentials in order to distinguish them from the scaled potentials which examine in Section 3.2
The following theorem characterizes the implicit bias of mirror flow for least squares regression with wide neural networks with ReLU activation (s) max(0, s).
Theorem 2 (Implicit bias of mirror flow for wide ReLU network). Consider a two layer ReLU network ① with d 2 1 input units and n hidden units, where we assume n is sufficiently large. Consider parameter initialization (2) and, specifically, let pB(b) denote the density function for random variable B. Consider any finite training data set {(aa y1) }mm that satisfies I; $ 1j when
i + j and {[]=a||2]]? ] , C. supp(ps). and consider mirror flow 44) with potential function +
satisfying Assumption 7 and learning rate = (1/nn) Then, there exist constants C1, > 0 such that with high probability over the random parameter initialization, for any t t2 0,
[[(()) - م| S SCnn-122
lim n-000
(((t)-HHOO)a) |H(()- - 22))2 0, llf(() - y1l2 S 1
yll2. (6)
Moreover, letting 0(20) -om O(1), we have for any given Rd, that lim (f 0 ((c)) f(x,, OCF (oo) O, where (cF( denotes the limiting point of gradient with (0) |1 on the same training data and initial parameter.
flow, flow,
i.e. mirror flow
Assuming univariate input data, i.e., d = L, we have for any given x E R. that lim (f 0(co))
fxx 0) h(x)) =0, 0, where F((.) is the solution to the following variational problem:
min hE-RELU
(h) + 92(h) + Ga(h)
s.t. Mx.)) - y; f(x,, vi E [m]]
(h) = [ !amppips) (6(x)
(h)(
de
where
G2(h) = (h'(+20) h'( 2²
(s(h) E[B2]
(suppive)
h" h"(x)|=/dz (th)a)d=>2h(0)=
Here FRELU (h(x) = a(ww b) [ - + E C(supp(a)) is uniformly )ontinuous),
11(+00)) = 1om (x),), and N1 (x0- = lim - (x)) In addition, if is supported on 50
[- for some B > 0, then G3(h)= E[6 (+(00+00) h1' (-oo) - ( + h( h--B)))
To interpret the result, (6) means that mirror flow exhibits lazy training and the network behaves like a kernel model with kernel (0). The theorem indicates that mirror flow has the same implicit bias as ordinary gradient flow, and the reason behind this is that, loosely speaking, since the parameters stay nearly unchanged during training, the effect of the potential reduces to the effect of a quadratic form, i.e., (0) = |[0||2 Further below we discuss scaled potentials breaking outside this regime.
To interpret the implicit bias described in 7 we discuss the effect of the different functionals on the solution function: (1) Gi serves as a curvature penalty, which can be viewed as the P& 1 -weighted L norm of the second derivative of h. Note in our setting pg corresponds to the density of B/W. which is the breakpoint of a ReI.U in the network. (2) G2 favors opposite slopes at positive and negative infinity. (3) The effect of G3 is less clear under general sub- Gaussian initializations. In the case where B is compactly supported, G3 can be interpreted x-axis (specifically, h(B) + h(-B) B> > 0), then G3 encourages h (++oo)= > h' (oo) and thus favors a U-shaped function; if on the contrary h is below the r-axis, G3 favors inverted U-shaped function.
follows: if h is generally above the
The functional G1 is familiar from the literature on space adaptive interpolating splines and in particular natural cubic splines when ps is constant (see, e.g.. .hhlbee al. 2016). However, with G2 and G3. the solution to 7 has regularized boundary conditions. In Figure [ we illustrate numerically that gradient flow returns natural cubic splines only on very specific training data.
5

---

Published as a conference paper at ICLR 2025
Remark (Relaxed potential assumption). Assumption I requires the potential & to be "centered" the initial parameter 0. In Appendix A.4 we show Theorem also holds for potentials of the form (0) = >> given the initial input biases are sampled from a bounded distribution.
Remark (Gradient flow with reparametrizationn.. As observed by Li et al. (2022), mirror flow is equivalent to gradient flow with commuting reparameterization. In Appendix D we use our results to characterize the implicit bias of gradient flow with particular reparametrizations.
Remark (Absolute value activations). For networks with absolute value activation (s)) = |s|, a similar result to Theorem hoolds, which we formally present in Theorem 24 G2 and G2 disappear in the objective functional.
The only change is that
Remark 6 (Skip connections). For networks with skip connections, i.e., fakip (I. (0, u. v)) f(z,0) + uzz + U with u, E I. the skip connections can modify the value of G2 and (s(kki)p) without affecting 91 (fakip). Therefore, if we assume that the skip connections have a much faster training dynamics than the other parameter then the solution function biased towards solving FRaao the result obtained by Jin & Montufaa (2023) for univariate networks with skip connections.
G1 (h) subject to (xx)) + u., + u ys for i E [m]] By this we recover
Remark 7 (Natural gradient descent). Theorem 2 holds for hypentropy potential pa for any 3 and in particular, in the limit as B -> 0+. This is consistent with a result by Zhang et al. 20199)
stating NGD converges the same point as ordinary gradient descent.
3.2 IMPLICIT BIAS OF MIRROR FLOW WITH SCALED POTENTIALS
We introduce scaled potentials and show that these induce rich class of implicit biases.
Assumption 8. Assume the potential function cb R - IR satisfies: (i) & can be written in the form of (0) N Ei=1 (n(Ox - 0z)), where @ is real-valued function on RR (ū) & takes the form of <xx) 14 (vx)) + wo²). where w > 0 and 1b (TR) is convex function on RR
In Assumption 8, the factor 1/n2 prevents the largest eigenvalue of the Hessian 2 from diverging as n - oc. The form of & can be interpreted convex combination of an arbitrary smooth, convex function ub and the quadratic function x². A further discussion is presented in Section 3.3.
The following result characterizes the implicit bias of mirror flow with scaled potentials for wide networks with absolute value activations. We consider absolute value activations because, compared to ReLU they yield smaller function space that makes the problem more tractable.
Theorem 9 (Implicit bias of scaled mirror flow for wide absolute value network). Consider two layer network ① with absolute value activation, d 2 1 input units and n hidden units, where we assume is sufficiently large. Consider parameter initialization (2) and, specifically, let pB(b) denote the density function for random variable B. Additionally, assume B is compactly supported. Given any finite training data set {(Ii, ys) with I; # I.j when i }j aand {[|raa] [|m=all2?7= ym C supp(ps), consider mirror flow (4) with a potential satisfying Assumptions [ and learning rate 1] ((1/n). Specifically, assume the potential is induced by a univariate function 6(2) T (v(2) + wz²). There exists wn > 0 (depending on ih the training data and the initialization distribution) such that for any w >> wo, there exist constants C1, C2 > O such that with high probability over the random parameter initialization, for any 2 O.
(8) Moreover, assuming univariate input data, i.e., d = and letring 0(0c) lim @tt), we have for any given a E R that limm-soe (fx 0(oo)) - f(x, 0) h(x) = O, where h(()) the solution to the following variational problem:
[[0(t) - 01 S Cin [[ - y112 S -e--e--0 -000)2)11(())))) - -y12
Do 2pp(a)) Po )
h" hx(x)
min
s.t. h(T,) = y: - f(x, Vi [m].
(9)
NETATb supp(rs)
Here Di denotes the Bregman divergence on IR induced by the univariate function o, and Fabe
{f(x)) = a(w, k.u.b 6) wx bldp: a E ((supp(H)) a is even and uniformly continuous).
'This is can achieved by setting the learning rate for skip connections significantly larger than the learning rate of the other network parameters.

---

Published as a conference paper at ICLR 2025
Note, given any non-hhmogeneous function o, the objective functional in 99 is not homogeneous and therefore it cannot be expressed by any function norm (including RKHS norm). Also note, even though mirror flow displays lazy training, the kernel matrix (t) does not necessarily stay close to (0)) due to the scaling factor n. inside the potential. This implies that training may not be in the kernel regime (see numerical illustration in Figure
The implicit bias described in ( can be interpreted as follows: the parameter initialization distribution determines ps and thus the strength of the penalization of the second derivative of the learned function at different locations of the input space; on the other hand, the potential function e determines the strength of the penalization of the different magnitudes of the second derivative. In particular, we choose o = 1 (|z]|" + x²), then smaller allows relatively extreme magnitudes for the second derivative, whereas a larger p allows for more moderate values. This is illustrated in Figure 2
Theorem 9 relies on the assumption that w in Ø T+ (v + =.r.) is larger than threshold wp. In Proposition 30, we give a closed-form formula for this threshold. Further, numerical experiments in Section indicate the theorem holds even when w takes relatively small values, for instance, w 1.
Remark 10 (Scaled natural gradient descent). In Theorem 38 we show Theorem also holds for scaled hypentropy potentials (0) = 2/2 (n 0x)) when 3 sufficiently large, and therefore provides insights into scaled NGD. Empirically we observe that as 3 - 0+, networks trained with scaled hypentropy converge to the solution to problem (9) with ((x) = |x| (see Figure
3.3 GENERAL FRAMEWORK FOR THE PROOF OF THE MAIN RESULTS
We present general framework to derive our main results, which builds on the strategy of Jin & Montúfar (2023) with a few key differences. First, we examine the training dynamics of mirror flow, for which the analysis is more intricate than that of gradient flow due to the need to account for the evolution of the inverse Hessian of the potential throughout training. In particular, in the scaled potential case, we carefully analyze how the Hessian's properties, such as the decay rate of its smallest eigenvalue, affect the parameter trajectories. Second, we solve the minimal representation cost problem to derive the function description of the implicit bias. This allows us to remove previous requirement to linearly adjust the data when the network does not include skip connections. The detailed proofs for Theorems and 1aa are presented Appendix A and Appendix respectively.
Linearization The first step of the proof is to show that, under mirror flow, functions obtained by training the network are close to those obtained by training only the output weights. Let ((1) = vec( (0), b(0), @(t), d(0)) denote the parameter under the update rule where W, b, d are kept at their initial values and
dt (V(===>(V'(((()))Vaaa
a(0) = a(O).
(10)
Here, we write (() for the potential function applied to the output weights. For unscaled potentials,
(ā) Ei=1 oāāz - āx); for scaled potentials, (@)) n "n d(n(āz - ax)). Let ft) = [f(11, (tt)) fam. O(t))] E R The evolution of ftt can be given by
d i(() -m)0(())00))))) - v). Ī(() f(0),
dt
(11)
where i1(1)
We aim to show that for any fixed x e Rd. lim sup [f 0(t)) - f(x,((t)
For unscaled potentials, we adapt an established approach from gradient flow analysis (Du et al. 2019
Lee et al. 2019). Specifically, we show [I0(t) O1loo and |0(()) 01 are bounded by O 1/2)
(Proposition 14 Then we bound (t) (t)lla with the continuity of the Hessian map 2-)
in Assumption which gives the desired bound [xx ((t) f(x 0(t) (Proposition 16). For
scaled potentials, we use similar strategy to show [[0(t)= 01 and [|0(t)) - 0I1 are bounded by O(n -1) (Proposition 27). Then we show []n((()) - n can be bounded when the smallest and largest eigenvalues of (0) have mild decay and growth rates, respectively, ||0||2 increases,
which allows to control [xx 0()) [xx (()) (Proposition 29
7

---

Published as a conference paper at ICLR 2025
Implicit bias in parameter space The second step is to characterize the implicit bias for wide networks in their parameter space. In view of the previous step, it suffices to characterize the bias for networks where only the output weights are trained. Since the network is linear with respect to the output weights, we can use the following result on the implicit bias of mirror flow for linear models.
Theorem 11 (Gunasekar et al. 20] If mirror flow (T)) converges to zero loss, then the final parameter is the solution to following constrained optimization problem:
min D& (4 â) s.t. fx,, 0) yk, Vi E |m]] aER"
(12)
We reformulate (12) in a way that allows us to consider the infinite width limit, i.e., - oo. Let Nn denote the empirical distribution on Rd x IR induced by the sampled input weights and biases {(wx, bR) ] Consider a function a: suрр(μ) - IR whose value encodes the difference of the output weight from its initialization for a hidden unit given by the argument, i.c., a(tx, 6%) nāā; - 0s), vk E [n]. Then the difference of the network output at given x E IR from its initialization is gn (Iz, a) fax f) f(x, 9) f a(w, b)o(w" x b)ddn, Noticing Mn, converges weakly to /4, we write the difference between the infinitely wide networks before and after training as:
g(r, a)) = f a(w, (lo,bb)(("== b)dj, a E.
(13)
Here θ C. ((supp())) is chosen as the parameter space for the model (3).. Our goal is identify a cost functional on the parameter space Cp: θ + IR such that, as the width of the network goes to infinity, the solution (2) converges to the solution to the following problem
min Ca(aa) ag
s.t. g(z,, a) y/s ( 0), Vi E [m]]
(14)
For unscaled potentials, noticing the final output weight returned by mirror flow, i.e., the minimizer of (12). has 1/2) bounded difference from its initialization, have the following approximation
for the objective function of (2) when n tends to infinity F
T (0) 2 (āx 0z))² (0) 2n
f(a(tu ())
D& D â)
(15)
where we consider the Taylor approximation of á at the origin. Based on this, for unscaled potentials we show that in the infinite width limit, problem (12) can be described by problem (14) with
Co (a) = (a(w, dd), regardless of the choice of & (Proposition 17
In the case of scaled potentials, the objective function in 12 can be reformulated as
D2 à) = n7 k=1
Σ (o(n(āx - ax)) - x(0) (0) - -a)) n.
_
Da(a(w, b),0) dpa-
(16)
Based on this, for scaled potentials we show that, as n tends to infinity, problem 122 can be described by problem (14) with (a) = f (arww,
(o, b), 0)dz (Proposition 32
Function space description of the implicit bias The last step of our proof is to translate the implicit bias in parameter space, as described by problem (14,, to the function space of univariate networks, F (M((r) = (ur, (wx b)dw: a E θ). We do so by "pulling back" the cost functional Co from the parameter space to the function space, via the map RR: F + O defined by
(h) arg min (a) s.t. g( a) = h(x), Vx E R. a
(17)
The optimization problem in (17) is known as the minimal representation cost problem, as it seeks the parameter a with the minimal cost Co (a) such that g(-,a) represents given function h. In the following result, we show that for general cost functionals, the map R.p allows one to translate problem (14) to a variational problem on the function space. The proof is presented in Appendix A..
(A) Eا Law(,, bx)) for any measurable set A C Rd x H. where lA is the indicator of A. "The convergence is in the following sense: let be the solution to (12) and o,, E O be continuous function satisfying ā. (r0x, by) = ná@2 - dx) for all k E [m). Let ā be the solution to [1.. Then for any given x E Rd gn (z, a.) converges in probability to g(2, a) as n > co.
'Here we write xn "2 yn to signify that lim Un Un 0.
8

---

Published as a conference paper at ICLR 2025
Proposition 12. Assume that θ is a vector subspace in ((supp(μ)) and C2 is strictly convex functional on O. Consider the map R2 defined in (17). If ā is the solution to problem (14). then g(+āā) solves the following variational problem:
min Co o (h) s.t. - y: [II, vi E [m] KEF
(18)
With Proposition 12 to obtain function description of the implicit bias, it suffices characterize the exact form of the map Ra, which is equivalent to solving the minimal representation cost problem. For unscaled potentials and ReLU networks, we set ReLU as the space of uniformly continuous functions on supp(a) and our approach builds on the following key observations:
(I) every a RRELU has unique even-odd decomposition, a at + a where at +
Even
{ar E ÓRELU: is }ven) and a E Goad = fa E ORELU: a is odd
(2) under the cost functional Co (a) = [(((w, ([) d the Pythagorean- typa equality holds: (a
Co (a+ + Co (a Va E ORELU-
(3) the even-odd decomposition in the parameter space is equivalent to decomposing any infinitely wide ReLU network as the sum of an infinitely wide network with absolute value activation and an affine function (this has been previously observed by Ongic et al. 2020):
f aa(w, /a(w.b))u b)lux= - b)+ = f xx* <w. b)
we
du + / (w,, (w,
WT
du.
(19)
2
Based on these, we solve the minimal representation cost problem (17) for FReLU by decompos- ing h as in (99 and then solving two simpler problems on FAbs and F FAffine (Proposition 21). where FAffine is the space affine functions on I.. We show C. = Gi on FAbs and C. R. = G2+ G3
on FAmie and thus the implicit bias on FReLU can be described by G1 + G2 + G9 (Proposition
For the case of scaled potentials, where the cost functional is (a) = Da b), O)dμ, the
decomposition technique discussed above no longer applies since the Pythagorean equality generally does not holds. Therefore, solving the minimal representation cost problem remains challenging (see more details in Appendix [).T). To mitigate this, restricted our attention to networks with absolute value activation functions, and solve problem (17) on the subspace FAbs (Proposition 33))
EXPERIMENTS
Mirror flow with unscaled potentials
Approximation error
Effect of position the data
63
dilowal
fildee! dilde
dilow
Error
-22
-2
-33-
-14
-0.5
Figure 1: Illustration of Theorem Left: ReLU networks with 4860 hidden neurons, uniformly initialized input biases and zero initialized output weights and biases, trained with mirror flow common data set using unscaled potentials: 61 = x², 62 |x] + x²² and d3 = x + x². Middle: e between the solution to the variational problem and the networks trained using mirror descent with Φ1, d2 and Φ3. against the network width. Right: ReLU networks trained with gradient descent on five different data sets, each obtained by translating the same data set along the y-axis.
We numerically illustrate Theorem 2 in Figure In agreement with the theory, the left panel shows that mirror flow with different unscaled potentials have the same implicit bias. The middle panel verifies that the solution to the variational problem captures the networks obtained by mirror descent training, with decreasing error as the number of neurons increases. Notice the errors across different
Recall function on (-1,1) [- Bs, Bo] is called even function if a(w, b) a(-w,-b) for all (w,, b) in its domain; and an odd function if (( b) = -a(-w, -b) for all (we, b) in its domain.

---

Published as a conference paper at ICLR 2025
potentials are nearly equal. In fact, we find that the parameter trajectories for different potentials also overlap (see Figure The right panel demonstrates that applying vertical translations to the training data results in much different training outcomes. Only on specific training data, the trained network is equal to the natural cubic spline interpolation of the data (shown as a dashed gray curve). Meanwhile, networks have opposite slopes positive and negative infinity, which illustrates the effect of functional G2 and that when the data lies above the æ-axis, networks have a U-shape and when the data is below the r-axis, they have inverted U-shape, which illustrates the effect of G3.
scaled potentials
Aparaximation errar
Magnitude second derivative
Fid
85-
64-
62-
04-
Densilye
-10
13
80
8:20 xf
30
Left: absolute value networks with 4860 hidden neurons, uniformly initialized input biases and zero initialized output weights and biases, trained with mirror descent on common data set using scaled potentials: Φ1 = x², Φ2 |z| + x², and Ø3 x² + zz For comparison we also plot a network with Gaussian initialized input biases, trained with gradient descent. Middle: L between the solution to the variational problems and networks trained using mirror descent for scaled potentials Φ1, Φ2 and Φ3,, against the network width. Right: the distribution of the magnitude of the second derivative of the solutions to the variational problems for &1 (blue), 2 (orange), and Φ3 (green). The inset shows the second derivatives over the input domain.
Figure 2: Illustration of Theorem 9
We illustrate Theorem numerically in Figure 2 The left panel illustrates the effect of the scaled potentials and parameter initialization on the implicit bias of mirror descent. When fixing the potential, network with Gaussian initialization (dashed blue curve) learns function whose curvature is more concentrated around the origin, compared to network with uniform initialization (solid blue curve). When fixing the initialization, different potentials result in functions with different curvature magnitudes (solid curves in different colors). The middle plot verifies that the solution to the variational problem captures the trained network solutions as the network width increases. The right panel further illustrates how the potentials in the left panel lead to functions with different curvature magnitudes. We see Φ allows occurrences of more extreme small and large magnitudes, whereas Φeg allows the second derivative to concentrate more on moderate values.
5 CONCLUSION
We obtained function space description of the implicit bias of mirror descent in wide neural networks. In the infinite width limit, the implicit bias hecomes independent of the potential and thus equivalent to that of gradient descent. On the other hand, for scaled potentials strongly depends on the potential and in general cannot be captured as a minimum norm in a RKHS. takeaway message 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.
found that the implicit bias
Limitations and future work Our variational characterization of the implicit bias only applies to univariate, shallow networks with unit-norm initialized input weights. Characterizing the implicit bias for deep nonlinear networks remains an open challenge even for ordinary gradient descent. Extending the results to multivariate networks with general input weights initialization requires stronger theoretical tools (see Appendix E.2 for details). We considered separable potentials which allow for convenient definition in the infinite width limit and we did not cover potentials defined on a proper subset of the parameter space, such as the negative entropy. Extending our results to more general potentials could be interesting direction for future work. Other future research directions include deriving closed-form solutions for the variational problems and investigating the rate of convergence the number of neurons increases.
10

---

Published as a conference paper at ICLR 2025
Reproducibility statement Code to reproduce our experiments is made available at https
//github. httpa//aooaaaaaaaaaaaaaaaaaoooaa baorratr--baa-mmrrr---r-eeeeee
ACKNOWLEDGMENT
This project has been supported by NSP CCF-2212520. GM has also been supported by NSF DMS-2145630, DFG SPP 2298 project 464109215, and BMBF in DAAD project 57616814.
REFERENCES
J Harold Ahlberg, Edwin Norman Nilson, and Joseph Leonard Walsh. The Theory of Splines and Their Applications: Mathematics in Science and Engineering: A Series of Monographs and Textbooks, Vol. 38, volume 38. Elsevier, 2016.
Navid Azizan and Babak Hassibi. Stochastic gradient/mirror descent: Minimax optimality and implicit regularization. In International Conference on Learming Representations, 2019. URL
net/ Fo
Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent on overparameterrize nonlinear models. IEEE Transactions on Neural Networks and Learning Systems, 33(1)):7717- 7727, 2022. doi 1011199/
Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference on Machine Learning, pp. (68-477. PMLR, 2021.
Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-175, 2003.
Etienne Boursier and Nicolas Flammarion. Penalising the biases in norm regularisation enforces sparsity. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
https: http::////peeevveev net/ Eorum (iceummmmdd--ee Jiteetiddddaaaii.e 7DAC
Etienne Boursier, Loucas Pillaud- Pillaud-Vivien, and Nicolas Flammarion. Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs. In Advances in Neu- ral Information Processing Systems, volume 35, pp. 20105-20118 18. Curran Associates, Inc., 2022. URL
www.ipeurri..e.. httppppperellleee//eer///222 pdff
Lev M Bregman. The relaxation method of finding the common point of convex sets and its application the solution of problems convex programming. USSR computational mathematics and mathematical physics, 7(3):200-217, 1967.
Dmitry Chistikov, Matthias Englert, and Ranko Lazic. Learning a neuron by a shallow ReLU network: Dynamics and implicit bias for correlated inputs. In seveth- Conference Neural Information Processing Systems, 2023. URL httpps://ppezeeiew...e
1d=xgY40c01E7
Iénaïc Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 1305-1338. PMLR, 2020. URL mlr wtt.mmtw.....i.//./ htw.122/chiiattaa.m html
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. In Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL /tL//tttps///poceeeennng wttpeediiiisssusuiuiee www.:ss.eeerrrr........e...e...s wctpppppppppppppp.....ee...p files/paper/
Hung-Hse Hsu Chou, Johannes Maly, and Dominik Stöger. How to induce regularization generalized linear models: A guide reparametrizing gradient flow. arXiv preprint arXiv:2308.044: 0492/, 2023.
11

---

Published as a conference paper at ICLR 2025
Steven Diamond and Stephen Boyd. CVXPY: A Python-embedded modeling language for convex optimization. Journal of Machine Learning Research, 2016. URL / ettanford.edu/
+beyd/papers/pddt hepppy_ppppr_ppppp paf To appear.
Manfredo Perdigao Do Carmo and J Flaherty Francis. Riemannian Geometry, volume Springer, 1992.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations, 2019. URL httptsss/oppeeeeriiieee net. http://wwww..eeerroooddddeeeeee.
Spencer Frei, Gal Vardi, Peter Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky reLU networks trained on high-dimensional data. In The Eleventh International Conference Learning
Representations, 2023. URL httph: hattpss:/openreviiers net/ http:/ewwnnne///ommmmmmddaaappaa ISEWW
Udaya Ghai, Elad Hazan, and Yoram Singer. Exponentiated gradient meets gradient descent. In Proceedings of the 31st International Conference on Algorithmic Learning Theory, volume 117 of Proceedings of Machine Learning Research, pp. 386 407. PMLR, 2020. URL https://
proceedings . .ml..rr....... wtt.pprrssssvss/ppepepp
html
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2), 2021.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 1832-1841. PMLR, 2018a. URL http:: httt:://proceedingg. .ml
press/v80/ hwtp////uuaaeaaarr. Ba. html
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018b.
Jakob Heiss, Josef Teichmann, and Hanna Wutte. How implicit regularization of ReLU neural networks characterizes the learned function -pat I: the 1-D case two layers with random first layer. arXiv preprint arXiv:1911.022. 1910999,,,,, 2019.
Sigurdur Helgason. The Radon Transform, volume 5. Springer Science Business Media, 1999.
Jean Jacod and Michael Sørensen. A review of asymptotic theory of estimating functions. Statistical Inference for Stochastic Processes, 21:415 2018.
Arthur Jacot. Implicit bias of large depth networks: notion of rank for nonlinear functions. In The Eleventh International Conference on Learning Representations, 2023. URL https
/ //openreview.oo .nt. /tte/efoonaaeeee httt:///ormmiddddood 6idicc 0B-a
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and general- ization in neural networks. In Advances in Neural Information Processing Systems, volume 31. Cur- ran Associates, Inc., 2018. URL. https:tttt//tttt http::///////cceeeeeeeeoeec ngs neur ips wwwppppppp....eeeeeeeeeeeeeeeeep firf_fieee_s es/ paper/2018/f11 e 5a 4bebel faleee22aaaafccce115122222222222
Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106. 15933, 2021.
Sine R Jensen. Sufficient conditions for the inversion formula for the k-plane Radon transform in R..
Mathematica Scandinavica, pp. 207-226, 2004.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019. URL URL/prtppe//ooeeeevviee
net / www.mmiiddoe..............
Hui Jin and Guido Montúfar. Implicit bias of gradient descent for mean squared error regression with two-layer wide neural networks. Journal of Machine Learning Research, 24(137):1-97, 2023. U http:// / // wwo.g//appss//////222
0832 .html
12

---

Published as a conference paper at ICLR 2025
Jianfa Lai, Manyun Xu, Rui Chen, and Qian Lin. Generalization ability of wide neural networks on R. Xiviv preprint arX1v:2302.05933, 2023.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. In Advances in Neural Information Processing Systems, volume 32. Cur-
wwtppppppppppppppppppppppp.p files/ +ttpa///wwoo.aa/////aooaaaoaaoaa
Associates, Inc., 2019. URL http://///...tte///ppaaaaaaaaaea hatpp///0000111err....o.33.66362
Zhiyuan Li, Tianhao Wang, Jason D. Lee, and Sanjeev Arora. Implicit bias of gradient descent on
reparametrized models: On equivalence mirror descent. In Advances in Neural Information Processing Systems, 2022.
http://wwwwweeeeuuuuummmamaaaaaa wttprrummmmmkikkaaaxc....x. zOV
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020. URL http:ttss///opeeeeiiies
net id- 1
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory two-layers neural networks: dimension-free bounds and kernel limit. In Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 2388-24644 PMLR, 2019. /ttpstts::://///tt httpttpp:///rooceediingn
mlr www.prrsseeeee. w99///e11111ee htt.//mmm1a.a.h.m.11...
Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of min- ima stability: A view from function space. In Advances Neural Informa- tion Processing Systems, volume 34, pp. 17749-17761. Curran Associates, Inc.. 2021. URL 111e//44bba333333eeeee111.bbbbb neurips. http://sss.uuppps..eceeppaaee
cc/paper files halles/pape:/221//
Arkadij Semenovič Nemirovskij and David Borisovich Yudin. Problem Complexity and Method Efficiency in Optimization. A Wilee-Intersciencee publication. Wiley, 1983. URL https://
books , gwogle..ooooooooooooooooo com tw http:/occkkkkcss (ttockkkkiid-e00AAAAAAAAA..
Whitney K Newey and Daniel McFadden. Large sample estimation and hypothesis testing. Handbook of econometrics,
1994.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. In 3rd International Conference on Learning Representations, ICI. 2015, San Diego, CA, USA, May 7-9, 2015, Workshop Track Proceedings, 2015. URL /tL:ttp: /tt.t:://axxx.ooot/ org/ abs 1412.41111120 6614
Behnam Neyshabur, Ryota Tomioka, Ruslan Salakhutdinov, and Nathan Srebro. Geometry of optimization and implicit regularization in deep learning. arXiv preprint arXiv:1705.0307,, 2017.
Greg Ongie, Rebecca Willett, Daniel Soudry, and Nathan Srebro. A function space view of bounded norm infinite width ReLU nets: The multivariate case. In International Conference on Learning Representations, 2020. URL htth::t:::/ httttts:://peereiieee..... http:://wpwpwww.peeteteeeeeeeeea wttp::/wwwwwwweeeoeeeeeee.eee.ee
xHKDH
Suzanna Parkinson, Greg Ongie, and Rebecca Willett. ReLU neural networks with linear layers are biased towards single- and multi-index models. arXiv preprint arX5,:2305.15598, 2024.
Scott Pesme, Radu-Alexandru Dragomir, and Nicolas Flammarion. Implicit bias of mirror flow on separable data. arXiv preprint arXiv:2406. 12763, 2024.
Dan Qiao, Kaiqi Zhang, Esha Singh, Daniel Soudry, and Yu-Xiang Wang. Stable minima can- not overfit in univariate ReLU networks: Generalization by large step sizes. arXiv preprint arXiv:2406.06833 2024.
Justin Sahs, Ryan Pyle, Aneel Damaraju, Josue Ortega Caro, Onur Tavaslioglu, Andy Lu, Fabio Anselmi, and Ankit Patel. Shallow univariate ReLU networks as splines: Initialization, loss surface, hessian, and gradient flow dynamics. Frontiers in artificial intelligence, 5:889981, 2022.
13

---

Published as a conference paper at ICLR 2025
Pedro Savarese, Itay Evron, Daniel Soudry, and Nathan Srebro. How do infinite width bounded norm networks look in function space? In Proceedings of the Thirty-Secone Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 2667-26999 PMLR, 2019.
URI wtt.htttts///proceeeeee ngs ml www.pressvvvv wwt 9a html
Alexander Shevchenko, Vyacheslav Kungurtsev, and Marco Mondelli. Mean-field analysis of piecewise linear solutions for wide relu networks. Journal of Machine Learning Research, 23(130); 1-55, 2022. URL
html
Hirohiko Shima and Katsumi Yagi. Geometry of Hessian manifolds. Differential geometry and its applications, 7(3):277-290, 1997.
Donald Solmon. Asymptotic formulas for the dual Radon transform and applications. Mathematis- che Zeitschrift, 195:321-343, 1987.
Haoyuan Sun, Kwangjun Ahn, Christos Thrampoulidis, and Navid Azizan. Mirror descent maximizes generalized margin and can be implemented efficiently. In Advances in Neural Information Processing Systems, 2022. [tt.pttts://ppeneeiiie.
fett/formm?cno www.:ummmddddooooeoeeommemm
Haoyuan Sun, Khashayar Gatmiry, Kwangjun Ahn, and Navid Azizan. A unified approach to controlling implicit regularization via mirror descent. Journal of Machine Learning Research, 24 139,) 2023. URL /tttt//::/ wwtpppappg///aeee/..o
0836. html
Martin Wainwright. High-dimensional Statistics: Non-asymptotic Viewpoint, volume 48. Cam- bridge university press, 2019.
Francis Williams, Matthew Trager. Daniele Panozzo, Claudio Silva, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate ReLU networks. In Ad- vances Neural Information Processing Systems, volume 32. Curran Associates, Inc.,
2019. URL. httpp:://proceeedingss wtw.::/ssssnniss.ronos.occco cc/paper fi esspaper/2019/ [1222211111911eee777//0002222222
Fan Wu and Patrick Rebeschini. Implicit regularization matrix sensing via mirror descent. In Advances in Neural Information Processing Systems, volume 34, pp. 20558-20570. Curran Asso- ciates, Inc., 2021. URL /ttpt///ttss//ppooeeeeeee.ererrr
firrr_iiii__ fww.fffaperrrrrrere pdf
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
Guodong Zhang, James Martens, and Roger B Grosse. Fast convergence of natural gradient descent for over par-merramtter neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Yaoyu Zhang, Zhi-Qin John Xu, Tao Luo, and Zheng Ma. A type of generalization error induced by initialization deep neural networks. In Proceedings of The First Mathematical and Scientific Machine Learning Conference, volume 107 of Proceedings of Machine Learning Research, pp. 144 164, PMLR, 2020. /ttppp/p/pppeeeuuu....p.psss///p
14

---

Published as a conference paper at ICLR 2025
APPENDIX
The appendix is organized into the following sections.
Appendix
Proof of Theorem 2
Appendix
Proof of Theorem
Appendix
Discussion on parametrization and initialization
Appendix
Gradient flow with reparametrizztion
Appendix
Minimal representation cost problem in general settings
Appendix
Riemannian geometry of the parameter space (O, V2Φ)
Appendix
Experiment details
Appendix Additional numerical experiments
A PROOF OF THEOREM 2
In this section, we present the proof of Theorem We follow the strategy that we introduced in Section 3.3 and break down the main theorem into several intermediate results. Specifically, in Proposition 14 we show mirror flow training converges to zero training error and is in the lazy regime. In Proposition 17 we show mirror flow with unscaled potentials has the same implicit bias as ordinary we obtain the function space
gradient flow the network width tends to infinity. In Proposition 22 description of the implicit bias.
We use plim, Xn = X or Xn - X to denote that random variables X,, converge to X in probability - We the notations Op and Og to denote the standard mathematical orders in probability, i.e., X, = Op(an) for every & > 0, there exist Me, Nu > such that P([X,) MMaa) > 1 - for every No; and Xn, (p(a,) if for every > 0, there exists n, such that P([Xnl < Ean) > - e for every n >> fg..
A.1 LINEARIZATION
We show that under mirror flow with unscaled potentials, the network obtained by training all parameters can be well-approximated by that obtained by training only the output weights.
We first show that the initial kernel matrices (0), (0) converge to positive definite matrix the width of the network increases.
Lemma 13. Consider network [1] with ReLU or absolute value activation function and initialization 22. Assume that the potential function satisfies Assumption 7 and that the training data (ae)) satisfies v, & Hy when i + j. and [1
m supp(ps). Then there exists a constant À > which is determined by the activation function, the training data, and the initialization distribution, such that plom
Amiu (0)) = plim, Awin (H
(0) Ao-
Proof. [Proof of Lemma
3 Note for i.j [m]] we have
11++1
H2y(0) n
Vo. f (a 8) dy(0) f(z;()))))
k=1
a (wk in ) (1i x. bz) n²"(()
no"(0)
o x; ba (x) 1).
no"
Noticing (ax ]ee()) are independent samples from the random variable J. A and by sub- Gaussann concentration bound, we have that maxse[m] l s CV/(log n)/n holds with high probability for some constant C > O. It then follows that:
plim (1) = Th) E [o(WT x, - B)o
x, - -B]]
15

---

and Published as a conference paper at ICLR 2025
Define H (d"(0))='C G where Gi = Blo(WT xz B)o 2j B)].
Then we have
plim flim-->o 1-000 H(0) = H In the meantime, we have
(0) = Vax foy; ("(ax) V f(ay()) E H
which means (0) also converges to 4d in probability. Now it suffices to show G is positive definite and let Xn = Amin (.). Consider the Hilbert space of squared integrable functions:
H. {h: supp(u) + R: B[(h(W, B)) < oo}
with inner-product (h, h' = (W, (S]](,W,)'' (W,, B)]. For each 2; E R. define h,(w, b) = o(w"r; b). Since (W, B) is sub-Gaussian and hú is uniformly continuous, h; H. Note G is the Gram matrix of [--) Hence, G > 0 is equivalent to {h+] are linearly independent in H.
in H. which means - C,h,(w, b) = 0 almost everywhere. Without loss I| 2 O. We aim to show e₁ = 0. Notice that, for 0. In partieular, hn is not differentiable at (z1/||21le/) ll,, [|æ1]|2). Now for all j > 1, by noticing that ||,||| s || and œ & Iy we have,
Assume Ei=1 c,a,
of generality, assume | 2 | 2 2... |mmml any [m]] h; is not differentiable (we, b) if and only if w xr;
F1 -
|| = ||ا cos( (((()(z11 x3)) - ||, <0
|
where L(x,, (T)) denotes the angle between x and zy. This implies that h2 is differentiable at (x1//|<z) (12, || (2) and so Ejfı cyhy. Therefore, to make Σ. ca,, differentiable everywhere, we
have G O. With the same strategy, we can show Cj for 2.3...., m 1, since ||=r|| > 0. then follows that all coefficients are zero, which completes the proof.
LI
In the next proposition, we show that the model output ft) converges exponentially fast y and the parameters have O, (n bounded Loe distance from the initialzation.
Proposition 14. Consider singl--hidden-layer univariate network (1 with ReLU or absolute value activation function and initialization (2). Assume the number of hidden neurons is sufficiently
large. Given any finite training data set {I,, yy) bm that satisfies o Iz when i + j. and {||=:||]]= C supp(ps), consider mirror flow 44) with the potential satisfying Assumptions] and learning rate Mu/ /n/n with no > O. Then there exist constants C1 C2 > 0 such that with high probability over initialization the following holds:
sup 10(t) 01100 S Cin 20000 120
1/2 lim sup]| (0)122 = 0, inf Amiu (H( > C2. 120
120
Moreover:
[I((1) - y112 S eexp(= - vl1,, AI 2 0.
Specifically, the constants C₁ and C2 can be chosen
C (0) + 1)K() - vllp C2 = 40)
Au
(20)
No
where Ao = plim Amin Jo) (of 3 max (2)) m is the size of the training data set, B2 max,e[m] || and K a constant determined by B and A.
Proof. [Proof of Proposition 14 To ease the notation, we let J(L) Jo and Q(I) (V2)((()))- For i [m]] and &s E [n]. let (1) a(wx(U (1)² I; bx(()) and (4) = o' (t) x; - bb(()).
By Lemma 13 for sufficiently large n and with high probability, we have Amin (H()) > A2/(22"(0) According to Assumption [] we have 0 < 3 oc.
Consider
T(t) max VTilla(t) a(0)
kE(n] max |[ wx (0)loo, n b()) (0)),,,,nd(( xld()))))) d(0)).
s= E (0, oo): -(1) Ci).
16

---

Published as a conference paper at ICLR 2025
Let 1 = inf S. Clearly, 0 < t1. We now aim prove that 1 +00. Assume, for the sake of contradiction, that 11 < to0. Since (t)) is continuous, we have T(1 = C1.
In the following let denote any point in [0, f1]. Noticing that rtt) S C₁ implies |[((t) ((0) s Cin-1/2 we have for large enough
1
[12(t)l2 Diag( d'(O,(t) 02 (0) ) max ("(O (0)
max xc[-1,1] (²
B.
ke(3+
(21)
Next we bound the difference between H (t) and (O). By concentration inequality for sub-Gaussian random variables (see e.g., Wainwright 2019), there exists K > 1 such that with high probability,
[b(0) I ((( 12 S K KVñ, [|a(0)] ||a(0)ll2 K. Then we have for any 2.j E [m]] h E [n].
(ta((t) - Imx (t) xd + [bx(t) S (B₂ + 1)2K
and
(t) (zx(()] < [xx(t) wx (0) 10 (t) - bba S vd(B, 1)C
For i E [m] we have
Hty(t) - (0) 1 - 12 + 13 11.
(22)
where
ou (t)
Tth (0)o (0)
T1
8" (1) - (0))
(²(0)
p ax (110)
12
("(,,r wk, (0))
(0())
(23)
au (t)² // (t)
ag (0) o (0)o
13
n
d' (t) 6y (0))
(0
1
11 n (d((t)) d(0)) e²(0)
Notice that according to Assumption
(h(-) is continuous and hence is uniformly continuous
[-1, 1J. Then by noticing [[((() 0(0) Icc ( -1/2), we have for large enough
max kE (3n+11
(²(0) (1)))
(24)
0% (0))
It immediately follows that I4 (1).
For I1, notice that, according to the law of large numbers.
plim
Σ
loways = Flo(W" BB - B)}.
RE(M]
Therefore, have 1/ n
EkEin Jankass = (1). Then have
h <
[oae()) ox)()) (oxx(()))) + lou (0)llasx()) (0)
("(ax(t) (0)) )7(o) 1)
1
+ (=0(00))1)))000
< n
Σ
(O,(n ()),-))) + (0,,n + o1) 7
Σ
kE(m]
kE[n]
(1).
17

---

Published as a conference paper at ICLR 2025
For 12, notice that 2² is Lipschitz continuous [-2 Therefore, have
max kEIn]
(ax(t)² an (0))) (000))))))))))) (00)l - (p()).
Then we have
B
I2 -
R,
+ (0)
6" ("(wx, wk, wx oy(0) + (=x)))==x(00)1)))
B 23 2/ (0))l + oz Bn op(1))
n
(1).
Using a technique similar to that applied for 12 above, we have 13 = 0p(1). Therefore, with union bound, we have
[1(()) (())l22 < (1((() - H(O)EF S
E
(1,(()) (46))--H;(00)) H1,( (1).
(25)
(JE[m]
Then with large enough we can bound the least eigenvalue of (t) as
Ao
Amin (H())) 2 Amin (H()) - (1) H())]l2 40)((0)
(26)
Therefore we have
7n [I)(1)) =11
dt. [lf(()) y112 S
(0)
and it follows that
70
[I((1) - 2112 S exp((
(())(())))) 2112.
(27)
(0"(0)
Notice that
[Nazf(t))l S JallJauffle fl S Vm ww (t)l2 (+)) S 2K + 1) mn.
m
(28)
Then with large enough n and for all E [n). we have
d dt
ax(())
Yo
(()lalll( v112
(29)
K
7o
2K(B, + +) mm -mmp-epp------------ no 3
t)lf((0) yll2.
40"(0)
Hence,
aa(()) (dm ()K] 8112 -1/2
(ax(1)) a0 (0))
Cin -1/2
dt
Ao
(30)
Note that with large enough n,
|e,, f(t)ll2 S (mBrlau(t) < 2KVmBB. e (m) [d]:
1 (1)l22 < (m)an(1) < 2K Vm, Vk [m]:
(31)
[|aaf(t) vm.
Comparing (3T) to (28) and using similar computations as in 29 and 30 we conclude that
T(11) < C1, which yields contradiction. Thus ( oo and T(t) C1 for all > 0. Meanwhile, notice that (26) and (27) hold for all 2 0, which gives all the claimed results.
We point out that Proposition 14 can be extended to the where only output weights trained, without requiring additional efforts. We present the results in the following proposition and omit the proof for brevity.
18

---

Published as a conference paper at ICLR 2025
Proposition 15. Assume the same assumptions in Proposition 14 and consider the mirror flow of
training only the output weights of the networks, as described in (10p. Let C1, C2 be the constants defined in 20). Then with high probability over initialization the following holds:
sup ([(()) - 011m0 S <Cin 120
lim sup (1) - (0) || 0,
n- 120
and
ff(()) 2ll2 S exp((>nC2 v122 A2 0.
Next we show that the network obtained by training only the output weights can be well-approximated by the network obtained by training all parameters.
Proposition 16. Under the same assumption in Proposition given any E Rd,, with high probability over initialization,
lim sup [z 0()) (()) 120
f(x, ((x))(()) 0.
Proof. [Proof of Proposition
16 For E [m] and E [n], let o' (1) = o' d'(wxx(1) bb())). Under initialization 2 have ||a(0]]
(1) and thus
HH,(0) H2y
5 no" (0)
la (00(0))0d(0000 (0) +| n²"(()
kE(m
B² +
(32)
[a(0]l2+
no"
Op(n
Note the size (0) and (0) is independent n.. Therefore, [((()) = o1).. According to Proposition 14 and Proposition 15 we have
sup - (1)l2 sup [H()) - ((0)ll2 sup x20
o1).
120
120
(33)
In the following, we use the notation = sup:20 (t)) (t) Now we show ft) remains
close to ((). Let (()) = f(t) 1000-ffo (() and C2 be the constant in Proposition 14 Consider the function
vu() = |r()]]2 defined on f. 2 0. Clearly, u(t) is continuous. Since u(() = Of and only if r(t) = O. we have that u(t) is differentiable whenever u(t) O. Note by Proposition 14, we have that, with high probability, Amin (t)) 2 C2 for all t 2 0. Then when u(t) * O, we have
d (e Conot "(t))) = dt dt
d
Canot r(t)l3
2e Camot (())) d (e
di
2e22 (C2nor(1) - -100) H()(((() + + ++)((() - -v) (())(((1) -9) -H(0 < 2mp 2 2C2not ( (t)
2np 202no (r(1)
((()) -9)
< 20p 2C
u( (t)) ((())aa))(())) y112.
(34)
Meanwhile, we have
(e (e 26 Cono ml)) d (eCono dt
dt
and hence by
34
d (e (Canou((t) S noe O2nat | ((t)leel(((t) - 8112-
dt
19

---

Published as a conference paper at ICLR 2025
Now for any t 2 0, consider t' = inffs E [0,t]: r(s) + 0}. Then we have
Canmou u(() ' no (2m0 -()) (())l22l yl|2ds.
Hence for all 2 0, we have
||r()]l2 se Canot [ No00 (s))laal))(s) yll2ds
S noe Canot [ ff(0) - yl22ds
(35)
noe Canot [f(0) - yll
Given x E RR, by the chain rule we have:
dt f(x (0,,0(1)) (100))))00(())( -
(36)
where (tt, 1 -.l_
and
d dt (zx,((1) (,.0))))=m(((.)(00)))) y),
(37)
where & (
Similar (22) and (23). one can show suP:20 [ = -1) and sup |
H 2)l2) (1). Meanwhile, similar to 32 one can show [0, x) HO, x))22 = (1).
Therefore, we have
A₂ (((,x)) Ht, x))12 = op))
(38)
Notice that for both ReLU and absolute value activation, we have 1a(2)] S | and [o'(2)] s 1. Then for large enough we have
Jo Je
Ww (0) x; bxx))))
au (0 1))
(39)
2 21| (0) + 20) Iax (0) +1)
R= <
2B b(0) (B) + 1)]1a
Using concentration inequality for sub-Gaussian random variables, have |b0)] and (0) can be bounded by C'vn with some constant C > 0. Therefore, by (39 we have [ Op(vn). Note similar bound applies to [|Ja/(x; 0)]|2. Then for large enough have
1 [afs(z 0)112 *^(0)
(0, x)I22 <
|1af(00)l = O1)))
20

---

Published as a conference paper at ICLR 2025
Then using (35) and 38 for t 2 0 we have that
d
fzz,
Jo ds (f(z, f(z, ((,)
1 noll(s - y) m) y) ds
< no (1) x)) - 22s
+ I (H( Hss, - yl12ds
$ O,(1)A / Camu sds + Ap) / c-Conosss
s 0 + Op(1)Ar (I),
which gives the claimed results.
A.2 IMPLICIT BIAS IN PARAMETER SPACE
As we discussed in Section 3.3 the final parameter of the networks whose only output weights are trained is the solution to the following problem:
min aER
Do ( a())
s.l. Σ (n((uc(()) (x)] yt, Vi [mm]
(40)
kE(M)
In the following result, we characterize the limit of the solution to 40 when n tends to infinity. Proposition 17. Let go and defined as follows:
In (x, / l(ww,)) b)(w b)+ dun
(41)
a(z,a)= = / (aw, b)lw - -6]+ du.
Assume à is the solution to (40) where the potential T satisfies Assumption ! Let ORALU = fa E C(supp())): is uniformly continuous). Assume an +RRELU satisfies that ā, (w (0), by (0)) n(ā 0g (0) for all l E [n]. Assume ā is the solution of the following problem:
/ ((a(w.)) ((d(w,)))2ap du,, s.t. g(z(, = yys f 0), Vi E [m].
min agGRaLt
(42)
Then given any x E R4,, with high probability over initialization, lim (x, ā..) - g(x, ã)]
Proof. [Proof of Proposition 17 Recall (0), by (0) denote the sampled initial input weights and biases, and μ4n denotes the associated empirical measure. We use the notation ||x for the (pmn)-smmmnn in OReLU, i.e., for a E +ReLU.
1 Σ (0) n
||al
We write OReLU/1-1 for the quotient space of +RELU by the subspace of functions a with ||a]]n = 0. In the sequel, will use the notation represent the corresponding element (equivalence class) in ORA.U/ /I-lll if there is no ambiguity. Let )) denote the inner-product in OROLU/IIH II-.. defined as: for a, E OReLU//HH..
(0,0") = Σ a (0) (x))
21

---

Published as a conference paper at ICLR 2025
Note equipped with (+, .) OReLU/1 ||s. becomes a finite dimensional Hilbert space and is isometric to IR".
Notice that for any a R" and a that satisfies (x) n(aa axx(O)) for & E [n], have
Da a(0)) = Σ
daax ax(0)) <(0) (({ an (o))
kE(m]
a(w,b)) b) ) d'(0) alw, b) djp.
= n Φ
n
Define (,(w,b) b) = lw x, - b1+ E ORELU. Since a is the solution to 40 we have that a, solves
minaEGRaLo [o) s.t. (a, oi)n y/ fx; 0), i [m]:
where Ln (a) 2 n² /0l a(w,b) o(0) ø' d'(0) <x( b)
(43)
dun.
n
At the same time, let a₂ be the solution to the following problem:
min (a) s.t. (a, (s) = -(x,, 0), E [m],
where (a) - n² f (0(0) alw, b) ddn.
(44)
n
Note the objectives and constraints in problem ORLLU/I-II- Specifically, for any a, a' E OELLL if ||o a'I1n 0, then and a' are equivalent with respect to feasibility and objective function value in both problems. Therefore, we shall address these two problems OReLU/ |-]|,. Note Ln and Ln are strictly convex in OReLU/I-II. Hence, on and a' are unique in the L (1,,) sense.
and problem 44) are essentially defined in
Consider the Taylor expansion with Peano remainder: o(=) = 6(0) + +00))++ 20 R(2)=² where lim R(2) Then we have
[L)(a) - Ln(aa)= / R( a(w, b) )a (w, b)dunl
"
(c) (0))
< max R kE [m)
||a:|
15 ān and 0', are determined by the limit points of mirror flow 10) with potential (0) Σn (02 0%) and x00) Ex 0z)2 respectively. Specifically, we have that maX 0e (w (0). by (0))/n maxse(n) au (0) O, -1/2) and that ||ā.,l1 nlla (0) = (1). Similarly, we have maXkeju] a', (w,(O), by (0))/n = Op(n 11) and ||ā ll? =
By Theorem 11 and Proposition
(1). Therefore, have
[L)( (ān)). (L.((r)))) -= o
(45)
Using 45 and noting the optimality of ā. and α',, we have
(a,)) < Ln (ā.) + 01()) S Ln(d')) + (1) s L2 (d,) + (1).
(46)
Notice that, in the sense of Fréchet differentialbiliity L, is twice differentiable and d" (0) convex. Thus we have
6²(0)
12 14,(a))) + (VI, on on), +
2 a a'
(47)
where VI denotes the Fréchet differential. Note (I, = =0"(0)a', Then combining (46) and gives
(o"())"", an -,,) + (0)
aa - 22,12 < (1).
4
22

---

Published as a conference paper at ICLR 2025
By the first-order optimality condition for problem 44). 0' be written linear combination of {ar ym On the other hand, by the constraints we have for all E [m]]
(ūnn O1) y: [ (()) (y) [(z,, ((0))) 0.
Therefore, (dn an - dn)n = 0 and hence
||ān a'zlln = 0)p())
For given data æ E Rd, notice that
(0) bba < ||| 1 [b(()] O2(1).
lo x b) S
1=1
Therefore, we have
(gn an) - 9n 9n (())
= ((n(ws b) - (w, b), oww 6))ml
I lān 2 x b)1lm
(48)
= ol()
Now it remains to show that 19, (z, a',) g( a)) op(1). By the optimality condition for problem a', can be chosen as follows
%% (w, b)
o), (w,b) supp(4);;
where the Lagrangian multipliers (A, (16 ) are determined by the following random linear system
x (m) oi o3), = Vj. j = 1....,m. m.
(49)
Note that a', is uniformly continuous. Similarly, by the optimality condition for problem
42 can be given by
m
ã [. b)
Σ (w, b).
where the Lagrangian multipliers (M)) are determined by the following determinant linear system
Noi σ3) Vis j Lm.
(50)
Here (5-) denotes the usual inner product in (n), i.e. (J.g) = fg du. Note that {ax) [em) are all continuous functions and that y is compactly supported. Then by law of large number, we have (a. (r).o,), converges probability (o,, o,)) for all i, E [m]] which means the coefficients in
system 49 converges in probability to the coefficients in system (50). It then follows that
[|A0)) - A|l = (1).
(51)
Given fixed E Rd, recall that
In ()( ((w:b)(w² b)+ dumi
g(x, 6) = / a(w, b) b) To x - bl.
Noting that ā and the ReLU activation function is uniformly continuous and that μ is sub-Gaussian, by law of large number we have
(g)( ((x, (1).
(52)
23

---

Published as a conference paper at ICLR 2025
On the other side, we have
/ Σ Eie 1!th oos(w;)) (w" [b) dynn
(x, - In go dn))
$ || 1lo / Eo((w,b) [w- x - 61+
Again by law of large number,
flu== b]+ (w, dun & fwwT (0-2-b],- b1+
m
Σ f=1
m
of w. b) du.
Hence, with large enough and with high probability, f'w" b|+ Σド= or (w, b) cle, can be
bounded by a finite number independent of n. Therefore, by 51 have
9 (z. @) 9n (z, d',' o1))
Finally, according to (48) and 52 we obtain
[g(x, (i) - In (a,, an) S (g(a,, a) In (r, ai] 19n - go (2, d')) + (x, gn an)l
(1),
which concludes the proof.
Note the above proof can be directly applied to networks with absolute value activations. Then Proposition [7 for networks with absolute value activations is stated as follows.
Proposition 18 (Proposition 17 for networks absolute value activations). Let In and g defined as follows:
In (2, a)
a(w, b) "0l""== djt
(53)
g(I, a))
a(w, b)|w" To bl dp.
For potential t that satisfies Assumption 7
assume à is the solution to the following problem:
min (e(ā,a())) s.t.
Σ
āu wa (0) (0) x, bx() = yvs Vi E [m]]
aCR"
kEin]
Let Aas = {a C(supp(a)): : a is even and uniformly continuous) and on E @ be any continuous function that satisfies on (o)) n(Øs ax())). Assume is the solution of the following problem:
/ (a(w,, b)) s.t. g( a) = 31 f(x 0), Vi [m],
min aE
Then given any x E IRd, with high probability over initialization, lim Ign an) g(x, a)
A.3 FUNCTION DESCRIPTION THE IMPLICIT BIAS
Assume the input dimension is one, i.e., d 1. Next aim characterize the solution the function space:
4>
in
b)(wxr - d/p: a E ORRLU }
FRELU =
Recall the following result from classical analysis.
Lemma 19. Assume that (1) is continuous function defined neighborhood U of zo R: (2) f is continuously differentiable on U - {2o}; and (3) lim - f'(x) = lim, limz-1z7 f'(x) = k for constant E TR. Then differentiable Do and f'(xx) = k.
24

---

Published as a conference paper at ICLR 2025
Proof. (Proof of Lemma 19
Consider > zo and x U. Since f is continuous on [ro z| and continuously differentiable on (Io, 2), by mean value theorem there exists x' (To, such that
(x) f(xo) f'')). I Zo
Considering the limit process x + 20. we know f is right differentiable at zo and
(x0) lim f'(x') k.
x'
Similarly, we know f' is left differentiable at zo as well and
(ro) (2o) k.
This implies f is differentiable at zu and fl(xo) = k.
In the following result, characterize the analytical properties of functions in FRELU- Proposition 20. Assume h E FReLU. Then (1) h is continuously differentiable
IR; (2)
h0(-20) = lim (x) and (oo) = limm-++a< (((x) are well-defined; (3) h is twice con- finuously differentiable almost everywhere on IR; (4) supp(h") C supp(ps)); and (5) Jg h² (b)]jjáb well-defindd.
Proof. [Proof of Lemma 20) Assume (xx) takes the form of h(x) f aw,, b) - bj. Notice that a and ReLU activation function are uniformly continuous and that y is sub-Gaussian. Hence, h(x)) well-defined for all x R.
By our assumption (W, B) is symmetric and B has continuous density function. Hence, we have Pb(b) Pw.s(1,1))) (1 + Pw.s(-1, = pw.a(-1, (-1, bb) + Pw.B(1, bb) = PB((b),
which implies ps is an even function.
Since ps is a continuous on R, (supp(ps)) i.e., the interior of supp(ps), an open set. Hence, (supp(ps)) can be represented by countable union of disjoint open intervals. Without loss of generality, we assume supp(ps) [m,M] where m, M E (00, -00} U.. R. We have
1 / a(1, ((1,b)p((b][z db + 2 h(x)
1 MM a(-1, bb)p((bb db.
/
2
first show continuously differentiable on IR. For x E (supp(ps))) we have that
h(x) / _" ((b)))) - b)db + 2
f M a(-1, ((-1,-b)ps(b)( æ)db.
By Leibniz integral rule, differentiating fizz) with respect to x gives
1 h'(x) 2 /" a(1, a(L)b))a(b)db 2
[ M
a( (b)p
(54)
Note (4) implies that h is continuous on (supp(ps))). M is a finite number, we examine the continuity of h M. For x > M, we have
1 M
Mxx)
a(1,b)ps(V)r= b)db.
Again by Leibniz integral rule we have
h'(x) = / M
((b)d))
(55)
Clearly, (a:) is continuous when x: > M. According to 54 and (55) we have lim ((x)) = lim h((x). By Lemma (19). h is differentiable at M and h is continuous M.
Therefore, h (x) continuous when x > M. Noticing a similar discussion can be applied to m, we
conclude that h/ is continuous R.
25

---

Published as a conference paper at ICLR 2025
Next, we show (oo) and (-oo) are well-defined. Note by (54) and 55 have
1 M
('(o))
a(1,6)pp (b)d(bbb.)
(56)
2
This integral always well-defined since uniformly continuous and ps (b) has heavier tail than e-b² Similarly, we have
1 - / a(-1, (b)=,(d N' (-(-) =
(57)
which implies (-() is also well-defined.
We then examine the second-order differentiability of h. For x (supp(pa)) differentiating (54) gives:
1 1" (x) 2 [1 + (000)880)(2220011) -0--),-))s(()))
1
(58)
=ps(b)a((L1)
which implies h is twice continuously differentiable (supp(ps)) For R (supp(ps)) notice such a exists only if m or M a finite number. For the case of M < oc, by (53). we have h²(x)= when
M. Similarly, if m -00, have h((x) 0 when m.
Finally, m = - oo or M oo, by 58 and by noticing that a uniformly continuous and that B is Sabb-Gasssian, we' ve just shown h" is compactly supported. Therefore, h² jbjdb is also well-defined. By this, we conclude the proof.
have f (b) db is well-defined. If both and M are finite numbers, then
Recall that, for any ReLU, a can be uniquely decomposed into the sum of even function and an odd functior 6 a(w, b) = a (w. b) + (w, b). where a b) = (a(w, ww b) +a(-w, bb))/2 and a (w, b) (a(w, b) - a(-w, w. /2) In the following result, we present the closed-form solution for the minimal representation cost problem (17) with the cost functional described in problem (42)
Proposition 21. Given h E FReLU, consider the minimal representation cost problem
/ ((((w))) d)) s.t. h(x) = g(x. a.,, Vx R, min oEGRa
(59)
where g(,, a) is defined in 41). The solution to problem (59) denoted by R(h), is given by
(b)
(R(h))+(1,B)
Pb(b)
(60)
JR (bbbbb/ 2h(0)
(R(1)) (1,b)
+ 1/(00) + (-(0))
E[B²
where h1(+00)) = -iim (x) and h (-co) = lim l()x))
Proof. [Proof of Proposition 21 Fix an E FReLU. We now show that the constraints g(:, = h(-))
in problem (59)
equivalent to
h²(b)
(1,b) Pis(b)
(1, h1(+00) + h'(-oo),
(61)
a (1, / h² hb(b)]bldb (h()).
We first show 61 is sufficient for the equality o(; a) h(.)) For E ONeLU satisfying (61) consider the function
s s(z) = h(x)) g(. a).
Recall function on (-1,1) B,, Ba] is called even function if a(w,)) a(-w, -b) for all (m0, b) in its domain; and odd function a(10, = -a( b) for all (we, b) in domain.
26

---

Published as a conference paper at ICLR 2025
Clearly, s(z) FRELU- By Proposition 20 for & supp(ps). we have s²(x) 0 0. For x E supp(ps), applying the computation given in (58). we have
<"(x) = h+) pb(b)a (1,6) =0.
This implies s' (x) is piece-wise constant function on R. According to Proposition 20 (z) is continuous. Hence, s' szxx) constant on R. Using similar computations as presented in (56) and (57) we have that lim 0, (x. a) + lima-+ dzg(z, a.) J a (1, b)dus. Therefore,
s'( (-o0) n+(+00) n'( M((-)) -
(1,b)dug 0.
Therefore, s((x) for all R and s(x) is constant function. By 66 have
g(0, a))
a(w, b)l +
bl
-b + (w: b)
at (w, b)
dj
2
2
dus + / a (1,1, b1, dpg
(62)
at (1,
2
16
db / (1,(1)) b
1²(b)
(a=(1b))
JR
2
h(O).
which implies s 0 and hence h(x) g(2, a) for E.R
We now show (61 is also necessary for the equality g(: = h(-). Assume g(,aa) = -h-).. It's
clear from (56). (57). and (58) that must satisfy the first and second equations in (61p. Noting
that 900 a)) h() and according to (62). we have a must satisfy the third equation in 61 as well.
Therefore, ( a:) h(-) equivalent (61) and the minimal representational cost problem (59
equivalent to the following problem:
/ a2, b)dp s.t. satisfies 61
min aEQRaU
(63)
We now explicitly solve problem
(63 To simplify the notation, we let
Sh h h'(+o0) + (oo-,))
(64)
C h (bbbb]| (0.
Since the measure μ is symmetric, the following Pythagorean-type equality holds for any a in R
f (a(u, 6))
/ (a+ b)) + f(a (a-(w, (w, d)
/ (1,b)))) / ((a (a
Note in 61 (at(())) is completely determined. Hence, solve 663 it suffices solve
/ (a (w, d) s.t. [f
a (1,b)duu Sn,
min E odd
If a (1,,,))ddp = Ch-
(65)
The first order optimality condition for the above question can be written as
2x (1, b)
X Azb = 0,
(66)
where ^1, A2 E IR are Lagrangian multipliers. Plugging in (66) to the constraints in 65 gives
E[B]
A1 +
X/ Sh.
2
E
E(B²2
Ch.
27

---

Published as a conference paper at ICLR 2025
Since the measure is symmetric, the density function ps must be function R. This implies [[ = O. Therefore, A1 2S,, and X2 / Finally, R(h), which is the solution to 663 is determined by (R(A))+(1, = (B(6)
(b)
Ch
(R(h)) (1,
+ Sh
E
which concludes the proof.
Finally, with Proposition 12 and Proposition 21 we can write the implicit bias in the function space.
Proposition 22. Assume solves 42 then h(x) J a(w, bw 6]+ solves the following variational problem
min NEFRELU
(1) + 92(h) + (s(h)
s.t. h(x,) - ys -((x,, 0), vi E [m].
(i) S... _B, pT(
Bo (x))
dr
where G2(h) = (h'(+00) + h2(-2)))
Ga(h) ETB²] h"(b)]bldb (((6))(db2hh00)
Proof. [Proof of Proposition 22 By Proposition 12 h(-) solves the following problem
/ (R(N)(w, b)) dp s.t. h((;)) = y1 - [(xs, 0) E [m]]
min hEFRatu
Notice that
/ (R(h)(w,b))
(w, (w. / ((R(A))* = ((R())* (1, + ((R(h)) (1) 2 dus.
((R(h))T( (w, dd
(67)
((R()))((1))))
Recall Sn, and O₁ defined in (64). Then plugging in 60 to 67
gives:
/
()) =
h1 (x)
dØB /
+
Sh
dus
PB (x)
Bn (h" (h"(x)) (z)))
Br
B,
dx
x² x²
Sh __B, B,
dus
B, PF(
[
B.
Ba (h"(r))²
C2
dr
+
S2.
pb(x)
E[
which concludes the proof.
The above results can be summarized by the following commutative diagram:
RRLLU
+
Godd
[Ro
R
Co
F
FAba
Co
FAttine
Ca
(68)
G1
G2- Ga
R..
We state Proposition 21 for networks with absolute value activations below. The proof can be viewed special case of the proof of Proposition 21 and thus is omitted.
28

---

Published as a conference paper at ICLR 2025
Proposition 23 (Proposition 21 for networks with absolute value activations). Let FAbs
{h(x) = - b)|ur - bbdμ: a E C(supp(u)) a is even and uniformly continuous). Given h E FAbs, consider the minimal representation cost problem
min e
(a(w, dl s.t. h(x) = g( a), Vx E E.
(69)
The solution to problem
denoted by R(h), is given by
(b)
R(h) b) = (p
With Proposition 14, Proposition I6 Proposition [8] and Proposition 23 we obtain Theorem for networks with absolute value functions, which is stated below.
Theorem 24 (Implicit bias of mirror flow for wide univariate absolute value network). Consider a two layer absolute value network (T) with d input units and n hidden units, where we assume is sufficiently large. Consider parameter initialization (2) and. specifically, let ps(b) denote the density function for random variable Consider any finite training data set ((r,, ys) m that satisfies supp(ps). and consider mirror flow (4) with a potential function satisfying Assumption and learning rate n = ((/nn). Then, there exist constants C1, C2 > such that with high probability over the random parameter initialization, for any t 1 O,
Ia F 2 when i + j and {1[z; | mm
[10(t) - 011o0 S SCin
lim (H( H()) 0, (tf 2ll2 S - (l.f(()))) - y112-
1-000
Moreover, letting ((00) lim O()) we have for any
aorany given R2,, that limmn-soo (f (cc))
i.e. mirror flow f(z, 0GF (po)) = O, where (CF( denotes the limiting point of gradient flow, flow, with (0) |0.12. on the same training data and initial parameter.
Assuming univariate input data, i.e., d = 1, we have for any given æ E R. that limon (f ((00))= f(x 0) -](=)]== O, where T((.) is the solution to the following variational problem:
(x)) dr s.t. h(x;) yi -(x, 0), Vi [m]. REF supp(ps) (B(
min [
Here
= (h(x)) = a(w, b)lw soc--blot bldp: E C(supp(u)), is even and uniformly continuous) is the space of functions that can be represented by an infinitely wide absolute value network with an even and uniformly continuous output weight function a.
Before ending this section, we present the proof of Proposition 12
Proof. [Proof of Proposition 12 Note R. is well-defined map since Cp is strictly convex on
O. By assumption, problem 14 has a unique solution. Notice (g(((()))
feasible for (14).
Hence, by the optimality of ã, we have (g ,)) ā, i.e. & is the parameter with the minimal
cost for a), Now assume i solves (17). By noticing ạ(;,a) is feasible for (18). we have
R. (h) < (g),R((() (a.). Since R.g. (f) is also feasible for problem 14 R (h) and therefore h(()) = g(),ā).
A.4 THEOREM 2 FOR UNCENTERED POTENTIAL
In this subsection, we characterize the implicit bias of mirror flow with "uncentered" and unscaled potential functions that satisfying the following assumptions.
Assumption 25. The potential function tb 333+1 + R satisfies: (i) can be written in the 3n+1 ((0x), where & is a real-valued function on R; (ii) de is twice continuously differentiable; (ii) (2)((0) is positive definite for all 0 E R33+11
form (()) k=1
Theorem 26 (Theorem 2 for uncentered potentials). Consider a two layer ReLU network ① with d input units and hidden units, where we assume is sufficiently large. Consider parameter initialization (2) and assume B is compactly supported on [-B. B, B) for some > O. Let (b( denote the density function for random variable B. Consider any finite training data set ( vi) vh) m
29

---

Published as a conference paper at ICLR 2025
that satisfies 2; # x, when i j and {1|a:1211" supp(pe), and consider mirror flow 4 with a potential function & satisfying Assumption 25 and learning rate n, = ((1/n). Then, there exist constants C1, C2 0 such that with high probability over the random parameter initialization, for any 2 0,
[|O(t) Ol S 1/n
lim ( o, (t u1l2 SE -o n00>+ 1112. (70)
Moreover, letting ((00) lim f(t). we have for any given R that limn (f ((0)) fæx BGF (oo) = O. where (c) denotes the limiting point of gradient flow, flow, i.e. mirror flow with (0) |0 on the same training data and initial parameter.
Assuming univariate input data, Le., d = 1, we have for any given x E R, that lim (f ((0o))
f(z 0) h(x)l = O, where h(()) is the solution to the following variational problem:
min bEFRELU
(1(h)) + G2(h) + Gs(h)
s.t. h(x;) = yı - fxx, 0), vi E [m]]
(1(h) [ !aupptee) PB( dr
(x))
(71)
where
92(h) (W0+)0) + 11 (-o0))
Gs(h)
- h1' h'(-)o)) (+00)++0))
(h(B) h(-B)))²
E|
Here FReLU = (h(x) = a(w, ww - b a E C(supp(()), a is uniformly continuous), h'(+000) lim h((x), and ho)- = lim hz).).
The proof of Theorem mirrors the proof of Theorem 2 which we discussed in previous subsections. For brevity, below we only provide the proof sketch and highlight the differences brought by the change in assumptions.
First, we show Lemma holds for uncentered potentials. Notice that
3m++1
Hoj(0)
Vo, (x;;
Vo. (zz;;
("(Ox)
k=1
5 ( v2 bo - ba) + (dd
n
("(a))
k=1
1
dg (w) x; - bx) x. bu)
Fi 8d (wx,r)
("(bs)
By sub-Gaussian concentration inequality, have max låvl = O0(1////p)) (2ogg)))) (loggn/mm By noticing that max) a (w) bx) = (1) since W and B are bounded, we have
1 o (wT xi b1)o - bb)) 0, (0)
(ak)
which implies (i)(0) P Ho. The remainder analysis proceeds as in the proof of Lemma 13
Next, show Proposition 14 and Proposition (onniitioonnnoolddd 16]hold. Note the key ingredients which needed in proving these two results and which relate to the potentials are: (i) there exists a constant 3 > 0 such that with high probability, ||Q(0)]22 3 and [[Q(1)]2 < 3 if [[((I) ((0) lo o, and (1) if [|@(t) ((0) = (1), where we use the same notation as in Proposition 14] We verify these two properties for uncentered potentials. Since is compactly
(ii) [[Q(t) Q(0)]|2
supported, we have with high probability, maxxe(3n++ 10x (0) < C' for some constant C' independent
of n. Since 1/0"(-) () is uniformly continuous on 2C, 2C', we have
[IQ(t) Q(0) max
= (1).
keE 1]
(*(((tt)
("(Ox(0))
30

---

Published as a conference paper at ICLR 2025
given [0(t) (c0) Op(1). Meanwhile, notice that
max ke be[3n+11] (0)) + +1 ("(O((t))
1
[[(()]l2
max kE(3++1] (0)((tt)
('(0x
By the continuity d" we have [[Q(t)](2 ||2 can be bounded by some constant 3 > with high
probability, given [[0(t) 0(0) (1).
Then, we show Proposition 17 holds. Notice that for uncentered potentials,
D& a()) = Σ
(o(ax) s(au (0)) - ('(ax - (a))))
kE|m]
8 (0))))))) + R(ax - ax (0)) (ax - ax (0))²
2
kEIn]
where R captures the error produced by Taylor approximation and satisfies lim R(=) = 0.
Consider L, (a) f v(w, (a(ww, b)) dun, where v(wx(0), (0), by (0) (0)) + aa 0% (0))
for k [n]. Hence, if solves (40), .n. which satisfies ā,, (0), by (0)) ag ax (0) for E [n].
solves the following problem:
min (a) s.t. (a, 0) = yi - f oe DE
Let 0% be the solution the following problem:
minageR.Le (a) s.t. (a, 21 [ax, 0), i [m],
where (a) f
o²(0)
(a(w, (o(w;b))'r ddj..
2
Notice that
[n((a) - Lnaaa) =1/ (u(w,))
_0((0) -0) a (w, b)dual
2
6" (a))
(0(0) R(as ax (x)) ||a1
< max kEim
2
Following the same analysis in the proof of Proposition
17 have with high probability Ozn 2 = Op(1), and |0, O₁(1).
maXse[и) 0A - (0) = (1), maxke[n] de - (0) = 0p (1),
Meanwhile, noticing that maXke(m) (ax Op( (llooll n)/n) holds with high probability and by the continuity of d", we have maxke[m] 10" (ax (0)) 60()) (1). Therefore, we have
(ā,) - (a,))., (ā,, =1,
The remainder analysis proceeds as in the proof of Proposition 17
So far, have shown that the training converges and the implicit bias in the infinite width limit can be exactly described as Proposition 17 Then Proposition 22 directly applies and gives the claimed results.
training error, the lazy training occurs,
B PROOF OF THEOREM 9
In this section, we present the proof Theorem We follow the strategy that we introduced in Section 3.3 and break down the main Theorem into several intermediate results. Specifically, in Proposition 27 we show the mirror flow training converges zero training error and is in the lazy regime. In Proposition 32 we characterize the implicit bias of mirror flow with scaled potentials in implicit bias.
parameter space. In Proposition 37 we obtain the function space description
We use plim, X,, X or X, P X to denote that random variables X,, converge X in probability. We use the notations O, and op to denote the standard mathematical orders in probability (see Appendix A for detailed description).
31

---

Published as a conference paper at ICLR 2025
B.1 LINEARIZATION
We show that for mirror flow with scaled potentials, if the Hessian of the potential function is flat enough, the mirror flow converges to zero training error and in the lazy regime (Proposition (7). and the network obtained by training all parameters can be well-approximated by that obtained by training only output weights (Proposition 29 Further, we show that if the potential is induced univariate function of the form 6(2)
(+(2) + wz²) with convex function s and positive
constant in Assumption 8 one can always choose sufficiently large w to obtain the sufficient "flatness" (Proposition 30).
Note that under mirror flow with scaled potential, the evolution of the network output is governed by the following matrix:
(l))
where, by Assumption g the inverse Hessian matrix takes the form of
(V22(0(())~ Diag ("(n(0x 0))
It's clear that the behavior of the Hessian map is determined by the behavior of Φ. In the following analysis, we consider the following functions to measure the flatness of the Hessian map: for C 2 0,
K1 (C)
(x)) K2(C)= = sup =E|-C,C1
inf
(d" ka(C)= = sup =E(-C.C)
(x))
=E|-c,0]
Note function Hi (C) characterizes the smallest possible eigenvalue of (V when satisfies
||0 - 0 < Cn-1, Similarly, characterizes the largest possible eigenvalue. Hence, k1 /1/M2 captures the largest possible condition number of the inverse Hessian matrix. By definition, Ki /22 is monotonically decreasing. In the following Proposition, moderate rate such that takes a larger value than L/C for a particular constant L determined by the training data and the network architecture at some point C, the mirror flow converges to zero training error and it is in the lazy training regime.
show that if K1 K1/K₂ decays with
Proposition 27. Consider a singlehhideen-layer univariate network with ReLU or absolute value activation function, d input units and hidden units. Assume is sufficiently large. Consider parameter initialization (2). Assume B has a compact support. Given any finite training data set that satisfies I4 F Øj when i 4 j and {[]=;]1a) ) C supp(ps), consider mirror flow (4p with learning rute 1 n,,/n, yoo > O and potential function t. Assume &b can be written in the form of (0) 1 Eznt1 ((n(0,- 6k)), where & is a real-valued C -mooohh function on IR and has strictly positive second derivatives everywhere. Assume there exists constant C₁ > O satisfying
(1 ye) ) ==
Ky (1())) (1)
Sym(Ba + ()) yll2 CI'
(72)
Aa
where Ao plim Amin (o) m is the size of the training data set, Ba maxie[m] |2,112, and K is a constant determined by B and A. Then with high probability over random initialization.
(1
sup @i(() - 0l1oc Cin inf Amin (H())) 120
t20
4
Moreover,
(I)(L) y112 < Sexp (-
Ki (Ci)do
nol ) y112, 14 0.
4
Proof. [Proof of Proposition 27 To ease the notation, let ((t) (t) and (V²) For i E [m] and l E [n], let au (t) = o (wx(t)" x; by (t) and d' (t) = (w) x3 (t)).
Assume C1 satisfies (72). Consider the map :: [0, oc) - R. r(t) n|((()) Ollocc Consider
t = inf 2 0: r(t)= > ).). Clearly, t > 0. Now we aim to show t = +oo. Assume t < co for the sake of contradiction. Since r(.) is continuous, 7(4) C₁ and (()) < Ci for / [0,4]]
32

---

Published as a conference paper at ICLR 2025
In the following, let be any point in [0, 11). Since n n1l@(t) -11100 C1 we have
1
min (Q(()) -min 6" (m(((((())
(C1 0.
By Lemma 13 with sufficiently large n and with high probability, Amin ( {(0))(()))) > As For i [m]] we have
T
(n-' (n='J(()J(tt)), - (n =0))0"))))) 13.
(73)
where
1 I1
n
12 = |of Tij
(1) 1::
(74)
13
a (10)
Since A is uub-Gaussian random variable, and B are bounded random variable, there exists such that with high probability,
[16(0)loo s K, |la(0)llo & ||a(0)ll2 s K.
(75)
Then we have for any i.j E [m]. k E [n].
|oxx(1)) S (xx(1)" [b) S (B, + 1)2K,
and
(t) (t)(()) [x) wx (0) 10 (bx -b((0) s 1)C
Therefore, we have
a
I - 0000000-000000000000000000000000 -0xx(0))
Σ 02(1))0(())+0(0000pp)
kEjn]
= (1).
Noticing |a(0) (1), for In have
B
1
Σ
d' -c (nt))=axx(0)))) (0)9))
T1
kE[n]
B² 21
+ Σ o2(1))
kE[n)
(1).
Using a technique similar that applied for 12 above, we have 13 op(1). By applying union bound we have (()) n ((0) (1). Then for large enough n, we have
Amin (n -(())((()))) 2 0o/4. Therefore, we have
Ri x1(Ci)An
Amin (H())) 2 Amin (n-(-(](][(t)))) Aanin (Q(1))
(76)
Then we have
Ki (C())
lf((t) y11 < -210
[[((t) 21
dt
4
and
[[f(t) 211 S exp(
Kı
(C
Not) (l(()) -1122
(77)
4
33

---

Published as a conference paper at ICLR 2025
Notice that,
[a (t)l22 S VmilJay fl (m(BBllwxxtlla (bx( & 2/mK +
(78)
Then with large enough and for all [n]. have
d
ax((t)
-
7o
dt
7
(79)
11o
m2(Ci)
2/ + 1) exp(- xi(Ci)Ao
no/) (ol)1l((0) - yl12.
n
4
and hence
(ax(1) an (0)) - {
d
(s)
dt
SNa m + 1)) ulla
(80)
(C
$ Cin
Similarly, by noticing that with large enough n,
|| Jwn, f()lll & /a & m,mB, Ak ( [d];
1o (t)]l2 < (mlax(t) s 22mK, Ak E [n];
(81)
[f)f(()) = m.
Comparing (8 to (78) and using similar computations as in 779 and 80). conclude that
r(t2) < C1 which yields a contradiction. It then follows that 776 and hold for 2 0. This concludes the proof.
We point out that Proposition 27 can be extended to the case where only output weights are trained, without requiring any additional efforts. We present the results in the following proposition.
Proposition 28. Assume the same assumptions in Proposition 27 Assume constant C1 > 0 satisfies Then with high probability over random initialization, for any E [0, oo).
K5 (C
(0(1) - 01 S C II5()) - y11 < exp(
no mo))) - yl12.
4
In the following result we show that training only the output weights well approximates training all parameters.
Proposition 29. Assume the same assumptions as in Proposition 27 C which satisfies (72) and
Assume there exists a constant
(2((1) (C1) + 10m 10m 1)²
3 (1)
R2 (21
(Aa)
(CC)))ssCC)
16(B, + 1)K m f(0) y112 (82)
Then for any E I and with high probability over initialization,
lim sup [z 6(t) f(x ((t))l2 0.
n=+00 120
Proof. [Proof of Proposition 29
We define the following function defined for 2
((t) n 3/2 ||a(t) - act)lloo.
For any fixed C >>0 0, consider f inf(1 0: <(t) C). Now aim to prove ( +o0. For the sake contradiction, assume ( < +00. Then we have for f E [0, ((t) & C and in particular, ((ti) C.
Assume that C₁ > 0 satisfies 72) and 82 According to Proposition 27 and Proposition 28 we
have sup:zo (0()) - 0 |oo < C1 and sup, [])) - 0] < C1. By assumption, (ø"(x)
(x)
is Lipschitz continuous on |-O,, C1). Let L denote the corresponding Lipschitz constant. Note that
(x"(x)) (ox())2xs()))))
L < sup
(0(x))
YI
sup
(x)(
# C1,Ci)
xE(Cr.C:]
34

---

Published as a conference paper at ICLR 2025
Next we bound the difference between (t) and (t). To simplify the notation, we let 0 (t) o (wx(()) x, (t)) for i E [m] and & E [n]. For i, E [ma] we have the following decomposition:
HH H.)(t)l - 4 + 1e 13.
(83)
where
Jsh (0)0 (0)
11
6 (t(((x((t) - ax (0))) ("(m(āx(t) a& (0))
71
OK (1() ("(((wx) Wk.
o
wk.r (00) (0))) 1)+
(84)
1
(a(bb(t)-bb(00
13 =
(n(d())) d(o)))
Since A is uub-Gaussian and B is bounded, there exists K > such that with high probability,
[((0)]l S K, ||a(0)]lo S ||(())l2 - K.
(85)
Noticing [lA(t) -f1oo Cin we have for large enough and for any i.j E [m], E [n], 2 0.
(axx(1)) (1x + [bx( (B2 + 1)2K
and
(tx(()) (1) < (() (0) (xx))))b((0)) S vd(B, 1)CCn
Recall x2((i)) sup (x)) -1 By noticing [la(() - a(()]lo Cn/3/2, we have for t [0,t]]
S n
(oa(()) (((0) + -x0)))/))()00000))((((())))
1
bas (0)lo (0)(( (((ax((1) (0)())())()((((((n((((()) )'((āx(1))
(0p(n"))n((()
(22))) + (n(2n))n2(1)1 (K 11
4K 1) 1/2 Op
In the meantime, note that [|a())]22 + ||a(() a(0)]l2 = (1(). Then we have that for E (0,4].
a(())() = O(n-1(()))))))))11)1)))))))())) Therefore, [|a())]12
Ig [B2 + (1)] n
(n(mm")))(()))) (G))
I3 s (nn(())²((G))
Then with large enough we have for all i [m]] that
sup (t)(()) (t) < (B2 +E(0.11
and hence
sup :<(0.t)]
[H(t) -
+ (B.
(86)
Recall (Ci) infeer-ci, (o"" Let r()) = f(t) j(1) and C2 = (C)) /4.
Consider the function u(t) = [t) defined on E [0, t.]. Clearly, u(t) is continuous on [0, t... Since u(t) if and only if rt)) O, have that u(t) is differentiable whenever u(t) * 0. In
35

---

Published as a conference paper at ICLR 2025
particular, when u(t) + O, notice that
d dt
(e (c u(1)) = (lo (1))
d
dt
= 2e Cenot rtt) (c
dt
2e2 (C2nor(1) (()o((C))r)) no (())(((( y) + + +) F(0((())) v)
2noe 20 (r(t) (H((t) (()) y) -100((1(1--222))()
S 2npe 2anot (I)) (()) v)
I 2np 2Onto (t) - H(OO))21(()) - yl12.
(87)
where in the first inequality we use the property Amiu (H())) > C2 for any 2 0 from Proposition Meanwhile, have
C2n0 u(1) d (eCrno (CComo u(1)² 2e
dt
and hence by
187
(e u(t)) S noe c20001 - ((OIla))(() - y/12-
dt
Now consider e t.. Let t' inf E [0,t]: r(s) + 0}. Then we have
u() < I'' noe (1 (+))l2l) yll2ds.
Hence for all E [0, 41. we have
(r(t)lla S: Cc (())1221F(s) - y/l2ds
< e (5mM L) J. [lf(0) y
(88)
< e Cano 5m
() 2/12/n - 1/2
We have that for all & E [n]] and with sufficiently large n,
(ax (t) ax(())
70
Jan f(1) (f(() y) ("(nā((t)) Jo, f(0)= (F(1) (F() 3)
8" (manx (t)) 1
[|Jar (tt]l2l1f(t) vll2
d" (nax
("(nāz(t))
1
f(t Jo. f(0)llll)))) vll
("(nā((t)) 1
|| (0))==)))))))(
8" ("(nāx(t))
36

---

Published as a conference paper at ICLR 2025
Noticing [[((t) 20 Cnn and using 88 we have that
(ax(t) (t) ax({
dt
[CCn 11/2-1122 1/2 /m + (bx (t)) (I - y122ee -aano
Canot
+ n2 (2(())) /m (0) + (t) - by (0)1) -11((0) - -y112 C2nott
xx2(Ci)
(B (0)]22 (0)1) Vhe Com k 2B -)) y1122n 1/2)
m f(() - ull [LC- 1)m 3/2 (c()
exp(C2not)
+ +5xg(Ci)(BBB + 1) k}*K3/1///3 nomLLnn 3/2
BLC(B, + 1)Kno f() 8112 (2 5wz (1)) + K) 7omm) -/2
exp(C2not)
Recall L S ( and C2 = (x1(C2)Ao)44, Then t E [0, and l E [n] we have
(t) (t!l
2 5k2 (C) + K2 K2 2 noms ds -3/2
<2LC(BB 1)) ylla
exp(C2n10)
=2LC(B, + 1)K f(0) 2112 8K (Ci)Ao + (0n)(C + k) -3/2
(Ni(C))Ao)²
(Ci)A + 10m BN + 122( 16(B, + (mllmllff 3112 On 3/2
(C,)²
(2o)2
CCn33/2
where the last inequality comes from the fact that C1 satisfies (82). Hence, we have <tt) < C for
t E t1). Note this holds at t = t1, which yields a contradiction. Hence, t = +00. By (88) we
have that all 2
(I - F1)122 = o1(()e -Crottt
(89)
Next we fix the constants Cj and C' and aim to show f(t) agrees with f(- ((,)((t)) on any given test data x Rdd For fix x E Rdd consider
f(x,0(1) 0(1) -o0l1(1 - y)
((t,
and
ह. ((1)((()) - - (t) y)
((t,x)) V
Using the fact that [|a(t) - a(t)l S Cn-3/2 3/2 and applying a decomposition similar to we have
83 and
84
sup
H (t, x)) 11 (t, x))l2 1//
(00)
120
For large enough n, by 85 and [[((t) - 01100 < 11 we have
|)f((1)33 = Σ
osk (1) + (11 +1))
(91)
5m+
2|| (1)12 + 2/6 (1) (a
= =0,(n,,
37

---

Published as a conference paper at ICLR 2025
and hence [t)l2 = Op(vn). Note similar bound applies to f enough n we have for t 2 0,
(2)) ((2) Then with large
[1(1, - n
[ (l)).
Finally, using 89 and (90). we have for 0.
(z, ((t) - (t))l
d (((,, f(x, (ss))
ds
(f y) no yy
$ !. H. ()). 2) f(s)ll2
+ H() - yll2ds
$ Op(1)o, Cano sds + Op(n 1/2) /o 0C
llf(0) u1ße Camaa
y
(1)),
which gives the claimed results.
Next we show that, if the potential & satisfies Assumption 8 one can select a sufficiently large w to ensure that the conditions in Proposition 27 and 29 are satisfied.
Proposition 30. Consider potential function that satisfies the assumptions in Proposition Further, assume that (xx) = w+I (M(x)) + wx²) with convex function vp E (RR) and positive
constant i > 0. Let D1, D2 be constants defined as follows:
D1 8ym(BB (f yll
D2 10mm10m + 1)?
An
Consider functionn Çi and (2 defined as follows:
(CC) = SUPse[-c.c] (2) -D
AC > D1;
2(C - D2)
(2(C) sup|2/<o co 14b" (2)) Dr (No(1+ sup (2)) + D2 ( sup (=)))) 4A 0. EEC
Ao
|=KC
(92)
Then there exists Ci > that satisfies
and
whenever w > wo with wa defined as follows:
wp maxt 2
1 + inf G(("), C = supp(CC>>>> Dr: GCC) < + inf
(CC)).
O'SD:
C>D,
Proof. [Proof of Proposition
30
Since J > 0, we have that for E [-C,
w +1
W +1
w w+1
sup sup=e[-C.o1 (2) + 2w
0²(2)
(2) + 2w
2n
Hence we can take
+1
xC(C)= supse[-c,cT (0 2 2w
K2(C) 2w
Then we have
(C)
+1
w +1
2
(C) supse[-c.c 4 (2) 2w
2w
sUPse[-C.c] v/" + 2n
38

---

Published as a conference paper at ICLR 2025
Then 72 can be reformulated as
2w
- D1.
(93)
PPPEe(C,cy (2) + 2w
When C' > Di, we can rearrange (93) as follows
sup -u - DI ((()
(94)
2(C' D1)
Note (I(C) is a continuous function on C > D1. This implies whenever w satisfies
w inf <(C")
(95)
CD>D;
there exists C' > D, such that (94) holds and thus 72 holds. On the other side, notice that
SuPzEl-C.CT (x)]
Ka (C) = sup
(x)
=EI-c,01
w
Then 82
can be reformulated as
sup -clc (=)]
sup -. (x)) D2(11
sup suPsE(-C,o (11++))---000((111)))
do
2w
2w
2w
2D
(96)
w > 1/2, we have
sup sup:el SaPPE(-Co C,O (2)
< +
sup
(2)
2w
sE(-C,0]
Therefore, 966 holds if the following holds
w sup =|ec (2)l D1
(20(1+ + sup (²)))) - Dz
(+)))) 2 (e(C). (97)
sup
Ao
|=|00
|s|e
Now, assume satisfies
w > max + inf G(("), C = sup > Dr: (iCC) <1 inf (i(C)), ctD.,
Since w + inf Co (C"), then for any C that satisfies
Ce fC> D1: (CC) 1+ inf ((C"))
(98)
o's.
we have (94) holds and thus
72 holds. At the same time, since Çe is non-decreasing, we have for
any C that satisfies (98).
w > (2(C*) 2 (2(C),
which implies 82 also holds. This concludes the proof.
In the following result, show that when é is the hypentropy potential Ps, the conditions in Proposition 27 and aree satisfied with sufficiently large B.
Proposition 31. Consider a potential function that satisfies the assumptions in Proposition 27
Further, assume that o(z) pa(x)) is the hypentropy function with sufficiently large B. Then there exists C₁ > 0 that satisfies (22) and (8))
Proof. Notice that
(3(x)
1
(x)
x
Va2 + B2²
(w2 + 3213/2
Therefore, we can take K1, K2 and K3 follows:
Ki (C) = B, (2(C) = VC2 +
B kx(C) (C + +/2233/2
39

---

Published as a conference paper at ICLR 2025
Therefore, given any fixed C' > 0, we have
Cki (C)
CB
lim 8-x
lim
C.
CC(CC
80000
VC² B
Therefore, by selecting C₁ = 2 28
+1 f(o) ulla
we have that for sufficiently large 3,
Ao
(11) C1 (2())))
8 8vm(B, + 1) - -lll
2
A
This implies 72 holds. Now fix C1. Since K2 (2(C))Na((1) = VCT+ we have Xi (ri -
O as B + co. Meanwhile, notice that (1)
+ 1 as B + xc. It then follows that
(C)
lim (2( (Ao
K2 (C1)
10m (B, 1)²
Kog (P((C)) K1 (Ci) 9 = 0.
8000
K1 (Ci)
Therefore, with large enough 3, we have (82 also holds for C₁ This concludes the proof.
B.2 IMPLICIT BIAS IN PARAMETER SPACE
By Theorem II the final parameter of the network with absolute value activation function whose only output weights are trained is the solution to following problem:
min (ā (0))
s.t. Σ
du ww (0)".a x; (b(( yt. Vi E [m]]
GER"
(99)
kE|m]
Similar to the case of ReLU networks, we can write infinitely wide absolute-value networks as follows:
(2r.a)=((() g(z, = f a(w:,))w" b))w T bl dp.
Note the above integral always annihilates the odd component of a since |w -b is an even function of (w,, b) and μ symmetric. Therefore, for networks with absolute value activation functions, we lose no generality to restrict the parameter space to the space of even and continuous functions:
A {a E ORELU: is even and uniformly continuous).
(100)
In the following result, we characterize the limit of the solution to 99 when n tends infinity. Proposition 32. Let Jn and g be defined as follows:
/ a(w, b)lw"* Tx duns
In. (x(x, a)
(101)
/ a(w, ((w.b)lw" w)]=T== du.
g(2, at)
Assume à the solution to (99 Assume the potential function can be written in the form of (0) = n 3n++ o(n(fz 9x)), where & real-valued C -mooohh function on and has strictly positive second derivatives everywhere. Let Abs be defined in (I00) and o, E θ Abs be any function that satisfies o (w(w (0), by (0)) = n((a (x(()). Assume a is the solution of the following problem:
/ (),a(w,b),0)) s.t. g( a) Vi -(x,, 0), vi [m].
min ac REGADS
(102)
Then given any x E IR,, with high probability over initialization, lim (z, a..) - g(r, ã) = 0.
We will need the following two lemmas for the proof. We omit the proof of the Lemmas to maintain clarity.
Lemma 33 (Jacod & Sørensen 2018 Theorem 2.5). For n E N, let = Gn(,, 2) be a random
".' function defined on T where (z1. 2n) consists of independent samples from a pre-specified random variable 2g ~ Z. Let G(.) be a deterministic RmD -alued function defined on R.m. Assume there exist X R.m and neighborhood U of X such that the following hold:
40

---

Published as a conference paper at ICLR 2025
i) (A( 0 as + co, and G(X) =
ii For any 2, (.) and G(-) are continuously differentiable on R and as n + 00
(A) & 0.
sup
XEU
iii The matrix (x) is non-singular.
Then there exists sequence (X()) InEN such that Gn((A and
|| 2110o -
Lemma (Newey & McFadden
1994 Lemma 2.4). Let a(.;z) be a random real-valued function
on compact set U C R Let 21..... 22 be independent samples from pre-specified random variable Let pn denote the empirical measure associated with samples {= and μ denote the probability measure of Z. Assume the following hold:
i) For any 2. al.,2) is continuous.
ii There exists d(z) such that aa>, < d(z) holds for all E and Ez [d(()) < oc.
Then
sup / a(A)z)ddum / a(l, 2(), 4 0. XEE
We now present the proof of Proposition 32
Proof. [Proof of Proposition 32 To ease the notation, for k E [n], let 25 (0), (0)) Rd+1 denote the imput weight and bias associated with the k--th neuron at initialization. For j [m], define (2(2) = (w, b) = w 2; b]. Let Z denote the random vector (W,B), and 0s = y - fx,, 0). Notice that for any R" and that satisfies a bx n(ax ax(()) for & E [m]] have
Da a()))
Σ (d(n(ax- au (0))) <0) no (0) aax()=))()))) an))
n2 kE(M)
1 / (a(w,, b)) x(0) d' ('())a(w, b)dj,,
n
1
/ Do(a(w, b), 0)dpa-
n
Therefore, since à is the solution to 199) we have that on solves:
min Ob,, Ab.
Da (w, b). 0)d s.t. (n(I,, an) Dr. ie [m].
(103)
The Lagrangian of problem 103 can be written as
L(an, A) = / D Do(an b), 0)dpa
Σ x (ʃ (f an())) - ū.).
m
1=1
Here A (AI,..., Am...,,.m
Am) )T E R" is the vector Lagrangian multipliers. The optimality condition
V a. L = 0 implies that
0 ("(),(z))) (²(0) Σ A=))(Ek). Ak E [n].
By Assumption, (²(x) > 0 everywhere and & (R). Therefore, o' is invertible and -d')-1 is
continuously differentiable. Therefore, on must satisfy:
m
= (d)
Σ x (2((2)) Vk [n],
41

---

Published as a conference paper at ICLR 2025
where (m) is determined by the following equations
f (0)) + Σ (o 0z(2)) 02(2) dun = Us. Vi E [m))
(104)
J=1
Problem [033 strictly convex in the quotient space of Abs by the subspace of functions with zero (2()))) Semi-nome Hence, 0g is unique in the (2(p)) sense. Note that (o') is injective, and that by Lemma 13 for large enough n and with high probability, the matrix (a) (a, (2x(2))) (2x)),, is full
row-rank. Therefore, the Lagrangian multipliers ( is uniquely determined by (104).
According to [04 can select Ox
(6))
(8(0)) + Σ 1 5
))()))
V2 E supp(u)-
(105)
Note supp(u) is a compact set and (dd)-1 is continuous, we have on defined above even and uniformly continuous and therefore lies in A
Let = (#1,..., Zn =T. We define a random Rm -vlued function as follows:
(A, &
)5))(()) X 02(2) dun - Ds. i 1,..., m.
j=1
Then by (104). (m) is root of Ga. and, as we have shown, it is the unique root.
Similarly, according the optimality condition, the solution â problem (102) can be given by
((=) = (d') (') (1), Vz supp(u),
(106)
where À is root of the following deterministic -valude function:
(G(N)), a a f(0) + Σ )ooz(z) 0z(2) du Di, i= 1=....,,. m.
1=1
Since problem 102 is strictly convex in Abs. ā is unique, At the same time, note that (Ø) is injective, and that { are linearly independent in (²()),, as have shown in the proof of Proposition 13 Hence, À is the unique root of G.
Now we aim to prove that (m) converges to Ā in probability. Notice that, by law of large number, G (A) converges to G(A) in probability for any fixed X. This implies that
GT E. G(+) 0.
Since & (R), know for any fixed z, (), x) and G(A) are continuously differentiable in A. To ease the notation, let E(x) = (zr). Then we have that
(JAGUCA),, = 5 (8(00) + (I.G(N), = =
=55(
(v)(()) (((())))(())0((3)) dum.
o3)())) du.
Let U be compact neighborhood of Å. Notice that μ is compactly supported and g is continuous. Then there exists a constant such that
('(6'(0) + Σ {ai(z)) SK, V. supp(μ).
42

---

Published as a conference paper at ICLR 2025
Define dzy(2) = (o)(2)),(z). Then we have that
|< (d(1))) (0)
Σ (
()) (2)))11)))))))
AA E U.
1=1
Meanwhile, it's clear that for any i.j E [m]. there exists polynomial (z) of z, whose degree is
independent of n, such that [1)(2) < (2)(2)] holds for any E supp(u). Recall is bounded random variable. Therefore, Ez (di co. Then, by Lemma 34 we have
(JAGU(A)) (JAG(N)
sup XeU
Notice the size of the matrix JxG, and JxG is independent of n. With union bound we then have
[JA JAG(A)le & 0.
sup
xEU
Next show Jx)(A), given below, is non-singular,
(JAG(T) - - = Σ [xor(t)oo(())oy(e dp.
To this end, we consider the map ),)) Lz(H) x L2(1) + R defined by
m (1.g) = / E (0(00)) Σ (Io(()) (==)(()))))))))))))))))))))))))) du, vf.g E L2(n).
Clearly, (:. .. is symmetric and linear. For any f E L2(n) with f 0, i.e. u([=: f(z) + 03) > 0. Then there exists compact set A C supp(u) such that ((A (f()) + 0}} > 0. Since {othemm are continuous and {At) are fixed, the set
A' E) CR
is non empty and bounded. Then by noticing E' is continuous and g'(x) > 0, Vo E IR, we have
(f.f) f. = E' ((((0))))) + Σ ) P(x)
1=1
2 inf (x(x) zEA'
(f()))
0.
Therefore, (+ is well-defined inner product on L²(p)) Noting that JxG(X) can be the viewed as the Gram matrix of (or)?" with respect to (:.))., and that {a) Im are linearly independent in L (u), have JxG()) is non-singular. Since ((n) and Ā is the unique root of G,((+) and G(-), respectively, by Lemma we have
(A(m) - -1loo = (1).
(107)
Now consider fixed R. We aim show that Ig a)) - gg(x, a)] (1). Recall that
gn gn( a)) / a(w, b)l djm
g(x, @)) / ā(w, (lw.))lWT= bb dp.
Since ā and the absolute value functions are continuous and μ is compactly supported, by law of large number, we have
(g)(x, ((x,a)) (1).
(108)
43

---

Published as a conference paper at ICLR 2025
Hence, it remains to show (I, 9n a)] oy(1). Notice that À is fixed point in R and ( converges Ā in probability, and that y is compact Then there exists constant K' independent of n such that with high probability,
m
6'(0)
+
N ( b) ('(0) + Σ (1, S K', v(w,b) E supp(u).
Notice (6')-1 is continuously differentiable on IR and hence is Lipschitz continuous on [- K'). Let L denote the corresponding Lipschitz coefficient. Then we have that
(n, an) 9n. a)l
< an b) a(w, b) law ()oo - bb dun
m $ / (0)- -1 (d((0) Σ ow: b)) (0)- Σ
m
Too (w. b) o0 b) dp
m
$ / L Σ (A o,) (w,
x ((
dun
1=1
< / L LI( - 11100 Σ
(ow' [ow"2zz- 1z- b)) aovot x - b) dum
j=l
s(z / olwo x - b)
Σ
alw 2j - b) dpmm) |1 - :11
1=1
By the law of large numbers,
f JolwTo x - b) Σ
oow x. din Ps f(ww"= -o [ww"s b) dp-
==1
Hence f x x 6) Σz= ow x, - b) dun can be bounded by a finite number independent of n.
Then by (107). we have
Ign (x, ūn) (2, o,,(
(109)
Finally, combine (08) and (109)
we have
In (z, a.) - g a)l (g,( an) - 2n(2, + @) - g(x,ā)) = =1))),
which concludes the proof.
B.3 FUNCTION DESCRIPTION OF THE IMPLICIT BIAS
Assume the input dimension is one, i.e., d = 1. Now we study the solution to problem 1022 in the function space
FABS = ()) f
b)lur a
bldp: +
In the following result, we show that FAbs is subspace of FROLM which can be obtained by eliminating non-zero affine functions in FReLU-
Proposition 35. Let FAmue denote the space of affine functions on R. Then FRELU Fabs e FAfine- Moreover. if h FAbe, then (h) = 0 and Gs(h) = O, where G2 and G3 are defined in 7
Proof. [Proof of Proposition 35 We first show both FAbs and FAftine are subspaces of FRELU-
Assume h(z) a(w, b)lwrr w.r bldj FAbe- Notice the ReLU activation function can be
decomposed as [z]+ = z/2 /2|/2. Since a is an even function, we have for all x E IR,
Mxx= / (ww ww bldy 20(w, [w b)
44

---

Published as a conference paper at ICLR 2025
Since 2x ORELU, FAbs C FReLU. Now assume h(æ) + with E R. Consider an defined by
2w
(n, = 11
b+ 2z;
on = -Ocn -1, bb).
Note 0% is odd function. Then have for any R,
On an b) 1 on (n, dus
on (n b)
x- -- = (1 b)
PB (b)ddb
(b(b))db
2
2
zpB =pb(b)db) x
w
PB
Pb(6)b2
E
2 + W.
Hence, FAttine C FRELU.
Next show FAbs FAftine {0), i.e. the only affine function in FAbs is the zero function.
Assume f(()) = a(w, (/w bldμ E FAbs is an affine function. For x E (supp(ps)) by
Leibniz integral rule have
h h((x)
(110)
Since h is an affine function, we have 2a(1, x)5 <)) 0. Now since pB(x) 0, a(1,z) a(-1, -x) = O, which implies h(x) = 0 for E R.
Next we show that for any h h + h2. This can be directly verified by the following computation: for any E OReLU.
FROLU there exists h E FAbs and h2 E FAtine such that
/ a(w, b)lrr b)+ / aa( (w, wr bl du +
/ <(wwbb)
Wx - dp
2
-f a (w, aatw,,
WE bl dut
/ a (1,
dus x- / a (1,bb)
duB.
2
2
Therefore, we have FRELU = Abbs + FAftine.
We now show G2 and Ga vanish FAbs. Without loss of generality, assume supp(ps) [m, M] where m, M are finite numbers. Assume h(x) = =f(w a(w, b)lwx bld E FAbs Then for x > M, we have
h(x) = Im MM
a(1,b)(z bb)p
and
1'(x) I/ MM
a(1, (()))).
Similarly, for x m we have
PM h(x) = ( a(1,b)(-x +b)p
and
(()x) = - M a(1, (b)d
Therefore, (+(00) + 1'(-x)) = and (2) = O. To show (h) 0, consider the following function
s(x) h(x) / an b)lwx bld,,
where OE O Abs defined by
(b)
ay (1, 2ps(b)
45

---

Published as a conference paper at ICLR 2025
It is clear that s(x) E FAbs FReluu By Proposition 20 s" vanishes when supp(ps). For I E (supp(pa))), similar to computation in (110) we also have s"(x) O. Therefore, s' is a piece-wise affine function. According to Proposition 20 s' is continuous. Therefore, s' is constant function. By what we just shown above, the only constant function in FAts the zero function. Therefore, h(()) g(,, 04k) for all R. In particular, have that
h1" 1²(b) b db. 2
(0) =000=9(000)0000 an) = an (w, b)l - bldj supp(ps)
Therefore, (ampp(re) N R² (b)66/b - -2h(0) and (a(h) = With this, we conclude the proof.
In the following result, we solve the minimal representation cost problem (17 with the cost functional
described in problem (TO2), i.e. Ca(a)= (a((u b), 0) du.
Proposition 36. Let be as defined in (10T). Given E FAbs, consider the following minimal
representation cost problem
min a
Do(a(w, b), 0)du s.t. g(,,a)) a- = h(x), Vo E I.
(111)
Then the solution to problem
denoted by, Ro (h), is given by
Ro (1), = (b(b))
(112)
(p
Proof. [Proof of Proposition 36) Note suffices to show that the constraint g(, = -h(-) in problem
111 is equivalent
(b(b)
a(1, b) 2pg(b)
(113)
If h(-)) g(,, a.), then by calculation in 110 have must satisfy 113). This implies IB is
necessary for the equality a( = h1(-). Assume a satisfies (13)). Consider the following function
s(x) h(x)) alw, b) wax bld
As we shown in the proof for Proposition 35 s(x) = 0 on R. Therefore, the solution to 111 is given by (112) and we conclude the proof.
Now we can write down the implicit bias problem in the function space. The following Proposition is a direct result of Proposition 12and Proposition 336]
Proposition 37. Assume ā solves problem (102). Then Mxx = g(x, a), as defined in 101 solves the following variational problem
[
Do hx(x) 2pp(x)) 0) PB (r))d s.t. h(aa)) f(x - 0), Vi E [m].
min hEF
(114)
In the following result, show that Theorem 9 also holds for hypentropy potentials pم(x) J (aosinh(x//1) /2² - 32 when 3 is sufficiently large.
Theorem 38 (Implicit bias of mirror flow with scaled hypentropy). Assume the same assumptions
pa(n(0--0)) where as in Theorem except that the potential is given by: 4(0) = n "(0)===><n<[B+ [3]43311]
Pa is the hypentropy potential. Assume is sufficiently large. Then with high probability over the random initialization, there exist constant C1, C2 > 0 such that for any 2
[[@(t) - 01 S Cin -1 ft) - y S no (0) y112
Moreover, assuming univariate input data, i.e., d = and letting ((oo) lim, 0(t), have for any given that lim,- (sx (ooo0))o)))))))))))))())))))(( ) fax 0) T((x)) O, where π((.) the solution to the following variational problem:
h²(z))
(p)(x)" () (z))d s.t. h(xx) = y f(,, 0), Vi E [m].
min REFALe supplpee)
Dos
Here Dp, denotes the Bregman divergence on TR induced by the hypentropy potential and is given by Dep (z, x (arcsinh(æ//) arcsinh(y/8)) V2² + 82 / B
46

---

Published as a conference paper at ICLR 2025
Proof. [Proof of Theorem 38 By Proposition 31 with sufficiently large B. Proposition 27 and 29 holds for Φ = pg. Note in Proposition 32 we only require function ø to be C -mooth and to have positive second derivatives everywhere, which is satisfied by the hypentropy function. Also, notice Proposition does not use information on the potential function. Therefore, Theorem holds in this case and we conclude the proof.
C DISCUSSION ON PARAMETRIZATION AND INITIALIZATION
In this section, we comment the parametrization ① and initialization scheme 2 that we considered this work, and compare these to NTK parametrization and Anti-Symmetrical Initialization.
Unit-norm initialized input weights initialization 2 we require the input weights to be sampled from the unit sphere gd-1, This choice aims to make the minimal representation cost problem (17) more tractable. Specifically, by this, for univariate input data we have supp(μ) = [-1,1] x R and the parameter space for infinitely wide ReLU networks, i.c. the space of uniformly continuous functions on supp(u), becomes much smaller than the case when input weights are sampled from general sub- Gaussinn distribution and suрр(μ) = R x R. We note this is a common setting in literature on representation cost for overparameterized networks (Savarese et al.) 2019 Ongie et al.) 2020). In Appendix [.2] further discuss the challenge of solving the minimal representation cost problem for input weights sampled from general distributions.
Dependence of lazy training on parametrization and initialization Lazy training is known to be dependent on the way the model is parameterized and how the parameters are initialized (Chizat et al. 2019). In this work, we consider the parametrization (T) and initialization (2) for shallow ReLU networks. However, there are other settings that also lead to lazy training in gradient descent. Generally, consider a two-layer ReLU network of the following form:
f(x) = 1 8n Σ
x b]+ + d,
where 8, is scaling factor dependent on the number of hidden units n. According to Williams et al. (2019), for gradient descent training, there two ways to enter the lazy regime: (1) initializing the input weights and biases significantly larger than output weights; and (2) setting s, = o(n) and considering - cc. Infinitely wide networks with standard parametrization, which consider in this work, aligns with the second approach, as 8n is set to 1.
Standard vs. NTK parametrization
When 8n is set to vū, we obtain the NTK parametrization (Jacot et al. 2018) in which the networks are parametrized as follows:
1 Σ [x)
ano(w x by) +
1=1
Under NTK parametrization, a typical initialization scheme is to independently sample all parameters from standard Gaussian distribution or, more generally, from sub-Gaussian distribution. Note that, compared to standard parametrization in ( and (2). the factor 1/vn appears before the sum, rather than as a scaling factor in the initialization distribution of the output weights. As result, not only the output weights are normalized at initialization, their gradients are also normalized throughout training. It is known that, although gradient descent under both settings in the kernel regime (lazy regime), but they correspond to training different kernel models and yield different implicit biases (Williams et al. 2019) Ghorbani al. (02T). Specifically, noted by JJin Montáfar (202)), under NTK parametrization one can not approximate the training dynamics of the full model by that of the model whose only output weights are trained. This implics the approach in this work, in particular the formulation of problem 14 can not be directly applied to networks with NTK parametrization.
Anti-Sy mini-Symmetrical Initialization Anti-Symmetrical Initialization (ASI) ZZang et al. 2020) is common technique in gradient descent analysis to ensure zero function output at initialization In the study of implicit bias of gradient descent, Jin & Montúfar (2023) utilize ASI base case and then extend to other initializations. However, we observe that for mirror flow, ASI could in contrast greatly complicate the implicit bias problem. In the following briefly discuss why this happens.
47

---

Published as a conference paper at ICLR 2025
To achieve zero function output initialization, ASI creates duplicate hidden units for the network with flipped output weight signs. Specifically, a network with ASI is described as follows:
0) v2
ax axwwg x - bul+
d' (((()))
(115)
V2
Here = (e b, a, W b', a') and the initial parameter ((0) is set as
( - (0) = W, 60) - 6(0) = b, (0)) - a d0) à,
where W, and á are as described 2 Assume one can simplify the implicit bias of training
model (115) to that of training the following linear model did for networks without ASI:
1
f(z, (a,
as de)o ba).
(116)
The corresponding mirror flow is given by
d (a(t), a'))) -n((2 Jaaa (f(t) - y)
(117)
dt
where (tt = (f(x1, (a(t), a' a'tt))..... f(Im. (a(t), a' a((t)))) E R.. Assume (117) is able to attain global optimality. Then by Theorem [1 the implicit bias is described by the following problem:
min (a.a')ER2n
Do ((a, a') (a, a)) s.l. [xx, (a, a')) = yi, Vi E [m].
(118)
However, observe that model (166 is "essentially" parameterized by a' rather than by (a, a'), meaning that the function properties of f are tied to a - a' instead of (a, a'). A gap thus appears: the implicit bias in parameter space describes the behavior of (a, a'), whereas the function properties are related to a - a'. To fill this gap, one can either solve the minimal representation cost problem with respect a cost function defined on (a, a') given in (118), which can be quite challenging, or to characterize a(x)) - a' (oc) returned by (117). The latter is also challenging, as a(I) a'(1) evolves according the following differential equation:
d (a(t) al (t)
(V"))))))((((((()))))))))))))))) - yy)
(119)
dt Noted by Azulay et al. 2021 possible way characterize a(2)) a'(oc) is solve the following equation for V::
(V
= v2y (av(a(() - -o'())
(120)
such a V exists, we have de (VW((())-a(())) = -n,) - yy- It follows that a(x))
a' dd(oo) is the solution to the following problem:
min (a - a') s.t. f(z,, (') = yı, vi [m].
a a-a'
In the special case of gradient flow, i.e., F(a) ||a.l we have (a = / for any E IR.. Therefore, (02 has a simple solution. However, it is known that 120 has a solution only
((V2 (at) )) + (V2 (a)(()))) -1)- satisfies the Hessian-man condition (see, e.g., Azulay et al. 202T. This is a very special property and typically does not hold for most potentials.
D GRADIENT FLOW WITH REPARAMETRIZATION
In this section, we apply our results to study the implicit bias of gradient flow with a particular class of reparametrization. Recall in this work, we focus on the following mirror flow with respect to potential & for minimizing loss function L:
(tt = - -V00))
(121)
dt
Consider reparametrization map
G(v) and the gradient flow in r-space for minimizing L G:
dt = -V)
(122)
48

---

Published as a conference paper at ICLR 2025
By the chain rule, the dynamics of (t) G(v(t)) in the 8-space induced by 122 is given by
d dt (()) = [U00=====(((((((((((())))))11110
(123)
where (G()tt)) denotes the Jacobian matrix of map G at v(t). Notably, the two dynamics, [2T and (23), are identical if the following equality holds for any 0 = G(v):
(V2)(0)) = =](()))G()))
(124)
Based on this, we can directly apply our characterization of the implicit bias to gradient flow with a reparametrization that permits a potential function satisfying [24].
Corollary 39 (Implicit bias of gradient flow with reparametrization). Consider two-layer ReLU network (I) with one input unit and n hidden units, where we assume n sufficiently large. Consider parameter initialization (2). Consider any finite training data set {(z4, y))
) that satisfies i= C supp(ps). Consider bijective reparametrization map G: 33++1 + 1n, : G(v) = 0 and the gradient flow in v-space (122) with n ((1/n). Assume
Ta # Tÿ. when i + j and {
there exists a function 4: R33++ + such that & satisfies Assumption | and the equality (124) holds for any 8 = G(v). Then letting v((o) lime-toc v(t), we have for any given x E R. that
lim,-roo (ar, G(v((c))) f(x, h(x)l = 0, where h(.) is the solution variational problem
similar result holds for scaled potentials and absolute value networks by Theorem 9 which we omit for brevity. We now give conerete examples of reparamettizations that satisfy 124 Corollary 39 Consider the map G determined by its inverse map G-1(()) V as
(z) dz, k l.... 3n + 1, On
(G (O)
where @ E (TR) satisfies (x)x) > 0 for all R. Notice that JG~ Diag( /²²(0x)) and by the inverse function theorem, JG(D) (0) Diag (0x) 112). It then it follows
that (124) holds for ((f) = Σ [x(33+1]] ((x)) In the special case when (xx) 12 x + 2 we have
1 2
(G)
x2 dx
0
+1+
arcsinh 0,, k= 1,.., 3n 1.
We note that Li ct al. 2022) established the equivalence between mirror flow and broader class of reparametrizations, which the authors refer to as commutative reparametrizations. We did not include discussion on these reparametrizations and leave it as a future research direction.
E MINIMAL REPRESENTATION COST PROBLEM IN GENERAL SETTINGS
In this section, we discuss solving the minimal representation cost problems for ReLU networks with respect general cost function and for multivariate networks with general parameter initialization.
E. REPRESENTATION COST FOR RELU NETWORK AND GENERAL COST FUNCTION
Consider the following the minimal representation cost problem for a given h E FReLU:
f
Do(a(w, b), 0)du s.t. h(x) / a(w, b) b|4 Vo R. (125)
arg arg.
As we show in Appendix
A.3 the constraints in problem 125 is equivalent
a (1,b)) = h²(b)
PB(
(1, (1b)) = 11' (+000) + (((-o0),
(126)
(1, (D(b))bj (0)
49

---

Published as a conference paper at ICLR 2025
For simplicity, assume the function satisfies <6()) 0 and (0( 0. Then the objective function in (125) can be rewritten as follows:
f Do(a(w, b), 0)du
/ ((((1,))) b)) ((a((((,)))lp)
/
δ( ((a (1, + (1,b) + (1, - - -1) (1)
Then by 126) solving problem 125) is equivalent solving the following problem
(b) (b)
+ (11bb) + Φ (-b)
arg min
pp((-b a (1, (1. dus
E+RaL odd
[ a (1, (1,b)dμg (+000) + 11((00),
(127)
s.l.
a (1, (1,b)6dug = - b|b6 2h(0).
In the special case of 6(x) x², where the subspace of even functions is orthogonal to the subspace of odd functions, the objective functional in problem (127) simplifies to [(aaa(, ddu... This allows one to explicitly solve the problem. For general Φ, however, obtaining a closed-form solution to problem (127) becomes much more challenging. We emphasize that, for any scaled potential that admits closed-form solution to problem (127), can obtain the function description of the corresponding implicit bias with the strategy we introduced in Section 3.3
E.2 REPRESENTATION COST FOR MULTIVARIATE NETWORKS AND GENERAL INITIALIZATION
Let d > 2 denote the input dimension for the network. Assume (W, B) is random vector Rd x R, with distribution i and density function pw,8- Consider the space of infinitely wide ReLU networks with d-dimensional inputs
Flel {0(z,0)= alw.,
f f) /aww;))/'> dj+ e ReLU
where (d) HeeLU = fa E ((supp(H)):: : a uniformly continuous). 3. For given h E F(d) ReLU. consider the minimal representation cost problem with the squared L norm cost:
arg min ae J [a(w d) s.t. h(x) ax,
(128)
For univariate networks, solve this problem by translating the constraints, h(x) = g(x a), to more manageable equations as presented (126). However, this strategy could be more challenging for multivariate networks with general input weights initialization. We discuss this difficulty for multivariate networks with unit-norm initialized input weights and for univariate network with general input weights initialization separately.
For multivariate networks with unit-norm initialized input weights, i.e., supp(a) C Cd x R, let a be a candidate for Problem (128). i.e., h(x))= g(,, a), Vor Rd Determining the even component of a becomes quite challenging. We now present potential approach with the Radon transform R and the dual Radon transform R* (for formal definition see, c.g., Helgason 1999). Note that
Ah(x) f b)o(() b)dj R² (opw.s)(x).
(129)
Assuming that apw.s is in the Schwartz space S(Sd-1 x R) then by the inversion formula for the dual Radon transform (Solmon 1987 Theorem 8.1) ) one has:
a (w. b) -R((- # - (ww 2(2 (p.))--pw..Bw.b
ax S(gd-1 x RR if aww, b) is -mooohh and all its partial derivatives decreases faster than (1 + + for any A: 0.
50

---

Published as a conference paper at ICLR 2025
However, the assumption apw,8 does not generally hold. In particular, consider â the solution to following problem:
/ (a(w,))) s.t. g(1; o.) = Vi E [m]. (d) Rebd:
arg min
(130)
a
By analogy with Proposition 17 flow training for infinitely wide ReLU networks the multivariate setting. By the first-order optimality condition for problem 130) & takes the following form:
one can expect á to characterize the parameter obtained by mirror
77
a(w, 6)
Σ
> 2z 6]+1
where [Ai are Lagrangian multipliers to be determined. It is clear that either à or à does not have continuous partial derivatives, implying that 6pw. does not lie in S($4-1 x R. even if Pw.8 does. A similar observation has been made by Ongie et al. (2020), who note the difficulty in explicitly characterizing the representation cost with the L norm in the multivariate setting. Some studies have relaxed the Schwartz space condition for the inversion formula of the Radon transform (Helgason]
1999 Theorem 6.2, Ch. I: J:Jeenenn 2004 Theorem 5.1), which recovers a function f IRd using RC)). However, to the best of our knowledge, corresponding relaxation has been established for the inversion formula of the dual Radon transform, which recovers a function a on gd-1 x E using R. R² (a). The latter necessary for characterizing the solution problem (128) with the relation given in 129).
For univariate ReLU networks with general input weight initializations, assume (W,B) a random vector on R x R. Let h(x)- g(,, a)), Va E R. When determining the even component of one can differentiate hxx) twice and obtain the following equality
[ a (w, (w, wx)w² wur)w"pw.be pw. (ww,(w, wa)dw lx))..
(131)
Then the minimal representation cost problem the subspace of even functions can be written
/ ((a d)) s.t. a + satisfies [31
arg min, o Eennau even
(132)
For unit-norm initialized input weights, i.e., W = Unif((-1,1)) the constraint 131 becomes a discrete sum and it in fact completely determines a + Under general distribution for input weights, solving ([32) can be challenging. We note that the unit-norm input weights constraint also appears in other works representation cost for overparameterized networks (Savarese et al., 2019 Ongie et al. 2020). Jin & Montúfar (2023) considered general initialization but at the same time they required either intractable data adjustments or otherwise adding skip connections the network.
F RIEMANNIAN GEOMETRY OF THE PARAMETER SPACE (θ, V2
this section, we briefly discuss the Riemannian geometry induced by the metric tensor 2 on parameter space.
The Riemannian curvature tensor a fundamental concept in differential geometry, describing how the space curves different points and along various directions. For detailed introduction and rigorous definition of the Riemannian curvature tensor, we refer the reader to Do Carmo & Flaherty Francis (1992). Next compute the Riemannian curvature tensor of the manifold TR V*Φ), where T is a potential function satisfying Assumption
With the trivial coordinate chart id: 3 I 31 the metric tensor 2 takes the form of diagonal matrix diag( (()))). The Riemannian curvature is determined by the Christoffel symbols for the connection (see, e.g.. S. Carmo & Flaherty Francis 1992), which can be computed
51

---

Published as a conference paper at ICLR 2025
as follows:
n 1/
km (d:gjm + +Oggmi - Imiis)
2 k=1
29
kk (aigzt + +0ggk+ - -Orgts)
(133)
(20"(0+))** (0x) i=j= k:
otherwise.
Then the components of the Riemannian curvature tensor is given by:
jk i
r
jk T.m + 0 0 jk =
(134)
jm
This implies the curvature tensor is everywhere zero and the Riemannian manifold (θ, V2 is flat everywhere. We remark that manifold equipped with a metric tensor of the form 4 for some general function & is commonly known a Hessian manifold. We refer the reader to /Shia Yagil (1997) for a broader discussion on the geometry on Hessian manifolds.
G EXPERIMENT DETAILS
In this section, we provide details for experiments in Section 4
Training data For the experiments with unscaled potentials in Figure I we use the follow- ing training data set: -0.15,, (-0.2, -0.15), (0, 0.15), (0.2. 0.15), (1, 0.15))-- In the right panel of Figure D we apply the following shift values to the y-coordinates of all the data points: 2.5, 1, -0.5, 22 3.5.. For the experiments with scaled potentials in Figure 2 we con- sider a different training data set in order to highlight the effect of different scaled potentials: {(-1, 0.15), (0.35, 0.15), (0.65, 0.115, 0.15)). Note the z-coordinates of all the training data points are contained 1, 1].
Mirror descent training We use full-batch mirror descent to train the network. For mirror descent with unscaled potentials, as the number of hidden units varies, we set the learning rate as 17 1/n across all three potentials ø1, 22 and Φ3. For mirror descent with scaled potentials, the learning rate for Φi and d2 is set 1/n and the learning rate for dy is set 2/n. We observe that for scaled potentials, a larger value in Φ(x) = | + x leads to slower convergence and thus we use a larger learning rate for Φ3. During each training iteration, we manually compute the inverse Hessian matrix, multiply it by the back-propagater gradients, and then update the parameters using the preconditioned gradients. The stopping criterion for the training is that the training loss is lower than the threshold value 10-7. In the middle panel of Figures I and 2 random initializations for the network. The training is implemented by Py PoTchc version 2.0.0.
we repeat the training five times with different
Numerical solution to the variational problem We solve the variational problem using discretiza- tion. We discretize the interval [-1.5, 1.5) evenly with 501 points -1.5 to < t1 <... < t500 = 1.5 and then set the vector be optimized as u = [ho. hy..... h200, N k4.]] where h2 represents the function value h((t,) for i = o,..., h-__ represents the slope h' (-(o) and h4 represents h' o00)) As we shown in Proposition 20 when B is compactly support on 1- 1]. we have (' h' (a) for any < -1 and h' o0(+oo)= h((x)) for any x > 1. Let (I4, Vk.) for k = l...., m denote the training data points. s For k = 1,..., m, let iq = arg fo....
I4 - z& record the index for the nearest discretization point to za. To solve variational problem (7). we consider the following two sets of constraints applying to u: (1) hs. y/: for k = 1..... m: (2) (h2- - hx-1)/(tz tv) = h'_ for ie I 11 and (hg hx-1)/(tt
to) h', for l is. The second set of constraints is used to ensure h lies in FRELU- We estimate the second derivative of h by finite difference method, i.e. by (h1+1 2h, + h4-1)/(4 )2 for i 1 ... 499. Noting G3 has closed-form formula under uniformly sampled input bias in Theorem we estimate the objective functionals as follows:
499
(1)) = Σ (h to)
1=1
(2((1) (h'4 + +22
(a( = 3((14 - -h-)) )- - (Nii + 1),m)))
52

---

Published as a conference paper at ICLR 2025
Then we optimize minimize the above objective function. To solve variational problem P consider additional sets of constraints (1) G2(h) O: and (2) 93 (h) O. This is to ensure Abe according to Proposition 35 Then we optimize u to minimize G2. We implement the optimization by CVXPY (Diamond & Boyd 2016) version 1.4.1. To estimate the L -eror between the solution the variational and the trained networks, we evaluate the trained network on [t+] 1500 ==0
denoted by {f:) and compute the Loo norm of the difference vector [/o ho...., ... -500 h500
H ADDITIONAL NUMERICAL EXPERIMENTS
In this section, we provide additional numerical experiments on mirror descent training.
Training trajectory in parameter space The left panel in Figure 3 illustrates that, training wide ReLU networks using mirror descent with different unscaled potentials not only have the same implicit bias, they also yield very close parameter trajectories. The right panel in Figure illustraas that, in contrast, training wide networks using mirror descent with different scaled potentials converge to different solution points and have significantly different trajectories.
0.003
4.00
0.015
8()0)
dulx) m
#:00 - +x?
9g0x) +xx
0.002
$,000 xxr
+-x?
((s))
0.005
o))
0(cc
0.000
((m))
0.000
-0.001
-0.005
-0.002
0
-0.010
0.01
Figure 3: Left: 2D PCA representation of parameter trajectories under mirror descent with unscaled potentials Φ1 = x², Φ - | + x², and Φ3 = x² + x², for ReLU networks with 4860 hidden units. Right: 2D PCA representation of parameter trajectories under mirror descent with scaled potentials Φ1. 62, and Φ3. for networks with absolute value activations and 4860 hidden units.
Lazy regime and kernel regime In Figure 4 we examine whether mirror descent with different unscaled potentials: Φ1 = x² d2 = | + x² and db = x² +x2 is in the lazy regime and whether it is in the kernel regime. According the plot,as the network width n increases, both the distance between the final parameter ((oo) and initial parameter 0, and the distance between the final kernel matrix H (co) and initial kernel matrix H (0) decrease to zero. This verifies that mirror descent with unscaled potentials is in both lazy regime and kernel regime, as shown in Theorem 2
In Figure we examine whether mirror descent with scaled potentials falls into lazy regime and whether it falls into kernel regime. Specifically, we observe that [10(co) - 011 decreases zero as n increases, which verifies that the training is in lazy regime as we shown in Theorem 9 In the right panel of Figure 5
that for potentials having & different from x² the kernel ma- trix moves significantly from its initial value during training, verifying our claim that networks trained by mirror descent with diverse scaled potentials do not fall in the kernel regime. Note di is 2-homogeneous. Therefore it induces the same potential function Tb under Assumption [ and Assumption 8 Eke(n) Φ1 (n(02 - 0x)) Σigin) &1 (0, Oz). This explains why mirror descent with the scaled potential induced by Φ1 is also in the kernel regime.
Effect of bp+lz potential We examine how scaled potentials of the form d(x) |>:| + 0.1.² affect the trained networks. In the left panel of Figure 6, we compare the solution to the variational problem ( with a constant ps and o(x) |z| 0.1x² for different p values, versus B-spline interpolations of the data with varying degrees. Interestingly, as p approaches 1, the variational solution behaves more like linear interpolation. This trend gets clearer in the right panel, which shows the 2D PCA representation of these interpolation functions. Here we directly consider the variational solution for two reasons: (1) these solutions well approximate absolute value networks trained by
53

---

Published as a conference paper at ICLR 2025
10a
dulel +if debel +x?
10
[|0t0] -011-
[|H(s)) -All
10m1
10"
4:00
6s50) m
10
Number of hidden neurons
Number of hidden neurons
Figure 4: Left: foo norm of the difference between the final parameter and initial parameter ReLU networks trained by mirror descent with different unscaled potentials: Φ1 x² の = | x²
and 03 = x² + x² plotted against the number of hidden units. Right: spectral norm (2-norm) of the difference between the final kernel matrix and the initial kernel matrix from the same training instances as the left panel, plotted against the number of hidden units.
+ Oibd
dabel +
109
4 dulel =x?
10-1
-81--
[[H(w)) -112
10-2
10-3
10
10
10
10:4
#:10)
$516) m Jof
*x0x)=x
10a
10a
10
10'
Number of hidden neurons
Number of hidden neurons
Figure 5: Left panel: for norm of the difference between the final parameter and initial parameter of absolute value networks trained by mirror descent with different scaled potentials: Φi = x² d2 = | + _², and db = 'r + x² plotted against the number of hidden units. Right panel: spectral norm (2-norm) of the difference between the final kernel matrix and initial kernel matrix from the same training instances the left panel, plotted against the number of hidden units. The inset shows the same data with the y-axis in linear scale instead of log-scale.
mirror descent with scaled potentials, as shown in Theorem 9 and (2) we observe that when p takes large value, the mirror descent converges in very slow rate and solving the variational problem is more efficient. Also note, potentials ((x) = |z| + x² with < 2 are not covered by Assumption g and hence the trajectory
as é does have a second derivative
can be ill-defined.
Effect of hypentropy potentials We examine the effect of hypentropy potentials ps(x)
arcsinh (z/ (z/6)
/ 1ي 32 Ghai et al. 2020 on the implicit bias of mirror descent training. The + 0+* their second derivatives, (x(()) 22
motivation behind these potentials is that,
converges to the second derivative of the negative entropy. Moreover, the Bregman divergence induced by the hypentropy function interpolates the squared Fuclidean distance and the relative entropy or Kullback-Leible (KL) divergence, as | varies from infinity to zero (Ghai et al. (020).. In agreement with this property, the left panel Figure 7 we observe that when takes large value, the trained networks behaves like the variational solution with o(x) = x², which is closely related to training
under Euclidean distance. Interestingly, as tends to zero, the trained networks behaves more like the variational solution with d(x) = | This pattern is also shown in the right panel, where we plot the 2D PCA representations for these interpolation functions.
Trained network vs. width Figure g compares ReLU networks of different width, trained by mirror descent with different unscaled potentials, against the solution the variational problem (7)
54

---

Published as a conference paper at ICLR 2025
1.0
Unear
4.0
Quadratic
Quadrabic
37
Linear
3.7
0.8
06
Quintic
Cubic
0.6
3.0
Quintic
3.0
2.7
02-
244
P
2.1
2.
0.2
-02-
14
-4
6.0
-06
Figure 6: Left: comparison of multiple interpolating functions on common data (black crossing points): B-spline interpolations with degree 1 (blue), 2 (green), 3 (purple), and (red); solutions to the variational problem with constant pg and 6(x) |x 1 + 0.1x where p ranges from 1.1 to 4.0 with the corresponding color indicated in the colorbar. Right: 2D PCA representation of the interpolation functions in the left panel.
Bo
x-
8.
6 Hi
0.015
$ x
0.3
54
0.010
5.4
0= [x|
0.2
0.005
37
28
2.8
0.0
11
-0.010
0.11
-0.015
Figure 7: Left: comparison of multiple interpolating functions on common data set (black crossing points): solutions to the variational problem with constant ps and d(x) |ar| for p = 2, 3, 4, and wide absolute value networks with uniformly initialized input biases, unit-norm initialized input weights and zero-initialized output weights and biases, trained by mirror flow with scaled hypentropy potentials dg where 3 ranges from 0.2 to 8.0 with the corresponding color indicated in the colorbar. Right: 2D PCA representation of the interpolation functions in the left panel.
In agreement with Theorem 2 and Figure [ (middle panel), the network function converges the solution to the variational problem the width increases. A further experiment is conducted for absolute value networks and scaled mirror descent, with results shown in Figure g In agreement with [ we observe that the network function trained with different scaled potentials converges to the solution to the corresponding variational problem as the width increases.
Theorem
55

---

Published as a conference paper at ICLR 2025
n 20
2 540
n 4860
Bm-1
R0
R/
(GD)
R1.8-111
R. B000
R.00-1)
A29
Figure 8: Comparison of ReLU networks of different width, trained by mirror descent with unscaled potentials (solid blue). against the solution to the variational problem (dashed orange). Networks of widths n = 30, 540, 4860 (columns) are trained using three different unscaled potentials: Ø1 x², s2 = |x|3 + x², and d3 .A + x² (rows). All networks have uniformly initialized input biases, unit-norm initialized input weights and zero-initialized output weights and biases.
n= 20
n 540
n 4860
R. 8--1
8:81=1)
d (GD)
6: Bcol
8:80111
0:.0011
A..0-1)
Figure 9: Comparison of absolute value networks different width, trained by mirror descent with scaled potentials (solid blue). against solutions the variational problem (dashed orange). Networks of widths n = 30, 540, 4860 (columns) are trained using three different unscaled potentials: Φ1 = x², 6ं = | + x²2 and d3 A² + x² (rows). All networks have uniformly initialized input biases. unit-norm initialized input weights and zero-initialized output weights and biases.
56


Reviews:
[{"review_id": "ckjzcM6sAN", "reviewer": "Reviewer_eX1T", "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.", "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\n\n* 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CqUlIhu3Lu", "reviewer": "Reviewer_QTJu", "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.", "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. \n\n2) Why is the mirror flow update scaled by $\\eta$? Do similar results hold if $\\eta$ is dropped?\n\n3) Is the restriction to the univariate case only needed for the variational characterization of the asymptotic solution? \n(i.e. can the univariate requirement be dropped for establishing lazy training results and risk minimization?)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Characterizing the asymptotic solution of mirror flow via a variational problem is interesting.", "weaknesses": "The work only considers univariate models. \n\n\nminor comment:\nConsidering the work only deals with mirror flow, I believe the title should be remedied to \"Implicit bias of Mirror Flow ...\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jt7u5nZti5", "reviewer": "Reviewer_8kZ2", "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.", "questions": "In both Theorem 2 and Theorem 8:\n“with high probability over the random parameter initialization,\nthere 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 …”?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "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.\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)