# Deconstructing the Goldilocks Zone of Neural Network Initialization

Artem Vysogorets  Anna Dawid  Julia Kempe

###### Abstract

The second-order properties of the training loss have a massive impact on the
optimization dynamics of deep learning models. Fort & Scherlis (2019)
discovered that a high positive curvature and local convexity of the loss
Hessian are associated with highly trainable initial points located in a
region coined the “Goldilocks zone”. Only a handful of subsequent studies
touched upon this relationship, so it remains largely unexplained. In this
paper, we present a rigorous and comprehensive analysis of the Goldilocks zone
for homogeneous neural networks. In particular, we derive the fundamental
condition resulting in non-zero positive curvature of the loss Hessian and
argue that it is only incidentally related to the initialization norm,
contrary to prior beliefs. Further, we relate high positive curvature to model
confidence, low initial loss, and a previously unknown type of vanishing
cross-entropy loss gradient. To understand the importance of positive
curvature for trainability of deep networks, we optimize both fully-connected
and convolutional architectures outside the Goldilocks zone and analyze the
emergent behaviors. We find that strong model performance is not necessarily
aligned with the Goldilocks zone, which questions the practical significance
of this concept.

Goldilocks Zone, Initialization, Curvature

##  1 Introduction

Every neural network gives rise to a high-dimensional optimization space
spanned by its trainable parameters. The complex geometry of these spaces and
embedded loss landscapes has been an area of prolific research since the
inception of machine learning models. The training loss Hessian and its
properties have received a lot of attention, as they offer key insights into
generalization (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017),
convergence speed (Becker et al., 1988), and broader optimization dynamics
(Jastrzębski et al., 2020; Cohen et al., 2021).

Figure 1: The Goldilocks zone is an area of high positive curvature of the
Hessian. Left: Originally, the Goldilocks zone is observed for a narrow range
of initialization scales. Middle: Setting the appropriate softmax temperature
$T$ allows for high positive curvature at initialization of any norm. Right:
Recreating the Goldilocks zone at an unscaled initialization just by varying
$T$.

While the Hessian is extensively studied throughout training and at
convergence, fewer works focus on the initialization stage. Recently, Fort &
Scherlis (2019) discovered the _Goldilocks zone_ —a region of the optimization
space marked by abnormally high values of positive curvature
$\text{Tr}(H)/\lVert H\rVert_{F}$ and local convexity of the loss Hessian $H$
(see Figure 1). The anomalous readings of these metrics are recorded at a
certain distance from the origin of the configuration space where some widely
used initialization schemes such as Xavier (Glorot & Bengio, 2010) and Kaiming
(He et al., 2015) are found. Thus, the Goldilocks zone is believed to be a
hollow centered spherical shell that contains a high density of suitable
initial points. Indeed, an appropriate parameter norm is crucial to avoid
exploding and vanishing signals, and it seems reasonable that initializations
of extreme norm (outside the Goldilocks zone) might suffer from this notorious
issue. However, internal covariate shift is largely solved in practice by
BatchNorm (Ioffe & Szegedy, 2015), and it can be directly accounted for in
some special cases. For example, given an $L$-layer homogeneous network
$f_{\alpha\theta}$ with an $\alpha$-scaled initialization (Dinh et al., 2017),
the appropriate logit variance and gradient norm can be restored by applying a
carefully selected softmax temperature $T=\alpha^{L}$ and learning rate
$\eta=\mathcal{O}(\alpha^{2})$. These adjustments ensure that
$f_{\alpha\theta}$ follows exactly the same training trajectory and has the
same initial positive curvature as $f_{\theta}$ (Figure 1). Thus, contrary to
the original claims by Fort & Scherlis (2019), the Goldilocks zone cannot be
characterized by the initialization norm alone and is not even a subset of the
configuration space. Instead, we find a more fundamental condition governing
the positive curvature of the Hessian in Section 3.

Previous works also attempted to describe the Goldilocks zone more formally.
Fort & Ganguli (2019) develop an abstraction, the random logit model, in which
they show that the positive curvature of cross-entropy loss vanishes with
increased logit variance. Cohen et al. (2021) attribute the decrease in cross-
entropy sharpness during optimization to the collapse of the softmax
distribution, which also has implications for the Goldilocks zone. However,
the phenomenon of positive curvature—its relation to the initialization norm
and trainability—remains unexplained. In this work, focusing on homogeneous
networks, we derive a comprehensive description of the Goldilocks zone from
the Gauss-Newton decomposition (Section 3) and formally associate high
positive curvature with certain properties of the network (Section 4). In
particular, we show that near-zero curvature results from a relatively low
spectral norm of the G-term in the Hessian decomposition due to either
saturated softmax or vanishing logit gradients, which naturally arise for
high- and low-norm initializations, respectively. Inside the Goldilocks zone,
we prove that the highest positive curvature is observed for networks with low
confidence, which in turn is associated with low initial loss and, for
balanced datasets, with vanishing expected loss gradient.

When the conditions governing the positive curvature of the loss are
established, we inquire about their relation to model trainability and
optimization dynamics. Fort & Scherlis (2019) find that optimization on random
subspaces of the parameter space converges only if they intersect the
Goldilocks zone. For unconstrained optimization, Gur-Ari et al. (2018)
discovered that gradients are mainly confined to a low-rank subspace spanned
by the Hessian top eigenvectors and vanish along flatter directions. Since the
excess of positive curvature manifests a larger separation between bulk and
outlier eigenvalues, it should then be associated with a more robust top-
eigenspace and, intuitively, a more informative training signal. Motivated to
unveil this relationship, we use gradient descent to optimize homogeneous
networks using a wide spectrum of initialization norms and learning rates and
taxonomize the emergent behaviors. Interestingly, we find that successful
training is not necessarily correlated with the Goldilocks zone. We
demonstrate setups where the slightest increase in the initialization norm of
LeNet-5 leads to degenerate learning despite happening well within boundaries
of the Goldilocks zone. These dynamics are marked by an increasing amount of
zero logits and, to the best of our knowledge, we are the first to report this
behavior.

### Contributions.

This paper conducts an extensive study of the Goldilocks zone of homogeneous
neural networks, both analytically and empirically.

  * •

In Section 3, we demonstrate that the Goldilocks zone cannot be characterized
by the initialization norm alone, as was previously claimed by Fort & Scherlis
(2019). Instead, we derive a more fundamental condition resulting in non-zero
positive curvature and find that it disappears due to saturated softmax on one
end and vanishing logit gradients on the other.

  * •

In Section 4, we closely study the interior of the Goldilocks zone and
analytically associate high positive curvature with low model confidence, low
initial loss, and low cross-entropy gradient norm.

  * •

In Section 5, we report the evolution and performance of scaled homogeneous
networks when optimized by gradient descent both inside and outside the
Goldilocks zone. Our investigation shows that positive curvature is a poor
estimator of the initialization trainability and exhibits a range of
interesting effects for initializations on the edge.

All of our experiments are implemented in PyTorch (Paszke et al., 2017) and
run on an internal cluster with NVIDIA RTX8000 and A100 GPUs. We use an open-
source implementation of stochastic power iteration with deflation to compute
the top eigenvalues and eigenvectors of the Hessian matrix (Golmant et al.,
2018). Our code is available here: github.com/avysogorets/goldilocks-zone.

##  2 Preliminaries & Notation

We begin by introducing the technical scope of this study, notation, and the
essential background. We consider a standard $K$-way classification problem
$\mathcal{D}=\\{(X^{\mu},y^{\mu})\\}_{\mu=1}^{N}$ with targets
$y^{\mu}\in[K]$. A neural network $f_{\theta}$, parameterized by a vector
$\theta\in\mathbb{R}^{P}$, computes $K$ logits $\\{z_{k}\\}_{k=1}^{K}$
associated to a probability distribution $p$ via the softmax function
$\sigma_{T}$ with a temperature parameter $T$:

$\displaystyle
p_{k}=[\sigma_{T}(z)]_{k}=\frac{\text{exp}[\frac{1}{T}z_{k}]}{\sum_{c=1}^{K}\text{exp}[\frac{1}{T}z_{c}]}\,.$
(1)

By default, $T=1$, in which case we refer to softmax simply as $\sigma$. The
corresponding cross-entropy loss is $\ell(p,y)=-\log p_{y}$. We restrict our
analysis to homogeneous models satisfying
$f_{\alpha\theta}=\alpha^{L}f_{\theta}$ for any scalar $\alpha>0$ where $L$ is
the number of layers (weight matrices) in $f$. This is a rather technical
assumption: homogeneous models are widely used in practice and include ReLU
networks without biases. The Hessian matrix $H\in\mathbb{R}^{P\times P}$ holds
the second derivatives of the loss at $\theta$:
$H_{ij}=\partial^{2}\ell/\partial\theta_{i}\partial\theta_{j}$. In principle,
all of the above quantities depend on one particular or a batch of inputs
$(X^{\mu},y^{\mu})$; we use the superscript $\mu$ to make this dependence
explicit where needed.

