Momentum › A Worked Example: Tracking Velocity Through an Update

To solidify the mechanics, let us trace through a few steps of momentum by hand on a simple two-parameter problem. Consider a loss function $\mathcal{L}(\theta_1, \theta_2) = 25\theta_1^2 + \theta_2^2$ (a narrow valley with curvature ratio 25:1 between the two dimensions). Starting at $(\theta_1, \theta_2) = (1.0, 0.5)$ with learning rate $\alpha = 0.01$ and momentum $\beta = 0.9$ (PyTorch formulation), initial velocity $\mathbf{v}_0 = (0, 0)$.

The gradient at position $(\theta_1, \theta_2)$ is:

$$\nabla \mathcal{L} = (50\theta_1, \, 2\theta_2)$$

Step 1: Gradient at $(1.0, 0.5)$ is $(50, 1.0)$.

$$\mathbf{v}_1 = 0.9 \cdot (0, 0) + (50, 1.0) = (50, 1.0)$$ $$\theta_1 = (1.0, 0.5) - 0.01 \cdot (50, 1.0) = (0.5, 0.49)$$

Step 2: Gradient at $(0.5, 0.49)$ is $(25, 0.98)$.

$$\mathbf{v}_2 = 0.9 \cdot (50, 1.0) + (25, 0.98) = (45 + 25, \, 0.9 + 0.98) = (70, 1.88)$$ $$\theta_2 = (0.5, 0.49) - 0.01 \cdot (70, 1.88) = (-0.2, 0.471)$$

Notice what happened: the velocity in the $\theta_1$ direction (the steep dimension) jumped from 50 to 70, even though the step reduced $\theta_1$ by half. The velocity has overshot and is now pointing in the wrong direction. This will cause oscillations in $\theta_1$ before damping out.

Step 3: Gradient at $(-0.2, 0.471)$ is $(-10, 0.942)$.

$$\mathbf{v}_3 = 0.9 \cdot (70, 1.88) + (-10, 0.942) = (63 - 10, \, 1.692 + 0.942) = (53, 2.634)$$

The $\theta_1$ velocity is shrinking (from 70 to 53) because the gradient flipped sign. Meanwhile $\theta_2$ velocity is growing from 1.0 to 1.88 to 2.634, because the gradient in $\theta_2$ has been consistently positive. After many steps, the $\theta_1$ velocity will oscillate and decay while $\theta_2$ velocity grows steadily. This is precisely the dampening-oscillations, accelerating-consistent-directions behavior described earlier, made concrete in numbers.