Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Chao Ma, Daniel Kunin, Lei Wu, Lexing Ying

TL;DR

This paper addresses the inadequacy of local quadratic approximations for neural network loss landscapes, revealing a multiscale structure that governs optimization beyond small neighborhoods. By visualizing training trajectories and formulating simplified subquadratic and multiscale models, it links subquadratic near minima to the Edge of Stability (EoS) and separate-scale landscapes to learning rate decay (LRD). The authors provide theoretical analyses for gradient-descent dynamics in one and higher dimensions, show how subquadratic growth explains non-divergent behavior at large learning rates, and demonstrate how LR decay can navigate multiple scales to improve convergence and generalization. They also construct simple neural-network-inspired examples showing that non-convexity and non-uniform training data can generate multiscale losses, offering practical guidance for optimization and LR schedules in deep learning.

Abstract

A quadratic approximation of neural network loss landscapes has been extensively used to study the optimization process of these networks. Though, it usually holds in a very small neighborhood of the minimum, it cannot explain many phenomena observed during the optimization process. In this work, we study the structure of neural network loss functions and its implication on optimization in a region beyond the reach of a good quadratic approximation. Numerically, we observe that neural network loss functions possesses a multiscale structure, manifested in two ways: (1) in a neighborhood of minima, the loss mixes a continuum of scales and grows subquadratically, and (2) in a larger region, the loss shows several separate scales clearly. Using the subquadratic growth, we are able to explain the Edge of Stability phenomenon [5] observed for the gradient descent (GD) method. Using the separate scales, we explain the working mechanism of learning rate decay by simple examples. Finally, we study the origin of the multiscale structure and propose that the non-convexity of the models and the non-uniformity of training data is one of the causes. By constructing a two-layer neural network problem we show that training data with different magnitudes give rise to different scales of the loss function, producing subquadratic growth and multiple separate scales.

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

TL;DR

This paper addresses the inadequacy of local quadratic approximations for neural network loss landscapes, revealing a multiscale structure that governs optimization beyond small neighborhoods. By visualizing training trajectories and formulating simplified subquadratic and multiscale models, it links subquadratic near minima to the Edge of Stability (EoS) and separate-scale landscapes to learning rate decay (LRD). The authors provide theoretical analyses for gradient-descent dynamics in one and higher dimensions, show how subquadratic growth explains non-divergent behavior at large learning rates, and demonstrate how LR decay can navigate multiple scales to improve convergence and generalization. They also construct simple neural-network-inspired examples showing that non-convexity and non-uniform training data can generate multiscale losses, offering practical guidance for optimization and LR schedules in deep learning.

Abstract

A quadratic approximation of neural network loss landscapes has been extensively used to study the optimization process of these networks. Though, it usually holds in a very small neighborhood of the minimum, it cannot explain many phenomena observed during the optimization process. In this work, we study the structure of neural network loss functions and its implication on optimization in a region beyond the reach of a good quadratic approximation. Numerically, we observe that neural network loss functions possesses a multiscale structure, manifested in two ways: (1) in a neighborhood of minima, the loss mixes a continuum of scales and grows subquadratically, and (2) in a larger region, the loss shows several separate scales clearly. Using the subquadratic growth, we are able to explain the Edge of Stability phenomenon [5] observed for the gradient descent (GD) method. Using the separate scales, we explain the working mechanism of learning rate decay by simple examples. Finally, we study the origin of the multiscale structure and propose that the non-convexity of the models and the non-uniformity of training data is one of the causes. By constructing a two-layer neural network problem we show that training data with different magnitudes give rise to different scales of the loss function, producing subquadratic growth and multiple separate scales.
Paper Structure (19 sections, 4 theorems, 41 equations, 7 figures)

This paper contains 19 sections, 4 theorems, 41 equations, 7 figures.

Table of Contents

  1. Introduction
  2. Related works.
  3. Organization of the paper.
  4. Two empirical observations
  5. The edge of stability phenomenon
  6. The effect of learning rate decay
  7. Loss landscape around training trajectory
  8. The subquadratic growth
  9. The separate scales structure
  10. The subquadratic property and the edge of stability
  11. A one-dimensional analysis
  12. A multi-dimensional analysis
  13. A more general setting
  14. What happens after the edge of stability
  15. The separate scales and learning rate decay
  16. ...and 4 more sections

