# Charting the Topography of the Neural Network Landscape with Thermal-Like
Noise

Théo Jules<EMAIL_ADDRESS>Raymond and Beverly Sackler School of
Physics and Astronomy, Tel Aviv University, Ramat Aviv, Tel Aviv, 69978,
Israel  Gal Brener  Raymond and Beverly Sackler School of Physics and
Astronomy, Tel Aviv University, Ramat Aviv, Tel Aviv, 69978, Israel  Tal
Kachman  Department of Artificial Intelligence, Radboud University, Thomas van
Aquinostraat 4, 6525 GD Nijmegen  Noam Levi  Raymond and Beverly Sackler
School of Physics and Astronomy, Tel Aviv University, Ramat Aviv, Tel Aviv,
69978, Israel  Yohai Bar-Sinai<EMAIL_ADDRESS>Raymond and Beverly
Sackler School of Physics and Astronomy, Tel Aviv University, Ramat Aviv, Tel
Aviv, 69978, Israel  The Center for Physics and Chemistry of Living Systems,
Tel Aviv University, Tel Aviv 69978, Israel

###### Abstract

The training of neural networks is a complex, high-dimensional, non-convex and
noisy optimization problem whose theoretical understanding is interesting both
from an applicative perspective and for fundamental reasons. A core challenge
is to understand the geometry and topography of the landscape that guides the
optimization. In this work, we employ standard Statistical Mechanics methods,
namely, phase-space exploration using Langevin dynamics, to study this
landscape for an over-parameterized fully connected network performing a
classification task on random data. Analyzing the fluctuation statistics, in
analogy to thermal dynamics at a constant temperature, we infer a clear
geometric description of the low-loss region. We find that it is a low-
dimensional manifold whose dimension can be readily obtained from the
fluctuations. Furthermore, this dimension is controlled by the number of data
points that reside near the classification decision boundary. Importantly, we
find that a quadratic approximation of the loss near the minimum is
fundamentally inadequate due to the exponential nature of the decision
boundary and the flatness of the low-loss region. This causes the dynamics to
sample regions with higher curvature at higher temperatures, while producing
quadratic-like statistics at any given temperature. We explain this behavior
by a simplified loss model which is analytically tractable and reproduces the
observed fluctuation statistics.

The optimization of neural networks lies at the core of modern learning
methodology, with the goal of minimizing a loss function that quantifies model
performance. Naturally, the landscape of the loss function plays a critical
role in guiding the optimization process and its properties are closely linked
to its performance and generalization capacities [1, 2]. However, the high
dimensionality of the parameter space, the non-convexity of the loss function,
and the presence of various sources of noise make it challenging to
characterize its geometry [3, 4] and subsequently to analyze the optimization
process over this complicated landscape.

Previous works have studied the topography of the loss landscape and found a
number of interesting features. Firstly, it was established that there exists
a wealth of global minima, all connected by low-loss paths, a phenomenon
referred to as Linear Mode Connectivity [5, 6, 7, 8, 9, 10, 11]. In the final
stages of training the network explores this low-loss region and gradient
descent predominantly occurs within a small subspace of weight space [12, 13,
14]. In addition, it was seen that the curvature of the explored region
sharpens progressively and depends on the learning rate through a feedback
mechanism termed “Edge of Stability” [15, 16, 17].

In this work we study the low loss region by injecting noise in a controlled
manner during training. Many previous works have studied the importance of
noise in the optimization process, modeling it as a stochastic process. Noise
sources might include sampling noise in the estimation of the gradient [18,
19, 20], the numerical discretization of gradient flow [21], noisy data [22,
23], stochastic regularization schemes [24] or other sources. Each such noise
source gives rise to different noise properties, which qualitatively affect
the optimization dynamics [25, 26].

We take a different approach than those described above: we do not use noise
to mimic noisy training dynamics, but rather as a probe that allows inferring
quantitative geometrical insights about the loss landscape [22, 23]. This is
done using standard tools of statistical physics to analyze loss fluctuations,
and ensuring that the thermal noise is the only noise source in the system so
the stochasticity is completely known.

To study the local landscape, we let the system evolve, starting at the
minimum, under over-damped Langevin dynamics, defined by the stochastic
differential equation

$\displaystyle\mathop{}\\!\mathrm{d}\theta_{t}$
$\displaystyle=-\nabla_{\theta}\mathcal{L}(\theta_{t})\mathop{}\\!\mathrm{d}t+\sqrt{2T}\mathop{}\\!\mathrm{d}W_{t},$
(1)

