Momentum / Nesterov

Why Momentum / Nesterov Matters

Plain gradient descent uses the current slope and forgets how it arrived. On a long, narrow loss valley, this can mean alternating across the steep direction while making slow progress along the shallow one. Momentum keeps a velocity vector so directions that persist over several updates can accumulate. Nesterov momentum measures the next gradient at a predicted position obtained from that velocity.

These methods are useful in smooth numerical optimization and in training models with full or mini-batch gradients. They add one persistent vector to the optimizer state. That small change makes the optimization a dynamical system: the same parameter vector can move differently depending on its velocity.

You will implement heavy-ball momentum, abbreviated HB, and a constant-coefficient look-ahead Nesterov form, abbreviated NAG. Both programs use an analytic diagonal quadratic, a fixed starting point, and zero initial velocity. This exposes every operation without an automatic-differentiation library or randomness. The short comparison illustrates the mechanics; it is not a claim that either method wins on every objective or every iteration. The displacement-velocity conventions follow the equations in Sutskever et al., Section 2.

Core Concepts and Terminology

We minimize a differentiable scalar function $f:\mathbb R^d\to\mathbb R$. Let $\mathbf x_t$ be the parameters after $t$ updates, $\mathbf v_t$ the stored displacement velocity, $\eta\gt0$ the learning rate, and $0\leq\beta\lt1$ the momentum coefficient. Initialize $\mathbf v_0=\mathbf0$. First form the retained motion $\mathbf m_t=\beta\mathbf v_t$.

Heavy ball evaluates the gradient at the old position:

$$\mathbf y_t=\mathbf x_t,\qquad \mathbf g_t=\nabla f(\mathbf y_t),$$ $$\mathbf v_{t+1}=\mathbf m_t-\eta\mathbf g_t,\qquad \mathbf x_{t+1}=\mathbf x_t+\mathbf v_{t+1}.$$

Nesterov changes the probe point, using the old velocity:

$$\mathbf y_t=\mathbf x_t+\mathbf m_t,\qquad \mathbf g_t=\nabla f(\mathbf y_t),$$ $$\mathbf v_{t+1}=\mathbf m_t-\eta\mathbf g_t,\qquad \mathbf x_{t+1}=\mathbf x_t+\mathbf v_{t+1}.$$

For Nesterov, $\mathbf x_{t+1}=\mathbf y_t-\eta\nabla f(\mathbf y_t)$ in real arithmetic. This alternative is a complete update, not a second update to perform after adding the new velocity. The programs retain the explicitly separated velocity form to make state and rounding order reproducible.

Velocity already includes learning-rate scaling: it has the same units as the parameters. Another common notation stores a positive gradient accumulator $\mathbf b_{t+1}=\beta\mathbf b_t+\mathbf g_t$ and writes $\mathbf x_{t+1}=\mathbf x_t-\eta\mathbf b_{t+1}$. For heavy ball with constant $\eta$, initialization zero, and no dampening, this corresponds to $\mathbf v_t=-\eta\mathbf b_t$. Changing learning rates or initialization rules changes that correspondence. Here there is no factor $1-\beta$ multiplying the gradient. A normalized exponential moving average is a different convention and requires a corresponding learning-rate interpretation.

The teaching objective is

$$f(\mathbf x)=\frac12\sum_{j=1}^d a_j(x_j-c_j)^2,\qquad [\nabla f(\mathbf x)]_j=a_j(x_j-c_j),\qquad a_j\geq0.$$

The curvatures are the diagonal entries of its Hessian. Write $L=\max_j a_j$. If all $a_j\gt0$, the unique minimizer is $\mathbf c$, the strong-convexity parameter is $\mu=\min_j a_j$, and the condition number is $\kappa=L/\mu$. A zero curvature makes that coordinate flat: its value need not reach $c_j$ to minimize $f$. Starting with zero velocity keeps that coordinate fixed.

The solver returns the final parameters, velocity, objective, counts, and snapshots of all states. It performs exactly $K$ updates and labels no state “converged.” For $L\gt0$, it conservatively requires $0\lt\eta\leq1/L$; if $L=0$, it skips this reciprocal test. This is a teaching guard for this quadratic family, rather than the largest possible stable step for every momentum method.

