Standard library · AI

ai.optim

Neural network optimisers: SGD, Adam, RMSProp and learning rate schedules.

import ai.optim

On this page

Choosing an optimiser
Learning rate schedules
Clipping gradients
The whole recipe on a real model
Types
Functions

An optimiser turns gradients into better parameters. This package has the three that cover almost all neural network training (SGD with optional momentum, Adam and AdamW, RMSProp), gradient clipping, and learning rate schedules. It works on the parameters of an nn ParamStore whose gradients autograd has filled in. For classical models that train themselves, such as linear regression, trees or k-means, use ml instead.

Every training loop has the same five steps:

tape.reset() and store.zero_grad()
run the model and compute the loss
autograd.Backward(tape, loss)
optionally ClipGradNorm, and set lr from a schedule
opt.step(store)

Things that trip people up

step does not clear the gradients. Gradients add up, so forgetting zero_grad makes every step use the sum of all earlier ones.
Momentum and Adam keep one state tensor per parameter, matched by position in the store and created on the first step. Use one optimiser per store and add every parameter before training starts.
weight_decay applies to every parameter in the store, biases and norm scales included.
The defaults differ: lr is 0.01 for SGD and RMSProp but 0.001 for Adam. The updates are computed in f32, like the tensors.
Optimisers with state hold memory; call free when training is done.

Choosing an optimiser

The three optimisers here minimize the same bowl, x² + 50y², from the same start for forty steps. The bowl is fifty times steeper in one direction than the other, which is what real losses look like. Plain SGD must keep its learning rate below the limit the steep direction allows (here anything at or above 0.02 never settles), so it crawls along the flat one. Momentum keeps a running velocity, carrying the step along the flat direction while the bouncing in the steep one cancels out.

optim-guide-compare.salam

import ai.optim
import ai.nn
import ai.autograd
import ai.tensor
import conv

func run(kind: int): f64:
    mut store := nn.ParamStore {}
    start := tensor.Full(tensor.Shape1(2), 1.0)
    handle := store.add("xy", start)
    mut scale_data := Vector {} as Vector<f32>
    scale_data.push(1.0 as f32)
    scale_data.push(50.0 as f32)
    mut scale_shape := tensor.Shape1(2)
    mut scale := tensor.FromArray(scale_shape, scale_data)
    mut plain := optim.SGD { lr = 0.019 }
    mut heavy := optim.SGD { lr = 0.01, momentum = 0.9 }
    mut adam := optim.Adam { lr = 0.1 }
    mut tape := autograd.NewTape(1)
    mut last := 0.0
    repeat 40:
        tape.reset()
        store.zero_grad()
        p := nn.Bind(tape, store, handle)
        loss := autograd.Sum(tape, autograd.Mul(tape, autograd.Constant(tape, scale), autograd.Square(tape, p)))
        autograd.Backward(tape, loss)
        last = autograd.Value(tape, loss).at(0)
        if kind == 0: plain.step(store) end
        if kind == 1: heavy.step(store) end
        if kind == 2: adam.step(store) end
    end
    tape.free()
    plain.free()
    heavy.free()
    adam.free()
    store.free()
    ret last
end

func main:
    names := ["SGD", "SGD with momentum", "Adam"]
    repeat 3 in kind:
        println names[kind], conv.FormatFloatPrec(run(kind), 4)
    end
end

OutputSGD 0.0622 SGD with momentum 0.0063 Adam 0.5352

Adam scales each weight's step by its own gradient history, so it needs the least tuning and is the usual first choice for a network. On this tiny problem its steps of about lr make it hover near the minimum instead of settling, which is a good reminder to lower the rate as training goes on. NewAdamW is Adam with decoupled weight decay, the common choice for transformers.

Learning rate schedules

A schedule is just a function from the step number to a learning rate; you assign the result to the optimiser's lr field before each step. WarmupCosine climbs linearly for the warmup steps, then follows half a cosine down to the minimum. StepLR multiplies the rate by gamma every so many steps.

optim-guide-schedule.salam

import ai.optim
import conv

func main:
    base := 0.01
    total := 100 as i64
    warmup := 10 as i64
    println "step  warmup+cosine  step decay"
    each step in [0, 4, 9, 10, 30, 55, 80, 99]:
        at := step as i64
        smooth := optim.WarmupCosine(base, at, warmup, total, 0.0001)
        stairs := optim.StepLR(base, at, 30, 0.5)
        println step, conv.FormatFloatPrec(smooth, 5), conv.FormatFloatPrec(stairs, 5)
    end
