Loss-Scale Dynamics and Spike Detection in Mixed Precision

Reduced-precision arithmetic is attractive for two separate reasons and breaks for two separate reasons. The arithmetic itself is faster and the operands are half the width, so a matrix multiply over them costs a fraction of the time and a fraction of the memory. It breaks because a 16-bit format’s exponent range is far narrower than its significand, so products that are perfectly representable in magnitude-independent terms overflow or underflow purely because of scale. Loss scaling is the standard remedy, and its interesting behaviour is in the parts that are usually hidden: how the scale is chosen, when it is reduced, and what a skipped step does to the optimizer’s accumulated state.

The range problem, stated in numbers

For the IEEE 754 binary16 format, the exponent field has five bits, which gives a maximum finite value of $65504$ and a smallest positive subnormal of about $6\times10^{-8}$. Normal numbers begin around $6\times10^{-5}$, and the significand has eleven bits including the implicit one, for a unit roundoff near $5\times10^{-4}$.

Three quantization effects follow. Values outside the range become infinity or zero, which is an unrecoverable loss of information. Values below the normal threshold lose precision gradually and are either represented with fewer bits or flushed to zero depending on the hardware’s denormal support — the mode is a hardware property, not a language one. And values within range round to eleven bits, which for an accumulated sum of many terms produces a relative error per addition of order $10^{-3}$; that is acceptable for a matrix multiply’s inputs but not for a running sum of gradient magnitudes over thousands of steps.

A deep network’s gradients span a wide range of magnitudes: the output head’s gradients can be large while the gradients of early layers, filtered through a long product of Jacobians, can be many orders of magnitude smaller. The backward pass’s intermediate values include products of gradients with activations, and a chain of multiplications by small numbers can drive values below the subnormal threshold before they are ever used.

Loss scaling addresses this by multiplying the loss by a constant $S$ before the backward pass. Every gradient in the chain is multiplied by $S$ by linearity, so a gradient that was $10^{-9}$ becomes representable if $S$ is around $10^{3}$ or larger. The gradients are divided by $S$ before the optimizer’s step, restoring the intended value.

The scale must be large and must not be too large

The lower bound on $S$ comes from the smallest gradient that must be represented, and the upper bound from the largest intermediate value, which must stay below $65504$. Both bounds move during training. Early on, gradients at the deepest layers can be large in absolute terms; later, they shrink as the loss approaches a minimum. A fixed scale chosen to fit the early phase wastes precision later, and a scale chosen for the late phase overflows early.

The intermediate products matter more than the final gradients. In a backward pass, a layer’s input gradient is the product of the upstream gradient and the layer’s Jacobian; both factors are representable at scale $S$ individually, and their product can overflow even when neither does. This is why overflow is frequently detected in a specific layer that has no obviously extreme gradient.

Dynamic scaling resolves both bounds by changing $S$ during the run and detecting overflow rather than predicting it.

Detecting an overflow without poisoning the weights

The detection works on the gradients themselves: after the backward pass, if any gradient element is non-finite, the step is discarded. The check must happen before the unscale and before any optimizer update, because an infinity in one element will otherwise propagate through the norm computation, through the clip, and into the moment estimates, from which recovery is impossible.

A step skipped for this reason leaves the parameters untouched and does not advance the optimizer’s step counter. Whether the second-moment estimate advances is the subtle part: if the step count used in bias correction is not deprecated alongside, the correction term $1-\beta_2^t$ diverges from the fraction actually used, and the estimate’s effective weight becomes wrong by a factor that grows with the number of skipped steps. Implementations that increment the step counter only on successful updates keep the bias correction consistent; implementations that increment unconditionally do not. The difference is invisible unless steps are skipped, and steps are skipped exactly when gradients are unstable, so the error appears at the worst moment.

def dynamic_scale_step(grads, scale, growth, backoff, max_scale):
    # grads is the scaled gradient buffer; inf/nan anywhere means overflow
    overflow = any(not (g == g and abs(g) != float("inf")) for g in grads)
    if overflow:
        return None, max(scale * backoff, 1.0)   # skip the update
    unscaled = [g / scale for g in grads]
    return unscaled, min(scale * growth, max_scale)

The growth window is what makes the scheme workable. Growing the scale every step oscillates: a step that happens not to overflow pushes the scale up, the next overflows, the scale drops, and the run spends most of its time at a scale that is one step behind the right one. The usual structure counts consecutive successful steps and grows the scale only after a window of them, so the scale settles at a value just below the overflow threshold and stays there. The backoff on overflow is immediate and larger than the growth factor, so a single overflow costs more than a window of successes gains; the asymmetry keeps the scheme stable against the runaway that a symmetric rule allows.