How Momentum / Nesterov Works, Step by Step

Inputs are the initial point, center vector, nonnegative curvature vector, learning rate $\eta$, coefficient $\beta$, update count $K$, and mode HB or NAG. Use the following order for every step:

1. Validate before computing. Check dimensions, finite numeric values, parameter bounds, integer update count, mode, and the reciprocal learning-rate guard. Copy the starting vector and initialize a zero velocity. Invalid inputs raise an exception; the caller receives no partial result.
2. Record step zero. Evaluate $f(\mathbf x_0)$ once. Its probe and gradient fields are empty because no update gradient has yet been requested. The empty fields print NONE, not a fabricated zero gradient.
3. Check the update cap. If $t=K$, return the final state and counters. A zero-step run therefore evaluates one starting loss and no gradients.
4. Freeze the old state and compute the carry. Set $\mathbf m_t=\beta\mathbf v_t$. Do not overwrite the old velocity before choosing the probe.
5. Choose exactly one probe. Heavy ball uses $\mathbf x_t$; Nesterov uses $\mathbf x_t+\mathbf m_t$. Obtain one full gradient there. Look-ahead does not inherently require a second gradient evaluation.
6. Construct a whole new state. For each coordinate, compute the correction $\eta g_{t,j}$, subtract it from the carry to obtain $v_{t+1,j}$, and add that new velocity to the old $x_{t,j}$. Do not mutate parameters while evaluating a gradient that may couple multiple coordinates.
7. Commit and monitor. Replace parameters and velocity together. Evaluate $f(\mathbf x_{t+1})$, store copies of the probe, gradient, and committed state, then repeat. The diagnostic objective does not accept, reject, or change an update.

A history row with iteration=k contains $\mathbf x_k$, $\mathbf v_k$, and $f(\mathbf x_k)$, but its nonempty probe and gradient describe the update into that row, namely $\mathbf y_{k-1}$ and $\mathbf g_{k-1}$. The gradient field is not a residual at $\mathbf x_k$.

There is intentionally no early exit for a zero probe gradient, one unchanged point, or increasing loss. For a production convergence test, evaluate the actual residual at the current parameters and consider both the residual and velocity, with appropriate absolute and relative tolerances. On a general nonconvex objective, even a small residual and velocity certify only approximate stationarity, not a global minimum.

flowchart TD
    A["Validate mode, shapes, ranges, and step size"] --> B["Copy x; set v to zero; evaluate loss and record start"]
    B --> C{"Completed K updates?"}
    C -->|Yes| Z["Return final x, v, loss, counts, and history"]
    C -->|No| D["Freeze old x and v; carry = beta times old v"]
    D --> E{"NAG mode?"}
    E -->|Yes| F["Probe = old x plus carry"]
    E -->|No| G["Probe = old x"]
    F --> H["Evaluate one full gradient at probe"]
    G --> H
    H --> I["New v = carry minus rate times gradient"]
    I --> J["New x = old x plus new v; commit both"]
    J --> L["Evaluate new loss; copy history row; count update"]
    L --> C

Worked Example with a Full Trace

Use the one-dimensional function $f(x)=x^2/2$, gradient $f'(x)=x$, and inputs $x_0=4$, $v_0=0$, $\eta=\beta=1/2$, $K=4$. There is no random seed because the function, gradients, and updates are deterministic. Initially $f(4)=8$. All intermediate state numbers below are dyadic rationals and exactly representable in binary64 arithmetic.

Heavy ball. Here the probe is always the old $x$.

