On the Implicit Bias of Initialization Shape: Beyond Infinitesimal Mirror Descent

Introduction

Gradient descent (GD) is the main optimization tool used in deep learning. A wealth of recent work has highlighted the key role of this specific algorithm in the generalization performance of the learned model, when it is over-parameterized. Namely, the solutions that gradient descent converges to do not merely minimize the training error, but rather reflect the specific implicit biases of the optimization algorithm.

In light of this role for GD, many works have attempted to precisely characterize the implicit bias of GD in over-parameterized models. Technically, these exact characterizations amount to identifying a function $Q(\mathbf{w})$ of the model parameters $\mathbf{w}$ such that GD converges to a minimizer (or, more generally, a stationary point) of $Q(\mathbf{w})$ under the constraint of having zero training error. The form of $Q(\mathbf{w})$ can depend on various hyper-parameters (e.g., initialization, architecture, depth) and its dependence sheds light on how these hyper-parameters affect the final solution. This approach worked very well in several regimes.

The first regime is the "Neural Tangent Kernel" (NTK) regime, which arises in networks that have an unrealistically large width Du et al. (2019); Jacot et al. (2018); Nguyen (2021) or initialization scale Chizat et al. (2019). In this regime, networks converge to a linear predictor where the features are not learned, but determined by the initialization (via the so-called “Tangent Kernel”), and in this case $Q(\mathbf{w})$ is just the RKHS norm for the linear predictor. Therefore, it is not surprising that models trained in this regime typically do not achieve state-of-the-art empirical performance in challenging datasets where deep networks perform well. Accordingly, this regime is typically considered to be less useful for explaining the success of deep learning.

The second regime is the diametrically opposed “rich” regime, which was analyzed specifically for classification problems with vanishing loss Lyu & Li (2020b); Chizat & Bach (2020). In this regime, the parameters converge to a stationary point (or sometimes a global minimum) of the optimization problem for minimizing $Q(\mathbf{w})=||\mathbf{w}||^{2}$ subject to margin constraints. This has been shown, under various assumptions, for linear neural networks Gunasekar et al. (2018b); Ji & Telgarsky (2019) and non-linear neural networks Nacson et al. (2019); Lyu & Li (2020a); Chizat & Bach (2020). This regime is arguably more closely related to the performance of practical neural networks but, as Moroshko et al. (2020) show, reaching this regime requires unrealistically small loss values, even in toy problems.

Understanding the implicit bias in more realistic and practically relevant regimes remains challenging in models with more than one weight layer. Current results are restricted to very simple models such as diagonal linear neural networks with shared weights in regression Woodworth et al. (2020) and classification Moroshko et al. (2020), as well as generalized tensor formulations of networks Yun et al. (2021). These results show exactly how the initialization scale determines the implicit bias of the model. However, these models are quite limited. For example, when the weights in different layers are shared, we cannot understand how the relative scale between layers affects the implicit bias.

Extending these exact results to a more realistic architectures is a considerable technical challenge. In fact, recent work has provided negative results with the square loss, for ReLU networks (even with a single neuron) Vardi & Shamir (2020) and for matrix factorization Razin & Cohen (2020); Li et al. (2021). Thus, finding scenarios where such a characterization of the implicit bias is possible and deriving its exact form is an open question, which we address here, making progress towards more realistic models.

Previous work (Woodworth et al., 2020; Gunasekar et al., 2018a; Yun et al., 2021; Vaskevicius et al., 2019; Amid & Warmuth, 2020a; b) that analyzes the exact implicit bias in such scenarios mostly focuses on least squares regression. All these analyses can be shown to be equivalent to expressing the dynamics of the predictor (which is induced by gradient flow on the model parameters) as Infinitesimal Mirror Descent (IMD), where the implicit bias then follows from Gunasekar et al. (2018a). This approach severely limits the model class that we can analyze because it is not always clear how to express the predictor dynamics as infinitesimal mirror descent. In fact, we can verify this is impossible to do even for basic models such as linear fully connected networks.

In this work, we sidestep the above difficulty by developing a new method for characterizing the implicit bias and we apply it to obtain several new results:

We identify degrees of freedom that allow us to modify the dynamics of the model so that it can be understood as infinitesimal mirror descent, without changing its implicit bias. In some cases, we show that this modification is equivalent to a non-linear “time-warping” (see Section 5).