### Positive curvature.

In their seminal work, Fort & Scherlis (2019) define the Goldilocks zone
rather informally as a region of high (non-zero) positive curvature of the
loss Hessian, which can be unmistakably identified in Figure 1. Positive
curvature of the loss Hessian is defined by

$\displaystyle\frac{\text{Tr}(H)}{\lVert
H\rVert_{F}}=\frac{\sum_{i=1}^{P}\lambda_{i}}{\sqrt{\sum_{i=1}^{P}\lambda_{i}^{2}}}\,,$
(2)

where $\\{\lambda_{i}\\}_{i=1}^{P}$ are the eigenvalues of $H$. Another metric
used by Fort & Scherlis (2019) is local convexity, which refers to the
fraction of positive eigenvalues. These two metrics are intimately related and
can detect the Goldilocks zone independently of each other, so we focus only
on positive curvature throughout this paper. Since computing the full Hessian
is intractable for most modern architectures, Fort & Scherlis (2019)
substitute it with a projection $H_{d}=R^{T}HR$ onto a low-rank random
subspace with basis defined by a sparse matrix $R\in\mathbb{R}^{P\times d}$
with orthogonal columns. Given an initialization $\theta$, this amounts to
computing the Hessian of the training loss of the model parameterized by
$R\hat{\theta}+\theta$ with respect to latent $d$-dimensional parameters
$\hat{\theta}$ at the origin of the chosen low-rank space, which is much more
accessible. We adopt the same strategy and assume that the first- and second-
order derivatives with respect to $\theta$ more generally represent
derivatives with respect to _trainable_ $d$-dimensional parameters, which can
be latent parameters $\hat{\theta}$ or the original model weights if we let
$d=P$ and $R=I_{P}$. This technical nuisance affects none of our analyses, but
we add further comments as it becomes necessary.

##  3 Revisiting the Goldilocks Zone

Fort & Scherlis (2019) introduced the Goldilocks zone as a region of the
parameter space with an excess of positive curvature and local convexity of
the loss function, as defined in Section 2. Starting from a gold standard
Kaiming initialization $\theta_{0}$, Fort & Scherlis (2019) record these
statistics over a ray $\\{\alpha\theta_{0}\colon\alpha>0\\}$ and read
unusually high values when $\alpha$ falls within a relatively narrow range
centered around $\alpha=1$ (Figure 1 left). Moreover, they find that SGD
constrained to a random subspace is successful only if it intersects the
Goldilocks zone. On the basis of these observations, the authors suggest that
the Goldilocks zone is a thick, hollow spherical shell about the origin in the
configuration space, which is densely populated with initial points amenable
for training.

A simple example shows that, strictly speaking, this visually appealing
representation cannot be true. For homogeneous models, the $\alpha$-scale
transformation $\theta\rightarrow\alpha\theta$ with $\alpha>0$ effectively
scales the underlying configuration space by $\alpha$ (Dinh et al., 2017). We
shall derive next that the cross-entropy loss landscape scales together with
the configuration space when the softmax temperature satisfies $T=\alpha^{L}$,
restoring the positive curvature of the scaled model $f_{\alpha\theta}$ to
that of the original model $f_{\theta}$ (see the middle plot in Figure 1).
Thus, we argue that initialization norm has a coincidental relationship to the
Goldilocks zone, calling for an alternative, analytically driven
characterization.

### Gradients and Hessian of scaled models.

We begin by simplifying the notation for clarity of presentation. Denote a
scaled model $f_{\alpha\theta}$ by $f^{\prime}$ and adopt a similar notation
for all of its attributes ($\theta^{\prime}=\alpha\theta$, etc.). The chain
rule allows us to express gradients of the cross-entropy loss as

$\displaystyle\frac{\partial\ell}{\partial\theta}=\frac{1}{T}\left[-\frac{\partial
z_{y}}{\partial\theta}+\sum_{k=1}^{K}\sigma_{T}(z_{k})\frac{\partial
z_{k}}{\partial\theta}\right].$ (3)

By virtue of homogeneity, the logit gradients of the scaled model $f^{\prime}$
satisfy $\partial z_{k}^{\prime}/\partial\theta^{\prime}=\alpha^{L-1}\partial
z_{k}/\partial\theta$. For the special case of $T=\alpha^{L}$, we have
$\sigma_{\alpha^{L}}(z^{\prime}_{k})=\sigma(z_{k})$, giving
$\nabla_{\theta^{\prime}}\ell={\alpha}^{-1}\nabla_{\theta}{\ell}$. This tells
us that the $\alpha$-scaled model $f^{\prime}$ follows exactly the same
optimization trajectory as $f$ if the ratio of their respective learning rates
is $\alpha^{2}$. This factor ensures equal update norms relative to parameter
norms across the two models. To derive a similar relationship for the Hessians
of $f^{\prime}$ and $f$, we turn to the Gauss-Newton decomposition (Sagun et
al., 2016; Fort & Ganguli, 2019; Papyan, 2020). For a single training sample,
we have:

$\displaystyle H_{ij}=\underbrace{\sum_{k=1}^{K}\sum_{c=1}^{K}\frac{\partial
z_{k}}{\partial\theta_{i}}\left[\nabla^{2}_{z}\ell\right]_{kc}\frac{\partial
z_{c}}{\partial\theta_{j}}}_{\text{G-term
}(G_{*})}+\underbrace{\sum_{k=1}^{K}\left[\nabla_{z}\ell\right]_{k}\frac{\partial^{2}z_{k}}{\partial\theta_{j}\partial\theta_{i}}}_{\text{H-term
}(H_{*})}\,,$ (4)

where $\nabla^{2}_{z}\ell=\text{diag}(p)-pp^{\top}$ is the Hessian of the loss
with respect to the logits and $\nabla_{z}\ell=p-y$ where $p$ is the softmax
output and $y$ is a one-hot encoded tagret (Cohen et al., 2021). To emphasize
the dependence of the G-term and H-term on $p$ we will sometimes refer to them
as $G_{*}(p)$ and $H_{*}(p)$, respectively. Note that
$G_{*}(p)=J^{\top}[\text{diag}(p)-pp^{\top}]J$ where $J$ is the Jacobian
matrix. For the $\alpha$-scaled model $f^{\prime}$ used together with softmax
$\sigma_{T}$, we similarly get
$\nabla^{2}_{z^{\prime}}\ell^{\prime}=T^{-2}\left[\text{diag}(p^{\prime})-p^{\prime}p^{\prime\top}\right]$
and $\nabla_{z^{\prime}}\ell^{\prime}=T^{-1}(p^{\prime}-y)$. Combining this
with homogeneity of gradients, we derive the Gauss-Newton decomposition of the
Hessian of $f^{\prime}$:

$\displaystyle H^{\prime}$
$\displaystyle=G^{\prime}_{*}+H^{\prime}_{*}=\frac{\alpha^{2L-2}}{T^{2}}G_{*}(p^{\prime})+\frac{\alpha^{L-2}}{T}H_{*}(p^{\prime}).$
(5)

Linearity of differentiation ensures that this equation is valid for the full-
batch Hessian, too. Indeed, the full G-term and the full H-term are just the
averages of the per-sample quantities $G^{\mu}_{*}$ and $H^{\mu}_{*}$ defined
in Equation 4, respectively. Thus, we will abuse this notation and refer to
their full counterparts in the same way where appropriate. In this case, the
probability vectors $p$ and $p^{\prime}$ can be viewed as matrices.

Returning to our discussion on the shape of the Goldilocks zone, we remark
that $T=\alpha^{L}$ yields $H^{\prime}=\alpha^{-2}H$ (as this temperature
ensures $p=p^{\prime}$). Since the measures of positive curvature and local
convexity are robust to scaling of the Hessian matrix, we conclude that
_initialization of any norm can be in the Goldilocks zone provided an
appropriate softmax temperature_ (see Figure 1). Now that the connection
between the initialization norm and the Goldilocks zone is compromised, we
leverage Equation 5 to establish the more fundamental principles governing
positive curvature in neural networks.

### Gauss-Newton decomposition.

The Gauss-Newton decomposition in Equation 4 is a common entry point for many
studies on the Hessian of large neural networks. The Hessian exhibits a “bulk-
outlier” eigenspectrum with the majority of eigenvalues small and clustered
around zero and only a handful of large positive outliers (Sagun et al., 2016;
Gur-Ari et al., 2018; Ghorbani et al., 2019). This decomposition is inherited
from the individual spectra of $G_{*}$ and $H_{*}$ with the top and the bulk
eigenvalues attributed to these two terms, respectively. We present a survey
of works concerned with this phenomenon in Appendix A.

### Rediscovering the Goldilocks zone.

