Beyond the Quadratic Approximation: the Multiscale Structure of Neural Network Loss Landscapes

Introduction

Despite the well-known nonconvexity of neural network’s loss functions, utilizing local quadratic approximations around (global or local) minimum has been a fruitful approach to study the optimization behaviors for deep neural networks. For example, global convergence of the gradient descent (GD) method can be established for many neural network models in the so-called “lazy training” (or neural tangent kernel (NTK)) regime, where the training trajectory stays within a region with good quadratic approximation . Additionally, quadratic approximations around global minima can be used to explain the preference of stochastic gradient descent (SGD) for flat minima and certain properties of the limiting dynamics .

However, regardless of the theoretical simplicity brought by local quadratic approximation, empirical study has shown that the effect of higher-order information is far from negligible in most stages of the training and for most reasonable choices of hyperparameters. One reflection of the important role played by non-quadratic terms is the Edge of Stability (EoS) phenomenon . The EoS phenomenon shows that GD can always find and stay in the sharpest region that it can be stable. As a sharp contrast, in a quadratic landscape GD either converges or blows up in most cases. Another relevant observation is the effect of learning rate decay. After a learning rate decay, the trajectory changes its moving direction and converges to a different solution. This cannot be explained by a quadratic approximation of the loss. Including non-quadratic and nonconvex information of the loss function into the study of neural network optimization is necessary to acquire more realistic understanding of the behavior of the optimizers, even locally around the minimum.

In this work, motivated by the two problems mentioned above, we study the behavior of optimizers on neural network loss functions beyond local quadratic approximation. We start from empirical observations by visualizing the loss landscape around the training trajectory. Then, from the observations we extract relevant simplified models and theoretically study the optimization dynamics on these simplified problems. Specifically, we obtain two typical observations: (1) around global minimum, the loss grows slower than a quadratic function, which we name the subquadratic growth; (2) in a larger region, the loss function shows distinct scales, which we name the separate scales structure. For the former, we propose to take a subquadratic function, a function that gets flatter when the parameters get farther from the minimum, as local approximation of the loss function. We study the behavior of GD minimizing this function and explain the mechanism behind the EoS phenomenon. We also consider a minima manifold with subquadratic landscape in the directions that are orthogonal to the manifold, and study the motion of the GD along the manifold driven by flatness after reaching EoS. For the latter observation (the separate scales structure), we consider a landscape with several valleys in different scales and explain the working mechanism of learning rate decay. Our theoretical studies, though not directly conducted on real neural network loss functions, help us build insights on what is happening during the training process.

In addition to characterizing the optimization behavior, we are also concerned with the origin of the observed properties of neural network loss functions. We understand both subquadratic growth and separate scales as manifestation of multiscale structures—a continuum of scales for subquadratic growth and finite scales for separate scales structure. By a construction, we show that the multiscale structure can be caused by the non-convexity of the models and the non-uniformity of the training data. Our construction is simple with a two-layer neural network model with only a few neurons. Despite its simplicity, it can already reveal the origin of complicated loss landscapes for neural network models. It also justifies that the simplified problems studied in our theoretical analysis are strongly connected with real neural network loss functions.

We visualize the loss landscape of neural networks in a region that cannot be approximated by the second-order Taylor polynomial at the minimum. We observe the multiscale structure of the loss functions, exhibited in two ways: subquadratic growth near minima and separate scales in larger regions.

Using the subquadratic growth of the loss functions, we theoretically explain the edge of stability phenomenon of GD.

By the separate scales structure of the loss functions, we provide detailed understandings for the working principle and necessity of learning rate decays during the optimization process, even for the deterministic GD algorithm.

We give a simple, yet neural network relevant, construction in which both the subquadratic behavior and separate scales structure happens for the loss landscape. The construction shows that such properties can be caused by the non-uniformity of the training data.

