Z-Loss Guardrail

In a Mixture-of-Experts router, an auxiliary 'z-loss' penalizes the squared log-sum-exp of the router logits before softmax. What training failure does this primarily prevent?

Gradient leakage through the top-k selection's discrete argmax operation

Router logits growing large in magnitude, causing softmax saturation and numerical instability (e.g. overflow in bf16)

Experts receiving too few tokens and collapsing to unused capacity

The load-balancing loss dominating the language modeling loss during early training