Our approach facilitates the analysis of a strictly more general model class. This allows us to investigate the exact implicit bias for models that could not be analyzed using previous techniques. Specific examples include diagonal networks with untied weights, fully connected two-layerBy ”two-layers” we mean two weight layers. linear networks with vanishing initialization, and a two-layer single leaky ReLU neuron (see Sections 4, 6, and 8 respectively).

Our improved methodology is another step in the path toward analyzing the implicit bias in more realistic and complex models. Also, by being able to handle models with additional complexities, it already allows us to extend the scope of phenomena we can understand, shedding light on the importance of the initialization structure to implicit bias. For example,

We show that the ratio between weights in different layers at initialization (the initialization “shape”) has a marked effect on the learned model. We find how this property affects the final implicit bias (see Section 7).

We prove that balanced initialization in diagonal linear nets improves convergence to the “rich regime”, when the scale of the initialization vanishes (see Section 7.1).

Taken together, our analysis and results show the potential of our approach for discovering new implicit biases, and the insights these can provide about the effect of initialization on learned models.

In what follows, Sections 4-6 present derivations of implicit biases for several models of interest, and Section 7 uses these results to study the effect of initialization shape and scale on the learned models.

Preliminaries and Setup

using gradient descent with infinitesimally small stepsize (i.e., gradient flow)

We focus on overparameterized models, where there are many solutions that achieve zero training loss, and assume that the loss is indeed (globally) minimized by gradient flow.

Background: Deriving the Implicit Bias Using Infinitesimal Mirror Descent

We begin by describing the crux of current approaches to implicit bias analysis, and in Section 5 describe our “warping” approach that significantly extends these.

We focus on linear models that can be written as

which is the KKT stationarity condition. Thus, in this case, it is possible to find the $Q$-function by solving the differential equation

The aforementioned papers now proceed to solve the differential equation $\mathbf{H}=\nabla^{2}Q$ for $Q$. However, this proof strategy fundamentally relies on this differential equation having a solution, i.e., on $\mathbf{H}$ being a Hessian map. We emphasize that $\mathbf{H}$ being a Hessian map is a very special property, which does not hold for general positive definite matrix-valued functions.Indeed, Gunasekar et al. (2020) show that the innocent-looking $\mathbf{w}\mapsto I+\mathbf{w}\mathbf{w}^{\top}$ is provably not the Hessian of any function, which can be confirmed by checking the condition Eq. (6). Indeed, Eq. (5) only has a solution if $\mathbf{H}$ satisfies the Hessian-map condition (e.g., see Gunasekar et al., 2020)

As we discuss in Section 6, this condition is not met for natural models like fully connected linear neural networks, and therefore a new approach is needed.

Comparing Eqs. (3) and (7), we see that the infinitesimal mirror descent view is equivalent to the approach we have described, with $\psi$ corresponding exactly to $Q$.

Although it may have been presented in different ways, these analysis techniques have formed the basis for all of the existing exactThere are some statistical (i.e. non-exact) results for matrix factorization with vanishing initialization under certain data assumptions Li et al. (2018). characterizations of implicit bias for linear models with square loss (outside of the NTK regime) that we are aware of (e.g. Gunasekar et al. (2017); Woodworth et al. (2020); Amid & Warmuth (2020a); Moroshko et al. (2020)). In Section 5 we show how to extend this analysis to cases where $\mathbf{H}$ is not a Hessian map.

Diagonal Linear Networks

All previous analyses of the exact implicit bias for linear models with square loss (outside of the NTK regime) are limited to cases where the different layers share weights. In this section, we will remove this assumption, which allows us to analyze the effect of the relative scales of initialization between different layers in Section 7.1. To begin, we examine a two-layer “diagonal linear network” with untied weights

Previous Results: Woodworth et al. (2020); Moroshko et al. (2020) analyzed these models for the special case of shared weights where $\mathbf{u}_{+}=\mathbf{v}_{+}$ and $\mathbf{u}_{-}=\mathbf{v}_{-}$, corresponding to the model

Both of these works focused on unbiased initialization, i.e., $\mathbf{u}_{+}(0)=\mathbf{u}_{-}(0)=\alpha\mathbf{u}$ (for some fixed $\mathbf{u}$). In Yun et al. (2021) these results were generalized to a tensor formulation, yet one which does not allow untied weights (as in Eq. (8)).

