Losses and optimizers

In short: every loss is built from autograd ops (two hand-written nodes where composing ops would get the gradient wrong) and ends in a shared Reduce step: mean, sum or none. Classification losses take logits, so log-softmax and log-sigmoid never overflow; cross-entropy adds ignored rows and label smoothing. There are regression losses (MSE, MAE, Huber), pair losses (contrastive, InfoNCE) and the two VAE pieces. Optimizers share one base class that skips parameters without a gradient, keeps per-parameter state in named slots for checkpoints, and calls one virtual Update per parameter: SGD with momentum and Nesterov, Adam, AdamW and RMSProp. Schedules are pure functions of the step; clipping and gradient accumulation are small helpers around the same loop.

The concepts are on /ml-models: loss functions, entropy and KL, variational autoencoders, optimizers, learning-rate schedules, the training loop and gradient accumulation.

One training step

The optimizer header gives the loop: ZeroGrad(), forward, loss.Backward(), Step() (optimizer.h L4–9). The optimizer holds the same Variables as the model, which share storage with it, so updating them in place updates the model.

Reductions

Every loss ends in Reduce(per_element, reduction, count = 0). kMean, the default, averages so the scale ignores batch size; kSum adds; kNone returns the per-element losses. When count is positive, kMean divides the sum by it instead of by the element count, which is how cross-entropy averages over the rows it did not ignore (reduction.cc L8–19).

Classification losses

These take logits, not probabilities, so the log of a softmax or sigmoid is computed in a form that cannot overflow or take log(0) (classification.h L4–8).

Why binary cross-entropy is a hand-written node

Composing max, abs and log1p from autograd ops gets the gradient wrong at x = 0, where abs has a kink that cancels in the true derivative. So SigmoidCrossEntropy is a Function whose forward loop is the stable formula and whose backward is σ(x) − y directly; the target gets no gradient (L29–56). A test checks that gradient at zero (GradientIsSigmoidMinusTargetIncludingAtZero). How a custom Function plugs into the graph is on Autograd.

Cross-entropy step by step

1. For each row, check the target is in range or equal to ignore_index, and record the flat index row·C + target (0 for an ignored row) and a keep mask of 1 or 0 (L80–96).
2. Take LogSoftmax over the class axis, flatten it, and gather the target entries with IndexSelect; negate (L98–101).
3. With smoothing ε, the loss against the uniform part is −mean_c log p_c, so each row becomes (1 − ε)·nll − ε·mean_c log p_c (L102–106).
4. Multiply by the mask and reduce, passing the number of kept rows (at least 1) as the mean's count (L107–108).

Worked example. Three rows with targets [2, −100, 0]: the middle row gathers index 0, is zeroed by the mask, and the mean divides by 2, not 3. With label smoothing 0.1 and C = 4, the effective target for a row whose label is 2 is 0.025 on classes 0, 1 and 3 and 0.925 on class 2.

Regression losses

Each treats the target as a constant and checks that the shapes match (L12–17). Huber is written with a clamp rather than a branch: for |r| = 3 and δ = 1, q = 1 and the loss is 0.5 + 2 = 2.5, the linear tail; for |r| = 0.5, q = 0.5 and the loss is 0.125, half the squared error.

Pair losses

Both add 10⁻¹² inside a square root, which keeps the gradient finite for a zero vector or a pair at distance zero (L23–26, L34–35). InfoNCE reuses cross-entropy, so it always takes the mean over the batch.

Variational pieces

A variational autoencoder's encoder outputs a mean and a log variance per latent dimension, defining q(z | x) = N(mean, exp(log_var)); training maximizes the ELBO, a reconstruction term estimated from one sample minus KL(q ‖ N(0, I)) (variational.h L9–14). rml provides the two pieces the autoencoder lacks:

The Optimizer base

NamedState is what a checkpoint stores under optim. keys next to the model's model. keys, so training resumes exactly where it stopped (checkpoint.h L22–32; see Model files). An Adam optimizer over two parameters that has stepped once saves step, exp_avg.0, exp_avg.1, exp_avg_sq.0 and exp_avg_sq.1.

The optimizers

The headers say why each exists. Momentum averages gradients over about 1/(1 − momentum) steps, damping oscillation across a ravine and speeding travel along it (sgd.h L7–13). Adam is momentum with a per-parameter step size; its L2 weight decay is weakened by the adaptive scaling for parameters with large gradients, which is why AdamW decouples it and why transformer training uses AdamW (adam.h L7–14). RMSProp divides each gradient by a running RMS so every parameter moves at a similar rate (rmsprop.h L4–11). Constructors check their ranges: Nesterov needs momentum > 0 and no dampening, Adam's betas lie in [0, 1), RMSProp's alpha too, and every eps is positive.

Worked example. Adam's first step: with m and v starting at zero, the corrected m̂ = g and v̂ = g², so the step is lr·g / (|g| + eps), almost exactly lr·sign(g). A parameter at 1 with gradient 100 and another with gradient −0.001 both move by 0.01 at lr = 0.01, to 0.99 and 1.01; that is the test FirstStepHasMagnitudeLr (optimizer_test.cc L84–92).

Learning-rate schedules

Schedules are pure functions of the step, 0 for the first, and the caller writes optimizer.set_lr(Schedule(step, ...)) before each Step() (schedule.h L4–6). Both curved schedules use one helper that moves from a start to an end value along half a cosine (schedule.cc L13–17).

Clipping and accumulation

Who uses it

rml/io's checkpoint is the one library user of rml/optim, through Optimizer* arguments that may be null to save only the model. The training demos in examples/ are the rest: the same nineteen directories that use rml/nn pick their losses and optimizers here. The inference stack never trains, so nothing on The inference stack uses either module.

Tests