Ster/Docs

Differentiable pass

Ster's decoder can run in two modes. Inference uses Candle's fused rope, softmax, and RMS-norm kernels; training swaps exactly those three for composed equivalents, because the fused ones have no backward pass.

Overview#

Everything in Ster except training is forward-only: it reads hidden states, fits a direction over them, and adds that direction back during decode. None of that needs a gradient with respect to a weight. Training does, and it cannot get one from the kernels the decode loop uses.

Why it has to exist#

Ster's decode loop leans on three fused Candle kernels, and every one of them ends in an apply_op*_no_bwd call, which records no node the autograd tape can walk back through. A loss built over them has no gradient at all — not a wrong gradient, no gradient:

Fused kernelComposed equivalentWhy the composition differentiates
candle_nn::rotary_emb::ropea hand-composed rotate-half ropebuilt from narrow, broadcast_mul, sub, add, and cat, each of which records a backward node.
candle_nn::ops::softmax_last_dimcandle_nn::ops::softmaxthe same softmax spelled out of max_keepdim, broadcast_sub, exp, sum_keepdim, and broadcast_div.
candle_nn::ops::rms_normRmsNorm::forward_diffthe same normalisation built from sqr, sum_keepdim, broadcast_div, and broadcast_mul.

Nothing else in the decoder changes. The gated feed-forward block needs no swap, because silu and the elementwise product already backpropagate.

Inference keeps the fused kernels#

The swap is selected per call, so inference never pays for it: ster generate, ster extract, ster vector train, ster vector optimize, and ster vector evaluate all run Pass::Inference and keep the fused kernels. That is deliberate — the decode loop is where Ster's speed lives, and a composed softmax or RMS-norm in the token loop would tax every generated token to serve a code path generation never enters.

The two passes once also decided what the forward returned, and no longer do. Which kernels run, whether the attached adapters apply, and how much of the vocabulary projection is computed are three independent choices. Decoding takes the last position's hidden state into the LM head, which is all sampling needs; anything that scores a sequence needs every position; and a reward model needs no vocabulary projection at all, because its head reads the residual stream. Separating them is what lets a frozen reference be scored with the fused kernels over a whole sequence — the combination the old pairing could not express, and the one preference optimization and `ster tune evaluate` both want.