where $\theta\in\mathbb{R}^{N}$ is the vector of the neural weights and
biases, $\mathcal{L}$ is the loss function (to be specified below), $T$ the
exploration temperature and $W_{t}$ is a standard $N$-dimensional Wiener
process. In terms of statistical physics, this is analogous to a system whose
phase space coordinates are $\theta$ and which is described by a Hamiltonian
$\mathcal{L}(\theta)$ in contact with a thermal bath at temperature $T$. As is
well known [27], the long time limit of the probability distribution of
$\theta$ is a Boltzmann distribution, $p(\theta)\propto
e^{-\mathcal{L}(\theta)/T}$, which balances between the gradient and the
random noise terms in Eq. (1).

Specifically, we explore the topography of the loss function in the vicinity
of a typical minimum, for a simple fully connected network performing a
classification task of random data in the over-parameterized regime. Our
analysis shows that, for the networks that we studied, the minimum is
constrained only in a small number of directions in weight-space, as was
previously observed in various contexts and is generally expected in the over-
parameterized regime [7, 28, 2, 29, 30, 12, 13, 14]. Furthermore, and inline
with previous studies, we find that at a given exploration temperature the
fluctuations behave as if $\mathcal{L}$ is effectively quadratic, with $N_{c}$
independent degrees of freedom with non-vanishing stiffness. In other words,
$N_{c}$ is the co-dimension of the low-loss manifold in the vicinity of the
minimum, which our method allows to measure directly.

However, contrary to previous works and quite counter-intuitively, we show
that this picture does not arise from a simple quadratic approximation of
$\mathcal{L}$ around its minimum, as one might naïvely interpret these
observations. Instead, we find that the stiffness associated with the $N_{c}$
constrained eigendirections depends linearly on $T$ over many orders of
magnitude, which is a distinctly nonlinear feature. As we explain below, this
dependence stems from the exponential nature of the “confining walls”
surrounding the low-loss region, and the flatness of the landscape far from
these walls. This exponential nature is also what gives rise to the seemingly
quadratic properties of the loss fluctuations, but this happens through a
delicate balance between the exponential walls and the noise, which cannot be
captured with a model of a quadratic loss function.

##  I Exact predictions for a quadratic loss

Before describing our results, it would be useful to remind the reader what
they would expect to observe in the case of a positive-definite quadratic loss
function, $\mathcal{L}=\sum_{i=1}^{N_{c}}\frac{1}{2}k_{i}\Theta_{i}^{2}\ $,
where $\\{\Theta_{i}\\}$ are the coefficients of the Hessian’s eigenvectors
and $\\{k_{i}\\}$ are their associated stiffnesses. $N_{c}$ is the number of
dimensions with non-vanishing stiffness. Plugging this into Eq. (1) yields a
multivariate Ornstein-Uhlenbeck process which is fully tractable analytically
[31]. We briefly summarize here the main results, whose derivations can be
found in the supplementary information.

First, the fluctuations of $\mathcal{L}$ follow a $\Gamma$-distribution

$\displaystyle P(\mathcal{L};\alpha,\beta)$
$\displaystyle=\frac{\beta^{\alpha}\mathcal{L}^{\alpha-1}}{\Gamma(\alpha)}\exp(-\beta\mathcal{L}),$
(2)

where $\alpha=N_{c}/2$, $\beta=1/T$ and $\Gamma$ is the Gamma function.

Second, a direct corollary of Eq. (2) is that the mean and standard deviation
of $\mathcal{L}$ are both proportional to $T$:

$\displaystyle\begin{split}\mu_{\mathcal{L}}&=\left\langle\mathcal{L}\right\rangle=\tfrac{1}{2}N_{c}T,\\\
\sigma_{\mathcal{L}}^{2}&=\left\langle\mathcal{L}^{2}\right\rangle-\left\langle\mathcal{L}\right\rangle^{2}=\tfrac{1}{2}N_{c}T^{2}\
.\end{split}$ (3)

This result, a standard example of the equipartition theorem [32], means that
each eigendirection contributes $\frac{1}{2}T$ to the total loss, regardless
of its associated stiffness. The “heat capacity”
$C_{h}=\partial\mathcal{L}/\partial T$ simply equals $N_{c}/2$ and is
$T$-independent.

Lastly, in terms of dynamics, the evolution of each eigendirection is
uncorrelated from the other ones and shows an exponentially decaying
correlation. This is quantified by the two-point correlation