Our Results: In this work, we analyze the model (8) for the square loss and show how both the initialization scale and the initialization shape (see Section 7.1) affect the implicit bias. To find the implicit bias of this model, we show how to express the training dynamics of this model in the form Eq. (3), which enables the use of the IMD approach (Sec. 3).

To simplify the presentation, we focus on unbiased initialization, where $\mathbf{u}_{+}(0)=\mathbf{u}_{-}(0)$ and $\mathbf{v}_{+}(0)=\mathbf{v}_{-}(0)$, which allows scaling the initialization without scaling the output Chizat et al. (2019). See Appendix A for a more general result with any initialization.

and $\sqrt{k_{i}}=2\left(u_{+,i}^{2}\left(0\right)+v_{+,i}^{2}\left(0\right)\right)$.

The function $Q_{\boldsymbol{k}}(\mathbf{w})$ in (9) generalizes the implicit regularizer found by Woodworth et al. (2020) to two layers with untied parameters. As expected, Eq. (9) reduces to Woodworth et al. (2020) when $\mathbf{u}_{+}(0)=\mathbf{v}_{+}(0)$ and $\mathbf{u}_{-}(0)=\mathbf{v}_{-}(0)$. Unlike the previous result, $Q_{\boldsymbol{k}}(\mathbf{w})$ can be used to study how the relative magnitude of $\mathbf{u}$ versus $\mathbf{v}$ at initialization affects the implicit bias. We present this analysis in Section 7.1, and highlight how initialization scale and shape have separate effects on the resulting model.

Warping Infinitesimal Mirror Descent

Perhaps surprisingly, for the right choice of $g$, the differential equation $g(\mathbf{w})\mathbf{H}(\mathbf{w})=\nabla^{2}Q(\mathbf{w})$ can have a solution even when $\mathbf{H}(\mathbf{w})=\nabla^{2}Q(\mathbf{w})$ does not! When such a $g$ can be found, we can continue the analysis just as before,

The above approach can also be interpreted as a non-linear warping of the time axis. The key idea is that rescaling “time” for an ODE affects neither the set of points visited by the solution nor the eventual limit point. Our approach essentially finds a rescaling that yields dynamics that allow solving for $Q$.

Therefore, the set of points visited by $\mathbf{w}(t)$ and $\mathbf{w}\left(\tau(t)\right)$ are the same, and so are their limit points $\mathbf{w}(\infty)=\mathbf{w}(\tau(\infty))$. All that changes is the time at which these points are reached. Furthermore, since $\tau^{\prime}>0$, $\tau$ is invertible so, conversely, a solution for Eq. (14) can also be converted into a solution for Eq. (13) via the warping $\tau^{-1}$. In this way, we can interpret $g$ as a time warping function which transforms the ODE

Fully Connected Linear Networks

In this section we examine the class of fully connected linear networks of depth $2$, defined as

For this model, the Hessian-map condition (Eq. (6)) does not hold and thus our analysis uses the “warped IMD” technique described in Section 5. In addition, our analysis of the implicit bias employs the following balancedness properties for gradient flow shown by Du et al. (2018):

Theorem 2.1 of Du et al. (2018) states that

In addition, Theorem 2.2 (a stronger balancedness property for linear activations) of Du et al. (2018) states that

First, we derive the implicit bias for a fully connected single-neuron assuming $\delta_{i}\geq 0$ (which ensures that we can write the dynamics in the form (4) for invertible $\mathbf{H}$), and then expand our results to multi-neuron networks under more specific settings.

In order to extend this result beyond a single neuron we require additional conditions to be met. For a multi-neuron network, in contrast to the single neuron case, we cannot use globally the “time warping” technique since it requires multiplying each neuron by a different $g$ function. However, for the special case of strictly balanced initialization, $\mathbf{\Delta}=0$, we can extend this result to $m>1$.

The Effect of Initialization Shape and Scale

Chizat et al. (2019) identified the scale of the initialization as the crucial parameter for entering the NTK regime, and Woodworth et al. (2020) further characterized the transition between the NTK and rich regimes as a function of the initialization scale, and how this affects the generalization properties of the model. Both showed the close relation between the initialization scale and the model width.