1. From $(x_0,v_0)=(4,0)$: carry $m_0=0$, probe $y_0=4$, gradient $g_0=4$, correction $\eta g_0=2$. Thus $v_1=0-2=-2$, $x_1=4-2=2$, and $f(x_1)=2$.
2. From $(2,-2)$: $m_1=-1$, $y_1=2$, $g_1=2$, and $\eta g_1=1$. Thus $v_2=-1-1=-2$, $x_2=2-2=0$, and $f(x_2)=0$.
3. From $(0,-2)$: $m_2=-1$, $y_2=0$, $g_2=0$, and $\eta g_2=0$. Thus $v_3=-1$, $x_3=-1$, and $f(x_3)=1/2$. The stored velocity carries the method away from the exact minimizer.
4. From $(-1,-1)$: $m_3=-1/2$, $y_3=-1$, $g_3=-1$, and $\eta g_3=-1/2$. Thus $v_4=-1/2-(-1/2)=0$, $x_4=-1$, and $f(x_4)=1/2$. An unchanged point is not a fixed state: the next gradient is still $-1$.

Nesterov. Now the probe includes the old carry.

1. From $(4,0)$: $m_0=0$, $y_0=4+0=4$, $g_0=4$, and $\eta g_0=2$. Thus $v_1=-2$, $x_1=2$, and $f(x_1)=2$, matching the first heavy-ball update.
2. From $(2,-2)$: $m_1=-1$, $y_1=2-1=1$, $g_1=1$, and $\eta g_1=1/2$. Thus $v_2=-3/2$, $x_2=1/2$, and $f(x_2)=1/8$.
3. From $(1/2,-3/2)$: $m_2=-3/4$, $y_2=-1/4$, $g_2=-1/4$, and $\eta g_2=-1/8$. Thus $v_3=-3/4+1/8=-5/8$, $x_3=1/2-5/8=-1/8$, and $f(x_3)=1/128=0.0078125$.
4. From $(-1/8,-5/8)$: $m_3=-5/16$, $y_3=-7/16$, $g_3=-7/16$, and $\eta g_3=-7/32$. Thus $v_4=-5/16+7/32=-3/32$, $x_4=-7/32=-0.21875$, and $f(x_4)=49/2048=0.02392578125$.

Each run requests four gradients and five objective values, including the initial loss. Nesterov has the smaller final loss for these inputs, while heavy ball reaches zero earlier and then leaves it. Nesterov's own loss increases on step $4$; neither sequence is monotone. The solver returns its last state, not the best state seen so far.

Implementation in Python, C++, and Java

Each panel is a complete program with the same sample data and operation order. The callable momentum routine accepts vectors initial, center, weights, numbers rate, beta, an integer steps, and string mode. It supports $1\leq d\leq32$, $0\leq K\leq2000$, finite coordinates and centers with magnitude at most $10^6$, and finite curvatures in $[0,10^6]$. It requires $0\lt\eta\leq10^6$, $0\leq\beta\lt1$, and, when $L\gt0$, $\eta\leq1/L$ as checked in binary64. Python rejects booleans as numeric controls and requires an actual integer update count. The other languages express those types in the function signature.

The function always performs the requested number of updates. Zero curvature, a zero update count, a zero gradient, and a numerical no-op are supported cases. There is no callback objective, mini-batch sampling, learning-rate schedule, restart, line search, or hidden convergence flag. For another objective, replace the scalar value and gradient routines and reconsider the validation and stability assumptions.

Run Python 3 with python main.py. Compile C++17 with g++ -std=c++17 -ffp-contract=off main.cpp -o momentum and run ./momentum; avoid fast-math so operations retain the stated order. Save Java as Main.java, use Java 17 or later, and run javac Main.java then java Main. Java's decimal helper only formats the exact stored binary64 value to eight places with half-even rounding; optimization uses ordinary double arithmetic.

Output prints the objective at every recorded parameter vector. probe and g refer to the incoming update, so NONE on step zero is expected. The returned counters count full vector-gradient evaluations and scalar objective evaluations, not individual coordinates. All snapshots own their parameter and velocity arrays.

from dataclasses import dataclass
import math


@dataclass
class State:
    iteration: int
    point: list
    velocity: list
    probe: list
    gradient: list
    value: float


@dataclass
class Result:
    point: list
    velocity: list
    value: float
    steps: int
    gradient_evaluations: int
    value_evaluations: int
    history: list


def objective(point, center, weights):
    total = 0.0
    for x, c, a in zip(point, center, weights):
        delta = x - c
        square = delta * delta
        term = a * square
        total = total + term
    return 0.5 * total


