# Z-Loss Backward Geometry in Dense Output Heads and Sparse Routers

arXiv: https://arxiv.org/abs/2609.16179 · HTML: https://mlsanity.com/paper/2609.16179/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts.
- Builds on 33 of its 43 references that are in the map.
- First posted on arXiv 31 Jul 2026.

## Placed in

- Mixture-of-experts / [Expert routing strategies](https://mlsanity.com/moe/expert-routing-strategies/)

## Builds on (most cited first, 30 of 33)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings](https://mlsanity.com/paper/1909.00512/) (2019)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Regularizing Neural Networks by Penalizing Confident Output Distributions](https://mlsanity.com/paper/1701.06548/) (2017)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Representation Degeneration Problem in Training Natural Language Generation Models](https://mlsanity.com/paper/1907.12009/) (2019)
- [Mitigating Neural Network Overconfidence with Logit Normalization](https://mlsanity.com/paper/2205.09310/) (2022)
- [MegaBlocks: Efficient Sparse Training with Mixture-of-Experts](https://mlsanity.com/paper/2211.15841/) (2022)
- [Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts](https://mlsanity.com/paper/2408.15664/) (2024)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Transformers without Tears: Improving the Normalization of Self-Attention](https://mlsanity.com/paper/1910.05895/) (2019)
- [Breaking the Softmax Bottleneck: A High-Rank RNN Language Model](https://mlsanity.com/paper/1711.03953/) (2017)
- [Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling](https://mlsanity.com/paper/1611.01462/) (2016)
- [NormFormer: Improved Transformer Pretraining with Extra Normalization](https://mlsanity.com/paper/2110.09456/) (2021)

## Abstract (arXiv)

Z-loss has been widely applied to the logits of language-model output heads and sparse mixture-of-experts routers. Z-loss constrains the softmax log-normalizers of these output heads and routers, thereby limiting large-logit excursions, reducing finite-precision roundoff exposure, and avoiding training-loss divergence. These use cases arise in modern Transformer settings where large-vocabulary softmax heads, top-k routing, fused losses, and mixed-precision optimizers interact. Z-loss has typically been understood only as a scalar penalty on the log-normalizer. This paper instead analyzes Z-loss from a backward-pass perspective, focusing on the gradients produced by the Z-loss penalty. The logit-space gradient, which we call the backward source, is injected at the logit boundary of the Z-loss branch of backpropagation; consequently, the backward source's effect depends on the architecture and implementation through which the gradient is transported. We develop a backward-transport view for Z-loss that separates the source's scalar amplitude and softmax shape from the transport factors. These factors include common-shift coordinates, tied-embedding pathways, output-to-hidden gain, fused-loss source consistency, optimizer-facing updates, and top-k router reduction scale. These diagnostics show that nearly identical forward Z-loss values can coexist with distinct logit-space Z-loss gradients and, after architectural and optimizer transport, distinct parameter updates. The transport diagnostics also explain why raw-logit Z-loss can reduce scalar tails without changing output-to-hidden gain and why active-route reductions alter the effective router coefficient. Across evaluations of models in the GPT-2 and Pythia families on WikiText-103 and FineWeb-Edu, architecture-aware variants reduce backward-geometry tails while maintaining comparable validation perplexity in low-coefficient regimes.