However, we identify another hyper-parameter that controls this transition between NTK and rich regimes for two-layer models, the shape of the initialization, which describes the relative scale between different layers.

We first demonstrate this by using the example of two-layer diagonal linear networks described in Section 4.

We denote the per-neuron initialization shape $s_{i}$ and scale $\alpha_{i}$ as

We can notice from Theorem 1 that $\sqrt{k_{i}}=2\left(u_{+,i}^{2}\left(0\right)+v_{+,i}^{2}\left(0\right)\right)$ controls the transition between the NTK and rich regimes. Using the definitions of the initialization shape and scale we write

Since $-1<s_{i}<1$, we can more accurately say that $\hat{k}_{i}=\frac{\alpha_{i}}{1-s_{i}^{2}}$ is the factor controlling the transition.

For simplicity, we next assume that $\alpha_{i}=\alpha,\,\,\ s_{i}=s\,\,\,\forall i\in[d]$. We can notice that for $k\rightarrow\infty$, i.e. $\frac{\alpha}{1-s^{2}}\rightarrow\infty$ we get that

which is exactly the minimum RKHS norm with respect to the NTK at initialization. Therefore, $k\rightarrow\infty$ leads to the NTK regime. However, for $k\rightarrow 0$, i.e. $\frac{\alpha}{1-s^{2}}\rightarrow 0$ we get that

which describes the rich regime. The proof for the above two claims appears in Appendix E.

Therefore, both the initialization scale $\alpha$ and the initialization shape $s$ affect the transition between NTK and rich regimes. While $\alpha\rightarrow 0$ pushes to the rich regime, $|s|\rightarrow 1$ pushes towards the NTK regime. Since both limits can take place simultaneously, the regime we will converge to in this case is captured by the joint limit

Intuitively, when $\alpha\rightarrow 0$ faster than $s\rightarrow 1$ we will be in the rich regime, corresponding to $\hat{k}^{\star}=0$. However, when $s\rightarrow 1$ faster than $\alpha\rightarrow 0$ we will be in the NTK regime, corresponding to $\hat{k}^{\star}=\infty$. For any $0<\hat{k}^{\star}<\infty$ the $Q$-function in Eq. (9) captures the implicit bias.

Figure 7.1 demonstrates the interplay between the scale and the shape of initialization. See Section 9 for details. The figure shows the population error (i.e., test error) of the learned model for different choices of scale $\alpha$ and shape $s$. Since in this case the ground truth is a sparse regressor, low error corresponds to the rich regime whereas high error corresponds to the NTK regime. It can be seen that as the shape $s$ approaches 1, the model tends to converge to a solution in the NTK regime, or an intermediate regime even for very small initialization scales. These results give further credence to the idea that the learned model will perform best when trained with balanced initialization ($s=0$).

2 Fully Connected Linear Networks

Similarly to the diagonal model, we again define the initialization shape parameter $s$ and scale parameter $\alpha$ as

Note that Theorem 2 is correct for $0\leq s<1$ and any $\alpha>0$. We also employ the initialization orientation, defined as $\mathbf{u}=\frac{\mathbf{w}(0)}{\|\mathbf{w}(0)\|}$. Given $\alpha,s,\mathbf{u}$ we identify a few limit cases.

and it is easy to verify that the tangent kernel is given by $K(\mathbf{x},\mathbf{x}^{\prime})=\mathbf{x}^{\top}\mathbf{B}^{-1}\mathbf{x}^{\prime}$.

To sum-up, in order to achieve non-kernel bias for fully connected networks we prefer balanced initialization ($s\approx 0$). This observation is in line with our observation for diagonal models in Section 7.1.

Two-Layer Single Leaky ReLU Neuron

We further extend our analysis to the class of fully connected two-layer single neuron with Leaky ReLU activations, $\sigma(x)=\max(x,\rho x)$ for $\rho>0$. This is a first step in analyzing the implicit bias of practical non-linear fully connected models for regression with the square loss.

We follow Dutta et al. (2013) in the definition of KKT conditions for non-smooth optimization problem (see the definition in Appendix G).

For a single-neuron network with Leaky ReLU activation $\sigma$ of any slope $\rho>0$, and for any $\delta\geq 0$, assume $a(0)\mathbf{w}(0)\neq\mathbf{0}$. If the gradient flow solution $(a(\infty),\mathbf{w}(\infty))$ satisfies $a(\infty)\sigma(\mathbf{X}^{\top}\mathbf{w}(\infty))=\mathbf{y}$, then $(a(\infty),\mathbf{w}(\infty))$ satisfies the KKT conditions (according to definition 1) of the following optimization problem:

and $q_{\delta}(\mathbf{w})$ is identical to the definition given in Theorem 2.

Recently, Vardi & Shamir (2020) proved a negative result for depth 2 single ReLU neuron with the square loss. They showed that it is impossible to characterize the implicit regularization by any explicit function of the model parameters. We note that Theorem 4 does not contradict the result of Vardi & Shamir (2020) since it does not include the ReLU case ($\rho=0$).

Numerical Simulations Details

In order to study the effect of initialization over the implicit bias of gradient flow, we follow the sparse regression problem suggested by Woodworth et al. (2020), where $\mathbf{x}^{(1)},...,\mathbf{x}^{(N)}\sim\mathcal{N}(0,I)$ and $y^{(n)}\sim\mathcal{N}(\langle\beta^{*},\mathbf{x}^{(n)}\rangle,0.01)$ and $\beta^{*}$ is $r^{*}$-sparse, with non-zero entries equal to $1/\sqrt{r^{*}}$. For every $N\leq d$, gradient flow will generally reach a zero training error solution, however not all of these solutions will be the same, allowing us to explore the effect of initialization over the implicit bias.

See Figure 7.1 for results, and Section 7.1 for discussion.

Conclusion

Understanding generalization in deep learning requires understanding the implicit biases of gradient methods. Much remains to be understood about these, and even a complete understanding of linear networks is yet to be attained. Here we make progress in this direction by developing a new technique, which we apply to derive biases for diagonal and fully connected networks with independently trained layers (i.e., without shared weights). This allows us to study the effect of the initialization shape on implicit bias.

From a practical perspective it has been previously observed that balance plays an important role in initialization. For example, Xavier initialization Glorot & Bengio (2010) is roughly balanced by construction, and our results now provide additional theoretical support for the practical utility of this commonly used approach. We believe it is likely that further theoretical results like those presented here, can lead to improved initialization methods that lead to more effective convergence to rich regime solutions.

Acknowledgements

This research is supported by the European Research Council (ERC) under the European Unions Horizon 2020 research and innovation programme (grant ERC HOLI 819080) and by the Yandex Initiative in Machine Learning. The research of DS was supported by the Israel Science Foundation (grant No. 31/1031), by the Israel Inovation Authority (the Avatar Consortium), and by the Taub Foundation. BW is supported by a Google Research PhD Fellowship.

References

Appendix A Proof of Theorem 1

We examine a two-layer “diagonal linear network” with untied weights

The gradient flow dynamics of the parameters is given by:

We note that the quantity $u_{+,i}u_{-,i}+v_{+,i}v_{-,i}$ is conserved during training, since

Combining Eq. 16 and Eq. 17 we can write:

which can be easily shown since $\frac{d}{dt}\left(v_{+,i}^{2}-u_{+,i}^{2}\right)=0$ and $\frac{d}{dt}\left(v_{-,i}^{2}-u_{-,i}^{2}\right)=0$. So using Eq. 18 we can write:

Coming back to $u_{+,i}^{2}+v_{+,i}^{2}+u_{-,i}^{2}+v_{-,i}^{2}$ we have using Eq. A that:

We next turn to solving for this $q$, beginning with Eq. 20:

where $k_{i}\triangleq\left(\delta_{+,i}-\delta_{-,i}\right)^{2}+4c_{i}^{2}$.

Integrating the above, and using the constraint $q^{\prime}\left(0\right)=0$ we get:

Finally, we integrate again to obtain the desired $q$:

Therefore, we get that gradient flow satisfies the KKT conditions for minimizing this $Q$, which completes the proof. ∎

Appendix B Proof of Theorem 2

We start by examining a general multi-neuron fully connected linear network of depth $2$, reducing our claim at the end to the case of a network with a single hidden neuron ($m=1$).

The fully connected linear network of depth $2$ is defined as

The parameter gradient flow dynamics are given by:

Using Theorem 2.1 of Du et al. (2018) (stated in Section 6), we can write

where we again employed Theorem 2.1 of Du et al. (2018). Also, since