Many works that study the optimization behavior of neural networks resort to a local quadratic approximation of the loss function. This approach is equivalent with linearizing the optimization dynamics around the minima, or fixing the second-order derivatives. One line of works uses this idea to study the minima selection effect of optimizers by analyzing their linear stability in the quadratic approximation . Another notable series of works prove convergence of the GD dynamics for highly over-parameterized neural networks using the fact that the initialization is already in a region with good quadratic approximation . This technique is usually referred to as lazy training or Neural Tangent Kernel (NTK) . Besides direct analysis on (approximately) quadratic landscape, the Hessian of the loss function is widely used to study and characterize the landscape. For example, eigenvalues of the Hessian are used to measure the flatness of minima . It is also used to study the escaping of SGD from local minima .

The mechanism of SGD’s exploration among different minima is made clear in the recent work , which characterizes the movement of SGD iterators along the minima manifold. This picture of exploration along minima manifolds suits the neural network problem better than the exploration among isolated minima. Prior to , similar analysis has been conduct for SGD with label noise . In this paper, we also analyze the motion of optimizers along the minima manifold (see Section 4.4). Our theory is essentially different from these works, because we consider deterministic GD rather than SGD, and the motion along the manifold in our case is driven by an interaction of subquadratic growth and changing flatness, rather than the SGD noise.

The optimization on mulitscale objective function has also been studied. For example, in , a diffusion effect was derived from deterministic gradient descent due to the small scales of objective functions. In the very recent work , the authors generalize the edge of stability phenomenon into the concept “unstable convergence”, which happens when the objective function is complicated. Some examples studies therein is similar to the ones we study in this paper.

Organization of the paper.

The rest of the paper is organized as follows: In Section 2, we discuss two empirical observations that show rich behaviors of GD/SGD on neural network loss functions, that cannot be easily explained by quadratic approximation. Then, in Section 3 we visualize the loss landscape and summarize two aspects of the multiscale structures of the loss—subquadratic growth and separate scales—that may help explain the phenomena in Section 2. We theoretically explain the edge of stability phenomenon using the subquadratic growth in Section 4, and discuss how separate scales structure can help understand the behavior of learning rate decay in Section 5. In Section 6, we study the origin of the multiscale structure and construct simple examples showing the important role played by non-uniform training data. Finally, a summary and conclusions are given in Section 7.

Two empirical observations

In this section, we discuss two empirical observations that cannot be well explained by analyzing the optimizer on a quadratic approximation of the loss function.

The edge of stability (EoS) phenomenon is discussed in detail in . It is also observed in . The EoS states that when GD is used to train neural networks, the sharpness at the iterator (measured by the largest eigenvalue of the Hessian) tends to increase until it arrives at $2/\eta$, where $\eta$ is the learning rate. Note that $2/\eta$ is the largest sharpness that GD can be stable assuming a quadratic loss landscape. Surprisingly, even after the sharpness stabilizes, the training loss keeps decreasing. We show one example in Figure 1 (left). Extensive experiments are done in .

The EoS cannot be explained on quadratic loss functions. On a quadratic loss, GD either converges or blows up exponentially fast, unless the learning rate is exactly $2/\lambda$, where $\lambda$ is the largest eigenvalue of the Hessian. In later sections, we observe that the EoS is caused by loss landscapes that grow slower than quadratic functions around the minimum, which we call subquadratic growth. We then theoretically study how EoS happens, and why the loss value keeps decreasing after EoS, on simplified landscapes.

2 The effect of learning rate decay

Another observation that cannot be explained by quadratic approximation is the effect of learning rate decay (LRD). LRD not only helps find parameters with lower training loss, but also benefits generalization if used at a proper time. In Figure 1 (right), we show that doing LRD later gives better generalization performance, even though in both experiments the learning rate is decayed after the test accuracy is nearly stable and increasing very slowly. The figure also shows that the training loss decreases very slowly after decay.