def momentum(initial, center, weights, rate, beta, steps, mode):
    if type(steps) is not int or not 0 <= steps <= 2000:
        raise ValueError("steps must be an integer in [0,2000]")
    if mode not in ("HB", "NAG"):
        raise ValueError("mode must be HB or NAG")

    def number(x):
        return type(x) is int or (type(x) is float and math.isfinite(x))

    d = len(initial)
    if not 1 <= d <= 32 or len(center) != d or len(weights) != d:
        raise ValueError("invalid vector sizes")
    if any(not number(x) or abs(x) > 1000000 for row in [initial, center] for x in row):
        raise ValueError("invalid initial point or center")
    if any(not number(a) or not 0 <= a <= 1000000 for a in weights):
        raise ValueError("invalid curvature")
    if not number(rate) or not 0 < rate <= 1000000:
        raise ValueError("invalid learning rate")
    if not number(beta) or not 0 <= beta < 1:
        raise ValueError("invalid momentum coefficient")
    largest = float(max(weights))
    if largest > 0 and rate > 1.0 / largest:
        raise ValueError("learning rate exceeds conservative 1/L limit")
    point, origin, diagonal = [list(map(float, row)) for row in [initial, center, weights]]
    rate, beta = float(rate), float(beta)
    velocity = [0.0] * d
    value = objective(point, origin, diagonal)
    history = [State(0, point[:], velocity[:], [], [], value)]
    for k in range(1, steps + 1):
        carry = [beta * v for v in velocity]
        probe = [x + m for x, m in zip(point, carry)] if mode == "NAG" else point[:]
        gradient = []
        for y, c, a in zip(probe, origin, diagonal):
            delta = y - c
            gradient.append(a * delta)
        next_velocity, next_point = [], []
        for x, m, g in zip(point, carry, gradient):
            correction = rate * g
            v = m - correction
            next_velocity.append(v)
            next_point.append(x + v)
        point, velocity = next_point, next_velocity
        value = objective(point, origin, diagonal)
        history.append(State(k, point[:], velocity[:], probe[:], gradient[:], value))
    return Result(point[:], velocity[:], value, steps, steps, steps + 1, history)


def vector(xs):
    return "(" + ",".join(f"{x:.8f}" for x in xs) + ")" if xs else "NONE"


def main():
    for mode in ("HB", "NAG"):
        result = momentum([4], [0], [1], 0.5, 0.5, 4, mode)
        print(mode)
        for s in result.history:
            print(f"k={s.iteration} x={vector(s.point)} v={vector(s.velocity)} "
                  f"probe={vector(s.probe)} g={vector(s.gradient)} f={s.value:.8f}")
        print(f"steps={result.steps} gradients={result.gradient_evaluations} "
              f"losses={result.value_evaluations}")


if __name__ == "__main__":
    main()

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
using namespace std;

struct State {
    int iteration;
    vector<double> point, velocity, probe, gradient;
    double value;
};
struct Result {
    vector<double> point, velocity;
    double value;
    int steps, gradient_evaluations, value_evaluations;
    vector<State> history;
};

double objective(const vector<double>& point, const vector<double>& center,
                 const vector<double>& weights) {
    double total = 0.0;
    for (size_t j = 0; j < point.size(); ++j) {
        double delta = point[j] - center[j];
        double square = delta * delta;
        double term = weights[j] * square;
        total = total + term;
    }
    return 0.5 * total;
}