We are now in position to describe the exact conditions that result in an
excess of positive curvature—the hallmark of the Goldilocks zone. The
eigenstructures of the individual terms in the Gauss-Newton decomposition of
the Hessian suggest that the G-term is the one and only source of positive
curvature of the Hessian, and so

$\displaystyle\lVert G^{\prime}_{*}\rVert_{2}\gtrsim\lVert
H^{\prime}_{*}\rVert_{2}$ (6)

is a necessary and sufficient condition for the Goldilocks zone. We observe
this exact correspondence in Figure 2. Otherwise, when $\lVert
G^{\prime}_{*}\rVert_{2}$ is sufficiently small relative to $\lVert
H^{\prime}_{*}\rVert_{2}$, the bulk-like eigenspectrum of the H-term swallows
the outliers of the G-term and is inherited by the Hessian, resulting in a
near-zero positive curvature. This sudden change in the eigenstructure is
known as the BBP phase transition (Baik et al., 2005; Fort & Ganguli, 2019).
Normally (when $\alpha=1$ and $T=1$), we expect the above inequality to be
true. Figure 2 (bottom) confirms that the G-term has an edge over the H-term
within a neighborhood around the unaltered model (orange bar), which also
corresponds to the Goldilocks zone.

Therefore, it is left to understand when and why Equation 6 is violated.
Figure 2 reveals that this occurs at sufficiently low or sufficiently high
values of $\alpha$ given a fixed temperature $T=1$. For $\alpha>1$, the norm
of the G-term plummets to zero when the logit variance becomes sufficiently
large for the softmax output $p^{\prime}$ to collapse to a one-hot
distribution for all training samples. In this scenario,
$\text{diag}(p^{\prime})-p^{\prime}p^{\prime\top}$ is identically zero, as has
been observed by Cohen et al. (2021) in the context of progressive sharpening.
For us, this implies that
$G^{\prime}_{*}(p^{\prime})=J^{\prime\top}[\text{diag}(p^{\prime})-p^{\prime}p^{\prime\top}]J^{\prime}$
is identically zero as well, making the Hessian equal to the H-term. Fort &
Ganguli (2019) observed the causal relationship between the increased logit
variance and vanishing positive curvature but never explained it analytically.

Figure 2: Positive curvature (top) and spectral norm (bottom) of the Hessian,
G-term, and H-term across initialization scales. We computed these quantities
on a low-rank subspace with $d=50$. Left: LeNet-300-100 (fully-connected) on
FashionMNIST; Right: LeNet-5 (convolutional) on CIFAR-10.

For $\alpha<1$, on the other hand, the G-term does not completely vanish; in
fact, it achieves the highest positive curvature at low initialization scales
(see the top plots in Figure 2), suggesting that higher entropy predictions
should generally be associated with larger positive curvature. According to
Equation 5, $\lVert G^{\prime}_{*}\rVert_{2}=\mathcal{O}(\alpha^{2L-2})$ while
$\lVert H^{\prime}_{*}\rVert_{2}=\mathcal{O}(\alpha^{L-2})$. Thus, the G-term
simply decays faster than the H-term as $\alpha\rightarrow 0$, eventually
letting it dominate in the Gauss-Newton decomposition. The higher decay rate
of the G-term comes from the cross-class product of logit gradients that
vanish as $\mathcal{O}(\alpha^{L-1})$ each. Thus, we call this phenomenon
_vanishing logit gradient_ and emphasize that vanishing _loss_ gradient can in
fact co-occur with high positive curvature as we demonstrate in the next
section.

##  4 Features of the Goldilocks Zone

In this section, we inquire about the properties of the interior of the
Goldilocks zone. In particular, we will associate extreme values of positive
curvature with certain features of the network, initialization, and the data.
To this end, we study the spectral properties of the G-term that, provided
Equation 6 holds, transfer to the loss Hessian as well. Equation 5 reveals
that $G^{\prime}_{*}(p)\propto G_{*}(p)$ for the same $p$, so that structural
changes in the eigenspectrum of the G-term at scaled initializations are due
to the changing softmax output $p$ alone. Therefore, we can drop the
superscript and study the behavior of $G_{*}$ as a function of $p$, which is
the goal of this section. In particular, we will find that high positive
curvature is associated with low model confidence, which, in turn, leads to
vanishing expected loss gradients for balanced datasets and low initial loss.

At the basis of our analysis is the random model introduced by Fort & Ganguli
(2019). They observed that same-logit gradients cluster across training
samples, and that the $K$ corresponding logit gradient means $c_{k}$ are
nearly orthogonal to each other. In light of this observation, they model
$c_{k}\sim\mathcal{N}(0,\sigma_{c}^{2}I_{P})$ and, for every input $X^{\mu}$,
let the corresponding $k$-th logit gradient be
$\nabla_{\theta}f_{k}(X^{\mu})=c_{k}+e^{\mu}_{k}$ with iid residuals
$e_{k}^{\mu}\sim\mathcal{N}(0,\sigma_{e}^{2}I_{P})$. Recall from Section 2
that, in practice, we compute positive curvature of the Hessian projected to a
low-rank hyperplane given by an orthonormal basis $R\in\mathbb{R}^{P\times
d}$, so we actually care about gradients with respect to the $d$-dimensional
trainable parameters. Conveniently, the random model still holds in this low-
rank subspace since the corresponding gradients are just
$R^{\top}c_{k}+R^{\top}e_{k}^{\mu}$ where
$R^{\top}c_{k}\sim\mathcal{N}(0,\sigma_{c}^{2}I_{d})$ and
$R^{\top}e^{\mu}_{k}\sim\mathcal{N}(0,\sigma_{e}^{2}I_{d})$ as linear
transformations of isotropic Gaussian variables by $R^{\top}$ satisfying
$R^{\top}R=I_{d}$. Thus, we drop the transformation and simply assume
$c_{k}\sim\mathcal{N}(0,\sigma_{c}^{2}I_{d})$ and
$e^{\mu}_{k}\sim\mathcal{N}(0,\sigma_{e}^{2}I_{d})$. Note that, in particular,
the variance of logit gradients and the corresponding residuals does not
depend on $d$. We make two more simplifying assumptions and require $c_{k}$ to
be pairwise orthogonal and have equal squared length $\lVert
c_{k}\rVert^{2}=\mathbb{E}\lVert c_{k}\rVert^{2}=d\sigma^{2}_{c}$.

Figure 3: Positive curvature of the true G-term $G_{*}$ and the expected
G-term $\mathbb{E}G_{*}$ from Equation 8 computed using logits associated with
different initialization scales (model confidence). We used a low-rank
subspace with $d=50$. Error bands represent min/max across $3$ seeds. Left:
LeNet-300-100 on FashionMNIST; Right: LeNet-5 on CIFAR-10.

For a single training sample, let $C$ and $E$ be $k\times d$ matrices of logit
gradient means and the corresponding residuals (i.e., $C_{k}=c_{k}$ and
$E_{k}=e_{k}$), respectively; the Jacobian is then $J=C+E$ and so
$G_{*}=(C+E)^{\top}[\text{diag}(p)-pp^{\top}](C+E)$. Taking expectation over
data, we get

$\displaystyle\mathbb{E}_{e}G_{*}$
$\displaystyle=C^{\top}[\text{diag}(p)-pp^{\top}]C+(1-\lVert
p\rVert^{2})\sigma_{e}^{2}I_{d}$ $\displaystyle\equiv
d\sigma_{c}^{2}\begin{bmatrix}\text{diag}(p)-pp^{\top}&0\\\
0&0\end{bmatrix}+(1-\lVert p\rVert^{2})\sigma_{e}^{2}I_{d}$ (7)

where in the second line, we switch to the orthonormal basis of the normalized
logit gradient means $c_{k}$. We can now compute the positive curvature of
$\mathbb{E}_{e}G_{*}$ directly from its elements. Assuming that $p$ is not
one-hot, we get

$\displaystyle\frac{\text{Tr}(\mathbb{E}_{e}G_{*})}{\lVert\mathbb{E}_{e}G_{*}\rVert_{F}}=\frac{\sqrt{d}(\sigma_{c}^{2}+\sigma_{e}^{2})}{\sqrt{\sigma_{e}^{4}+2\sigma_{e}^{2}\sigma_{c}^{2}+d\sigma_{c}^{4}{\Gamma}_{p}^{-2}}}$
(8)

where $\Gamma_{p}$ is the positive curvature of $\text{diag}(p)-pp^{\top}$ as
defined in Equation 2. For the expected G-term computed on a batch of samples
with potentially different softmax probabilities $p^{\mu}$, $\Gamma_{p}$
computes positive curvature of the average matrix
$\frac{1}{B}\sum_{\mu=1}^{B}\text{diag}(p^{\mu})-p^{\mu}{p^{\mu}}^{\top}$.
Having empirically estimated values $\sigma_{c}^{2}$ and $\sigma_{e}^{2}$, we
use Equation 8 to compute the positive curvature of $\mathbb{E}_{e}G_{*}$ on
real softmax outputs $p$ found by scaling initialization, and find that it
approximates the positive curvature of the true G-term across varying model
confidence levels very well (Figure 3).