Since $\left\|\mathbf{w}_{i}(t)\right\|^{2}\geq 0$ we choose the (+) sign and obtain

Comparing the form above with the Hessian in Eq. 24 we require

Finally, for the case of a fully connected network with a single hidden neuron ($m=1$), the condition

which since $\nu^{(n)}$ has no dependency on the index $i$ is a valid KKT stationarity condition for the $q$ we found above. Therefore, the gradient flow satisfies the KKT conditions for minimizing the $q$ we have found. ∎

First, we show that Eq. 23 cannot take the form suggested by Eq. 4 (as in the standard IMD approach described in Section 3):

From Eq. 23 we get that $\mathbf{H}(\mathbf{w})$ takes the form

Suppose $\mathbf{H}(\mathbf{w})$ is indeed the Hessian of some $Q(\mathbf{w})$, then is must respect the Hessian-map condition (see Eq. 6) for any $\delta\geq 0$. Specifically, for $\delta=0$ we get

which does not satisfy the Hessian-map condition

Therefore, Eq. 23 cannot be solved using the standard IMD approach, and requires our suggested “warped IMD” technique (see Section 5).

We notice that $\hat{g}(x)$ is smooth and positive for $\forall x>0$, and since $\lim_{x\rightarrow 0^{+}}\hat{g}(x)=\sqrt{\delta_{i}}$ (see Lemma 3) it is also bounded for any finite $x$.

we see that $\hat{g}^{\prime}(x)>0$, $\forall x>0$ and so $\hat{g}(x)$ is monotonically increasing.

We denote $f(x)=\frac{1}{x}\sqrt{\sqrt{x^{2}+\frac{\delta^{2}}{4}}-\frac{\delta}{2}}$ and $h(x)=\frac{f(x)}{2\left(\frac{\delta}{2}+\sqrt{\frac{\delta^{2}}{4}+x^{2}}\right)\sqrt{\frac{\delta^{2}}{4}+x^{2}}}$.

Without loss of generality it is enough to observe the following settings:

Therefore, if $\forall x\,\,,\frac{f^{\prime}(x)}{x}=-h(x)$ we get that $\frac{\partial\mathbf{H}_{i,i}(\mathbf{w})}{\partial\mathbf{w}_{k}}=\frac{\partial\mathbf{H}_{i,k}(\mathbf{w})}{\partial\mathbf{w}_{i}}$.

Using the derivative of $f(x)$ we can write:

and so $g({\mathbf{w}})\mathbf{H}(\mathbf{w})$ respects the Hessian-map condition.

Appendix C Proof of Proposition 1

We recall that the fully connected linear network of depth $2$ is defined as

Returning to the dynamics of model parameters (Eq. 21) we have

By using the Woodbury matrix identity we can write

From Theorem 2.2 of Du et al. (2018) (stated in Section 6) we get that

For the case of strict balanced initialization we have $\mathbf{\Delta}=0$, and therefore

where in the last transition we used the Sherman-Morrison lemma. It follows that

Using Theorem 2.1 of Du et al. (2018) (stated in Section 6), we know that

Comparing the form above with the Hessian in Eq. 25 we require

Therefore, gradient flow satisfies the KKT conditions for minimizing this $q$. ∎

Appendix D Proof of Theorem 3

We recall the proof of Theorem 2 given in Appendix B.

When the linear term captured by $\mathbf{z}$ in the $q$ function is equal to zero, we have

Hence, gradient flow satisfies the KKT conditions for minimizing this $q$.

It follows that for a multi-neuron fully connected network with non-zero infinitesimal initialization,

Appendix E Characterization of the Implicit Bias Captured in Theorem 1

In this Appendix we provide a detailed characterization of the implicit bias for a diagonal linear network as described in Theorem 1,

For simplicity, we next assume $\alpha_{i}=\alpha,\,\,\ s_{i}=s\,\,\,\forall i\in[d]$.

We can notice that for $k\xrightarrow{}\infty$, i.e. $\frac{\alpha}{1-s^{2}}\xrightarrow{}\infty$ we get that:

Calculating the tangent kernel at the initialization we get

For the case of unbiased initialization ($u_{+,i}\left(0\right)=u_{-,i}\left(0\right),v_{+,i}\left(0\right)=v_{-,i}\left(0\right)$) we have