Result momentum(const vector<double>& initial, const vector<double>& center,
                const vector<double>& weights, double rate, double beta,
                int steps, const string& mode) {
    if (steps < 0 || steps > 2000 || (mode != "HB" && mode != "NAG"))
        throw invalid_argument("invalid steps or mode");
    size_t d = initial.size();
    if (d < 1 || d > 32 || center.size() != d || weights.size() != d)
        throw invalid_argument("invalid vector sizes");
    for (size_t j = 0; j < d; ++j) {
        if (!isfinite(initial[j]) || abs(initial[j]) > 1000000 ||
            !isfinite(center[j]) || abs(center[j]) > 1000000 ||
            !isfinite(weights[j]) || weights[j] < 0 || weights[j] > 1000000)
            throw invalid_argument("invalid coordinates or curvature");
    }
    if (!isfinite(rate) || rate <= 0 || rate > 1000000 ||
        !isfinite(beta) || beta < 0 || beta >= 1)
        throw invalid_argument("invalid learning rate or momentum");
    double largest = *max_element(weights.begin(), weights.end());
    if (largest > 0 && rate > 1.0 / largest)
        throw invalid_argument("learning rate exceeds conservative 1/L limit");
    vector<double> point = initial, velocity(d, 0.0);
    double value = objective(point, center, weights);
    vector<State> history;
    history.push_back(State{0, point, velocity, {}, {}, value});
    for (int k = 1; k <= steps; ++k) {
        vector<double> carry(d), probe(d), gradient(d), next_velocity(d), next_point(d);
        for (size_t j = 0; j < d; ++j) {
            carry[j] = beta * velocity[j];
            probe[j] = mode == "NAG" ? point[j] + carry[j] : point[j];
        }
        for (size_t j = 0; j < d; ++j) {
            double delta = probe[j] - center[j];
            gradient[j] = weights[j] * delta;
        }
        for (size_t j = 0; j < d; ++j) {
            double correction = rate * gradient[j];
            next_velocity[j] = carry[j] - correction;
            next_point[j] = point[j] + next_velocity[j];
        }
        point = next_point;
        velocity = next_velocity;
        value = objective(point, center, weights);
        history.push_back(State{k, point, velocity, probe, gradient, value});
    }
    return Result{point, velocity, value, steps, steps, steps + 1, history};
}

string vector_text(const vector<double>& xs) {
    if (xs.empty()) return "NONE";
    ostringstream out;
    out << fixed << setprecision(8) << '(';
    for (size_t j = 0; j < xs.size(); ++j) {
        if (j) out << ',';
        out << xs[j];
    }
    out << ')';
    return out.str();
}

int main() {
    for (const string& mode : {string("HB"), string("NAG")}) {
        Result r = momentum({4}, {0}, {1}, 0.5, 0.5, 4, mode);
        cout << mode << '\n' << fixed << setprecision(8);
        for (const auto& s : r.history)
            cout << "k=" << s.iteration << " x=" << vector_text(s.point)
                 << " v=" << vector_text(s.velocity) << " probe=" << vector_text(s.probe)
                 << " g=" << vector_text(s.gradient) << " f=" << s.value << '\n';
        cout << "steps=" << r.steps << " gradients=" << r.gradient_evaluations
             << " losses=" << r.value_evaluations << '\n';
    }
    return 0;
}

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.ArrayList;
import java.util.List;

public class Main {
    static class State {
        int iteration;
        double[] point, velocity, probe, gradient;
        double value;
        State(int k, double[] x, double[] v, double[] y, double[] g, double f) {
            iteration = k; point = x.clone(); velocity = v.clone();
            probe = y.clone(); gradient = g.clone(); value = f;
        }
    }
    static class Result {
        double[] point, velocity;
        double value;
        int steps, gradient_evaluations, value_evaluations;
        List<State> history;
        Result(double[] x, double[] v, double f, int k, List<State> rows) {
            point = x.clone(); velocity = v.clone(); value = f;
            steps = k; gradient_evaluations = k; value_evaluations = k + 1; history = rows;
        }
    }

    static double objective(double[] point, double[] center, double[] weights) {
        double total = 0.0;
        for (int j = 0; j < point.length; ++j) {
            double delta = point[j] - center[j];
            double square = delta * delta;
            double term = weights[j] * square;
            total = total + term;
        }
        return 0.5 * total;
    }