The explanation of the phenomena shown in LRD is beyond the reach of quadratic approximation and relies on more complicated structures of the loss. For quadratic loss function, LRD makes convergence slower, but the iterators will finally converge to a same solution and show the same generalization performance. In this paper, we will explain these observations using the separate scales structure of the loss functions.

Loss landscape around training trajectory

In this section, we visualize the landscape of neural network’s loss functions around the training trajectory. We observe the subquadratic growth and separate scales phenomena of the loss functions. Both these characteristics are aspects of the multiscale structure of the loss function—one with a continuum of scales and one with finite scales.

It is important to note that neural network’s loss landscape possesses very rich behaviors, and almost any curvature can be found somewhere in the parameter space . In this work, we are only concerned with the loss curvature around the trajectory of SGD or GD. It is widely known that these optimizers only explore a very confined but important region of the whole parameter space.

In the experiments shown in Figure 2, we train neural networks using GD until the loss stops decreasing, or decreases very slowly, in which case we suppose GD starts oscillating around some minima. Then, we pick a point on the GD trajectory and visualize the “gradient direction loss landscape” around this point—the loss landscape along the line going through this point and towards the gradient direction at this point. The gradient direction landscape is important because it is the landscape that GD sees. Experiments are conducted on VGG network, ResNet, and DenseNet, on CIFAR10 and CIFAR100 datasets. The results for gradient direction landscapes are shown in the first row of Figure 2. The figures show that around the minimum the gradient direction landscape is convex and grows slower than quadratic functions. This subquadratic growth is verified by the second-order finite differences shown in the second row of Figure 2.

This subquadratic growth of the landscape around minimum (at least along the gradient direction) explains the edge of stability phenomenon. The stable learning rate for GD increases as the parameters move close to the minimum. Hence, when the learning rate is not small enough, the iterator becomes unstable when it is too close to the minimum, and hence can only oscillate around the minimum at a certain level—but it may not blow up. Note that this subquadratic growth is not contradictory with the local quadratic approximation—the Taylor expansion of the loss function still holds locally, but the radius of this region is very small.

2 The separate scales structure

If the landscape is visualized in a larger domain, we will inevitably see non-convex behaviors of the loss functions. Some examples are shown in the left and middle panels of Figure 3. Here, we can observe another typical structure of neural network’s loss functions—separate scales. Figures in 3 show that the minimum lies in a small, sharp well located within a large, flat well. The separate scales structure can give richer behaviors in high dimensional spaces. For example, in the right panel of Figure 3 we indirectly show the multiscale loss by visualizing the principal moving directions of the model parameters before and after learning rate decay. In the figure the iterator goes along very different directions when the learning rate is decayed at different epochs, showing rich fine-scale structures of the loss function, which are different from the large-scale structure reflected by the blue curves picked by large learning rate. Later in Section 5 we show this high dimensional multiscale structure is an important cause of LRD’s complicated behaviors.

The subquadratic property and the edge of stability

In Section 3, we observed a subquadratic growth property of the loss landscape. This subquadratic growth makes it possible to explain the edge of stability phenomenon discussed in Section 2. In this section, using a simple problem inspired by the landscape curves shown in Figure 2 we reveal the mechanism of edge of stability—the iterator oscillates at a certain level related to the learning rate when the learning rate is too large for the optimizer to converge.

By classical theories of gradient descent, it is easy to show that the iteration (1) converges to the minimum $x=0$ if $\eta<\frac{2}{f^{\prime\prime}(0)}$. If $\eta>\frac{2}{f^{\prime\prime}(0)}$, instead, $x=0$ becomes an unstable stationary point for the dynamics. In this case, other than $x=0$, there is a 2-periodic solution for the GD: the iterator jumps between $x_{\eta}$ and $-x_{\eta}$, where $x_{\eta}$ satisfies $\eta f^{\prime}(x_{\eta})=2x_{\eta}$. We assume $x_{\eta}>0$ and denote this periodic solution by $\{\pm x_{\eta}\}$. The following simple lemma shows that $x_{\eta}$ exists for any $\eta>\frac{2}{f^{\prime\prime}(0)}$ as long as $f^{\prime\prime}(x)$ goes to as $|x|$ tends to infinity, and $x_{\eta}$ is increasing with respect to $\eta$.

