Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Stochastic Gradient Descent › Loss Landscape Visualization

Michael BrenndoerferApril 30, 20255 min read

Language AI HandbookAll chapters

Part 10 of 67Neural Network Foundations

Chapter 66Stochastic Gradient Descent › Loss Landscape Visualization

Reading section 10 of 16 · Loss Landscape Visualization · 5 minSection 10 of 16 · 5 min

Section 10 of 16 · 5 min read

Stochastic Gradient Descent › Loss Landscape VisualizationOne section · Sections in this chapter10 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Visualizing the loss landscape directly is difficult: it lives in $D$ -dimensional space, where $D$ can be billions. But we can create informative 2D projections.

The technique introduced by Li et al. (2018) in "Visualizing the Loss Landscape of Neural Nets" creates a 2D slice of the loss surface by choosing two random directions $\boldsymbol{\delta}$ and $\boldsymbol{\eta}$ in parameter space and evaluating the loss at:

\mathcal{L}(\boldsymbol{\theta}^* + \alpha \boldsymbol{\delta} + \beta \boldsymbol{\eta})

where:

$\boldsymbol{\theta}^*$ : the final trained parameters (the center of the visualization)
$\alpha, \beta$ : scalars that vary over a grid (e.g., both from $-1$ to $1$ in steps of 0.1)
$\boldsymbol{\delta}, \boldsymbol{\eta}$ : random unit vectors in parameter space (often constructed by sampling each component from a Gaussian and then normalizing)

The result is a 2D "slice" through the loss landscape centered on the solution. Sharp solutions produce highly curved, steep contours. Flat solutions produce gently curving, wide contours.

This visualization technique has been used to show that residual connections (skip connections, as in ResNets) dramatically smooth the loss landscape, making optimization much easier. Similar studies have shown that batch normalization also smooths the landscape. These empirical findings help explain why these architectural choices work so well.

To build intuition for how SGD moves across a loss surface, we can visualize optimizer trajectories on a simple 2D quadratic loss function. While real neural network losses are far more complex, the basic dynamics of learning rate effects and batch noise are captured clearly in this simplified setting.

Out[3]:

Visualization

2D contour plot with three optimizer paths showing different learning rate behaviors. — SGD trajectories on a 2D quadratic loss surface for three learning rates. The small learning rate (0.05) approaches the minimum cautiously. The medium learning rate (0.15) converges faster with damped oscillations, while the larger stable rate (0.18) repeatedly overshoots the steep direction before settling near the minimum.

1D loss curve plot showing a sharp narrow minimum and a flat wide minimum. — Comparison of sharp versus flat minima in 1D. The sharp minimum (narrow valley) is sensitive to small perturbations and can be easily escaped by SGD noise. The flat minimum (wide valley) is robust to noise, and SGD naturally spends more time here. Flat minima are associated with better generalization.

The trajectory plot on the left demonstrates a key challenge with SGD on ill-conditioned problems: on the elongated ellipsoidal loss surface, a large learning rate causes the optimizer to overshoot and oscillate across the steep dimension while making slow progress along the shallow one. A well-chosen learning rate finds the right balance. The right plot illustrates why flat minima are preferred: SGD noise (represented by the shaded band) easily pushes the optimizer out of the narrow sharp minimum, but stays contained within the flat minimum's wide basin.