    static Result momentum(double[] initial, double[] center, double[] weights,
                           double rate, double beta, int steps, String mode) {
        if (steps < 0 || steps > 2000 || (!"HB".equals(mode) && !"NAG".equals(mode)))
            throw new IllegalArgumentException("invalid steps or mode");
        int d = initial.length;
        if (d < 1 || d > 32 || center.length != d || weights.length != d)
            throw new IllegalArgumentException("invalid vector sizes");
        double largest = 0.0;
        for (int j = 0; j < d; ++j) {
            if (!Double.isFinite(initial[j]) || Math.abs(initial[j]) > 1000000 ||
                !Double.isFinite(center[j]) || Math.abs(center[j]) > 1000000 ||
                !Double.isFinite(weights[j]) || weights[j] < 0 || weights[j] > 1000000)
                throw new IllegalArgumentException("invalid coordinates or curvature");
            largest = Math.max(largest, weights[j]);
        }
        if (!Double.isFinite(rate) || rate <= 0 || rate > 1000000 ||
            !Double.isFinite(beta) || beta < 0 || beta >= 1)
            throw new IllegalArgumentException("invalid learning rate or momentum");
        if (largest > 0 && rate > 1.0 / largest)
            throw new IllegalArgumentException("learning rate exceeds conservative 1/L limit");
        double[] point = initial.clone(), velocity = new double[d];
        double value = objective(point, center, weights);
        List<State> history = new ArrayList<>();
        history.add(new State(0, point, velocity, new double[0], new double[0], value));
        for (int k = 1; k <= steps; ++k) {
            double[] carry = new double[d], probe = new double[d], gradient = new double[d];
            double[] nextVelocity = new double[d], nextPoint = new double[d];
            for (int j = 0; j < d; ++j) {
                carry[j] = beta * velocity[j];
                probe[j] = mode.equals("NAG") ? point[j] + carry[j] : point[j];
            }
            for (int j = 0; j < d; ++j) {
                double delta = probe[j] - center[j];
                gradient[j] = weights[j] * delta;
            }
            for (int j = 0; j < d; ++j) {
                double correction = rate * gradient[j];
                nextVelocity[j] = carry[j] - correction;
                nextPoint[j] = point[j] + nextVelocity[j];
            }
            point = nextPoint; velocity = nextVelocity;
            value = objective(point, center, weights);
            history.add(new State(k, point, velocity, probe, gradient, value));
        }
        return new Result(point, velocity, value, steps, history);
    }

    static String fixed8(double x) {
        String s = new BigDecimal(x).setScale(8, RoundingMode.HALF_EVEN).toPlainString();
        return Double.doubleToRawLongBits(x) < 0 && s.charAt(0) != '-' ? "-" + s : s;
    }
    static String vector(double[] xs) {
        if (xs.length == 0) return "NONE";
        StringBuilder s = new StringBuilder("(");
        for (int j = 0; j < xs.length; ++j) {
            if (j > 0) s.append(',');
            s.append(fixed8(xs[j]));
        }
        return s.append(')').toString();
    }
    public static void main(String[] args) {
        for (String mode : new String[]{"HB", "NAG"}) {
            Result r = momentum(new double[]{4}, new double[]{0}, new double[]{1}, 0.5, 0.5, 4, mode);
            System.out.println(mode);
            for (State s : r.history)
                System.out.println("k=" + s.iteration + " x=" + vector(s.point) +
                    " v=" + vector(s.velocity) + " probe=" + vector(s.probe) +
                    " g=" + vector(s.gradient) + " f=" + fixed8(s.value));
            System.out.println("steps=" + r.steps + " gradients=" + r.gradient_evaluations +
                " losses=" + r.value_evaluations);
        }
    }
}

Expected standard output, identical in all three languages:

• HB
• k=0 x=(4.00000000) v=(0.00000000) probe=NONE g=NONE f=8.00000000
• k=1 x=(2.00000000) v=(-2.00000000) probe=(4.00000000) g=(4.00000000) f=2.00000000
• k=2 x=(0.00000000) v=(-2.00000000) probe=(2.00000000) g=(2.00000000) f=0.00000000
• k=3 x=(-1.00000000) v=(-1.00000000) probe=(0.00000000) g=(0.00000000) f=0.50000000
• k=4 x=(-1.00000000) v=(0.00000000) probe=(-1.00000000) g=(-1.00000000) f=0.50000000
• steps=4 gradients=4 losses=5
• NAG
• k=0 x=(4.00000000) v=(0.00000000) probe=NONE g=NONE f=8.00000000
• k=1 x=(2.00000000) v=(-2.00000000) probe=(4.00000000) g=(4.00000000) f=2.00000000
• k=2 x=(0.50000000) v=(-1.50000000) probe=(1.00000000) g=(1.00000000) f=0.12500000
• k=3 x=(-0.12500000) v=(-0.62500000) probe=(-0.25000000) g=(-0.25000000) f=0.00781250
• k=4 x=(-0.21875000) v=(-0.09375000) probe=(-0.43750000) g=(-0.43750000) f=0.02392578
• steps=4 gradients=4 losses=5