Key Result

Lemma 1

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

Figures (7)

  • Figure 1: (left) One example of the edge of stability. We train a three-layer fully-connected neural network on a subset of CIFAR10. The orange curve shows that the sharpness of the landscape at the iterator first increases, and then stabilizes around $2/\eta$ (the dashed orange line), where $\eta$ is the learning rate. The blue curve shows the training loss keeps decreasing after the edge of stability is achieved. (right) Training loss and test accuracy curves for two experiments with learning rate decay at different times. The blue curves show an experiment with LRD at epoch 250 (the left red vertical line), and the orange curves show an experiment with LRD at epoch 500 (the right red vertical line). The two experiments have the same initialization. The learning rate is $0.1$ initially and dropped to $0.01$.
  • Figure 2: The loss landscape on gradient directions for GD when the iterator oscillates around the (local or global) minimum. The first row shows the loss, and the second row shows the second-order finite difference of the loss. Models and datasets: (left) VGG11 on CIFAR10; (middle) ResNet18 on CIFAR10; (right) DenseNet121 on CIFAR100, only 50 classes are taken with 100 images per class. Note that on the second row we compute the second order finite difference of vectors (of loss values we plot on the first row). It is proportional but not equivalent to the second-order derivative.
  • Figure 3: The landscape around a point on the GD trajectory of training a VGG11 network on CIFAR10. The left panel takes the gradient direction, and the middle panel takes a random direction. The range of visualization is larger than that in Figure \ref{['fig:flattening']}. We can see separate scales in the loss landscape. The right panel shows the parameter vector of a VGG11 model projected in a 3-D space. The results of $3$ experiments are shown, with learning rate decayed at epoch $250$, $300$, and $350$, respectively. The x, y, z directions are obtained by orthogonalizing the first leading principal component of the blue segment before LRD (parameters with large learning rate), the orange segment (parameters after LRD at epoch 250), and the red segment (parameters after LRD at epoch 350), respectively.
  • Figure 4: The trajectory and objective function values when minimizing subquadratic functions by GD. (left) A 1-D convex function $f_1(x)=(1+|x|)\ln(1+|x|)-|x|$. The initialization is $x_0=-5$. The learning rate is initially $5$, and decreased to $3$ and $1$ on iteration $100$ and $150$. (middle) A 1-D non-convex function $f_2$ given by \ref{['eqn:fcn_integral']} with $C=3$. The initialization is $x_0=-5$. The learning rate is initially $2$, and decreased to $0.5$ and $0.2$ at iterations $500$ and $700$. GD show chaotic behaviors in this case. (right) A 2-D non-convex function $f_3(x_1, x_2)=f_2(x_1)+f_2(x_1+x_2)+0.5f_1(\sqrt{x_1^2+x_2^2})$. The initialization is $(-3,3)$. The learning rate is initially $1$, and decreased to $0.3$ and $0.1$ at iterations $200$ and $400$. GD show chaotic behaviors.
  • Figure 5: The trajectory of GD on a 2-D function with a flattening valley and subquadratic growth on $x_2$ direction. The iterator bounces back and forth on the valley and moves to the vicinity of flatter minima. The middle and right panels show the values of $x_1$ and $|x_2|$ for both GD and the quasistatic dynamics \ref{['eqn:quasistatic3']}. The results show: (1) the dynamics \ref{['eqn:quasistatic3']} is a very good approximation of the real dynamics. (2) The flatness-driven motion gets increasingly slower and finally stops. Experiment details: the loss function is $f(x_1,x_2)=f_1(x_2/(1+0.01x_1))$, where $f_1$ is defined in the caption of Figure \ref{['fig:traj_loss_1d2d']}. The initialization is $(1,1)$. The learning rate is $3$.
  • ...and 2 more figures

Theorems & Definitions (10)

  • Lemma 1
  • proof
  • Theorem 1
  • proof
  • Theorem 2
  • Definition 1
  • Theorem 3
  • proof
  • Remark 1
  • Remark 2