If $\lim\limits_{|x|\rightarrow\infty}f^{\prime\prime}(x)=0$, $x_{\eta}$ exists for any $\eta>\frac{2}{f^{\prime\prime}(0)}$. Moreover, viewed as a function of $\eta$, $x_{\eta}$ is monotonically increasing.

Note that if $x_{\eta}$ exists for some $\eta>0$, we have $\eta f^{\prime}(x_{\eta})=2x_{\eta}$, which means

Hence, let $h(x)=f^{\prime}(x)/x$, it suffices to show that $h(x)$ is a decreasing function in $(0,\infty)$, and $\lim\limits_{x\rightarrow 0^{+}}h(x)=f^{\prime\prime}(0)$, and $\lim\limits_{x\rightarrow\infty}h(x)=0$.

We first show the monotonicity of $h$. Taking derivative of $h$, we have

Since $f^{\prime\prime}(x)$ is decreasing, we have

Therefore, $f^{\prime\prime}(x)x-f^{\prime}(x)\leq 0$, and hence $h^{\prime}(x)\leq 0$. This shows $h(x)$ is monotonically decreasing.

For the limits, writing $f^{\prime}(x)$ as integral of $f^{\prime\prime}(x)$, we have

By L’hopital’s rule we obtain $\lim\limits_{x\rightarrow 0^{+}}h(x)=f^{\prime\prime}(0)$. On the other side, by $\lim\limits_{x\rightarrow\infty}f^{\prime\prime}(x)=0$ we easily have $\lim\limits_{x\rightarrow\infty}h(x)=0$. ∎

Next, we consider the GD dynamics, and show that if the objective function is “strictly subquadratic”, i.e. $f^{\prime\prime}(x)$ is strictly decreasing as $|x|$ increases, then the GD iterator converges to the periodic solution $\{\pm x_{\eta}\}$ except a zero-measure set of $x_{0}$.

When $\eta<\frac{2}{f^{\prime\prime}(0)}$, the proof of convergence is standard. When $\eta=\frac{2}{f^{\prime\prime}(0)}$, note that $f^{\prime\prime}(x)<f^{\prime\prime}(0)$ for any $x\neq 0$. Then, for any $x\neq 0$, we have

This gives convergence of the GD trajectory to .

Next, we consider the case when $\eta>\frac{2}{f^{\prime\prime}(0)}$. Let $A$ be the set of those $x_{0}$ such that starting from these $x_{0}$ the GD will arrive at after some steps. Because in the current case is an unstable stationary point, it is easy to show that $A$ contains countable number of points and hence has zero Lebesgue measure . We ignore the detailed proof here.

Rewriting the left hand side of (2), we have

We study the right hand side of (3) in two cases:

Also considering $\int_{x_{\eta}}^{x_{k}}f^{\prime\prime}(t)dt>0$, we have

In this case, we rewrite the right hand side of (3) as

For the integral term, due to the monotonicity of $f^{\prime\prime}$, we have

Since $x_{k}=x_{\eta}$ is the trivial case, we finish showing (2). And the only way for equality to hold is $x_{k}=x_{\eta}$. Therefore, the GD trajectory converges to $\{\pm x_{\eta}\}$. ∎

Theorem 1 shows that subquadratic growth can cause the edge of stability phenomenon. When the learning rate is too big to converge, the GD does not blow up. Instead, it oscillates at a certain level related with the learning rate. During the oscillation, if the learning rate is dropped to a smaller value, the iterator will leave the current periodic solution and converge to a new periodic solution at a lower level. Results of numerical simulations are shown in Figure 4. Similar phenomena are observed in .