Therefore, using Lemma 4, we can see that $Q_{\boldsymbol{k}}(\mathbf{w})$ is the RKHS norm with respect to the NTK at initialization. Therefore, $k\xrightarrow{}\infty$ indeed describes the NTK regime. For $k\xrightarrow{}0$, i.e. $\frac{\alpha}{1-s^{2}}\xrightarrow{}0$ we get that:

and $k\xrightarrow{}0$ describes the rich regime Woodworth et al. (2020).

Appendix F Characterization of the Implicit Bias Captured in Theorem 2

In this Appendix we provide a detailed characterization of the implicit bias for a two-layer fully connected neural network with a single hidden neuron ($m=1$) described in Theorem 2,

Note that for the sake of simplicity the notations above are an abbreviated version of those found Theorem 2.

F.2 Other special cases

We are interested in cases where the higher order terms vanish. Since $0\leq s<1$, we only need to require

In follows that when $\frac{\left(1-s\right)^{2}}{\alpha}\ll 1$ we can approximate

Note that $\mathbf{B}^{-1}$ is related to the NTK at initialization, since it is easy to verify that

and the NTK at initialization is given by

Next, we discuss the cases when condition (26) holds.

In this case (26) holds and thus the implicit bias is given by

F.2.2 The case s→1→𝑠1s\rightarrow 1 for any α>0𝛼0\alpha>0

In this case (26) also holds and thus the implicit bias is given by

where $\mathbf{B}$ defined in (27). Since $s\rightarrow 1$ we get that $\mathbf{B}\rightarrow\mathbf{I}$ and thus

Appendix G Proof of Theorem 4

where $\partial^{o}$ is the local (Clarke’s) sub-differential.

We follow the lines of the proof for Theorem 2 given in Appendix B.

As we do in Appendix B, we start by examining a general multi-neuron fully connected network of depth $2$, reducing our claim at the end to the case of a network with a single hidden neuron ($m=1$).

The fully connected depth $2$ network with Leaky ReLU activations is defined as

where $\sigma$ is a leaky ReLU with parameter $\rho$,

The gradient inclusion parameter dynamics are

Using Theorem 2.1 of Du et al. (2018) (stated in Section 6), we can write

Using the Sherman Morisson Lemma, we have

Next we use the relation proven in Appendix B,

We take notice that this is made possible since for a Leaky ReLU slope $\rho>0$, we have that $c_{i}^{\left(n\right)}\left(\infty\right)>0$.

Finally, for the case of a fully connected network with a single hidden neuron ($m=1$), the condition

which since $\nu^{(n)}$ has no dependency on the index $i$ is a valid KKT stationarity condition for the $q$ we found above (according to definition 1, where we notice that the second KKT condition of complementary slackness is not needed for regression since we use an equality constraint).

Therefore, the gradient flow satisfies the KKT conditions for minimizing the $q$ we have found.

Additionally, from Eq. 31, using the chain rule, we get

which, together with the feasability of the solution are exactly the KKT conditions of this (non-convex, non-smooth) optimization problem

Appendix H Auxiliary Lemmas

$\delta=a^{2}\left(0\right)-\left\|\mathbf{w}\left(0\right)\right\|^{2}=\frac{4\alpha s}{1-s^{2}}$ .

The initialization scale $\alpha$, initialization shape $s$ and the balancedness factor $\delta$ satisfy:

be defined $\forall x>0$, and $\forall\delta\geq 0$. Then:

Let $\mathbf{A}$ be a positive definite matrix and $f\left(\mathbf{x}\right)$ a kernel predictor corresponding to a linear kernel $K\left(\mathbf{x},\mathbf{x}^{\prime}\right)=\mathbf{x}^{\top}\mathbf{A}\mathbf{x}^{\prime}$. Then

where $f\left(\mathbf{x}\right)=\mathbf{w}^{\top}\mathbf{x}$.

Write $K\left(\mathbf{x},\mathbf{x}^{\prime}\right)=\mathbf{x}^{\top}\mathbf{A}\mathbf{x}^{\prime}=\mathbf{x}^{\top}\mathbf{A}^{\frac{1}{2}}\mathbf{A}^{\frac{1}{2}}\mathbf{x}^{\prime}$, then $\phi\left(\mathbf{x}\right)=\mathbf{A}^{\frac{1}{2}}\mathbf{x}$ is the corresponding feature mapping and