Complexity, Correctness, and Invariants

Let $K$ be the requested number of updates and $d$ the dimension. Validation and initialization take $O(d)$. Every update evaluates one gradient, performs $O(d)$ vector arithmetic, evaluates one scalar objective for monitoring, and copies $O(d)$ state into history. For this diagonal quadratic both function routines cost $O(d)$, so best, average, and worst time are all $\Theta((K+1)d)$ for fixed valid $K,d$: no path exits early. The program makes exactly $K$ gradient calls and $K+1$ objective calls, including when the starting point is optimal.

Persistent optimizer state and temporary working vectors use $O(d)$ space. Retaining all snapshots raises the actual returned-storage requirement to $\Theta((K+1)d)$. Streaming or disabling the history gives $O(d)$ working space, but printed output and serialization still cost time. For a general objective with costs $C_g(d)$ and $C_f(d)$, total time is $O(d+C_f(d)+K(C_g(d)+C_f(d)+d))$. The constant full-loss monitoring policy here may be prohibitively expensive for a large training dataset.

State invariant and induction. Before update $t$, the stored pair is exactly $(\mathbf x_t,\mathbf v_t)$ from the chosen recurrence. Initially $\mathbf v_0=\mathbf0$, and the initial history row stores copies of the inputs and their objective. Assuming the invariant, the code forms $\beta\mathbf v_t$, selects the prescribed probe from the old state, and evaluates its gradient before mutation. Computing the new velocity and adding it once to the old parameters produces precisely the next recurrence state. Committing together and copying snapshots preserves the invariant. The counts follow because each loop performs one gradient call and one loss call. In floating-point arithmetic this means the specified, rounded sequence of operations; it does not imply exact real-arithmetic iterates on arbitrary inputs.

Why these guarded quadratics converge in exact arithmetic. Analyze one positive-curvature coordinate $a\gt0$, with error $e_t=x_t-c$. Set the fictitious $e_{-1}=e_0$ so $v_0=e_0-e_{-1}=0$. Eliminating velocity gives

$$e_{t+1}=(1+\beta-\eta a)e_t-\beta e_{t-1}\quad\text{(HB)},$$ $$e_{t+1}=q(1+\beta)e_t-q\beta e_{t-1},\qquad q=1-\eta a\quad\text{(NAG)}.$$

For a real characteristic polynomial $r^2-Ar+B$, both roots lie strictly inside the unit circle when $\lvert B\rvert\lt1$, $1-A+B\gt0$, and $1+A+B\gt0$. For heavy ball, $A=1+\beta-\eta a$ and $B=\beta$. Under $0\lt\eta a\leq1$ and $0\leq\beta\lt1$, these tests become $\beta\lt1$, $\eta a\gt0$, and $2+2\beta-\eta a\gt0$. For Nesterov, $A=q(1+\beta)$ and $B=q\beta$, with $0\leq q\lt1$; the tests become $q\beta\lt1$, $1-q\gt0$, and $1+q(1+2\beta)\gt0$. Thus every positive-curvature error tends to zero. A zero-curvature coordinate has gradient zero and, from zero initial velocity, remains fixed; it already contributes zero to the objective.

This establishes convergence of the infinite exact-arithmetic recurrence for this objective family under the stated mathematical guard. A finite run or rounded recurrence can stall or finish away from the minimum. Loss need not decrease at every iteration; the history above supplies counterexamples. Quadratic stability and implementation correctness do not establish a global convergence theorem for arbitrary nonlinear losses.

Acceleration needs assumptions and tuning. For smooth $\mu$-strongly convex deterministic objectives, a standard Nesterov choice is $\eta=1/L$ and $\beta=(\sqrt L-\sqrt\mu)/(\sqrt L+\sqrt\mu)$. Appropriate accelerated methods improve the condition-number dependence of iteration bounds from order $\kappa$ to order $\sqrt\kappa$, up to logarithmic accuracy factors. General smooth convex $O(1/t^2)$ guarantees use a suitable Nesterov parameter schedule; they do not follow from an arbitrary fixed $\beta$. Heavy ball's favorable quadratic tuning does not supply the same general strongly convex guarantee. Roger Grosse's momentum lecture discusses these distinctions. The fixed four-step demonstration is a recurrence trace, not a measurement of asymptotic rates.

Variants, Pitfalls, and Optimizations

Do not count carry twice. With displacement velocity, compute the new point as old point plus new velocity. Equivalently for Nesterov in exact arithmetic, use probe minus the gradient correction. Adding the new velocity to the probe is wrong because that velocity already contains the carry. The two correct algebraic forms can also round differently, so choose one operation order for cross-language reproduction.

Parameter stationarity and optimizer state are different. Heavy ball reaches the minimizer at step $2$ while its velocity is $-2$. It then moves away. On step $4$, its point is unchanged but its gradient is nonzero; step $5$ moves again. A proper fixed point of the full deterministic map must have $\mathbf v=\mathbf0$ and $\nabla f(\mathbf x)=\mathbf0$ for $\eta\gt0$. A small gradient at a look-ahead probe is not necessarily a small gradient at the current parameters. Finite precision adds underflow, cancellation, and numerical no-ops; the code counts those steps rather than declaring success.

Framework labels do not uniquely specify the state convention. PyTorch's SGD documentation describes a gradient-buffer implementation, its first-step buffer initialization, dampening, and a Nesterov update expressed at the stored parameters. Its stored parameter sequence should not be compared blindly with the explicit unshifted look-ahead sequence here. Relate the buffer and parameter variables first, and match learning-rate schedule, initialization, weight decay, and dampening. Restore optimizer buffers as well as parameters when resuming training.

Varying momentum and learning rates. A scheduled Nesterov method chooses coefficients to match a convergence theorem; a constant coefficient is a separate variant. Scheduling a learning rate in a gradient-buffer formulation changes how old gradients are scaled relative to storing historical displacements. Restarting can zero the velocity when a chosen diagnostic triggers, but it changes the recurrence and needs a stated rule and associated cost. Returning a best-so-far point can protect a reporting metric while requiring separate storage; it must not silently overwrite the actual optimizer state.

Stochastic gradients and regularization. Use one mini-batch at the chosen probe and average its gradients consistently. Noise can also accumulate through momentum, so a larger coefficient need not improve training. To evaluate an explicit look-ahead gradient with automatic differentiation, compute it at the temporary parameter values and restore or commit parameters carefully. Differentiating through the optimizer is another task. Coupled $L_2$ regularization adds a parameter-dependent term to the gradient; decoupled weight decay performs a distinct parameter update and is not interchangeable under momentum.

Conditioning and limits. Standardizing features or preconditioning a quadratic can reduce the curvature disparity before tuning momentum. A scalar preconditioned curvature analysis is then performed on the transformed system. The programs deliberately use the conservative $1/L$ guard; heavy ball has a larger stable interval on an individual quadratic mode, so a rejected rate is not a universal proof of divergence. Reusing a learning rate from plain gradient descent without checking the full momentum recurrence is also unsafe. Boundaries, nonsmooth terms, and constraints require suitable projection or proximal methods and their own analysis.

Practice Problems and Real-World Use

Use the sample trace as a small regression check before experimenting with a larger problem. Compare methods under the same objective, starting point, gradient-work budget, and reporting policy. Measure full-loss monitoring separately from training gradients when the dataset is large. For a convergence comparison, specify whether you report the final loss, the minimum observed loss, a current-point gradient norm, or elapsed time.

The following exercises build from arithmetic to optimizer engineering. The objectives and checks are explicit so that a passing trace tests a mathematical property, not just matching a particular implementation.