AI Concepts in Python

Neural Network Optimizers: SGD, Momentum and Adam in Python

An optimizer is the rule that turns gradients into weight updates. Plain gradient descent uses only the current gradient; momentum remembers past directions, and Adam also adapts the step size for each weight.

The concept in plain words

An optimizer is the rule that turns gradients into weight updates. Plain gradient descent uses only the current gradient; momentum remembers past directions, and Adam also adapts the step size for each weight.

Compare plain gradient descent, momentum and Adam by minimising the same function step by step.

Key ideas

  • Momentum keeps a velocity, so it rolls through flat areas and small bumps.
  • Adam keeps running averages of the gradient and of the squared gradient.
  • Bias correction (dividing by 1 − βᵗ) fixes those averages in the first steps.
  • Good weight initialization and mini-batches also make training faster and more stable.

When should you use it?

  • Understand what happens inside neural networks and AI libraries.
  • Learn the idea before you implement it in the matching challenges.
  • Prepare for machine learning courses, interviews and projects.

See it in action

Press Play to race three optimizers down the same bumpy valley. Try a larger learning rate to see momentum overshoot.

Example code

Run code →
main.py
import math

# Minimise f(x) = (x - 3)² with three optimizers. The gradient is 2(x - 3).
grad = lambda x: 2 * (x - 3)

def sgd(steps, lr=0.1):
    x = 0.0
    for _ in range(steps):
        x -= lr * grad(x)
    return x

def momentum(steps, lr=0.1, beta=0.9):
    x, v = 0.0, 0.0
    for _ in range(steps):
        v = beta * v + grad(x)        # remember the direction we were going
        x -= lr * v
    return x

def adam(steps, lr=0.3, b1=0.9, b2=0.999, eps=1e-8):
    x, m, v = 0.0, 0.0, 0.0
    for t in range(1, steps + 1):
        g = grad(x)
        m = b1 * m + (1 - b1) * g     # running mean of gradients
        v = b2 * v + (1 - b2) * g * g # running mean of squared gradients
        m_hat, v_hat = m / (1 - b1 ** t), v / (1 - b2 ** t)
        x -= lr * m_hat / (math.sqrt(v_hat) + eps)
    return x

for steps in (5, 20, 50):
    print(f"{steps:>2} steps  sgd={sgd(steps):.3f}  momentum={momentum(steps):.3f}  adam={adam(steps):.3f}")

Expected output

5 steps  sgd=2.017  momentum=4.741  adam=1.476
20 steps  sgd=2.965  momentum=4.048  adam=3.813
50 steps  sgd=3.000  momentum=3.091  adam=3.014

How it works

All three minimise (x − 3)², whose lowest point is x = 3. SGD moves steadily, momentum builds up speed and overshoots before settling, and Adam scales each step by running averages of the gradient and its square. After 50 steps all three are close to 3.

Change the values and run the program in the CodeUtility online Python compiler without installing Python locally.

Practice exercises

Modify the runnable example with the exercises below to build understanding beyond copying the result.

  1. Change the inputs or weights and predict the output before running.
  2. Print the intermediate values to follow each step of the calculation.
  3. Solve the linked challenges below to implement the idea yourself.
Run in Python IDE →

Practice with a challenge

Use what this lesson covers to solve a real problem. Each challenge runs in your browser with instant tests and an AI review.

All challenges