end

Outputstep warmup+cosine step decay 0 0.00100 0.01000 4 0.00500 0.01000 9 0.01000 0.01000 10 0.01000 0.01000 30 0.00884 0.00500 55 0.00505 0.00500 80 0.00126 0.00250 99 0.00010 0.00125

Warmup matters most for Adam and for large models: the first gradients are noisy and the moment estimates are still empty, so full sized steps at the start can throw the weights far off.

Clipping gradients

One bad batch can produce a huge gradient and undo many steps of training. ClipGradNorm measures the length of all gradients together and, when it is above the limit, scales them down to that length while keeping their direction. It returns the length before clipping, which is worth logging: a norm that keeps growing is an early sign the learning rate is too high. Call it after Backward and before step.

optim-guide-clip.salam

import ai.optim
import ai.nn
import ai.autograd
import ai.tensor
import conv

func main:
    mut store := nn.ParamStore {}
    handle := store.add("w", tensor.Full(tensor.Shape1(2), 1.0))
    mut tape := autograd.NewTape(1)
    mut sgd := optim.SGD { lr = 0.1 }
    repeat 3 in step:
        tape.reset()
        store.zero_grad()
        sgd.lr = optim.StepLR(0.1, step as i64, 1, 0.5)
        w := nn.Bind(tape, store, handle)
        loss := autograd.Sum(tape, autograd.Scale(tape, autograd.Square(tape, w), 500.0))
        autograd.Backward(tape, loss)
        before := optim.ClipGradNorm(store, 1.0)
        mut p := store.at(handle)
        println "step", step, "lr", sgd.lr, "norm", conv.FormatFloatPrec(before, 1), "clipped grad", conv.FormatFloatPrec(p.grad.at(0), 3), conv.FormatFloatPrec(p.grad.at(1), 3)
        sgd.step(store)
        println "  w is now", conv.FormatFloatPrec(p.value.at(0), 3), conv.FormatFloatPrec(p.value.at(1), 3)
    end
    tape.free()
    sgd.free()
    store.free()
end

Outputstep 0 lr 0.1 norm 1414.2 clipped grad 0.707 0.707 w is now 0.929 0.929 step 1 lr 0.05 norm 1314.2 clipped grad 0.707 0.707 w is now 0.894 0.894 step 2 lr 0.025 norm 1264.2 clipped grad 0.707 0.707 w is now 0.876 0.876

The loss is 500w² for each weight, so the raw gradient is about 1400 long, and clipping cuts it to exactly 1. The program also halves the learning rate every step with StepLR. ClipGradValue is the blunter tool: it caps every single gradient value, which changes the direction too.

The whole recipe on a real model

Here everything comes together on a small delivery-time model: a Linear layer from nn learns how many minutes each kilometre and each item adds. AdamW trains it, WarmupCosine sets the rate before every step, and ClipGradNorm guards against a bad start. The inputs and targets are divided by 10 first; optimisers work best when the numbers are near 1.

optim-guide-recipe.salam

import ai.optim
import ai.nn
import ai.autograd
import ai.tensor
import conv