$\displaystyle\chi_{g}(t)$
$\displaystyle=\sigma_{g}^{-2}\left[\left\langle\vphantom{\Big{[}}g(t_{0})g(t_{0}+t)\right\rangle-\mu_{g}^{2}\right]$
(4)

where $g$ is any time-dependent quantity. For a quadratic loss we have
$\chi_{\Theta_{i}}=\exp\left(-|t|/\tau_{i}\right)$ and the correlation time
$\tau_{i}$ is simply the inverse of the stiffness $\tau_{i}=1/k_{i}$. We note
that in these terms, the stiffness of the “soft directions” does not need to
strictly vanish – $k_{i}$ should only be low enough so that the correlation
time $\tau_{i}$ would be so long that the dynamics in this eigendirection
would not equilibrate during the simulation time. The auto-correlation of
$\mathcal{L}$ is a sum of such exponentials,
$\chi_{\mathcal{L}}=\sum_{i}e^{-k_{i}|t|}$.

##  II Numerical experiment

We consider a classification problem with $C=3$ classes using a multi-layer
perceptron [33], represented by the function
$f(x;\theta)\\!:\\!\mathbb{R}^{d}\to\mathbb{R}^{C}$. The network is trained on
a training dataset $\\{x^{i},y^{i}\\}_{i=1}^{D}$ where
$x^{i}\in\mathbb{R}^{d}$ are the inputs and $y^{i}\in\\{0,1\\}^{C}$ are one-
hot vectors indicating a randomly assigned correct class. The $\\{x_{i}\\}$
are drawn from a standard $d$-dimensional normal distribution. Full details
regarding the architecture of the network and the dataset are given in the
supplementary information. The network’s output is transformed to a
classification prediction via a softmax function. That is, the estimated
probability that an input $x^{i}$ belongs to class $k$ is

$\displaystyle
p_{k}(x^{i};\theta)=\frac{\exp(f(x^{i};\theta)_{k})}{\sum_{m=1}^{C}\exp(f(x^{i};\theta)_{m})}\
,$ (5)

where $f(\cdot)_{k}$ denotes the $k$-th entry in $f$. Finally, the loss is
taken to be the cross entropy between the predicted and true labels:

$\displaystyle\mathcal{L}$
$\displaystyle=\frac{1}{D}\sum_{i=1}^{D}\ell(x^{i},y^{i},\theta)\ ,$
$\displaystyle\ell$
$\displaystyle=-\sum_{k=1}^{C}y^{i}_{k}\log(p_{k}(x^{i};\theta)).$ (6)

Our main objective is to explore the topography of the loss function in the
vicinity of a typical minimum. To find such a minimum, we train the network
using the ADAM optimizer [34] for a predefined amount of epochs. Since the
problem is over-parameterized, after some training, the data is perfectly
fitted and the loss becomes essentially zero, up to numerical noise. This
stage is denoted as “Adam” in Fig. 1a.

To explore the vicinity of this minimum, we then let the system evolve under
Eq. (1) using the Euler-Maruyama discretization scheme [35],

$\displaystyle\theta_{s+1}$
$\displaystyle=\theta_{s}-\eta\nabla\mathcal{L}(\theta_{s})+\sqrt{2\eta
T}\xi_{s}\ ,$ (7)

where $s$ is the step number, $\eta=t_{s+1}-t_{s}$ is the discrete time step
and $\xi_{t}$ is a Gaussian random variable with zero mean and unit variance.
This exploration is denoted as “Langevin” in Fig. 1. It is seen that the loss
increases quickly before reaches a $T$-dependent steady state (“thermodynamic
equilibrium”).

We stress while that the parameter $\eta$ is reminiscent of the “learning
rate” in the machine learning literature, they are not exactly equivalent.
Importantly, in our formalism $\eta$ serves only as the time discretization
and appears explicitly in the noise term, whose $\sqrt{\eta}$ scaling is
necessary in order for the dynamics to converge to the Boltzmann distribution
in the limit $\eta\to 0$ [27]. We also note that the convergence of the
probability distribution $p(\theta)$ in the limit $\eta\to 0$ is a different
concept than the convergence of the gradient descent trajectory to that of
gradient flow [21]. As such, $\eta$ is not a parameter of our exploration
protocol but rather of the numerical implementation of Eq. (1), and meaningful
results should not depend on $\eta$.

##  III Results: Loss fluctuation statistics

We begin by inspecting the moments of the loss fluctuations,
$\mu_{\mathcal{L}}$ and $\sigma_{\mathcal{L}}$, shown in Fig. 1c. It is seen
that both of them scale linearly with $T$. First, we note that our
measurements of $\mu_{\mathcal{L}}$ at a given temperature are independent of
$\eta$, as expected. Furthermore, a basic prediction of statistical mechanics
relates the variance of $\mathcal{L}$ in equilibrium with the heat capacity,
namely $\sigma_{\mathcal{L}}^{2}=T^{2}C_{h}(T)$ [32]. In our case of a
$T$-independent heat capacity this relation reads
$\sigma_{\mathcal{L}}=\sqrt{C}_{h}T$, which is numerically verified in Fig.
1c. These results support our claim that the dynamics are thermally
equilibrated and follow Boltzmann statistics.

Going beyond the moments, Fig. 1b shows the full distribution of the loss
fluctuation, which are well described by a Gamma distribution. Fig. 1d shows
the distribution parameters $\alpha$ and $\beta$, defined in Eq. (2), which
are estimated from the empirical loss distributions using standard maximum
likelihood estimators. It is seen that the distribution parameter $\beta$
agrees with the exploration temperature $T$, i.e. $\beta T\approx 1$, over
several orders of magnitude in $T$ and independently of $\eta$. The number of
stiff dimensions, $N_{c}=2\alpha$, seems to weakly depend on the temperature,
decreasing as $T$ grows. Lastly, we note that that the linear dependence of
$\mu_{\mathcal{L}}$ and $\sigma_{\mathcal{L}}$ on $T$ is a property of the
low-loss region explored by the dynamics at low $T$, and it is not observed if
the thermal dynamics are started immediately after initializing the network.
This is shown explicitly in the supplementary information.

Figure 1:  (a) Observed loss dynamics during the exploration. First, the
network is trained using the ADAM algorithm (black line). Then, the learning
algorithm is changed to Eq. (1), where the noise amplitude is controlled by a
temperature-like parameter $T$ (colored lines). Each curve corresponds to a
different temperature, all using $\eta=10^{-2}$. (b) Distribution of the loss
fluctuation in steady state normalized by the temperature. For each
distribution, the dashed black line corresponds to a gamma distribution, cf
Eq. 2, whose parameters are found using maximum likelihood estimation. The
inset shows the same data in log-linear axes. (c) Temperature dependence of
$\mu_{\mathcal{L}}$ (circles) and $\sigma_{\mathcal{L}}$ (squares). Each point
corresponds to an average over multiple runs. The solid line shows a fit to
$\mu_{\mathcal{L}}=C_{h}T$. The dashed line shows the equilibrium prediction
$\sigma_{\mathcal{L}}=\sqrt{C_{h}}T$ with the obtained value of $C_{h}$. (d)
Corresponding parameters $\alpha$ and $\beta$ for the gamma distribution. The
symbols and error bars show the average and standard deviation, respectively,
over multiple runs.

All these observations are _quantitatively_ consistent with a picture of a
(locally) quadratic loss function. In other words, at each temperature we can
interpret the loss statistics as if they were generated by an effective
quadratic loss, which has a $T$-dependent number of stiff directions,
$N_{c}(T)=2\alpha(T)$. This number, $N_{c}\approx 20-60$, is significantly
lower than both the dimensionality of $\theta$ ($N=900$) and the number of
elements in the dataset, $D=300$. It is also much _larger_ than the number of
classes $C=3$, which was suggested by Fort et. al. [7] as the number of
outlying large Hessian eigenvalues.

We find that the effective dimension of the low loss manifold is directly
related to the number of points that lie close to the decision boundary. To
demonstrate this, we examine the loss $\mathcal{L}$ of Eq. 6 as a sum over the
losses of individual sample points $\mathcal{L}=D^{-1}\sum_{i}\ell_{i}$. We
find numerically that most of the sample points are well classified,
contributing negligibly to the total loss. A common way to quantify how many
points contribute non-negligibly is the ratio of the $L_{1}$ and $L_{2}$ norms
of the loss vector [36],

$\displaystyle\phi\left(\\{\ell_{i}\\}\right)=\frac{\left(\sum_{i=1}^{D}\ell_{i}\right)^{2}}{\sum_{i=1}^{D}\ell_{i}^{2}}\
,$ (8)

where $\ell_{i}$ is the contribution of the $i$-th example to the loss. $\phi$
is a measure of sparsity, which counts how many entries in $\ell_{i}$
contribute to its sum $d$. For instance, if
$\ell_{1}=\ell_{2}=\cdots=\ell_{k}$ and all other $\ell_{i}$ vanish then
$\phi=k$. We calculate $\phi$ for random snapshots of the network during the
dynamics, and plot the averaged results in Fig. 2a. It is seen that $\phi$,
the effective number of sample points pinning the decision boundary,
quantitatively agrees with $\alpha$, twice the effective number of constrained
dimension in weight space.

###  III.1 Temperature dependence

However, while the time-independent statistics suggest an effective quadratic
loss, the dynamic properties show that the picture is not as simple. Examining
again Fig. 1a, one may notice that that the temporal dynamics of $\mathcal{L}$
seem to slow down at lower temperatures. This is readily verified by looking
at the loss auto-correlation, cf. Eq. 4, which shows a distinct slowing down
at low $T$, as seen in Fig. 2a. To quantify this, we define the correlation
half-time $\tau_{\frac{1}{2}}$ as the lag time at which $\chi_{\mathcal{L}}$
decays to $\frac{1}{2}$. Plotting $\tau_{\frac{1}{2}}$ as a function of
temperature, cf. Fig. 2b, shows a clear dependence $\tau_{\frac{1}{2}}\propto
T^{-1}$.

The fact that $\tau_{\frac{1}{2}}$ scales as $T^{-1}$ raises three interesting
insights. First, and most importantly, it is inconsistent with a picture of
quadratic loss, which implies that the dynamic timescales are $T$-independent,
$\tau_{i}=k_{i}^{-1}$. In contrast, we observe that $\tau_{\frac{1}{2}}$
changes over 4 orders of magnitude with $T$.

Secondly, while the quadratic analogy might not hold, one may still relate the
temporal timescale with the local stiffness, i.e. $k\sim\tau^{-1}$. If this
scaling relation holds, we should expect the eigenvalues of the loss Hessian
to scale linearly with $T$. To test this, we measured the Hessian of the loss
at 1000 randomly selected points during the exploration at steady state and
calculated their eigenvalues using standard numerical procedures [37, 38]. The
distribution of these eigenvalues is plotted in Fig. 3a, clearly showing a
linear scaling with $T$.

These observations are manifestly inconsistent with a picture of an
effectively quadratic loss: in the quadratic picture $\mu_{\mathcal{L}}$
increases linearly with $T$ because the system climbs slightly higher up the
confining parabolic walls, whose stiffness is constant. Our observation
suggests that the picture is quite different: $\mu_{\mathcal{L}}$ increases in
tandem with the stiffness of the confining walls, and due to a delicate
balance the net result is indistinguishable from a quadratic picture, as far
as static properties are considered. Below we explain this balance and show
that it is related to the exponential nature of the confining walls.

Figure 2:  (a) The sparsity $\phi$, cf. Eq. 8, as a function of $T$. In gray
we overlay our estimations of $\alpha$, plotted in Fig. 1. It is seen that
$\phi$ quantitatively agrees with $\alpha$, twice the effective number of
constrained dimensions of the low loss manifold. (b) Temperature dependence of
$\tau_{\frac{1}{2}}$. The measurement was repeated over multiple runs, and the
plot shows the average (points) and maximum and minimum values (color
shading). The black line shows a power law dependence
$\tau_{\frac{1}{2}}=\frac{L^{2}}{T}$. (c) Autocorrelation of the loss (cf. Eq.
(4)) in steady-state. The correlation half-time $\tau_{\frac{1}{2}}$ is the
time for which $\chi_{\mathcal{L}}=0.5$. It is seen that the auto-correlation
decays logarithmically at large $\Delta t$. (d) The same data as in panel c,
plotted as a function of the rescaled time lag $2T\Delta t$. The curves for
different $T$ collapse to a single curve, except at high temperature and long
times.

Lastly, we remark that the relation $\tau_{\frac{1}{2}}\sim T^{-1}$ gives rise
to a distance scale $L$, defined by $L^{2}=T\tau_{\frac{1}{2}}$. $L$ is the
distance, in parameter space, that $\theta$ would diffuse over the time
$\tau_{\frac{1}{2}}$, if subject only to isotropic Gaussian noise. Since the
diffusion coefficient scales with $T$, $L$ is $T$-independent. Furthermore,
since $\tau_{\frac{1}{2}}$ is the correlation time, one can also interpret $L$
as a correlation length, or the distance that two nearby networks need to
diffuse away from each other in order for their loss to decorrelate, i.e.
produce significantly different predictions. Since this distance scale does
not depend on $T$, we conclude that it is an intrinsic property of the loss
landscape, i.e. a characteristic length scale in weight space.

To demonstrate the effect of this length scale, we performed another numerical
experiment: Starting from the minimum (the end of training phase I in Fig. 2)
we let the system diffuse freely, i.e. evolve in time according to Eq. (1) but
without the gradient term. This procedure samples points uniformly and
isotropically around the starting point. Indeed, Fig. 3b shows that for
distances smaller than $L$ the loss does note deviate significantly from its
minimum value. At larger distances, $\mathcal{L}$ changes by orders of
magnitude over a relatively small distance.

###  III.2 Summary of the numerical observations

We summarize here the main properties of the loss fluctuations in the vicinity
of the minimum, described above:

  1. (a)

Both $\mu_{\mathcal{L}}$ and $\sigma_{\mathcal{L}}$ scale linearly with the
temperature $T$, as one would expect from a quadratic loss, cf. Fig. 1c.

  2. (b)

Interpreting the fluctuations as if they were generated from a quadratic loss,
the effective number of degrees of freedom is found to be small and weakly
$T$-dependent, cf. Fig. 1d. In addition, it is closely related to the number
of sample points that lie close to the decision boundary, cf. Fig. 2a.

  3. (c)

The correlation time $\tau_{1/2}$ scales as $1/T$ and the Hessian eigenvalues
scale as $T$, which is inconsistent with a quadratic loss and gives rise to an
emergent $T$-independent length scale $L$, cf. Fig. 2 and Fig. 3.

Figure 3: (a) The cumulative distribution function of the Hessian eigenvalues
sampled during dynamics with $\eta=10^{-2}$, for various values of $T$. Very
small negative eigenvalues are excluded from this plot. It is seen that at
higher temperatures the network explores regions with larger eigenvalues.
Inset: the same data plotted as function of $\lambda/T$ shows a collapse of
the distributions, suggesting that the eigenvalues scale linearly with $T$.
(b) The loss as a function of distance in weight space during the exploration.
In the warm-colored curves show Langevin exploration (same color code as panel
a). The black line shows the behavior in the case of pure diffusion (without
gradient descent). The dashed line marks $L$, the characteristic distance in
weight distance obtained from Fig. 2c.

##  IV An analytical toy model

In order to explain our numerical observations, one need to inspect the cross
entropy loss Eq. (6). For simplicity, consider a network performing binary
classification on a single training example
$\\{x,y\\}\in\mathbb{R}\times\mathbb{R}$. Since the network is
overparameterized, the networks in the low-loss region that we explore
classify most of the training samples perfectly. These examples contribute
negligibly to the total gradient. However, some samples lie close to the
decision boundary. We focus on one such sample $\\{x,y\\}$ and assume without
loss of generality that the correct class is $y=1$. Taking a linear
approximation of $f$, the contribution of this sample to the loss is (see
supplementary material for derivation)

$\displaystyle\ell(x;\theta)$
$\displaystyle=\log\left(1+e^{f(x;\theta)}\right)\ ,$ $\displaystyle f$
$\displaystyle=\sum_{i}a_{i}\theta_{i}+b$ (9)

In this description, the only property of $\theta_{i}$ that affects the loss
is its projection on $a_{i}$, the direction in weight-space that moves the
decision boundary towards the sample point. Since all other directions in
weight space are irrelevant, we ignore them and examine a one-dimensional loss
function

$\displaystyle\ell_{1D}(\theta)=\log\left(1+e^{a\theta+b}\right)\approx
Be^{a\theta}.$ (10)

The approximation in Eq. 10 holds in the vicinity of the minimum because the
point is well classified and the exponent is expected to be small. We define
$B\equiv e^{b}$, and assume for concreteness that $a>0$.

Figure 4: The loss function $\ell_{1D}$, cf. Eq. 10, is plotted in blue. The
probability distribution of $\theta$, $p(\theta)\propto e^{-\ell_{1D}/T}$ is
shown for three temperatures. It is seen that the probability distributions
are qualitatively different from the probability distribution of generated by
a quadratic loss, $p_{Q}$, which is Gaussian. For comparison, we plot $p_{Q}$
obtained from a quadratic approximation at $T=10^{-1}$. For this figure we
chose $B=1$ and $\theta_{*}=20$.

The statistical mechanics of $\ell_{1D}$ can be obtained in closed form by
calculating the partition function
$Z(\beta)=\tfrac{1}{\theta_{0}}\int_{-\infty}^{\infty}e^{-\beta\ell_{1D}(x;\theta)}d\theta$,
where $\theta_{0}$ is a resolution scale required to ensure the partition
function is dimensionless. Formally, Eq. 10 is minimized at
$\theta\to-\infty$, which effectively sets the decision boundary at infinity
and prevents the integral which defines $Z$ from converging. To avoid this
unphysical behavior we impose a hard cut-off at $\theta=-\theta_{*}$, where
$\theta_{*}>0$, which would realistically arise when the decision boundary
wanders far away and meets another sample point.

With this cutoff, the partition function $Z(\beta)$ can be obtained
analytically in closed form and consequently all other “thermodynamic”
quantities can be calculated (see supplementary information for the
derivations). The main finding is that this model reproduces the properties of
the loss fluctuations described above. Namely, in the limit $a\theta^{*}\gg 1$
and $T\ll 1$, both $\mu_{\mathcal{L}}$, $\sigma_{\mathcal{L}}$ and the average
curvature scale linearly with $T$, up to logarithmic corrections:

$\displaystyle\begin{split}\mu_{\ell_{1D}}&\simeq\frac{T}{a\theta_{*}-\gamma+\log(T/B)}\
,\\\
\sigma^{2}_{\ell_{1D}}&\simeq\frac{T^{2}\left(a\theta_{*}-\gamma+\log\left({T}/{B}\right)-1\right)}{\left(a\theta_{*}-\gamma+\log\left({T}/{B}\right)\right)^{2}}\
,\\\
H_{\ell_{1D}}&=\left\langle\nabla^{2}_{\theta}\ell_{1D}\right\rangle\simeq\frac{a^{2}T}{a\theta_{*}-\gamma+\log(T/B)}\
.\end{split}$ (11)

Here $\gamma\simeq 0.577$ is the Euler–Mascheroni constant. Finally, because
the loss is approximately exponential in $\theta$, it features an intrinsic
length scale $L\simeq a^{-1}$. We note that this length scale depends on the
gradient of the network and therefore in general might differ between two
different sample points that reside near the decision boundary.

In Fig. 4, we show the full loss given in Eq. 10, and the resulting
probability distribution $p(\theta)\propto e^{-\ell_{1D}/T}$ for various
temperatures. It is seen that, due to the flatness of the loss, $p(\theta)$ is
essentially constant at negative $\theta$ and drops sharply at the decision
boundary. As $T$ grows, the probability explores regions with higher loss and,
due to the exponential dependence on $\theta$, higher curvature. We compare
these results against a quadratic approximation for $\ell_{1D}$, expanded
around $\theta_{0}$ defined by $\ell_{1D}(\theta_{0})=\mu_{\ell_{1D}}(T)$. It
is seen that a quadratic loss is an extremely poor approximation in the low
temperature limit.

##  V Summary and conclusions

To summarize our findings, we have used Langevin dynamics to investigate the
geometry of the low-loss manifold of an overparameterized neural net. We find
that the fluctuation statistics of the loss are a powerful probe that allows
inferring geometrical insights about the loss topography. For the network
studied here – an overparameterized fully connected neural net performing a
classification task on randomly distributed data – the picture that emerges is
that in the low loss region, which is explored at low temperatures, most of
the sample points are well classified and do not contribute significantly to
the loss. However, a small number of sample points “pin” the decision
boundary, which fluctuates around them. At a given temperature, these
fluctuations have the same statistics as fluctuations produced by a quadratic
loss function, whose effective number of degrees of freedom is directly
related to the number of data points constraining the decision boundary and
can be immediately read off the fluctuation statistics.

However, we find that a quadratic description of the loss is fundamentally
inadequate: the effective stiffness scales linearly with $T$, and
correspondingly the characteristic time scale of loss fluctuations grows at
low temperatures as $1/T$. These observations cannot be reconciled with a
quadratic approximation of the loss. Rather, we suggest that this behavior is
due to the exponential nature of the cross-entropy loss in the low $T$ regime.
As we demonstrate analytically, an exponential loss function in 1D reproduces
the observed fluctuation statistics in the limit of low temperature. These
conclusions, of course, pertain to the simplified case studied here – a fully
connected network classifying random data. Understanding how they apply to
structured data or more complicated network architectures is left for future
studies.

##  VI Acknowledgements

We thank Nadav Cohen, Boaz Barak, Zohar Ringel and Stefano Recanatesi for
fruitful discussions. YBS was supported by research grant ISF 1907/22 and
Google Gift grant. NL would like to thank the Milner Foundation for the award
of a Milner Fellowship. TK would like to acknowledge Lineage logistics for
their funding. TJ was partly supported by the Raymond and Beverly Sackler
Post-Doctoral Scholarship.

## References

  * Dinh  _et al._ [2017] L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio, in  _Proceedings of the 34th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 70, edited by D. Precup and Y. W. Teh (PMLR, 2017) pp. 1019–1028.
  * Yang  _et al._ [2021] Y. Yang, L. Hodgkinson, R. Theisen, J. Zou, J. E. Gonzalez, K. Ramchandran, and M. W. Mahoney, arXiv:2107.11228 [cs.LG] (2021).
  * Li  _et al._ [2018a] H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein, in  _Advances in Neural Information Processing Systems_, Vol. 31, edited by S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (Curran Associates, Inc., 2018).
  * Czarnecki  _et al._ [2019] W. M. Czarnecki, S. Osindero, R. Pascanu, and M. Jaderberg, arXiv:1912.07559 (2019).
  * Draxler  _et al._ [2018] F. Draxler, K. Veschgini, M. Salmhofer, and F. Hamprecht, in  _Proceedings of the 35th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 80, edited by J. Dy and A. Krause (PMLR, 2018) pp. 1309–1318.
  * Garipov  _et al._ [2018] T. Garipov, P. Izmailov, D. Podoprikhin, D. P. Vetrov, and A. G. Wilson, in  _Advances in Neural Information Processing Systems_, Vol. 31, edited by S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (Curran Associates, Inc., 2018).
  * Fort and Jastrzebski [2019] S. Fort and S. Jastrzebski, Advances in Neural Information Processing Systems 32 (2019).
  * Fort and Ganguli [2019] S. Fort and S. Ganguli, arXiv:1910.05929 (2019).
  * Nguyen [2019] Q. Nguyen, arXiv:1901.07417 [cs.LG] (2019).
  * Cooper [2021] Y. Cooper, SIAM Journal on Mathematics of Data Science 3, 676 (2021).
  * Benton  _et al._ [2021] G. Benton, W. Maddox, S. Lotfi, and A. G. G. Wilson, in  _International Conference on Machine Learning_ (PMLR, 2021) pp. 769–779.
  * Gur-Ari  _et al._ [2018] G. Gur-Ari, D. A. Roberts, and E. Dyer, arXiv:1812.04754 [cs.LG] (2018).
  * Li  _et al._ [2018b] C. Li, H. Farkhoor, R. Liu, and J. Yosinski, arXiv:1804.08838 [cs.LG] (2018b).
  * Frankle  _et al._ [2020] J. Frankle, G. K. Dziugaite, D. Roy, and M. Carbin, in  _Proceedings of the 37th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 119, edited by H. D. III and A. Singh (PMLR, 2020) pp. 3259–3269.
  * Cohen  _et al._ [2021] J. M. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar, arXiv:2103.00065 [cs.LG] (2021).
  * Wang  _et al._ [2022] Z. Wang, Z. Li, and J. Li, in  _Advances in Neural Information Processing Systems_, edited by A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (2022).
  * Damian  _et al._ [2022] A. Damian, E. Nichani, and J. D. Lee, arXiv:2209.15594 [cs.LG] (2022).
  * Zhu  _et al._ [2018] Z. Zhu, J. Wu, B. Yu, L. Wu, and J. Ma, arXiv:1803.00195 [stat.ML] (2018).
  * Simsekli  _et al._ [2019] U. Simsekli, L. Sagun, and M. Gurbuzbalaban, in  _Proceedings of the 36th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 97, edited by K. Chaudhuri and R. Salakhutdinov (PMLR, 2019) pp. 5827–5837.
  * Xie  _et al._ [2020] Z. Xie, I. Sato, and M. Sugiyama, arXiv:2002.03495 [cs.LG] (2020).
  * Elkabetz and Cohen [2021] O. Elkabetz and N. Cohen, Advances in Neural Information Processing Systems 34, 4947 (2021).
  * Levi  _et al._ [2022a] N. Levi, I. Bloch, M. Freytsis, and T. Volansky, CoRR abs/2210.13599, 10.48550/arXiv.2210.13599 (2022a), arXiv:2210.13599 .
  * Levi  _et al._ [2022b] N. Levi, I. M. Bloch, M. Freytsis, and T. Volansky, CoRR abs/2210.15764, 10.48550/arXiv.2210.15764 (2022b), arXiv:2210.15764 .
  * Srivastava  _et al._ [2014] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov, The journal of machine learning research 15, 1929 (2014).
  * Li  _et al._ [2021a] Z. Li, S. Malladi, and S. Arora, in  _Advances in Neural Information Processing Systems_, Vol. 34, edited by M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan (Curran Associates, Inc., 2021) pp. 12712–12725.
  * Li  _et al._ [2021b] Z. Li, T. Wang, and S. Arora, arXiv:2110.06914 (2021b).
  * Tuckerman [2010] M. Tuckerman,  _Statistical mechanics: theory and molecular simulation_ (Oxford university press, 2010).
  * Brea  _et al._ [2019] J. Brea, B. Simsek, B. Illing, and W. Gerstner, arXiv:1907.02911 [cs.LG] (2019).
  * Simsek  _et al._ [2021] B. Simsek, F. Ged, A. Jacot, F. Spadaro, C. Hongler, W. Gerstner, and J. Brea, in  _International Conference on Machine Learning_ (PMLR, 2021) pp. 9722–9732.
  * Ainsworth  _et al._ [2022] S. K. Ainsworth, J. Hayase, and S. Srinivasa, arXiv:2209.04836 [cs.LG] (2022).
  * Gardiner  _et al._ [1985] C. W. Gardiner _et al._ ,  _Handbook of stochastic methods_ , Vol. 3 (springer Berlin, 1985).
  * Kardar [2007] M. Kardar,  _Statistical physics of particles_ (Cambridge University Press, 2007).
  * Murphy [2012] K. P. Murphy,  _Machine learning: a probabilistic perspective_ (MIT press, 2012).
  * Kingma and Ba [2014] D. P. Kingma and J. Ba, arXiv:1412.6980 [cs.LG] (2014).
  * Higham [2001] D. J. Higham, SIAM review 43, 525 (2001).
  * Hurley and Rickard [2009] N. Hurley and S. Rickard, IEEE Transactions on Information Theory 55, 4723 (2009).
  * Paszke  _et al._ [2019] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, in  _Advances in Neural Information Processing Systems 32_ (Curran Associates, Inc., 2019) pp. 8024–8035.
  * Harris  _et al._ [2020] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant, Nature 585, 357 (2020).