2 A multi-dimensional analysis

The problem analyzed above possesses an accurate mathematical characterization because the objective function considered is simple. It is a 1-D function, and though it enjoys subquadratic growth, it is still convex. For more general cases, e.g. non-convex high dimensional functions, the picture is much more complicated for at least two reasons: (1) there may be more than one periodic solutions, and the period of some solutions can be very long; (2) the dynamics can easily get chaotic when the learning rate is large. See Figure 4 for some experiments. Nevertheless, as long as there is subquadratic growth of the landscape around the minimum, the GD still does not blow up for large learning rates. In this and the next subsection, we make some extensions for the theory to consider high-dimensional/non-convex functions.

We first extend our analysis to a class of high dimensional functions that can be decomposed into a sum of subquadratic functions in eigendirections. Concretely, we consider functions like

The proof of Theorem 2 is a simple application of Theorem 1 on each of the components of $f$. This is possible because $\mathbf{p}_{1},...,\mathbf{p}_{n}$ are pairwise orthogonal. Note that Theorem 2 does not imply that GD will converge to a unique 2-periodic solution. Actually, there are exponentially many 2-periodic solutions due to the combination of signs of each $f_{i}$’s periodic solution.

3 A more general setting

Next, we extend our study to more general cases, including nonconvex functions. In this subsection, we consider a wide class of subquadratic functions and show that GD with large learning rate does not diverge on these functions.

For quadratic functions, the gradient grows linearly with the magnitude of the input. For functions that satisfies the definition above, $\|\nabla f(\mathbf{x})\|$ grows slower than $\|\mathbf{x}\|$. Thus, the function grows slower than a quadratic function. Compared with the condition in our 1-D example, Definition 1 is weaker, in the sense that $\|\nabla f(\mathbf{x})\|$ can decrease as $\|\mathbf{x}\|$ gets bigger. Hence, the function $f$ can be nonconvex. On the other hand, the condition on the inner product of the $\nabla f(\mathbf{x})$ and $\mathbf{x}-\mathbf{x}^{*}$ guarantees that the gradient always has a component pointing towards , which is the unique global minimum.

For functions satisfying the definition above, we can show that GD does not diverge with any learning rate.

Without loss of generality we assume $x^{*}=0$. We will use $B_{r}$ to denote the closed ball with radius $r$ and centered at the origin. Because $\lim\limits_{\|\mathbf{x}\|\rightarrow\infty}\|\nabla f(\mathbf{x})\|/\|\mathbf{x}\|=0$, we can find an $r_{1}$ such that for any $\|\mathbf{x}\|>r_{1}$ we have $\|\nabla f(\mathbf{x})\|<\frac{2c}{\eta}\|\mathbf{x}\|$, where $c$ is the constant in Definition 1. For such $\mathbf{x}$, we can easily verify that

Hence, for any $\mathbf{x}$ that satisfies $\|\mathbf{x}\|>r_{1}$, GD sends the iterator closer to the minimum. Equivalently speaking, only when the iterator is within $B_{r_{1}}$ can GD send the iterator to a farther (or with equal distance) location from the minimum. Now, consider the one step GD mapping

Since $\nabla f$ is continuous, $h$ is a continuous function. Hence, due to the compactness of $B_{r_{1}}$, there exists $r_{2}>0$ such that $\|h(\mathbf{x})\|\leq r_{2}$ for any $\mathbf{x}\in B_{r_{1}}$. Then, for any $\mathbf{x}\in B_{r_{2}}$, we always have $h(\mathbf{x})\in B_{r_{2}}$, i.e. $B_{r_{2}}$ is an invariant set for GD iterations. Note that $r_{2}$ only depends on $c$ and $\eta$.

We finish the proof by showing that the GD trajectory from any initialization $x_{0}$ will enter $B_{r_{2}}$. This is a natural result of (5). ∎