func main:
    km := [1.0, 2.5, 4.0, 5.5, 7.0, 8.5, 3.0, 6.0]
    items := [1.0, 3.0, 2.0, 5.0, 1.0, 4.0, 6.0, 2.0]
    minutes := [14.0, 22.5, 26.0, 37.5, 37.0, 48.5, 33.0, 36.0]
    mut x := tensor.Zeros2(8, 2)
    mut y := tensor.Zeros2(8, 1)
    repeat 8 in i:
        x.set_at2(i, 0, km[i] / 10.0)
        x.set_at2(i, 1, items[i] / 10.0)
        y.set_at2(i, 0, minutes[i] / 10.0)
    end
    mut rng := tensor.NewRng(1)
    mut store := nn.ParamStore {}
    model := nn.NewLinear(store, "eta", 2, 1, rng)
    mut opt := optim.NewAdamW(0.1, 0.01)
    mut tape := autograd.NewTape(1)
    total := 400 as i64
    repeat 400 in step:
        tape.reset()
        store.zero_grad()
        guess := model.forward(tape, store, autograd.Input(tape, x))
        loss := autograd.MSELoss(tape, guess, autograd.Constant(tape, y))
        autograd.Backward(tape, loss)
        norm := optim.ClipGradNorm(store, 5.0)
        opt.lr = optim.WarmupCosine(0.1, step as i64, 20, total, 0.001)
        opt.step(store)
        if step % 100 == 0 or step == 399:
            println "step", step, "lr", conv.FormatFloatPrec(opt.lr, 4), "loss", conv.FormatFloatPrec(autograd.Value(tape, loss).at(0), 4), "grad norm", conv.FormatFloatPrec(norm, 3)
        end
    end
    w := store.at(model.w).value
    b := store.at(model.b).value
    println "minutes =", conv.FormatFloatPrec(b.at(0) * 10.0, 1), "+", conv.FormatFloatPrec(w.at(0), 2), "per km +", conv.FormatFloatPrec(w.at(1), 2), "per item"
    tape.free()
    opt.free()
    store.free()
    x.free()
    y.free()
end

Outputstep 0 lr 0.0050 loss 15.9842 grad norm 8.950 step 100 lr 0.0896 loss 0.2397 grad norm 0.208 step 200 lr 0.0546 loss 0.0787 grad norm 0.109 step 300 lr 0.0170 loss 0.0456 grad norm 0.078 step 399 lr 0.0010 loss 0.0411 grad norm 0.073 minutes = 11.8 + 3.12 per km + 1.83 per item

The first gradient norm is almost 9, so the clip at 5 trims it, and after that the norm never comes near the limit. The rate climbs for 20 steps, peaks and falls to 0.001 by the last step, where the loss has flattened out. Because the inputs were scaled by the same factor as the target, the weights read directly as minutes per kilometre and per item; only the bias has to be multiplied back by 10.

Types

`Adam`

struct Adam

The Adam optimiser; a good default for most networks.

Name	Type	Description
`lr`	`f64 = 0.001`	The learning rate; 0.001 by default.
`beta1`	`f64 = 0.9`	The coefficient of the gradient average.
`beta2`	`f64 = 0.999`	The coefficient of the squared gradient average.
`eps`	`f64 = 1e-08`	A small number that keeps a division by zero from happening.
`weight_decay`	`f64 = 0`	Weight decay.
`decoupled`	`bool = false`	Apply the weight decay apart from the gradient, meaning AdamW.

ai-optim-Train.salam

import ai.optim
import ai.nn
import ai.autograd
import ai.tensor

func train(method: int, steps: int, target &: tensor.Tensor): f64:
    mut store := nn.ParamStore {}
    handle := store.add("w", tensor.Zeros1(4))
    mut adam := optim.Adam { lr = 0.1 }
    mut adamw := optim.NewAdamW(0.1, 0.01)
    mut rms := optim.RMSProp { lr = 0.05 }
    mut sgd := optim.SGD { lr = 0.05, momentum = 0.9 }
    mut tape := autograd.NewTape(1)
    repeat steps:
        tape.reset()
        store.zero_grad()
        w := nn.Bind(tape, store, handle)
        t := autograd.Input(tape, target)
        loss := autograd.Sum(tape, autograd.Square(tape, autograd.Sub(tape, w, t)))
        autograd.Backward(tape, loss)
        _ := optim.ClipGradNorm(store, 10.0)
        optim.ClipGradValue(store, 5.0)
        if method == 0:
            adam.step(store)
        else method == 1:
            adamw.step(store)
        else method == 2:
            rms.step(store)
        else:
            sgd.step(store)
        end
    end
    mut p := store.at(handle)
    result := p.value.at(3)
    tape.free()
    adam.free()
    adamw.free()
    rms.free()
    sgd.free()
    store.free()
    ret result
end

func main:
    mut target := tensor.Arange(4)
    each method in [0, 1, 2, 3]:
        rounded := (train(method, 400, target) * 100.0 + 0.5) as int
        println method, rounded
    end
    tensor.Free(target)
    println optim.StepLR(1.0, 25, 10, 0.5), optim.CosineLR(1.0, 50, 100, 0.0), optim.WarmupCosine(1.0, 0, 10, 110, 0.0)
end

Output0 300 1 299 2 300 3 300 0.25 0.5 0.1