Optimization Methods on an Energy Landscape

Abstract

This note is intended to build intuition for a set of standard optimization methods used in machine learning. In practice, using Adam or AdamW is largely a matter of calling an optimizer class and reading its specification in a framework's documentation; neither exercise reveals much about the qualitative behavior of the underlying algorithm, or why it can become trapped in a particular state of the loss surface. Here, each method is instead examined directly on three fixed, non-convex landscapes, chosen to isolate specific failure modes and structural bottlenecks.

1. The Energy Landscapes

The functions below serve as controlled test cases for isolating specific difficulties in finding optimal solutions. Realistic training scenarios are typically far more complex: the energy landscape is highly non-convex and irregular, and the dynamics of optimization over such surfaces are themselves a substantial area of study in their own right.

Himmelblau's function

Himmelblau's function[1] is a standard non-convex benchmark with four global minima of equal depth, marked as blue dots on the plot below.

It is used here because it isolates one specific question — which minimum a method finds — without the complication of any minimum being worse than another. A method that only follows the local slope cannot know which of the four basins is best; it can only find the one it happens to fall into. There is no shallow trap to escape from, and that is why the other two landscapes exist.

Rastrigin function

The Rastrigin function[2] adds a periodic term on top of a simple quadratic bowl:

The result is a surface tiled with local minima at a regular lattice of points, with a single global minimum at the origin; minima farther from the origin sit systematically higher, but nearby ones are close enough in depth to be difficult to distinguish from the true optimum by local information alone. Unlike Himmelblau's function, most of these local minima are strictly worse than the global one. The resulting surface is rugged enough that gradient-based methods struggle to navigate it reliably: as the later simulations show, a method frequently settles into one of the nearby local minima rather than reaching the global one.

Double Well (asymmetric)

The last landscape is not a named benchmark; it is built specifically for this series:

Neither of the first two landscapes is a good test of momentum specifically: Himmelblau has no shallow traps at all, and Rastrigin's traps sit too close together for a build-up of speed to matter much. This one has exactly two wells — a shallow local minimum and a much deeper global minimum — separated by a single barrier. The layout is such that a natural starting point rolls into the shallow well first.

2. Gradient Descent

The update rule can be derived directly, rather than just stated. Take a first-order Taylor expansion of \( E \) around the current point \( \boldsymbol{\theta}_t \), for a small step \( \Delta\boldsymbol{\theta} \):

To decrease \( E \) as much as possible for a step of fixed length \( \| \Delta\boldsymbol{\theta} \| = \eta \), we need the direction \( \Delta\boldsymbol{\theta} \) that minimizes \( \nabla E(\boldsymbol{\theta}_t)^{\!\top} \Delta\boldsymbol{\theta} \). By the Cauchy–Schwarz inequality,

with equality exactly when \( \Delta\boldsymbol{\theta} \) points opposite to the gradient. This gives the steepest local descent direction, \( \Delta\boldsymbol{\theta} = -\eta \, \nabla E(\boldsymbol{\theta}_t) / \| \nabla E(\boldsymbol{\theta}_t) \| \). In practice this normalized form is rarely used directly; dropping the normalization and letting \( \eta \) scale the raw gradient instead gives the familiar update rule:

The Taylor expansion is only accurate for small \( \Delta\boldsymbol{\theta} \), so this derivation already tells us why \( \eta \) has to be kept small enough — for larger steps, the linear approximation no longer holds and the resulting direction is no longer reliably a descent direction.

The update rule is applied iteratively from some initial point \( \boldsymbol{\theta}_0 = (x_0, y_0) \) until the gradient norm falls below a tolerance, or a maximum number of iterations is reached.

3. Stochastic Gradient Descent

In actual training, \( \nabla E(\boldsymbol{\theta}_t) \) is rarely available exactly: it is typically estimated from a random mini-batch of data, and that estimate carries noise relative to the true gradient. Stochastic Gradient Descent (SGD) is what Gradient Descent becomes once this estimation noise is taken into account:

There is no dataset behind \( E(x, y) \) here, so mini-batch noise cannot arise on its own — instead, the simulation below adds independent Gaussian noise \( \boldsymbol{\varepsilon}_t \) directly to the exact gradient at every step, which reproduces the same qualitative effect.

That noise works both ways. Because a noisy step is not guaranteed to be a descent direction, it can carry the iterate out of a shallow basin that plain Gradient Descent would be permanently stuck in — giving it a chance to find a better minimum nearby. The same noise also means SGD never truly settles: even once it reaches a minimum, it keeps jittering around it rather than coming to rest exactly on it.

4. Momentum, RMSProp, and Adam

A reasonable next guess is that averaging out that noise — giving each step some "momentum" from the recent past instead of reacting to a single noisy gradient — would already get most of the way to a modern optimizer. SGD with Momentum does exactly that: rather than moving by the (possibly noisy) gradient directly, it accumulates a velocity term over time and moves by that instead:

A consistent downhill direction builds up speed over several steps, the same way a heavy ball gains speed rolling down a slope. As a result, an iterate moving fast enough by the time it reaches a shallow dip can carry straight through it instead of stopping there, purely from accumulated inertia and with no noise required at all.

