Gradient descent
You will learn
How gradient descent walks down a loss, what momentum adds, and why too large a rate never settles.
Training is walking downhill on the loss. Gradient descent moves a weight by -lr times its gradient each step; in the spec the loss is (x - 3)^2 from x = 0, and lr 0.1 reaches 2.99996 in 50 steps. Momentum adds a share of the last move, which speeds up a long slope and overshoots a sharp bottom: at 0.9 the weight shoots past 3 to about 5.13 and rings down to 3.09 after 50 steps. Nesterov takes the gradient one step ahead and swings less, to about 4.25. Too large a rate never settles: at lr 1 the weight bounces between 0 and 6 for good. The spec sgd.t27 computes every step.
Try it
Find the largest learning rate that still settles without momentum; then add momentum 0.9 and count the swings.
Set the learning rate and momentum and watch a weight walk from 0 to the bottom of (x - 3)^2. Plain steps creep, momentum overshoots and rings, and a rate of 1 bounces forever.
specs/ml/optimizer/sgd.t27
// SPDX-License-Identifier: Apache-2.0
// t27/specs/
// Nesterov accelerated gradient (NAG) | φ² + 1/φ² = 3 | TRINITY
// #8980: init called `new f32[n]`; update read undeclared locals; dampening and
// nesterov were ignored; and weight decay was ADDED to the parameter
// (`- v + wd * p`), which grows the weights it should shrink. update is now
// PyTorch's SGD: g = grad + wd p; v = mu v + (1 - dampening) g;
// step = g + mu v (Nesterov) or v; p = p - lr step.
module Sgd;
use base::types;
use math::constants;
// ═══════════════════════════════════════════════════════════
// 1. Constants
// ═══════════════════════════════════════════════════════════
const DEFAULT_LR : f32 = 0.01;
const DEFAULT_MOMENTUM : f32 = 0.0;
// ═══════════════════════════════════════════════════════════
// 2. Types
// ═══════════════════════════════════════════════════════════
pub const SGDConfig = struct {
learning_rate : f32,
momentum : f32,
weight_decay : f32,
dampening : f32,
nesterov : bool,
};
pub const SGDState = struct {
velocity : []f32,
};
// ═══════════════════════════════════════════════════════════
// 3. Core Functions
// ═══════════════════════════════════════════════════════════
// init: the first param_count velocities start at zero (the caller owns the buffer)
fn init(state: SGDState, param_count: u32) -> void {
var i : usize = 0;
while (i < param_count) {
state.velocity[i] = 0.0;
i = i + 1;
}
}
fn update(params: []f32, grads: []const f32, state: SGDState, config: SGDConfig) -> void {
var i : usize = 0;
while (i < params.len) {
const g = grads[i] + config.weight_decay * params[i];
var step = g;
if (config.momentum > 0.0) {
state.velocity[i] = config.momentum * state.velocity[i] + (1.0 - config.dampening) * g;
if (config.nesterov) {
step = g + config.momentum * state.velocity[i];
} else {
step = state.velocity[i];
}
}
params[i] = params[i] - config.learning_rate * step;
i = i + 1;
}
}
fn cfg(lr: f32, momentum: f32, wd: f32, nesterov: bool) -> SGDConfig {
return SGDConfig{ .learning_rate = lr, .momentum = momentum, .weight_decay = wd, .dampening = 0.0, .nesterov = nesterov };
}
// Minimise f(x) = (x - 3)^2 (gradient 2 (x - 3)) from x = 0 for `steps` steps
fn descend(lr: f32, momentum: f32, nesterov: bool, steps: u32) -> f32 {
var params : [1]f32 = [0.0];
var velocity : [1]f32 = [5.0];
const state = SGDState{ .velocity = &velocity };
init(state, 1);
var n : u32 = 0;
while (n < steps) {
const grads : [1]f32 = [2.0 * (params[0] - 3.0)];
update(¶ms, &grads, state, cfg(lr, momentum, 0.0, nesterov));
n = n + 1;
}
return params[0];
}
fn init_zeroes(index: usize) -> f32 {
var velocity : [3]f32 = [7.0, 7.0, 7.0];
init(SGDState{ .velocity = &velocity }, 3);
return velocity[index];
}
// One step on params (1, 2), grads (0.1, 0.2), momentum 0.9, weight decay 0.001
fn one_step(index: usize) -> f32 {
var params : [2]f32 = [1.0, 2.0];
const grads : [2]f32 = [0.1, 0.2];
var velocity : [2]f32 = [0.0, 0.0];
update(¶ms, &grads, SGDState{ .velocity = &velocity }, cfg(0.01, 0.9, 0.001, false));
return params[index];
}
// A step with zero gradient and weight decay shrinks the weight toward 0
fn decay_step() -> f32 {
var params : [1]f32 = [2.0];
const grads : [1]f32 = [0.0];
var velocity : [1]f32 = [0.0];
update(¶ms, &grads, SGDState{ .velocity = &velocity }, cfg(0.1, 0.0, 0.5, false));
return params[0];
}
fn valid_input() -> f32 {
return DEFAULT_LR;
}
// ═══════════════════════════════════════════════════════════
// TDD: Tests (from .tri behaviors)
// ═══════════════════════════════════════════════════════════
test "init_basic_case"
given first = init_zeroes(0)
when last = init_zeroes(2)
then first == 0.0 and last == 0.0
// v = g + wd p (momentum starts at 0): p0 = 1 - 0.01 (0.1 + 0.001) = 0.99899
test "update_basic_case"
given p0 = one_step(0)
when p1 = one_step(1)
then p0 != 1.0 and p1 != 2.0 and @abs(p0 - 0.99899) < 0.000001 and @abs(p1 - (2.0 - 0.01 * 0.202)) < 0.000001
// Weight decay pulls the weight toward 0: 2 - 0.1 * 0.5 * 2 = 1.9
test "update_weight_decay_shrinks"
given result = decay_step()
then @abs(result - 1.9) < 0.000001
// Plain, momentum and Nesterov SGD all reach the minimum at 3
test "update_descends_a_quadratic"
given plain = descend(0.1, 0.0, false, 100)
when heavy = descend(0.05, 0.9, false, 300)
then @abs(plain - 3.0) < 0.0001 and @abs(heavy - 3.0) < 0.0001 and @abs(descend(0.05, 0.9, true, 300) - 3.0) < 0.0001
// ═══════════════════════════════════════════════════════════
// TDD: Invariants (from .tri constraints)
// ═══════════════════════════════════════════════════════════
// learning_rate > 0
invariant sgd_constraint_0
assert valid_input() > 0.0
// 0 <= momentum < 1
invariant sgd_constraint_1
assert DEFAULT_MOMENTUM >= 0.0 and DEFAULT_MOMENTUM < 1.0
// weight_decay >= 0: a non-negative decay never grows a weight with zero gradient
invariant sgd_constraint_2
assert decay_step() < 2.0
// params.size() == grads.size(): every parameter moves
invariant sgd_constraint_3
assert one_step(0) != 1.0 and one_step(1) != 2.0
// If momentum > 0: state.velocity.size() == params.size() -- init zeroes all of it
invariant sgd_constraint_4
assert init_zeroes(1) == 0.0
All lessons
Module 1 · One neuron
A neuron sums its inputs and bends the sum: ReLU, the sigmoid and tanh, and where each one stops passing a gradient back.
Module 2 · New bends
The smooth curves of current models: GELU with its dip, Swish with its knob, and ELU with a floor instead of a cliff.
Module 3 · A layer
What sits around the neurons: softmax turns scores into a choice, an embedding turns a word into numbers, dropout switches units off.
Module 4 · Pictures
Reading a picture in patches: a window averaged over its patch, max pooling, and batch normalisation of features on different scales.
Module 5 · The error
How wrong a guess is, as one number: mean squared error, cross-entropy for probabilities, and the Huber loss for outliers.
Module 6 · Downhill
Walking down the loss: gradient descent with momentum, RMSProp that ignores the gradient's size, and Adam that combines them.
Module 7 · Learning
Where the gradients come from and what they achieve: backpropagation through a layer, a network that learns XOR, and LAMB for large batches.
Module 8 · Memory and attention
Networks that read a sequence: a recurrent cell and its fading echo, the gates of an LSTM, and attention over every position at once.
Module 9 · The transformer and the reward
The parts of a transformer and of an agent: positions as clocks, one encoder block, and PPO, which learns from reward without lurching.