Theorem 3 characterizes the qualitative behavior of GD around a subquadratic minimum. With a certain learning rate, GD will oscillates in a learning rate dependent neighborhood of the minimum. Under the current conditions, we cannot fully characterize the trajectory—it may hit the global minimum in some step, or oscillates at a certain level, or oscillates chaotically in the neighborhood of the minimum. (though when $\eta$ is not too small, hitting or converging to the minimum is a zero measure event.) A typical behavior in 2-D space is shown in Figure 4.

The analysis in this section is based on the observation of the local landscape of neural network loss functions around minima. Though for the convenience of analysis we assume the objective function has a global subquadratic behavior, this is not true for neural network loss functions. The subquadratic growth will stop when the parameter is far enough from the minimum. We will address this issue in Section 5 when we study the separate scales structure of loss functions.

4 What happens after the edge of stability

For the objective functions we considered above, the GD iterator will oscillate around the unique minimum after arriving at the edge of stability. However, when training real neural network the iterator keeps moving and reducing the loss value even after reaching the EoS (also see Figure 1). This is mainly due to the over-parameterized nature of neural networks, which produces manifolds of minima instead of isolated minimum. Assume there is a manifold formed by global minima, taking a quasistatic approach in the direction tangent and orthogonal with the manifold, we can study how the GD iterator moves down the manifold and search for flat minima.

Consider a GD trajectory generated from $(\mathbf{x}_{0},y_{0})$ using learning rate $\eta$. In the quasistatic case, the update of $\mathbf{x}_{t}$ is

while $y_{t}$ follows the 2-periodic solution

By (7), we have $\eta f^{\prime}(h(\mathbf{x}_{t})y_{t})=2y_{t}/h(\mathbf{x}_{t})$. Substituting to (6) we obtain

Equation (8) shows that the motion of the GD iterator projected onto the manifold follows a GD of $\log h(\mathbf{x})$, and the speed of the dynamics is determined by $y_{t}$. Therefore, during the oscillation around the minima manifold, GD searches for flatter minimum by reducing the value of $h(\mathbf{x})$. An illustration of this effect is shown in Figure 5.

The idea of flatness driven motion along the manifold is similar to that in , but our result is essentially different. We treat GD instead of SGD, and in our case, the motion along the manifold is made possible by the subquadratic landscape around the minima, instead of the SGD noise. The two types of flatness driven motion have quite different behaviors in some situation. For instance, the SGD noise drives the iterator to the flattest minimum on the manifold, while for GD it will converge after finding a sufficiently flat minimum (relative to the learning rate). Also, if the manifold consists of interpolation solutions, SGD will not show flatness driven motion because the noise vanishes at the minima. However, in our analysis movement still exists as long as the minima have subquadratic property.

The quasistatic approach also serves as a handy tool to derive the noise-driven motion along manifold for SGD. As a simple illustration, consider an objective function $f(\mathbf{x},\bm{y})=\bm{y}^{T}H(\mathbf{x})\bm{y}$, which gives a quadratic approximation of a loss function with a global minima manifold $\{\bm{y}=0\}$. $H(\mathbf{x})$ always gives a positive definite matrix. Suppose an SGD is approximated by an SDE

Here we assume for convenience there is no noise along the manifold direction, and the noise on the $\bm{y}$ direction depends on the flatness of the minimum. Then, by assuming $\bm{y}$ is always at equilibrium, we first solve the dynamics of $\bm{y}$ fixing $\mathbf{x}=\mathbf{x}_{t}$. The equilibrium is

Plugging the equilibrium above into the dynamics of $\mathbf{x}$, and taking expectation over $\bm{y}$, we obtain the expected dynamics of $\mathbf{x}$ in quasistatic case:

which recovers the results in . This quasistatic approach can be easily adapted to other types of noise and other optimizers such as SGD with momentum.

