A visual exploration of optimization
One landscape, four ways down. Begin with the gradient, add memory, learn a scale, then learn the geometry itself.
The overview
Four ways to read a slope
- gradient
- update
- minimum
Start with the slope
Stochastic gradient descent
SGD trusts the current mini-batch gradient. It is cheap and honest—but one global learning rate must serve every direction.
The orange arrow is the raw gradient; the blue arrow is the step it produces.
Add memory
Momentum
An exponential moving average keeps useful velocity along the valley and cancels oscillation across its steep walls.
The blue update arrow carries memory: recent gradients vote on its direction, not only the current orange gradient.
Remember direction and scale
Adam
Adam combines momentum with a coordinate-wise second moment. Frequently large coordinates are divided down; quieter ones can move farther.
The ellipse is Adam’s diagonal preconditioner: axis-aligned, adaptive, and unable to rotate.
Learn the geometry
Shampoo
Diagonal scaling cannot see that two coordinates move together. Shampoo reshapes the same parameter vector as \(W=\operatorname{reshape}(\theta)\), then learns a matrix statistic for each tensor axis. Here, \(W\in\mathbb{R}^{2\times1}\) makes the preconditioner visible on the plane.
The tilted ellipse is the learned matrix preconditioner. It turns the gradient toward the valley.
Beyond the plane
What does 24-dimensional optimization look like?
Follow one \(6\times4\) tensor through its trajectory, loss, parameters, and learned geometry.
The matrices \(A\) and \(B\) couple rows and columns. Adam can scale each of the 24 entries; Shampoo can also learn which entries move together.
PCA of the loss geometry. Each state becomes \(Z=AWB\), where \(\mathcal{L}=\tfrac12\lVert Z\rVert_F^2\), and the two most informative directions form the view.
Every coordinate at once. The \(6\times4\) heatmap fades toward paper as the parameter tensor approaches zero.
Structure becomes visible. Shampoo’s row and column statistics reveal correlations that a diagonal second moment cannot represent.
Inside the update
What changes between a gradient and a step?
Compare four increasingly expressive views of gradient information: observation, memory, scale, and geometry.
Observe
A batch estimates the slope
Individual samples disagree. Their average reveals a useful direction without removing the noise entirely.
Samples form the cloud. The batch mean becomes \(g_t\).
Average through time
Momentum remembers a direction
Recent gradients receive the largest vote; older ones fade geometrically. Noise cancels while persistent motion remains.
Fading arrows are past gradients. The blue resultant is \(m_t\).
Normalize coordinates
Adam learns a diagonal metric
The second moment measures the typical magnitude of each coordinate. Dividing by it gives quieter coordinates room to move.
The raw vector is stretched independently along each axis. The metric stays axis-aligned.
Precondition correlations
Shampoo learns how axes interact
For a \(2\times2\) gradient matrix, the left factor rebalances rows and the right factor rebalances columns. Together they turn a skewed matrix update toward an even one.
Orange and blue arrows are the two columns. \(P_t^L\) acts on rows; \(P_t^R\) then mixes columns. A condition number near \(1\) means their shape is balanced.
- 01gradientsignal
- 02memorytime
- 03scalecoordinates
- 04geometrycorrelations