LAMB and large batches
You will learn
How LAMB scales each layer's step by its own weight norm, so one learning rate suits layers of every size.
Training with very large batches needs a large learning rate, and one rate can be too large for a small layer and too small for a big one. LAMB takes an Adam-style update for each layer and rescales it by the trust ratio, the layer's weight norm, capped at 10, over the update's norm, so every layer moves by the same fraction of itself. In the spec, the weights (1, 2) times a scale take a step whose norm is 0.001 times theirs: 0.00447 at scale 2, and once the weight norm passes 10, near scale 4.47, the step stays at 0.01. The spec lamb.t27 computes the update for one layer and for two at once.
Try it
Read the step at scales 1 and 2 and compare; then drag past the bend and explain why the line goes flat.
Scale a layer's weights and watch its LAMB step: it grows with the weight norm, so every layer moves by the same fraction of itself, until the norm passes 10 and the step stops growing.
specs/ml/optimizer/lamb.t27
// SPDX-License-Identifier: Apache-2.0
// t27/specs/
// LAMB parameter update with trust ratio clipping | φ² + 1/φ² = 3 | TRINITY
// #8980: built with `append` on `[]` and `range` / `len` / `pow`, so it never
// compiled. The math was also not LAMB: v_hat used beta1's bias correction, the
// "trust ratio" was ||g|| / ||v_hat|| and was computed, then never applied to the
// step, and m and v were never written back. It is now You et al. (2020):
// r = m_hat / (sqrt(v_hat) + eps) + wd * w, trust = min(||w||, clip) / ||r||,
// w = w - lr * trust * r,
// per layer, over buffers the caller owns.
module Lamb;
use base::types;
use math::constants;
use ml::optimizer::adam;
// ═══════════════════════════════════════════════════════════
// 1. Constants
// ═══════════════════════════════════════════════════════════
const DEFAULT_LR : f32 = 0.001;
const DEFAULT_BETA1 : f32 = 0.9;
const DEFAULT_BETA2 : f32 = 0.999;
const DEFAULT_EPSILON : f32 = 1e-6;
const DEFAULT_CLIP_THRESHOLD : f32 = 10.0;
// ═══════════════════════════════════════════════════════════
// 2. Types
// ═══════════════════════════════════════════════════════════
pub const LAMBConfig = struct {
learning_rate : f32,
beta1 : f32,
beta2 : f32,
epsilon : f32,
weight_decay : f32,
clip_threshold : f32,
};
pub const LAMBState = struct {
m : []f32,
v : []f32,
step : u32,
};
// ═══════════════════════════════════════════════════════════
// 3. Core Functions
// ═══════════════════════════════════════════════════════════
// m and v are the first num_params entries of the given buffers, zeroed; step 0
fn init_state(m: []f32, v: []f32, num_params: u32) -> LAMBState {
var i : usize = 0;
while (i < num_params) {
m[i] = 0.0;
v[i] = 0.0;
i = i + 1;
}
return LAMBState{ .m = m[0..num_params], .v = v[0..num_params], .step = 0 };
}
// powi (b^t for a whole t) comes from specs/ml/optimizer/adam.t27 through `use`
// Writes the moments for step (state.step + 1) into state.m / state.v and the update to
// ADD to params into `update`; returns update[0..n]
fn compute_layer_update(params: []const f32, grads: []const f32, state: LAMBState, config: LAMBConfig, update: []f32) -> []f32 {
const n = params.len;
const t = state.step + 1;
const c1 = 1.0 - powi(config.beta1, t);
const c2 = 1.0 - powi(config.beta2, t);
var w_norm : f32 = 0.0;
var r_norm : f32 = 0.0;
var i : usize = 0;
while (i < n) {
state.m[i] = config.beta1 * state.m[i] + (1.0 - config.beta1) * grads[i];
state.v[i] = config.beta2 * state.v[i] + (1.0 - config.beta2) * grads[i] * grads[i];
const r = (state.m[i] / c1) / (sqrt(state.v[i] / c2) + config.epsilon) + config.weight_decay * params[i];
update[i] = r;
w_norm = w_norm + params[i] * params[i];
r_norm = r_norm + r * r;
i = i + 1;
}
w_norm = sqrt(w_norm);
r_norm = sqrt(r_norm);
var trust : f32 = 1.0;
if (w_norm > 0.0 and r_norm > 0.0) {
trust = @min(w_norm, config.clip_threshold) / r_norm;
}
i = 0;
while (i < n) {
update[i] = -config.learning_rate * trust * update[i];
i = i + 1;
}
return update[0..n];
}
// One LAMB step over several layers packed back to back in `params` (layer k has
// sizes[k] entries), each with its own state; params change in place, each step + 1.
fn forward(params: []f32, grads: []const f32, sizes: []const usize, states: []LAMBState, config: LAMBConfig, scratch: []f32) -> void {
var start : usize = 0;
var k : usize = 0;
while (k < sizes.len) {
const end = start + sizes[k];
const upd = compute_layer_update(params[start..end], grads[start..end], states[k], config, scratch);
var j : usize = 0;
while (j < sizes[k]) {
params[start + j] = params[start + j] + upd[j];
j = j + 1;
}
states[k].step = states[k].step + 1;
start = end;
k = k + 1;
}
}
fn cfg(lr: f32) -> LAMBConfig {
return LAMBConfig{ .learning_rate = lr, .beta1 = DEFAULT_BETA1, .beta2 = DEFAULT_BETA2, .epsilon = DEFAULT_EPSILON, .weight_decay = 0.01, .clip_threshold = DEFAULT_CLIP_THRESHOLD };
}
// which: 0 = m.len, 1 = v.len, 2 = step, 3 = m[0], 4 = v[99]
fn init_probe(which: u8) -> f32 {
var m : [100]f32 = [9.0] ** 100;
var v : [100]f32 = [9.0] ** 100;
const st = init_state(&m, &v, 100);
if (which == 0) {
return @as(f32, @floatFromInt(st.m.len));
}
if (which == 1) {
return @as(f32, @floatFromInt(st.v.len));
}
if (which == 2) {
return @as(f32, @floatFromInt(st.step));
}
if (which == 3) {
return st.m[0];
}
return st.v[99];
}
// update[index] of the first step on params (1, 2, 3), grads (0.1, 0.2, 0.3); index 3 = its length
fn first_update(index: usize) -> f32 {
const params : [3]f32 = [1.0, 2.0, 3.0];
const grads : [3]f32 = [0.1, 0.2, 0.3];
var m : [3]f32 = undefined;
var v : [3]f32 = undefined;
var update : [3]f32 = [9.0] ** 3;
const u = compute_layer_update(¶ms, &grads, init_state(&m, &v, 3), cfg(0.001), &update);
if (index == 3) {
return @as(f32, @floatFromInt(u.len));
}
return u[index];
}
// The layer-wise property: scaling a layer's weights by 10 scales its step by 10
fn step_norm(scale: f32) -> f32 {
const params : [2]f32 = [1.0 * scale, 2.0 * scale];
const grads : [2]f32 = [0.1, 0.2];
var m : [2]f32 = undefined;
var v : [2]f32 = undefined;
var update : [2]f32 = undefined;
const u = compute_layer_update(¶ms, &grads, init_state(&m, &v, 2), cfg(0.001), &update);
return sqrt(u[0] * u[0] + u[1] * u[1]);
}
// Two layers (1, 2) and (3, 4), one step; which: 0..3 = new params, 4 = step of layer 1
fn two_layers(which: usize) -> f32 {
var params : [4]f32 = [1.0, 2.0, 3.0, 4.0];
const grads : [4]f32 = [0.1, 0.2, 0.3, 0.4];
const sizes : [2]usize = [2, 2];
var m : [4]f32 = undefined;
var v : [4]f32 = undefined;
var states : [2]LAMBState = [init_state(m[0..2], v[0..2], 2), init_state(m[2..4], v[2..4], 2)];
var scratch : [2]f32 = undefined;
forward(¶ms, &grads, &sizes, &states, cfg(0.001), &scratch);
if (which == 4) {
return @as(f32, @floatFromInt(states[1].step));
}
return params[which];
}
fn valid_input() -> f32 {
return DEFAULT_BETA1;
}
// ═══════════════════════════════════════════════════════════
// TDD: Tests (from .tri behaviors)
// ═══════════════════════════════════════════════════════════
test init_state_basic_case
given num_params = 100.0
when m_len = init_probe(0)
then m_len == num_params and init_probe(1) == num_params and init_probe(2) == 0.0 and init_probe(3) == 0.0 and init_probe(4) == 0.0
// Every update is negative (descent) and the step's norm is lr * min(||w||, 10) = 0.001 ||w||
test compute_layer_update_basic_case
given u0 = first_update(0)
when n = first_update(3)
then n == 3.0 and u0 < 0.0 and first_update(1) < 0.0 and first_update(2) < 0.0 and @abs(sqrt(u0 * u0 + first_update(1) * first_update(1) + first_update(2) * first_update(2)) - 0.001 * sqrt(@as(f32, 14.0))) < 0.0000001
// Each layer moves, and each layer's state advances
test forward_basic_case
given p0 = two_layers(0)
when p3 = two_layers(3)
then p0 != 1.0 and two_layers(1) != 2.0 and two_layers(2) != 3.0 and p3 != 4.0 and p0 < 1.0 and two_layers(4) == 1.0
// LAMB's point: the step is proportional to the layer's weight norm (up to the clip)
test compute_layer_update_scales_with_the_layer
given small = step_norm(1.0)
when large = step_norm(3.0)
then @abs(large / small - 3.0) < 0.001
// ═══════════════════════════════════════════════════════════
// TDD: Invariants (from .tri constraints)
// ═══════════════════════════════════════════════════════════
// 0 < beta1 < 1
invariant lamb_constraint_0
assert valid_input() > 0.0 and valid_input() < 1.0
// 0 < beta2 < 1
invariant lamb_constraint_1
assert DEFAULT_BETA2 > 0.0 and DEFAULT_BETA2 < 1.0
// epsilon > 0
invariant lamb_constraint_2
assert DEFAULT_EPSILON > 0.0
// clip_threshold > 0
invariant lamb_constraint_3
assert DEFAULT_CLIP_THRESHOLD > 0.0
// name: trust_ratio -- capped by clip_threshold: a layer of norm 50 steps as one of norm 10
invariant lamb_constraint_4
assert DEFAULT_CLIP_THRESHOLD == 10.0
// name: lamb_update -- bias correction makes the first m_hat equal g
invariant lamb_constraint_5
assert powi(DEFAULT_BETA1, 1) == DEFAULT_BETA1 and DEFAULT_LR > 0.0
// "You et al. (2020) - Large Batch Optimization for Deep Learning: Training BERT in 76 minutes"
invariant lamb_constraint_6
assert powi(0.5, 3) == 0.125
// "Used in BERT, T5, and other large-scale Transformer training"
invariant lamb_constraint_7
assert DEFAULT_BETA2 > DEFAULT_BETA1
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.