Figure 4: Left: The dependence of positive curvature on the dimension $d$ of
the low-rank hyperplane used to compute the projected G-term. Right:
$\Gamma_{p}^{-2}$ of the matrix $\text{diag}(p)-pp^{\top}$ where vectors $p$
are produced by scaling $30$ different logit sets by
$\alpha\in[10^{-2},10^{2}]$ ($30$ pink curves). The red curve corresponds to
$\Gamma_{p}^{-2}$ computed for the average matrix
$\frac{1}{30}\sum_{\mu=1}^{30}\text{diag}(p^{\mu})-p^{\mu}{p^{\mu}}^{\top}$.

### Positive curvature & dimension $d$.

We can now explain the dependence of positive curvature on the dimension $d$
of the hyperplane used to project the model parameters, which was first
noticed by Fort & Scherlis (2019), see the left plot in Figure 4. Indeed,
according to Equation 8, larger $d$ are associated with higher positive
curvature of the expected G-term given the same softmax outputs $p$.

### Positive curvature & model confidence.

Figure 2 shows that positive curvature of the G-term monotonically decreases,
reaching zero at some $\alpha>1$ due to a collapsing softmax. Since larger
$\alpha$ are associated with colder softmax distributions, it is natural to
hypothesize that positive curvature of the G-term is inversely related to
model confidence. Equation 8 reveals that positive curvature of the expected
G-term is directly related to that of $\text{diag}(p)-pp^{\top}$ and is larger
for smaller values of $\Gamma_{p}^{-2}$. Despite the seemingly nice form of
this matrix, $\Gamma_{p}$ is difficult to analyze algebraically; still, a few
general remarks shed light on the above relationship. First, Section 4 can be
used to show that $\mathbb{E}G_{*}$ has a “bi-level” eigenspectrum that yields
maximal positive curvature when the softmax output $p$ is uniform (Appendix
B). Second, as seen from Equation 2, positive curvature of any rank-$r$
positive semidefinite matrix is just the ratio of L1 and L2 norms of its $r$
non-zero eigenvalues, which is known to fall between $1$ and $\sqrt{r}$.
Therefore, $\Gamma_{p}\in[1,\sqrt{K-1}]$, and it achieves its maximal value
when $p$ is uniform. In Appendix C, we prove that $\Gamma_{p}\rightarrow 1$ as
$p$ collapses to a one-hot vector due to the increasing initialization scale
$\alpha$. Moreover, Figure 4 (right) reveals that $\Gamma_{p}$ is, in fact,
monotonic between its extreme values with respect to the entropy of $p$.
Overall, this suggests that $\Gamma^{-2}_{p}$ in Equation 8 grows with model
confidence, which supports our hypothesis.