The separate scales and learning rate decay

Learning rate decay is a widely adopted technique in training large scale neural networks, and has received much theoretical attention, too. Explanations of how LRD works include GD stability in different directions , SGD exploration , and pattern complexity . However, there are still some behaviors shown by training with LRD that cannot be well addressed by these explanations. For example, as shown in Figure 1, the generalization performance suffers if the learning rate is decayed too early. In this section, we build a simple loss function, inspired by the observations of separate scales structure, that can explain this behavior of learning rate decay.

with $k>1$. Finally, we build a multiscale landscape with two valleys by considering

where $\phi_{c}$ is a non-decreasing cutoff function that confines the value of $g_{2}$ within $[0,c]$. This makes the effect of $g_{2}$ local. The simplest choice of $\phi_{c}$ can be $\phi_{c}(z)=\min\{z,c\}$. An example landscape of $F$ is shown in Figure 6. The landscape has one large and flat valley and one small and sharp valley. The global minimum is at the origin, locating at the bottom of both valleys.

Now we can study the behavior of GD with learning rate decay in the landscape of $F$. Since the valley of $g_{2}$ is sharp, a small learning rate is necessary for GD to converge. On the other hand, when initialized far away from the global minimum, the iterator will first be attracted by the large valley. GD will first converge to the neighborhood of the large valley, then move along the valley to its flattest region (near the global minimum in this case) while bouncing between the valley walls via the mechanism discussed in Section 4.4. This moving process is faster with a larger learning rate. Therefore, an ideal strategy for learning rate decay is to use a large learning rate until the iterator moves to the vicinity of the small valley and then drop to a small learning rate to converge into the small valley. In this process, when the iterator is bouncing around and moving down the large valley, the loss value is not decreasing much. But that does not mean the learning rate can be decayed earlier. If it is decayed before the iterator is close enough to the small valley, it then has to move down the large valley using the small learning rate, which can cost much more time. In the extreme case, a small learning rate can cause convergence to a suboptimal minimum on the large valley.

An numerical example with the form (9) is given in Figure 6. In the experiment, we initialize GD at a point far from the small valley with a large learning rate. Afterwards, the learning rate is decayed by a same factor at different steps. Although the loss values when the learning rate is decayed are similar, the three trajectories take drastically different amount of time to converge.

The origin of the multiscale structure from training data

Both the subquadratic growth and separate scale loss structures can be understood as manifestations of multiscale structure—finite significant scales for separate scale loss and a continuum of scales for subquadratic growth. In this section, we study the origin of the loss’s multiscale structure. By a simple neural network based construction, we show that the non-convexity of the model and the multiscale structure of the training data together act as one cause of the multiscale loss.

Now we describe our construction. Consider the following two-layer neural network with $3$ neurons,

where $x$ is the input and $w$ is the parameter. Let $\sigma$ be the ReLU function. In this problem we assume the three neurons share the parameter $w$, so there is only one parameter. Suppose we have data $\{(x_{i},y_{i})\}_{i=1}^{n}$, with $x_{i}>0$ and $y_{i}=0$ for all $i$. The loss function is

It is easy to show that, for fixed $x$, $f(x,w)$ is the following piecewise linear function for $w$:

Consequently, each term in the loss function is

This is a function which is quadratic around , and takes constant when $w$ is away from . Moreover, the width of the quadratic segment depends on the magnitude of $x$. For bigger $x$, the quadratic part is narrow while sharp. For smaller $x$, it is wide and flat. From these properties, it is easy to show that the total loss $L$ gets sharper for $w$ closer to . When $x_{i}$’s vary a lot in their magnitudes, $L$ will have multi-scale structure—the sharpness increases by orders of magnitudes as $w$ moves towards 0.

To be clearer, assume without loss of generality that $x_{1}\leq x_{2}\leq\cdots\leq x_{n}$, then $L(w)$ is the following piecewise quadratic/constant function:

