Loss Scaling

Mixed-precision training makes neural networks faster by using low-precision numbers such as float16, but those numbers cannot represent very tiny values well. Loss scaling protects small gradients from disappearing before the optimiser can use them.

How it works
During backpropagation, gradients are proportional to the loss. Loss scaling multiplies the loss by a scale factor, such as 65,536, before backpropagation. This multiplies every gradient by the same factor, shifting values that would underflow to zero in float16 into its representable range. Just before the optimiser updates parameters, the training system divides—unscales—the gradients by that same factor. The intended parameter update is therefore unchanged compared with full-precision training; the scaling only makes the intermediate arithmetic numerically safe.

Dynamic scaling and safe updates
A fixed scale can be too small to help or too large, producing overflow: infinite or NaN gradients. Modern tools such as PyTorch’s torch.amp.GradScaler dynamically manage this balance:

• Scale the loss and run backward propagation.
• Check gradients for infinities or NaNs.
• Skip the optimiser step when overflow occurs, then lower the scale.
• Raise the scale gradually after many successful steps, preserving more small gradients.

Gradients must be unscaled before operations that depend on their true magnitude, especially gradient clipping. Clipping scaled gradients would make the clipping threshold meaningless.

Why it matters in practice
Without loss scaling, a deep model can appear to train but learn poorly because many low-precision gradients become exactly zero, particularly in earlier layers. With an excessively aggressive scale, a loss curve can stall because repeated overflows cause updates to be skipped. Loss scaling adds a small amount of checking and bookkeeping, but it enables the substantial memory savings and accelerator speed gains of mixed precision while retaining stable optimisation.