Skip to content
Last Dictionary
Plane
Training & Post-trainingDeep LearningLanguage Models & Transformers

z-loss

nounAlso: router z-loss, log-Z penalty, softmax normalizer loss

In plain wordsA small extra penalty added during training that stops the raw scores a language model produces before choosing a word from growing huge, which helps keep very large training runs from blowing up.

  1. An auxiliary loss proportional to the squared logarithm of the softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score. normalizer, (log Z)², where Z is the sum of the exponentiated logitslogit, n.: A raw score that a model gives each possible answer before the scores are turned into probabilities. In statistics, the 'log-odds': a probability rewritten on a scale running from minus to plus infinity.. Softmax probabilities do not change when every logit shifts by the same constant, so nothing in the ordinary cross-entropycross-entropy, n.: A score of how surprised a model's predictions are by what actually happens; lower is better. Language models are trained mainly by driving this number down on real text. objective stops logits from drifting to large magnitudes, where low-precision arithmetic loses accuracy and training can diverge. The z-loss pulls log Z toward zero, keeping logits well scaled and close to normalized log-probabilities. PaLM applied it to the output layer with a small coefficient (10⁻⁴); the router z-loss of ST-MoE applies the same idea to the gating logits of a mixture of expertsmixture of experts, n.: A way to build very large AI models cheaply: the model contains many specialist sub-networks, and for each word only a few are switched on, picked by a small dispatcher network. routerrouter, n.: In a mixture-of-experts model, the small component that decides which specialist sub-networks should handle each word, like a dispatcher assigning jobs to workers., whose instability had made sparse models hard to train. It is a standard mitigation for training instabilitytraining instability, n.: The tendency of very large training runs to go wrong — the error suddenly jumping or blowing up — which can waste weeks of expensive computing unless engineers prevent or repair it., often used alongside QK-normQK-norm, n.: A small change to how a language model pays attention — rescaling two sets of internal numbers before comparing them — that prevents a common way large models' training goes haywire..

Lz=λ (log⁡Z)2,Z=∑iezi\mathcal{L}_z = \lambda\,(\log Z)^2, \qquad Z = \sum_i e^{z_i}

Origin

The penalty existed as an option on the final softmax logits in Google's Mesh TensorFlow codebase before Barret Zoph and colleagues' ST-MoE paper (February 2022) adapted it into the router z-loss. Google's PaLM paper (Aakanksha Chowdhery and colleagues, April 2022) described it as an auxiliary loss keeping log Z close to zero, and its use in PaLM made it widely known.

Related

See also
  • training instabilitytraining instability, n.: The tendency of very large training runs to go wrong — the error suddenly jumping or blowing up — which can waste weeks of expensive computing unless engineers prevent or repair it.
  • loss spikeloss spike, n.: A sudden jump in a model's error partway through training. Sometimes the model recovers on its own; sometimes the run falls apart and engineers must rewind to an earlier saved copy.
  • softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score.
  • logitlogit, n.: A raw score that a model gives each possible answer before the scores are turned into probabilities. In statistics, the 'log-odds': a probability rewritten on a scale running from minus to plus infinity.
  • QK-normQK-norm, n.: A small change to how a language model pays attention — rescaling two sets of internal numbers before comparing them — that prevents a common way large models' training goes haywire.
  • routerrouter, n.: In a mixture-of-experts model, the small component that decides which specialist sub-networks should handle each word, like a dispatcher assigning jobs to workers.

Sources

  1. Chowdhery, A. et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311.
  2. Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N. & Fedus, W. (2022). ST-MoE: Designing Stable and Transferable Sparse Expert Models. arXiv:2202.08906. arxiv.org/abs/2202.08906

Mentioned in

training instabilitytraining instability, n.: The tendency of very large training runs to go wrong — the error suddenly jumping or blowing up — which can waste weeks of expensive computing unless engineers prevent or repair it., loss spikeloss spike, n.: A sudden jump in a model's error partway through training. Sometimes the model recovers on its own; sometimes the run falls apart and engineers must rewind to an earlier saved copy., QK-normQK-norm, n.: A small change to how a language model pays attention — rescaling two sets of internal numbers before comparing them — that prevents a common way large models' training goes haywire.