The cost of a skipped step

A skipped step is not free, and the cost has structure. For plain stochastic descent, skipping means the batch’s examples are processed and discarded, which for a batch of $B$ at a total budget of $N$ examples costs a fraction $B/N$ of the budget. Over a run with a moderate overflow rate, this is a small but real tax on compute efficiency.

The second cost is the state discontinuity. With momentum or adaptive second moments, the update after a skip uses statistics from before the skip and applies them to parameters that did not move. For Adam, the moment estimates $m$ and $v$ are unchanged by the skip only if the implementation does not merge the failing gradient into them, which is why the check must precede the merge. If it does not, the moments carry an infinite or NaN value and every subsequent step is NaN, which presents as a run that dies rather than one that loses a step.

The third cost interacts with clipping. As described in gradient clipping and norm scaling, the clip threshold applies to the unscaled gradient, so the order is: backward with the scaled loss, check finiteness, unscale, clip, step. Performing the check after the clip cannot detect an overflow that the clip has already converted to a finite value by multiplying by zero — an infinite norm produces a zero scale factor and a zero gradient, which is finite and would pass a finiteness check applied after the fact, while having destroyed the update’s information.

Interaction with the schedule and with adaptive methods

Loss scaling is a per-step transformation that depends on the gradient’s magnitude, and the schedule’s rate determines how much the gradient’s magnitude changes per step. A run whose rate decays by three orders of magnitude over training sees its gradients shrink correspondingly; the loss scale must grow to compensate, and the dynamic rule does that automatically, but the growth rate of the scale is bounded by the growth window. If the rate decays faster than the scale can grow, the late phase’s gradients fall below the representable range despite the scaling mechanism, and the symptom is a loss that stops improving without any overflow event.

The interaction with adaptive denominators is a scale invariance: multiplying all gradients by $S$ multiplies $m$ by $S$ and $v$ by $S^2$, so $m/\sqrt{v}$ is unchanged. The optimizer’s update is therefore invariant to the loss scale exactly, provided the unscale is applied before the moment estimates are updated and the $\epsilon$ is small relative to $\sqrt{v}$ at the scaled magnitude. This is a useful property: it means that changing the loss scale must not change the trajectory, and any observed dependence is a bug in the ordering. The $\epsilon$ is the exception — it does not scale, so a scale large enough to make $\sqrt{v}$ comparable to $\epsilon$ at the real magnitude will make the scaled estimate’s denominator dominated by $\sqrt{v}$ instead, changing the effective rate for small-gradient coordinates. The usual prescription of using a small $\epsilon$ with adaptive methods is partly a statement about loss scaling.

The warmup window described in learning rate warmup and early instability is where the scale is most likely to be wrong, since early gradients vary most between steps. A dynamic scheme converges to the right scale within a few dozen steps, which is comparable to the warmup length; a run with a short warmup and a fixed scale can overflow before the scale has adapted.

Where to keep master weights

The final structural question is which copy of the parameters the optimizer updates. Keeping a separate higher-precision copy — the master weights — and casting them to the reduced format for each forward and backward pass avoids the accumulation error that comes from keeping the parameters themselves in 16 bits: an update of $10^{-4}$ applied to a parameter of magnitude $1$ is representable in binary16 but rounds to the same value whenever the update is below half a unit in the last place at that magnitude, which is about $5\times10^{-4}$ relative. The parameter simply stops changing. With a master copy in single precision, the update accumulates in a format whose unit roundoff near one is about $6\times10^{-8}$, and the cast for the forward pass is a rounding of the input, not of the accumulated value.

The cost is memory: one extra copy of every weight, plus the update’s temporary. For a parameter count $P$, single-precision master weights occupy $4P$ bytes, while the two moments of an adaptive optimizer occupy $8P$ bytes and the gradients another $4P$; the reduced-precision copy for the forward pass adds $2P$. The accounting is what decides how much of the model fits on a device, and it is why the master copy and the optimizer state, not the parameters, set the memory budget for a training step.

The remaining decision is which operations run in reduced precision at all. Matrix multiplies and convolutions tolerate it because their inputs are rounded once and the accumulation happens in a wider format; reductions accumulate over many terms and lose precision proportional to the number of terms; and the loss computation, the normalization statistics, and the optimizer step require the wider format for the reasons above. The boundary between the two is a per-operation property, not a per-model one, which is why the mechanism is described as mixed precision rather than as reduced-precision training.