Figure 5: The effects of the average softmax output $\hat{Q}$ and the actual
target prior $Q$ on the batch gradient. Left: Cosine similarity of gradients
computed on different datasets (SVHN, CIFAR-10, and a randomly generated
dataset) of an unscaled LeNet-5. Right: Given a fixed unscaled initialization
of LeNet-5, we sample $2,000$ different class priors $Q$ uniformly from a
probability simplex $\Delta_{10}$. For each of them, we choose a subset of
CIFAR-10 with $5,000$, compute the cross-entropy gradient, and plot its norm
against $\lVert\hat{Q}-Q\rVert_{2}$. As we predicted, the norms follow a
linear trend with the predicted slope ($\sigma_{c}^{2}$ was estimated on the
entire CIFAR-10 dataset; $d=P=61,170$.

### Model confidence & vanishing gradients.

In Section 3, we saw that vanishing logit gradients associated with small-norm
network initializations diminish the role of the G-term in the Gauss-Newton
decomposition, leading to zero positive curvature of the Hessian. The term
“vanishing gradients” commonly refers to the condition of signal decay during
backpropagation to deeper layers, emerging either due to saturated activation
functions or small spectral radii of parameter matrices (Bengio et al., 1994;
Pascanu et al., 2013). The vanishing logit gradients observed at low
initialization scales of homogeneous networks fall in the latter category.
Using the random model of Fort & Ganguli (2019), we derive yet another,
previously unknown type of vanishing cross-entropy gradients caused by the
match between the average softmax output of the model
$\hat{Q}=\mathbb{E}p\approx\frac{1}{N}\sum_{\mu}p^{\mu}$ and the dataset prior
$Q=\mathbb{E}y\approx\frac{1}{N}\sum_{\mu}y^{\mu}$ where expectation is taken
over the data distribution. To this end, consider the expected batch cross-
entropy loss gradient:

$\displaystyle\mathbb{E}_{E}\left[\frac{\partial\ell}{\partial\theta}\right]$
$\displaystyle=\mathbb{E}_{E}\left[\frac{1}{B}\sum_{\mu=1}^{B}\sum_{k=1}^{K}(c_{k}+e_{k}^{\mu})(p_{k}^{\mu}-y_{k}^{\mu})\right]$
$\displaystyle=\sum_{k=1}^{K}c_{k}(\hat{Q}_{k}-Q_{k})$ (9)

By the orthogonality assumption on logit gradient means $c_{k}$, the norm of
this quantity is $d\sigma_{c}^{2}\lVert\hat{Q}-Q\rVert$. In Figure 5 (right),
we confirm this linear relationship on LeNet-5 shockingly well. Normally,
randomly-initialized models are biased and have non-uniform average priors
$\hat{Q}$, so that $\lVert\hat{Q}-Q\rVert>0$. In this case, we expect loss
gradients computed on two potentially unrelated but balanced datasets to have
considerable overlap as both should align with the non-zero expected gradient.
Indeed, Figure 5 (left) shows that gradients of LeNet-5 computed on CIFAR-10
and on random images drawn from a Gaussian have cosine similarity $>0.7$,
which is quite noteworthy for a 61,170-dimensional space. In constrast, for
low-confident models that always predict uniform distribution (e.g., when
$\alpha\ll 1$), logit gradients cancel each other out when computed on a large
balanced batch. Hence, rather counterintuitively, we derive that high positive
curvature is associated with low-norm cross-entropy gradients for balanced
datasets, as we verified in Figure 6.

### Model confidence & initial loss.

Fort & Scherlis (2019) observed that initializations with higher positive
curvature tend to also have lower initial loss, which is often considered
favorable for learning (Agarwala et al., 2023). Now, this relationship comes
naturally since both phenomena are associated with low-confidence models.
Indeed, the expected cross-entropy loss on samples from class $k$ is
$\mathbb{E}_{p|y=k}(-\log p_{k})$, which is lower bounded by
$-\log\mathbb{E}_{p|y=k}p_{k}$ by virtue of Jensen’s inequality. If the
softmax output $p$ is independent of the input’s class identity $y$—a
reasonable assumption for randomly-initialized models—then the lower bound is
just $-\log\mathbb{E}p_{k}=-\log\hat{Q}_{k}$. Thus, the total expected loss is
no less than $\sum_{k=1}^{K}(-\log\hat{Q}_{k})\geq K\log K$ with equality when
the expected model prediction $\hat{Q}$ is uniform. In principle, the model
does not have to always be uncertain for $\hat{Q}$ to be uniform; the
predictions can even be one-hot if all classes are equally-likely to be
chosen. However, this symmetry is practically unachievable for randomly-
initialized models, not to mention that the above bound becomes quite loose
when the model makes confident mistakes. The right plot in Figure 6 confirms
that lower initial losses correspond to higher average entropy of softmax
output and, hence, larger positive curvature.

### Summary.

In this Section, we scrutinized the interior of the Goldilocks zone by
presenting novel observations and justifying some existing claims made in
previous studies. In particular, we related high positive curvature to low
model confidence, vanishing expected full-batch gradients for balanced
datasets, and low initial loss, as is once again demonstrated in Figure 6.

Figure 6: The interplay between model confidence (avg. softmax entropy), loss
gradient norm, positive curvature, and the initial loss value. Given a single
balanced batch of $3,000$ CIFAR-10 images, we sampled $2,000$ Kaiming
initializations for LeNet-5 and computed the above statistics for each of
them.

##  5 Connections to Optimization

Previous works suggest that, in a standard setup (i.e., within the Goldilocks
zone), the top eigenspace of the training loss Hessian plays an important role
during optimization. Gur-Ari et al. (2018) demonstrate that gradients are
largely confined to that space, and Ben Arous et al. (2023) prove this
phenomenon for two-layer neural networks. Outside the Goldilocks zone, where
Equation 6 no longer holds, the loss curvature vanishes along all directions
and so gradient descent must behave differently. In this section, we
characterize the behavior of gradient descent outside the Goldilocks zone and
identify some failure modes that result in poor convergence or generalization
as shown in Figure 7.

Figure 7: Overview of the training outcomes observed in this study.
Homogeneous networks with large-norm initializations either diverge, develop
an increasing number of samples with zero logits, or learn in the lazy regime.
Homogeneous networks with small-norm initializations diverge as long as their
output remains uniform and either circle back to the Goldilocks zone and train
normally, or continue diverging with a potential to develop zero logits and
remain trapped in that regime.

### Experimental setup.

We optimize $\alpha$-scaled homogeneous networks using vanilla full-batch
gradient descent for 20,000 epochs. In particular, we use a fully-connected
LeNet-300-100 with two hidden layers on FashionMNIST and a convolutional
LeNet-5 with $4$ hidden layers on CIFAR-10 (LeCun et al., 1998), all
implemented in PyTorch (Paszke et al., 2017). We set softmax temperature to
$T=1$ to link the Goldilocks zone to the initialization norm. Recall from
Section 3 that logit gradients of the $\alpha$-scaled network $f^{\prime}$ are
$\alpha^{L-1}$ times the respective logit gradients of the unscaled model $f$,
so we multiply a preset learning rate $\eta_{0}$ by $\alpha^{2-L}$ to ensure
that updates are initially commensurate to the weights of $f^{\prime}$. In
fact, this correction factor ensures that $f^{\prime}$ has exactly the same
training dynamics as $f$ provided that logit gradients are combined using the
same softmax output $p=p^{\prime}$ to compute the update vector as indicated
in Equation 3. From this perspective, our experiments essentially constitute
an ablation study examining the effects of extreme model (un)confidence, which
corresponds to the exterior of the Goldilocks zone in our setup.

### Main results.

Our findings are summarized in Figure 7. As long as the softmax ouptut remains
uniform ($\alpha<1$), homogeneous models remain in the divergence regime
characterized by an increase of the parameter norm in accordance with previous
studies (Liu et al., 2023). Once the weights become sufficiently large for the
predictions to become non-uniform, the network circles back to the Goldilocks
zone and trains normally thereafter provided that the learning rate is
approximately admissible, i.e., $\eta=\eta_{0}\alpha^{2-L}<\lVert H\rVert_{2}$
(Lewkowycz et al., 2020; Cohen et al., 2021; Wang et al., 2022). Hence, in
Figure 8 (left) we observe a linear separation between trainable and divergent
setups that extends well into the low-norm initialization region and way
beyond the Goldilocks zone (see Appendix E for further discussion). When the
learning rate is too high, however, the network breezes past the Goldilocks
zone and either diverges to infinity or develops an increasing amount of zero
logits (middle column in Figure 8). This happens as some negative weights grow
much faster than the positive ones, promoting zero activations within the
network by virtue of ReLU.

Figure 8: Training statistics of LeNet-300-100 (top) and LeNet-5 (bottom)
optimized with gradient descent across initialization scales $\alpha$ and base
learning rates $\eta_{0}$. Left: maximum achieved training accuracy. Test
accuracy exhibits the same general patterns and is omitted. Middle: maximum
achieved parameter norm; azure values correspond to models that develop over
75% zero logits and never recover from this regime. Right: product of the
effective learning rate $\eta=\eta_{0}\alpha^{2-L}$ with the initial loss
curvature $\lVert H\rVert_{2}$; inadmissible configurations with values above
$2$ are shown in pink. The highlighted strip corresponds to the Goldilocks
zone where the G-term is dominant (cf. Figure 2).

Homogeneous networks with large-norm initializations ($\alpha>1$) exhibit a
more diverse collection of behaviors. Previous works suggest that these models
adhere to the lazy learning regime characterized by approximately linear
optimization dynamics (Chizat et al., 2019; Moroshko et al., 2020; Kumar et
al., 2023). We confirm this observation for LeNet-300-100 but not for LeNet-5.
Figure 9 shows that, while the training error of LeNet-300-100 reaches zero,
the test accuracy saturates at 75%, which is 10% lower than that of the
unscaled network, suggesting lazy learning. In contrast, even the slightest
increase in the initialization norm drives LeNet-5 to be completely
untrainable (Figure 8 bottom-left), which is not well-aligned with the
Goldilocks zone (Figure 8 bottom-right). In these cases, the network fails to
train and develops zero logits (Figure 8 bottom-center). Unlike the case with
$\alpha<1$, however, here zero logits often emerge when the weights are
completely balanced, presenting a rather mysterious phenomenon.

Liu et al. (2023) and our own experiments show that networks with large
initialization and no regularization neither diverge nor circle back to the
Goldilocks zone. Thus, as long as these models are inaccurate, their loss
remains extremely large and scales exponentially with $\alpha$ (Fort &
Scherlis, 2019). Thus, we hypothesize that, unable to learn meaningful
representations, gradient descent finetunes the weights of confident but
inaccurate networks to make logits of misclassified samples zero in an attempt
to reduce the exploding loss to just $\log(K)$—the value of cross-entropy
under uniform softmax distribution. In favor of this intuition, we discover
that models trained on random labels tend to develop zero logits as well. This
condition arises early in training and disappears, provided that the labels
remain fixed, as the network memorizes the correct output. On the other hand,
it only gets worse if labels are randomized on every iteration, in which case
the lowest possible loss is achieved by always making uniform predictions.

Figure 9: Training curves of LeNet-300-100 optimized with full-batch gradient
descent on FashionMNIST across different base learning rates $\eta_{0}$ (note
that the effective learning rate is $\eta_{0}\alpha^{2-L}$). Left: low-
confidence initialization with $\alpha=0.01$. Right: high-confidence
initialization with $\alpha=1,000$.

### Goldilocks zone & trainability.

The observations above reveal that homogeneous networks may not only converge
when initialized outside the Goldilocks zone but also fail to learn when
optimized well within its boundaries (Figure 8). The former occurs as the
model circles back to the Goldilocks zone ($\alpha\ll 1$) or when it learns
the lazy solution ($\alpha\gg 1$), and the latter—as it develops zero logits.
Indeed, LeNet-5 does not perform above random for any $\alpha>1.8$ and any
learning rate, whereas the corresponding Goldilocks zone spans
$\alpha\in[1,18]$. Lastly, Figure 8 simply illustrates that phase transitions
between different optimization regimes are not necessarily aligned with the
Goldilocks zone, which is naturally highlighted in the right plots. Thus, we
argue that positive curvature of the loss Hessian is not an appropriate metric
to evaluate the quality of initializations in terms of network trainability
and generalization.

##  6 Conclusion

This paper studies the Goldilocks zone of neural network initializations in
the context of homogeneous architectures, both analytically and empirically.
We demonstrate that the positive curvature of the training loss requires a
robust top eigenspace of the Hessian inherited from its positive semi-definite
component, the G-term and not a particular initialization norm per se.
Normally, the G-term dominates in the Gauss-Newton decomposition of the
Hessian but vanishes, for example, at extreme initialization scales or softmax
temperatures, which is formally shown in Equation 5. We study properties of
the G-term itself and relate high positive curvature to low model confidence,
low initial cross-entropy loss, and a previously unknown type of vanishing
gradient caused by the match between class priors found in the data and
computed by the model. Finally, we report the training behavior of homogeneous
networks optimized by gradient descent from a wide variety of initialization
norms and uncover interesting training modalities at the edges of the
Goldilocks zone (Figure 7). Furthermore, we find that model performance is
only vaguely related to the Goldilocks zone. For example, homogeneous networks
sometimes develop zero logits for an increasing number of inputs even when
initialized and trained within the Goldilocks zone; thus, we argue against
using positive curvature of the loss Hessian as a reliable estimate of the
model performance.

### Limitations & future work.

Most of our analysis requires neural networks to be homogeneous. Individual
examinations are warranted for architectures with residual connections or non-
homogeneous activation functions. Interestingly, our preliminary experiments
show that BatchNorm eliminates the Goldilocks zone altogether, showing no
excess of positive curvature even for normally initialized networks.
Furthermore, while we comprehensively studied the behaviors of scaled
homogeneous models when optimized by gradient descent, the conclusions might
change with adaptive algorithms such as AdaGrad or Adam as these optimizers
are known to effectively escape flat loss regions (Orvieto et al., 2022). We
note that existing studies propose logit normalization for calibration not
only after but also during training to improve convergence and generalization
(Wei et al., 2022; Agarwala et al., 2023). Given the connection between model
confidence and the Goldilocks zone unveiled in our study, future research may
investigate how these methods affect positive curvature or, more generally,
the eigenspectrum of the loss Hessian. Finally, our findings demonstrate the
need for alternative measures to evaluate the quality of network
initializations as positive curvature fails to faithfully capture model
performance after training.

## Acknowledgements

The authors thank Mahalakshmi Sabanayagam, Jonathan Niles-Weed, and Kyunghyun
Cho for numerous enlightening discussions. AV and JK were supported by the NSF
Award 1922658. The Flatiron Institute is a division of the Simons Foundation.
This work was supported in part through the NYU IT High Performance Computing
resources, services, and staff expertise.

## References

  * Agarwala et al. (2023) Agarwala, A., Schoenholz, S. S., Pennington, J., and Dauphin, Y.  Temperature check: theory and practice for training models with softmax-cross-entropy losses.  _Transactions on Machine Learning Research_ , 2023.  ISSN 2835-8856.
  * Allen-Zhu et al. (2019) Allen-Zhu, Z., Li, Y., and Song, Z.  A convergence theory for deep learning via over-parameterization.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 242–252. PMLR, 09–15 Jun 2019.
  * Baik et al. (2005) Baik, J., Ben Arous, G., and Péché, S.  Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices.  _The Annals of Probability_ , 33(5):1643 – 1697, 2005.
  * Becker et al. (1988) Becker, S., Le Cun, Y., et al.  Improving the convergence of back-propagation learning with second order methods.  In _Proceedings of the 1988 connectionist models summer school_ , pp. 29–37, 1988.
  * Ben Arous et al. (2023) Ben Arous, G., Gheissari, R., Huang, J., and Jagannath, A.  High-dimensional sgd aligns with emerging outlier eigenspaces, 2023.
  * Bengio et al. (1994) Bengio, Y., Simard, P., and Frasconi, P.  Learning long-term dependencies with gradient descent is difficult.  _IEEE Transactions on Neural Networks_ , 5(2):157–166, 1994.
  * Bunch et al. (1978) Bunch, J. R., Nielsen, C. P., and Sorensen, D. C.  Rank-one modification of the symmetric eigenproblem.  _Numerische Mathematik_ , 31(1):31–48, 1978.
  * Chizat et al. (2019) Chizat, L., Oyallon, E., and Bach, F.  On lazy training in differentiable programming.  _Advances in neural information processing systems_ , 32, 2019.
  * Cohen et al. (2021) Cohen, J., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A.  Gradient descent on neural networks typically occurs at the edge of stability.  In _International Conference on Learning Representations_ , 2021.
  * Dinh et al. (2017) Dinh, L., Pascanu, R., Bengio, S., and Bengio, Y.  Sharp minima can generalize for deep nets.  In Precup, D. and Teh, Y. W. (eds.), _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 1019–1028. PMLR, 06–11 Aug 2017.
  * Du et al. (2019) Du, S. S., Zhai, X., Póczos, B., and Singh, A.  Gradient descent provably optimizes over-parameterized neural networks.  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_ , 2019.
  * Fort & Ganguli (2019) Fort, S. and Ganguli, S.  Emergent properties of the local geometry of neural loss landscapes.  _CoRR_ , abs/1910.05929, 2019.
  * Fort & Scherlis (2019) Fort, S. and Scherlis, A.  The Goldilocks zone: Towards better understanding of neural network loss landscapes.  In _Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence_. AAAI Press, 2019.
  * Ghorbani et al. (2019) Ghorbani, B., Krishnan, S., and Xiao, Y.  An investigation into neural net optimization via Hessian eigenvalue density.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 2232–2241. PMLR, 09–15 Jun 2019.
  * Glorot & Bengio (2010) Glorot, X. and Bengio, Y.  Understanding the difficulty of training deep feedforward neural networks.  In Teh, Y. W. and Titterington, M. (eds.), _Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics_ , Proceedings of Machine Learning Research, pp. 249–256. PMLR, 13–15 May 2010.
  * Golmant et al. (2018) Golmant, N., Yao, Z., Gholami, A., Mahoney, M., and Gonzalez, J.  pytorch-hessian-eigenthings: efficient pytorch hessian eigendecomposition, October 2018.
  * Gur-Ari et al. (2018) Gur-Ari, G., Roberts, D. A., and Dyer, E.  Gradient descent happens in a tiny subspace.  _CoRR_ , abs/1812.04754, 2018.
  * He et al. (2015) He, K., Zhang, X., Ren, S., and Sun, J.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 1026–1034, 2015.
  * Hochreiter & Schmidhuber (1997) Hochreiter, S. and Schmidhuber, J.  Flat minima.  _Neural computation_ , 9(1):1–42, 1997.
  * Ioffe & Szegedy (2015) Ioffe, S. and Szegedy, C.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37_ , ICML’15, pp. 448–456. JMLR.org, 2015.
  * Jacot et al. (2018) Jacot, A., Gabriel, F., and Hongler, C.  Neural tangent kernel: Convergence and generalization in neural networks.  In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_. Curran Associates, Inc., 2018.
  * Jastrzębski et al. (2020) Jastrzębski, S., Szymczak, M., Fort, S., Arpit, D., Tabor, J., Cho, K., and Geras, K.  The break-even point on optimization trajectories of deep neural networks.  In _International Conference on Learning Representations, ICLR_ , 2020.
  * Keskar et al. (2017) Keskar, N., Nocedal, J., Tang, P., Mudigere, D., and Smelyanskiy, M.  On large-batch training for deep learning: Generalization gap and sharp minima.  In _5th International Conference on Learning Representations, ICLR_ , 2017.
  * Kumar et al. (2023) Kumar, T., Bordelon, B., Gershman, S. J., and Pehlevan, C.  Grokking as the transition from lazy to rich training dynamics.  _arXiv preprint arXiv:2310.06110_ , 2023.
  * LeCun et al. (1998) LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.  doi: 10.1109/5.726791.
  * Lewkowycz et al. (2020) Lewkowycz, A., Bahri, Y., Dyer, E., Sohl-Dickstein, J., and Gur-Ari, G.  The large learning rate phase of deep learning: the catapult mechanism, 2020.
  * Liu et al. (2020) Liu, S., Papailiopoulos, D., and Achlioptas, D.  Bad global minima exist and SGD can reach them.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Liu et al. (2023) Liu, Z., Michaud, E. J., and Tegmark, M.  Omnigrok: Grokking beyond algorithmic data.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Moroshko et al. (2020) Moroshko, E., Gunasekar, S., Woodworth, B., Lee, J., Srebro, N., and Soudry, D.  Implicit bias in deep linear classification: Initialization scale vs training accuracy.  In _NeurIPS 2020_. ACM, December 2020.
  * Orvieto et al. (2022) Orvieto, A., Kohler, J., Pavllo, D., Hofmann, T., and Lucchi, A.  Vanishing curvature in randomly initialized deep ReLU networks.  In Camps-Valls, G., Ruiz, F. J. R., and Valera, I. (eds.), _Proceedings of The 25th International Conference on Artificial Intelligence and Statistics_ , volume 151 of _Proceedings of Machine Learning Research_ , pp. 7942–7975. PMLR, 28–30 Mar 2022.
  * Papyan (2018) Papyan, V.  The full spectrum of deepnet Hessians at scale: Dynamics with SGD training and sample size.  _arXiv: Learning_ , 2018.
  * Papyan (2019) Papyan, V.  Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet Hessians.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 5012–5021. PMLR, 09–15 Jun 2019.
  * Papyan (2020) Papyan, V.  Traces of class/cross-class structure pervade deep learning spectra.  _Journal of Machine Learning Research_ , 21(252):1–64, 2020.
  * Pascanu et al. (2013) Pascanu, R., Mikolov, T., and Bengio, Y.  On the difficulty of training recurrent neural networks.  In Dasgupta, S. and McAllester, D. (eds.), _Proceedings of the 30th International Conference on Machine Learning_ , volume 28 of _Proceedings of Machine Learning Research_ , pp. 1310–1318, Atlanta, Georgia, USA, 17–19 Jun 2013. PMLR.
  * Paszke et al. (2017) Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A.  Automatic differentiation in pytorch.  2017\.
  * Pennington & Bahri (2017) Pennington, J. and Bahri, Y.  Geometry of neural network loss surfaces via random matrix theory.  In Precup, D. and Teh, Y. W. (eds.), _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 2798–2806. PMLR, 06–11 Aug 2017.
  * Perry et al. (2018) Perry, A., Wein, A. S., Bandeira, A. S., and Moitra, A.  Optimality and sub-optimality of PCA I: Spiked random matrix models.  _The Annals of Statistics_ , 46(5), oct 2018.
  * Sabanayagam et al. (2023) Sabanayagam, M., Behrens, F., Adomaityte, U., and Dawid, A.  Unveiling the Hessian’s connection to the decision boundary, 2023.
  * Sagun et al. (2016) Sagun, L., Bottou, L., and LeCun, Y.  Singularity of the hessian in deep learning.  _CoRR_ , abs/1611.07476, 2016.
  * Wang et al. (2022) Wang, Z., Li, Z., and Li, J.  Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability.  In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 9983–9994. Curran Associates, Inc., 2022.
  * Wei et al. (2022) Wei, H., Xie, R., Cheng, H., Feng, L., An, B., and Li, Y.  Mitigating neural network overconfidence with logit normalization.  In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 23631–23644. PMLR, 17–23 Jul 2022.

##  Appendix A Hessian Decomposition: Related Works

The Gauss-Newton decomposition in Equation 4 is a common entry point for many
studies on the Hessian of large neural networks. The Hessian exhibits a “bulk-
outlier” eigenspectrum with the majority of eigenvalues small and clustered
around zero and only a handful of large positive outliers (Sagun et al., 2016;
Gur-Ari et al., 2018; Ghorbani et al., 2019). This decomposition is inherited
from the individual spectra of $G_{*}$ and $H_{*}$ with the top and the bulk
eigenvalues attributed to these two terms, respectively. A few observations
offer an intuitive explanation for this claim.

First, the G-term can be rewritten in matrix form as
$G_{*}(p)=J^{\top}[\text{diag}(p)-pp^{\top}]J$ where $J$ is the logit
Jacobian. The matrix $\text{diag}(p)-pp^{\top}$ is positive semi-definite, so
$G_{*}$ has non-negative eigenvalues. At the same time, $G_{*}$ computed on a
single training sample is at most rank $K-1$, suggesting that the full G-term
should be low-rank, especially for overparameterized architectures and small
training datasets (Pennington & Bahri, 2017). Moreover, Fort & Ganguli (2019)
find empirically that logit gradients on same-class training samples
considerably overlap, which implies that eigenspaces of the G-terms of
different data samples align. In support of this claim, Papyan (2019) shows
that the top eigenvectors of the Hessian are attributable to high magnitude
gradient class means. This agrees with a common observation that the Hessian
has $K$ outlier eigenvalues. Recent studies argue that the corresponding
eigenvectors encode the decision boundary of the network, and that more
outliers emerge for more complex decision boundaries (e.g., after training
from an adversarial initialization) (Liu et al., 2020; Sabanayagam et al.,
2023). This observation can potentially correspond with a three-level
hierarchical decomposition of the G-term demonstrated by Papyan (2020), which
reveals additional clusters of smaller outliers.

The H-term, on the other hand, has a continuous bulk of small eigenvalues, as
independently verified by a number of studies (Papyan, 2018). Assuming that
$H_{*}$ is a Wigner random matrix, Pennington & Bahri (2017) find its limiting
spectral density to be consistent with experiments on randomly-initialized
ReLU networks without biases. They also note that the H-term has a block off-
diagonal structure for ReLU networks, which in particular implies that its
positive curvature is zero. In the context of Neural Tangent Kernel (NTK), a
kernel-based model for optimization dynamics, logits depend linearly on model
parameters, and so $H_{*}$ is identically zero (Jacot et al., 2018).

This pervasive analytical and empirical evidence suggests that, normally, the
G-term has a pronounced bulk+outlier eigenspectrum that dominates the bulk
eigenvalues found in the H-term. Pennington & Bahri (2017) note that the
eigenspaces of these two matrices are in a nearly generic position and do not
align in any special way. To this end, spiked matrix theory suggests that the
Hessian eigenspectrum should be well approximated by the combination of
individual eigenspectra of $G_{*}$ and $H_{*}$ (Perry et al., 2018). Since
outlier eigenvalues of the Hessian are particularly important, researchers
often drop the H-term from Equation 4 and focus on the G-term alone (Sagun et
al., 2016; Fort & Ganguli, 2019; Cohen et al., 2021; Wang et al., 2022).

##  Appendix B Eigenvalues of $\mathbb{E}G_{*}$

Under the assumptions of the random model on logit gradient clustering
described in Section 4, Section 4 offers a clear picture of the eigenstructure
of the expected G-term. In particular, it consists of a constant bulk plus
$K-1$ outliers:

$\displaystyle\lambda_{i}(\mathbb{E}_{e}G_{*})=\begin{cases}(1-\lVert
p\rVert^{2})\sigma_{e}^{2}+d\sigma_{c}^{2}\tilde{\lambda}_{i}&\text{ if }i\leq
K-1,\\\ (1-\lVert p\rVert^{2})\sigma_{e}^{2}&\text{ otherwise. }\end{cases}$

where $\tilde{\lambda}_{i}$ is the $i$-th largest eigenvalue of the matrix
$\text{diag}(p)-pp^{\top}$. Since this is a rank-$1$ update to a diagonal
matrix, its eigenvalues are interlaced with those of $\text{diag}(p)$; more
precisely, for all $i\in[K-1]$, $p_{k_{i}}\geq\tilde{\lambda}_{i}\geq
p_{k_{i+1}}$ where index $\\{k_{i}\\}_{i=1}^{K}$ sorts elements of $p$ in non-
increasing order (Bunch et al., 1978). This, for example, immediately proves
that when $p$ is uniform, all $K-1$ outlying eigenvalues of
$\mathbb{E}_{e}G_{*}$ are identically $1/K$ as they get squeezed between
values $p_{k}$. The same “bi-level” eigenstructure emerges when the expected
G-term is taken across multiple training samples, as long as their softmax
output remains uniform, as is the case for completely unconfident models
($\alpha\ll 1$).

##  Appendix C The Limiting Behavior of $\mathbf{\text{diag}(p)-pp^{\top}}$

Recall that the positive semi-definite matrix $\text{diag}(p)-pp^{\top}$ of
rank $K-1$ defined for some probability distribution $p$ over $K$ classes is
identically zero when $p$ is one-hot. In this section, we derive a few
limiting properties of this matrix as the distribution $p$ gets colder. To
this end, for any $\epsilon>0$ and $2\leq S\leq K$, let $p(S,\epsilon)$ be
defined as $p_{k}=\epsilon$ for $k\in[S-1]$, $p_{S}=1-(S-1)\epsilon$, and zero
otherwise. In other words, $p(S,\epsilon)$ has $S$ non-zero entries with $S-1$
of them equal to $\epsilon$.

### Positive curvature $\mathbf{\Gamma_{p}}$.

In Section 4, we defined $\Gamma_{p}$ to be positive curvature of the matrix
$\text{diag}(p)-pp^{\top}$. By inspecting the elements of this matrix, we find

$\displaystyle\Gamma^{2}_{p}=\left[\frac{\text{Tr}(\text{diag}(p)-pp^{\top})}{\lVert\text{diag}(p)-pp^{\top}\rVert_{F}}\right]^{2}=\frac{\left(\sum_{k=1}^{K}p_{k}-\sum_{k=1}^{K}p_{k}^{2}\right)^{2}}{\sum_{k=1}^{K}(p_{k}-p_{k}^{2})^{2}+\sum_{k\neq
c}p_{k}^{2}p_{c}^{2}}.$ (10)

Note that $\Gamma_{p}$ is undefined when $p$ is one-hot, but we can still
analyze its limiting behavior. Rewriting Equation 10 for a distribution
$p(S,\epsilon)$ defined above, we get

$\displaystyle\Gamma^{2}_{p}(S,\epsilon)=\frac{\epsilon^{2}(4S^{2}-8S+4)+\mathcal{O}(\epsilon^{3})}{\epsilon^{2}(S^{2}+S-2)+\mathcal{O}(\epsilon^{3})}\quad\Longrightarrow\quad\lim_{\epsilon\rightarrow
0}\Gamma^{2}_{p}(S,\epsilon)=4\frac{S-1}{S+2}.$ (11)

While this limit evaluates to $3$ when $S=10$ (the number of classes in the
datasets considered in this study), this is typically not how softmax output
$p$ collapses as we increase logit variance. Assuming no two logits are
identically the same, we should eventually get $p=p(2,\epsilon)$ with the two
largest logits associated with non-zero softmax probabilities. In this
realistic scenario, the limiting value of $\Gamma_{p}^{2}$ is $1$, as observed
in Figure 4. Note that some trajectories of $\Gamma_{p}^{-2}$ over the entropy
of the underlying softmax distribution $p$ seem to converge to a smaller
number before eventually reaching $1$. These trajectories correspond to logits
with several close high values, so that more than two entries in the softmax
output remain positive as we increase the initialization scale.

### The eigenspectrum.

As argued above, just before the softmax output $p$ succumbs to the increasing
logit variance and collapses to a one-hot vector, it necessarily has the form
$p(2,\epsilon)$. In this case, up to a rearrangement of some rows and columns,

$\displaystyle\text{diag}(p)-pp^{\top}=\begin{bmatrix}\epsilon&-\epsilon&\cdots&0\\\
-\epsilon&\epsilon&\cdots&0\\\ \vdots&\vdots&\ddots&\vdots\\\ 0&0&0&0\\\
\end{bmatrix}+\mathcal{O}(\epsilon^{2})I_{K}.$ (12)

It is easy to see that, neglecting the $\mathcal{O}(\epsilon^{2})$ terms, this
matrix has a single non-zero eigenvalue equal to $2\epsilon$. In Appendix D,
we present some interesting observations about the corresponding eigenvector.

Figure 10: The pre-collapse regime of LeNet-300-100 on FashionMNIST. Top row:
the top eigenvector of the first layer’s Hessian is reshaped into the weight
matrix dimensions ($300\times 784$), averaged across the hidden dimension
($300$), and reshaped as an input image ($28\times 28$). Note that the above
heatmaps have different scales. As we gently increase $\alpha$, only one input
sample (a picture of a bag in this case) remains uncertain, taking full
control of the G-term and representing the principle curvature direction for
the input layer. Bottom: The collapse of per-sample softmax entropies
$\mathbb{H}(p^{\mu})$ and the corresponding positive curvature of the loss
Hessian.

##  Appendix D The Pre-collapse Regime

Having examined positive curvature of low-confident models in Section 4, we
focus our analysis on the other end of the Goldilocks zone associated with
high model confidence. Recall that positive curvature of a positive semi-
definite matrix is lower bounded by $1$. Figure 2 illustrates this property
for the G-term: as we increase $\alpha$, $\text{Tr}(G_{*})/\lVert
G_{*}\rVert_{F}$ approaches $1$, remains level for a short time (we call this
_the pre-collapse regime_), and then suddenly drops to $0$. This property even
transfers to the full Hessian, as seen from Figure 2. Further, in the pre-
collapse regime of fully-connected architectures, we mysteriously observe that
the principle curvature direction of the first layer’s Hessian aligns with one
particular sample from the dataset (Figure 10). It turns out that we have all
the necessary tools to explain this surprising behavior. As we increase
$\alpha$, the G-terms $G_{*}^{\mu}$ of individual samples approach zero
because $\text{diag}(p^{\mu})-p^{\mu}{p^{\mu}}^{\top}$ does; however, they do
so at different rates. When $\alpha$ falls within a very specific narrow
range, there is only one sample $\mu_{0}$ (the least confidently classified
one) left with a large enough G-term, so that
$G_{*}=\frac{1}{N}\sum_{\mu}^{N}G_{*}^{\mu}\approx G_{*}^{\mu_{0}}$. At this
point, the top Hessian eigenspace comes from that of $G_{*}^{\mu_{0}}$ alone,
which is spanned by the corresponding logit gradients. Since logit gradients
with respect to parameters in the input layer are just the input vector
itself, it manifests as the top eigenvector of the Hessian. Furthermore, in
Appendix C, we show that $G_{*}(p)$ has only one non-zero eigenvalue when $p$
is on a knife edge from collapsing to a one-hot vector, so that positive
curvature of $G^{\mu_{0}}$ and, hence, of the full Hessian, approaches $1$ in
the pre-collapse regime.

##  Appendix E More Connections to Optimization

In this section, we elaborate on the phenomena summarized in Section 5. In
particular, we present more details about the optimization behaviors observed
when scaled homogeneous networks are trained by gradient descent.

Figure 11: Training dynamics of the USO model and networks with low
initialization norm ($\alpha<1$), all for LeNet-300-100 on FashionMNIST and
$\eta_{0}=0.0002$. Left: The divergence paths of networks with $\alpha<1$ that
initially have a uniform softmax output. The parameter magnitudes of these
networks are shown after rescaling back by $\alpha^{-1}$ for a fair comparison
with the USO model. The $\alpha$-scaled networks follow the same exact
divergence path as the USO model until their weights grow large enough to make
softmax output non-uniform. Right: Properties of the divergence (USO model):
cosine similarity between the current parameter vector $\theta_{t}$ and (1)
the initial point $\theta_{0}$, (2) current negative loss gradient $-g_{t}$.
The red curve represents the training error achieved when training
LeNet-300-100 initialized at the parameter vector of the current USO model
rescaled to an appropriate magnitude:
$\lVert\theta_{0}\rVert\theta_{t}\lVert\theta_{t}\rVert^{-1}$ where
$\theta_{0}$ is the original Kaiming initialization.

### Low-confidence initializations ($\alpha<1$).

Initially, the softmax output of these models is numerically uniform for
sufficiently small $\alpha$ due to vanishing logit variance. We observe that,
as long as it stays uniform, the model remains in a divergence regime
characterized by the increase of parameter norm $\lVert\theta_{t}\rVert$,
gradient norm $\lVert g_{t}\rVert$ and the alignment
$\cos(\theta_{t},-g_{t})$. Similarly, Liu et al. (2023) report parameter
growth during the early training stage of down-scaled networks (see Figure 1
in their paper). Once the parameter norm grows large enough for the prediction
to become non-uniform, networks either begin training normally or, if the
learning rate is too high, continue diverging. The left plots in Figure 8
reveal that the maximum admissible base learning rate $\eta_{0}$ that yields
convergence is linearly related to $\alpha$, forming a linear “staircase”-like
separation between trainable and non-trainable setups.

To understand this relationship, recall that classical optimization theory
suggests that, for quadratic loss functions, gradient-based methods converge
when $\eta<2/\lVert H\rVert_{2}$ and diverge otherwise (Wang et al., 2022).
This result was confirmed in deep learning with more complex loss functions,
although slightly larger learning rates are often admissible, too (Lewkowycz
et al., 2020; Cohen et al., 2021). Assuming that the loss curvature $\lVert
H\rVert_{2}$ is approximately constant when different $\alpha$-scaled networks
escape the divergence regime and re-enter the Goldilocks zone, this
requirement translates to $\eta_{0}\alpha^{-1}=2/\mathcal{O}(1)$, which
corresponds to the linear separation between trainable and divergent models
seen in the left plots in Figure 8. This extension of the convergence region
well into $\alpha<1$ suddenly ends at sufficiently small $\alpha$. In this
case, it is the limited number of epochs that prevents networks with smaller
initialization reach the Goldilocks zone. Still, we believe that this phase
transition should inevitably emerge at some, possibly smaller, value of
$\alpha$ due to the adverse effects of remaining in the divergence regime for
arbitrarily long.

To study these adverse effects, we propose a “uniform softmax output” (USO)
model, which is a normally initialized network that always receives updates as
if its softmax output is uniform, regardless of its actual predictions. In
other words, on every training iteration, we manually curate the update
direction by combining logit gradients using coefficients
$p_{k}=1/K-1\\{y=k\\}$. Equation 3 then implies that the USO model and the
down-scaled networks with uniform output have exactly the same evolution for
the same base initialization because they receive identical updates. Figure 11
(left) confirms that low-confidence networks evolve according to the USO model
until their parameters grow large enough for the softmax output to become non-
uniform, at which point they re-enter the Goldilocks zone. On the other hand,
the USO model never escapes the divergence regime by design, allowing us to
study the properties of scaled models with arbitrarily small $\alpha$ as they
continue diverging. Figure 11 (right) demonstrates that longer divergence is
associated with a more severe rotation of the parameter vector, which
significantly affects its trainability even after fixing its magnitude, as
evident from the training error shown in red. Thus, we expect that for
sufficiently small values of $\alpha>0$, scaled networks cannot train normally
by the time they reach the Goldilocks zone. Furthermore, longer divergence
entails larger gradients, so that a network may just overshoot the Goldilocks
zone altogether unless $\eta_{0}$ is appropriately adjusted.

Figure 12: Left: Evaluating the features learned by LeNet-300-100 at different
initialization scales. We train a linear classifier on data representations
extracted from the penultimate layer of scaled networks at convergence (solid)
as well as randomly initialized unscaled networks (dashed). Error bands
represent min/max across 3 seeds. Right: The proportion of training samples
with all logits equal zero as computed by LeNet-5 throughout optimization.

### High-confidence initializations ($\alpha>1$).

For large-norm initializations, homogeneous networks are overly confident and
compute one-hot softmax outputs for all input samples. While both
architectures considered in this study exhibit similar optimization patterns
for $\alpha<1$, their behavior is remarkably different in this scenario. For
LeNet-300-100, we find that predictions remain one-hot throughout training for
admissible step sizes where learning is possible. In fact, on every iteration,
each training sample requires the learning rate to fall into a very specific
narrow range for its softmax output to become non-degenerate, and these ranges
are wildly different for different samples. Therefore, it is not possible to
tune the step size to escape this extremely confident training regime.
Further, note from Figure 9 (left) and Figure 8 (top-left, top-right) that
LeNet-300-100 reaches zero training error as long as the learning rate is
admissible at initialization. These solutions, however, do not generalize
well; the right plot in Figure 9 shows that test accuracy of a trained
LeNet-300-100 saturates at 75% across all learning rates shown, which is 10%
lower than that of the unscaled network. This observation suggests that
LeNet-300-100 with high-norm initializations adhere to the lazy learning
regime in accordance with previous studies even though our optimization
methodologies differ (Chizat et al., 2019). To verify this claim, we train a
linear classifier on features of the trained $\alpha$-scaled LeNet-300-100 and
find them only marginally better than those extracted from a randomly
initialized network (see the left plot in Figure 12). At the same time,
features learned within the Goldilocks zone exhibit much better
generalizability. Thus, even in our setup, gradient descent applied to
networks with large initialization can be in the well-studied lazy regime.
(Allen-Zhu et al., 2019; Du et al., 2019).