If we have a continuum of $x$, sampled from a probability distribution $\mu$ supported on $(0,\infty)$, we can also write down the “population” loss function:

If the training data contains $3$ data points with different magnitude, the empirical loss function looks like the blue curve in Figure 7 (left). We obtain a loss with separate scales. On the other hand, if $\mu$ is a uniform distribution on $[0,C]$, then the population loss is

The curve is shown in orange in Figure 7 (left). Now, we obtain a loss with continuous scales and thus shows subquadratic growth. Similar to the population loss, if the training data set is very large, the empirical loss function will also have (nearly) continuous scales.

With similar approach, we can also construct multiscale losses in multi-dimensional spaces, especially multiscale valleys we studied in Section 5. Recall the definition of $f(x,w)$ above. To produce a 2-D loss landscape with two valleys, we just need to consider a model with a two-dimensional input $(x^{(1)},x^{(2)})$, two parameters $w_{1},w_{2}$, and a two-dimensional output:

Then, consider data $\{(x^{(1)}_{i},x^{(2)}_{i},y^{(1)}_{i},y^{(2)}_{i})\}_{i=1}^{n}$ with $y^{(1)}_{i}=y^{(2)}_{i}=0$ for all $i$. The loss function becomes

The loss function is a superposition of two valleys along $w_{1}$ and $w_{2}$ directions. If we assume $x^{(2)}$ is bigger than $x^{(1)}$ (e.g. $x^{(2)_{i}}=10x^{(1)}_{i}$ for any $i$), then the valley along $w_{1}$ direction (generated by landscape of $w_{2}$) is in a smaller scale than the other valley. An example using a continuous distribution of training data is given in Figure 7.

Discussion on homogeneity.

Many factors can contribute to the special structure of neural network’s loss functions. In the construction in this section, we focus on the distribution of training data. We show that if the training data are not well normalized, multiscale structure will appear in the loss landscape. Besides the training data, nonconvexity of the model and the loss also plays an important role in this example. If the loss for each data is convex or even quadratic, then the total loss as a sum of many single losses will not show very rich structures.

Although in practice the input data is always standardized before being fed into the network, we note that usual standardization applies a fixed transform to all data to achieve zero mean and identity variance, but does not eliminate the multiscale structure in the data. After standardization, the length of different data can still differ by several orders of magnitude. Hence, it is still possible for training data to cause multiscale structure in the loss function.

Finally, we show that for large deep neural networks, multiscale data can still cause multiscale loss due to the homogeneity of ReLU function. For example, consider an L-layer fully connected neural networks with ReLU activation function and without bias:

Let $l_{1}$ and $l_{2}$ be the losses of two input data $\mathbf{x}_{1}$ and $\mathbf{x}_{2}$, both with the same target $y$. Then, if $\mathbf{x}_{2}=k\mathbf{x}_{1}$, we can easily verify

i.e. the two losses have the same shape but different scales. Such relation is not unique. For instance, fixing $W_{2},...,W_{L}$, we have

In this case, $l_{2}$ is a scaling of $l_{1}$ only in the $W_{1}$ space. The second relation above is true even for networks with bias.

Summary

In this paper, we study the limitations of using the quadratic approximation for neural network’s loss functions and highlight the importance of a multiscale structure. Firstly, we empirically observe two manifestations of the multiscale structure—the subquadratic growth and the separate scales structure. These properties can explain some intriguing phenomena observed during the training process of neural networks. Specifically, we explain (1) the edge of stability phenomenon of GD using the subquadratic growth and (2) the behavior and effect of learning rate decay using the separate scales structure. Then, we study the origin of the multiscale structures, and show by constructive examples that non-convex models and non-uniform training data can lead to multiscale loss. It is worth noting that our study puts more focus on GD due to its simplicity. Extending the study to SGD is an important and meaningful direction for future work.

References