RMSProp[3] takes a different, independent idea: instead of accumulating the gradient itself, it maintains a per-dimension exponential moving average of the squared gradient \( \mathbf{s}_t \) as an estimate of each dimension's recent gradient magnitude. The update then divides the current gradient by the square root of this estimate before scaling by \( \eta \):

Here \( \epsilon \) is a small constant (typically around \( 10^{-8} \)) added purely to prevent division by zero in dimensions where \( \mathbf{s}_t \) is still close to zero; it plays no other role in the dynamics. Dividing by \( \sqrt{\mathbf{s}_t} \) means each dimension is scaled down in proportion to how large its gradient has typically been, independently of the others — but unlike momentum, there is no velocity term, so RMSProp has no memory of direction, only of magnitude. On its own, it does not carry through a local minimum the way momentum can.

Adam[4] (Adaptive Moment Estimation) combines both ideas rather than choosing between them: a per-dimension exponential moving average of the gradient itself \( \mathbf{m}_t \) and an exponential moving average of the squared gradient \( \mathbf{v}_t \), playing the same role as \( \mathbf{s}_t \) in RMSProp:

Both are initialized at zero, which biases them toward zero during the first few steps — the effect is more pronounced here than in RMSProp, since \( \beta_1 \) and \( \beta_2 \) are typically larger, giving more weight to the (empty) history; a standard correction removes this bias before the values are used:

As before, \( \epsilon \) only guards against division by zero and otherwise plays no role. Because \( \mathbf{m}_t \) has the same structure as the velocity term in the Momentum update above, Adam inherits the same inertia effect described there: it can carry through a shallow local minimum the same way Momentum can, on top of its per-dimension rescaling. Because of that rescaling, \( \eta \) here (and for RMSProp) operates on a different natural scale than in plain Gradient Descent, so it is treated as a separate setting below. The standard values \( \beta = 0.9 \) (RMSProp), \( \beta_1 = 0.9 \), \( \beta_2 = 0.999 \), \( \epsilon = 10^{-8} \) are used throughout.

5. Methods Sandbox

The controls below run any of the five methods on any of the three landscapes and animate the resulting path in both views at once, updating the status line with the current step as it runs. "Stop" halts a run in progress; "Show gradient" draws the current gradient direction on the 2D view, updating live while a run is in progress.

6. Comparing the Methods

Running plain Gradient Descent on \( E(x, y) \) makes its two defining limitations easy to see:

• Initialization determines the outcome: since the method only ever follows the local downhill direction, it converges to whichever of the four minima is "downhill" from the starting point. Two nearby starting points on opposite sides of a ridge can converge to entirely different basins.
• No mechanism to escape a basin: once inside the basin of attraction of a minimum, there is nothing in the update rule that could push the point back out, even if a better minimum is nearby.
• Sensitivity to the learning rate: too small an \( \eta \) makes convergence slow, particularly near the saddle region around the origin where the gradient is close to zero. Too large an \( \eta \) causes overshooting and oscillation around steep regions of the surface, and can diverge entirely.

On the Rastrigin function, the second point stops being a corner case and becomes the typical outcome: switch to it in the controls above and run Gradient Descent from a few different starting points. It converges quickly every time — just almost never to the global minimum at the origin, since it stops at whichever nearby local minimum it happens to fall into first. RMSProp, run on the same case, behaves no differently in this respect: its per-dimension rescaling changes how quickly it approaches a minimum, not whether it can leave one once inside.

Switching to Stochastic Gradient Descent, on the same Rastrigin landscape and the same starting point, changes this picture: run it a few times in a row and the endpoint is no longer the same every time. With enough noise \( \sigma \), some runs escape the first local minimum they reach and settle in a different, sometimes better, one — at the cost of never quite standing still once they get there, visibly jittering around the minimum instead. Momentum shows a related but distinct effect: on the Double Well, a run that reaches the shallow well with enough accumulated velocity carries straight through the barrier into the deep one, with no noise involved at all.

Adam, run on the same Rastrigin case, tends to sit somewhere between the noise-driven and velocity-driven pictures: its own accumulated gradient history smooths out the raw oscillation from the \( 20\pi\sin(2\pi x) \) term, so it moves more steadily than SGD does, but it is still following the same local landscape and can still settle into a local rather than the global minimum. The per-dimension rescaling changes how it gets there, not fundamentally whether it can see past a nearby basin.

Across all three landscapes, then, the same pattern recurs. Gradient Descent and RMSProp follow only the local gradient and commit to whichever basin they first encounter, with no way back out. SGD, Momentum, and Adam can each occasionally carry a run past a nearby basin — through noise, through inertia, or through both — but none of them does so reliably or by design: escaping a basin, where it happens at all, is a side effect of a mechanism built for a different purpose, not a deliberate search strategy.

References

1. Himmelblau, D. M. (1972). Applied Nonlinear Programming. McGraw-Hill.
2. Rastrigin, L. A. (1974). Systems of Extremal Control. Nauka.
3. Tieleman, T., & Hinton, G. (2012). Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning.
4. Kingma, D. P., & Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR.