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

# Z-Loss Backward Geometry in Dense Output Heads and Sparse Routers
Source: [https://arxiv.org/html/2609.16179](https://arxiv.org/html/2609.16179)
###### Abstract

Z\-loss has been widely applied to the logits of language\-model output heads and sparse mixture\-of\-experts \(MoE\) 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\-kkrouting, 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\-kkrouter 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\. The empirical suite includes continued pretraining, mixed\-precision and fused\-kernel audits, Adam\-state stress tests, and end\-to\-end MoE training\. In stress regimes, the same variants expose explicit trade\-offs between validation quality and both update\-tail behavior and mixed\-precision headroom\. Together, these results yield a transport\-aware framework for measuring, reporting, and intervening on the Z\-loss logit\-gradient source, its transport path, and the resulting update tails\.

## Introduction

Modern Transformer training often includes auxiliary loss terms intended to constrain final\-softmax and router log\-normalizers\. In dense output heads, the motivation is to limit output\-logit or common\-shift excursions and preserve finite\-precision numerical headroom\. In sparse routers, the motivation also includes reducing roundoff\-sensitive routing and training\-loss divergence\. A prominent example is the log\-normalizer penalty commonly called Z\-loss, which has been widely used in modern Transformer training\. Mesh TensorFlow and Pathways Language Model \(PaLM\) apply Z\-loss to final\-softmax logits\(Shazeeret al\.[2018](https://arxiv.org/html/2609.16179#bib.bib3); Chowdheryet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib7)\)\. PaLM reports that keeping the softmax log\-normalizer near the target increases training stability\(Chowdheryet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib7)\)\. Prior work accordingly uses logit\-control methods as training\-stability measures\. Router Z\-loss keeps gating logits small, reduces low\-precision roundoff error, and prevents training\-loss divergence\(Zophet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib6)\), while output\-embedding centering removes common\-shift degrees of freedom associated with output\-logit divergence\(Stollenwerket al\.[2026](https://arxiv.org/html/2609.16179#bib.bib9)\)\. Despite the broad use of Z\-loss, prior work has typically presented Z\-loss only as a scalar penalty on the log\-normalizer\.

This paper formulates Z\-loss as architecture\-dependent backward transport and uses this gradient perspective to analyze and diagnose how Z\-loss fundamentally affects training behavior\. We develop diagnostics for common backward\-transport phenomena involving scalar source amplitude, source shape, output\-to\-hidden gain, pathway coupling, and router effective scale\. Using the standard shift invariance of softmax, we derive common\-shift sensitivity of Z\-loss and output\-head centering that preserves cross\-entropy \(CE\)\. We also derive tied\-embedding gradient decomposition and top\-kkrouter scale distortion\. Our empirical evaluation uses pretrained Transformer language models\(Vaswaniet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib1)\)and text from WikiText\-103 and FineWeb\-Edu\. The model set includes four models from the Generative Pre\-trained Transformer 2 \(GPT\-2\) family: GPT\-2, DistilGPT\-2, GPT\-2 Medium, and GPT\-2 Large\. The same framework is tested with matched dense continued pretraining, effectively unclipped high\-coefficient stress sweeps, static mixed\-precision endpoints, low\-rank output\-to\-hidden gain spectral audits, optimizer\-facing update audits, and end\-to\-end top\-kkmixture\-of\-experts \(MoE\) language\-model training\. The implementation study additionally audits a row\-wise fused CE\+Z\-loss graphics processing unit \(GPU\) kernel implemented inTriton, a language and a compiler for writing tiled GPU kernels\(Tilletet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib26)\)\.

##### Contributions\.

Our primary contribution is a backward\-transport framework that treats the Z\-loss logit gradient as a source whose optimizer\-facing effect depends on architecture and implementation, rather than on the scalar penalty alone\. We derive diagnostics for common shift, tied pathways, output\-to\-hidden gain, fused\-source consistency, and router reduction scale, together with coordinate\-specific probes and interventions: centering, factorization, gain\-aware weighting, source\-consistent computation, and scale matching\. Matched pretrained diagnostics, continued pretraining, fused\-loss audits, stress endpoints, and MoE training show that these coordinates can behave differently\. Low\-coefficient settings can preserve validation perplexity \(PPL\) while stress settings reveal model\- and regime\-dependent trade\-offs among validation quality, update tails, and mixed\-precision headroom\. The resulting audit protocol therefore selects coefficients and interventions based jointly on transport diagnostics and validation quality instead of prescribing a universally optimal Z\-loss variant\. Figure[1](https://arxiv.org/html/2609.16179#Sx1.F1)summarizes the framework\.

This framing also identifies a contrast with the usual recipe\-level interpretation\. In that interpretation, Z\-loss is often treated as a scalar auxiliary term whose role is to makelog⁡Z\\log Zsmaller and therefore improve numerical stability\. Nevertheless, lowering the scalar log\-normalizer tail is not equivalent to lowering the transported update tail\. Raw\-logit Z\-loss can make the reported Z\-loss curve look successful while leaving output\-to\-hidden gain, tied\-path coupling, adaptive moment estimation \(Adam\) state pressure, or static loss\-scale headroom essentially unchanged or worse\. Conversely, centering, factorization, gain\-aware weighting, and router scale matching are not interchangeable ways to reduce one scalar; these methods are interventions on different coordinates of the same backward source\.

Forward scalarℒZ≔λ​\(log⁡Z−c\)2\\mathcal\{L\}\_\{Z\}\\coloneqq\\lambda\(\\log Z\-c\)^\{2\}Z\-loss logit gradient, the backward sourceδzZ=2​λ​\(log⁡Z−c\)​p\\delta\_\{z\}^\{Z\}=2\\lambda\(\\log Z\-c\)pImplementation reuselog\-sum\-exp statsΔsrc,Δθ\\Delta\_\{\\rm src\},\\Delta\_\{\\theta\}Common shiftz=z~\+μ​𝟏z=\\widetilde\{z\}\+\\mu\\mathbf\{1\}remove2​λ​μ​p2\\lambda\\mu pOutput and tied pathsδuZ≔WU⊤​δzZ\\delta\_\{u\}^\{Z\}\\coloneqq W\_\{U\}^\{\\top\}\\delta\_\{z\}^\{Z\}tiedEout\+EinE\_\{\\rm out\}\+E\_\{\\rm in\}Router scaletop\-kkreductionsλR​α\\lambda\_\{R\}\\alphaandNtl​αN\_\{\\rm tl\}\\alphaOptimizer\-facing updateslog\-normalizer and gain tails, Adam state, and half\-precision headroomtransport coordinates and diagnostics

Figure 1:Backward\-transport view of Z\-loss\. The scalar penalty creates the logit gradientδzZ\\delta\_\{z\}^\{Z\}, termed the backward source\. Architecture and implementation determine whether this vector is preserved, rescaled, amplified, or coupled into shared parameters before the transported source reaches optimizer\-facing updates\.

## Backward Transport View

This section develops the backward\-transport view by treating the Z\-loss logit gradient as the source injected into backpropagation\. We first specify the primary language\-modeling objective and auxiliary Z\-loss, then separate source construction from architectural transport, identify the coordinates and pathways that modify the backward source, and use these results to motivate the diagnostics and experiments that follow\. All proofs and more detailed theoretical developments are provided in the Appendix\.

### Language\-Modeling Objective and Z\-Loss Setup

Z\-loss is added to the standard language\-modeling loss\. For an autoregressive token sequencex1:Tx\_\{1:T\}, the primary objective is the CE for next\-token prediction,

ℒCE=−1T−1​∑t=1T−1log⁡pθ​\(xt\+1∣x≤t\)\.\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{CE\}\}=\-\\frac\{1\}\{T\-1\}\\sum\_\{t=1\}^\{T\-1\}\\log p\_\{\\theta\}\(x\_\{t\+1\}\\mid x\_\{\\leq t\}\)\.Applied to the deployed\-logit log\-normalizer, the Z\-loss term is added as an auxiliary regularizer:ℒtotal=ℒCE\+ℒZ\\mathcal\{L\}\_\{\\mathrm\{total\}\}=\\mathcal\{L\}\_\{\\mathrm\{CE\}\}\+\\mathcal\{L\}\_\{Z\}\. This setup also delineates the scope of our stability claims: Z\-loss directly constrains that log\-normalizer rather than output\-head precision or representation geometry\. Smaller logit excursions may reduce finite\-precision roundoff exposure, but any resulting changes in output\-embedding geometry, numerical headroom, or optimizer\-facing updates are downstream effects that must be evaluated separately\.

The penalty studied here has the form

ℒZ\\displaystyle\\mathcal\{L\}\_\{Z\}≔λ​\(log⁡Z−c\)2,\\displaystyle\\coloneqq\\lambda\(\\log Z\-c\)^\{2\},Z\\displaystyle Z≔∑i=1Vexp⁡zi,\\displaystyle\\coloneqq\\sum\_\{i=1\}^\{V\}\\exp z\_\{i\},λ\\displaystyle\\lambda≥0,\\displaystyle\\geq 0,\(1\)for a softmax\-axis logit vectorz∈ℝVz\\in\\mathbb\{R\}^\{V\}, log\-normalizer targetcc, and coefficientλ\\lambda\. The targetccis the zero\-penalty value oflog⁡Z\\log Z, andVVdenotes the size of the relevant softmax axis\. Prior final\-softmax and router formulations commonly useλ​\(log⁡Z\)2\\lambda\(\\log Z\)^\{2\}, corresponding toc=0c=0\(Shazeeret al\.[2018](https://arxiv.org/html/2609.16179#bib.bib3); Chowdheryet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib7); Zophet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib6)\)\. We instead study the target\-shifted generalization in Eq\.[1](https://arxiv.org/html/2609.16179#Sx2.E1)and usec=log⁡Vc=\\log Vunless stated otherwise;ccis fixed rather than tuned, while coefficient sweeps varyλ\\lambda\. This choice yields zero violation for the zero\-logit uniform vector\. For raw logits, settingc=log⁡Vc=\\log Vselects a common\-shift target, whereas for centered logits, the same setting defines the uniform reference for the remaining relative\-logit coordinate\. Accordingly, the Standard Z\-loss rows in our tables apply Eq\.[1](https://arxiv.org/html/2609.16179#Sx2.E1)to raw logits withc=log⁡Vc=\\log V; conventional formulations commonly usec=0c=0\.

This scalar definition of Z\-loss in Eq\.[1](https://arxiv.org/html/2609.16179#Sx2.E1)specifies the forward loss value\. However, the scalar definition alone cannot determine the resulting gradient\-based parameter update\. The update also depends on how the Z\-loss logit gradient, which we term the backward source in the next subsection, is transported through the architecture and training implementation\. A scalar audit can therefore certify the quantity being added to the objective while missing the vector actually injected into training\. This distinction matters precisely in the finite\-precision and training\-divergence settings where Z\-loss is used\. The reported penalty can remain unchanged across output\-head gauges, tied embeddings, fused low\-precision losses, adaptive optimizers, and top\-kkrouter reductions, even as these factors change where and how strongly the logit gradient reaches parameters\. A complete analysis must therefore identify this gradient and the transport path through which the gradient acts, in addition to measuring the Z\-loss value\.

### Z\-Loss Logit Gradient and Transport

Letz∈ℝVz\\in\\mathbb\{R\}^\{V\}be the output\-logit vector,p≔softmax⁡\(z\)p\\coloneqq\\operatorname\{softmax\}\(z\), andZ≔∑ieziZ\\coloneqq\\sum\_\{i\}e^\{z\_\{i\}\}\. Differentiating Eq\.[1](https://arxiv.org/html/2609.16179#Sx2.E1)gives

δzZ≔∂ℒZ∂z=2​λ​\(log⁡Z−c\)​p\.\\delta\_\{z\}^\{Z\}\\coloneqq\\frac\{\\partial\\mathcal\{L\}\_\{Z\}\}\{\\partial z\}=2\\lambda\(\\log Z\-c\)p\.\(2\)Throughout, source is shorthand for this Z\-loss logit gradient, equivalently the logit\-space adjointδzZ\\delta\_\{z\}^\{Z\}\. The logit\-space adjoint is a covector at the logit boundary\. The componentδz,iZ≔∂ℒZ/∂zi\\delta\_\{z,i\}^\{Z\}\\coloneqq\\partial\\mathcal\{L\}\_\{Z\}/\\partial z\_\{i\}represents the local sensitivity to logitziz\_\{i\}\. The term source reflects the adjoint’s role as the starting backward signal that the Z\-loss branch of reverse\-mode differentiation transports toward hidden states and parameters\. The adjoint is neither a data source nor a parameter gradient\. In a dense output head, the source has one entry per vocabulary item for each token; in an MoE router, the source has one entry per expert for each token\-layer routing decision\. Unless explicitly qualified as a total CE\+Z source, source refers only to the auxiliary Z\-loss contribution; the CE logit gradient is the separate vectorp−yp\-y\.

The source has two separable factors\. The first is a signed scalar coefficient2​λ​\(log⁡Z−c\)2\\lambda\(\\log Z\-c\), whose absolute value2​λ​\|log⁡Z−c\|2\\lambda\|\\log Z\-c\|we call the scalar source amplitude\. The second is the softmax\-shaped vectorpp, which we call the source shape\. Consequently, the full Euclidean source magnitude is‖δzZ‖2=2​λ​\|log⁡Z−c\|​‖p‖2\\left\\lVert\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}=2\\lambda\|\\log Z\-c\|\\left\\lVert p\\right\\rVert\_\{2\}\. For fixedλ\>0\\lambda\>0, this identity can also be written as

‖δzZ‖2=2​λ​ℒZ​‖p‖2,\\displaystyle\\left\\lVert\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}=2\\sqrt\{\\lambda\\mathcal\{L\}\_\{Z\}\}\\left\\lVert p\\right\\rVert\_\{2\},so the exact per\-instance source vanishes with the scalar penalty\. The scalar alone nevertheless leaves the source sign and softmax shape unspecified and does not determine amplification through the logit Jacobian\. When discussing gain, we view source shape through both the concentration‖p‖2\\left\\lVert p\\right\\rVert\_\{2\}and the normalized directionp/‖p‖2p/\\left\\lVert p\\right\\rVert\_\{2\}\. The downstream architecture then determines how this shaped source is transported\. For any parameter blockθ\\thetawith logit Jacobian

Jθ≔∂z∂vec⁡\(θ\),\\displaystyle J\_\{\\theta\}\\coloneqq\\frac\{\\partial z\}\{\\partial\\operatorname\{vec\}\(\\theta\)\},the Z\-loss gradient is

vec⁡\(∇θℒZ\)=Jθ⊤​δzZ\.\\operatorname\{vec\}\(\\nabla\_\{\\theta\}\\mathcal\{L\}\_\{Z\}\)=J\_\{\\theta\}^\{\\top\}\\delta\_\{z\}^\{Z\}\.\(3\)Thus, the backward terminology has a direct operational meaning\. The complete causal pipeline can be written compactly as

δ^zZ\\displaystyle\\widehat\{\\delta\}\_\{z\}^\{Z\}=2​λ​α​\(log⁡Z^−c\)​p^,\\displaystyle=2\\lambda\\alpha\(\\widehat\{\\log Z\}\-c\)\\widehat\{p\},\(4\)gθZ\\displaystyle g\_\{\\theta\}^\{Z\}=Jθ⊤​δ^zZ,\\displaystyle=J\_\{\\theta\}^\{\\top\}\\widehat\{\\delta\}\_\{z\}^\{Z\},\(θ\+,s\+\)\\displaystyle\(\\theta^\{\+\},s^\{\+\}\)=Optimizer⁡\(θ,s,gθCE\+gθZ\),\\displaystyle=\\operatorname\{Optimizer\}\(\\theta,s,g\_\{\\theta\}^\{\\mathrm\{CE\}\}\+g\_\{\\theta\}^\{Z\}\),where hats denote the statistics actually supplied by a fused or low\-precision implementation,α\\alphais the per\-instance reduction weight, andssis optimizer state\. Common shift andccaffect source amplitude; implementation precision affects the hatted source statistics; top\-kkreduction affectsα\\alpha; tied pathways and output\-to\-hidden gain enter throughJθ⊤J\_\{\\theta\}^\{\\top\}; and the optimizer transforms the total CE\+Z gradient in the final line\. Source construction, architectural transport, and optimizer transformation are distinct stages rather than interchangeable transport paths\. For a single exact softmax instance,α=1\\alpha=1,log⁡Z^=log⁡Z\\widehat\{\\log Z\}=\\log Z, andp^=p\\widehat\{p\}=p\.

In plain gradient descent with step sizeη\\eta, the Z\-loss contribution to the parameter step is

ΔZ​vec⁡\(θ\)≔−η​Jθ⊤​δzZ\.\\displaystyle\\Delta\_\{Z\}\\operatorname\{vec\}\(\\theta\)\\coloneqq\-\\eta J\_\{\\theta\}^\{\\top\}\\delta\_\{z\}^\{Z\}\.This transport distinction has three direct consequences\. First, implementations that agree on the forward scalar can transport different backward sources and therefore apply different parameter updates\. Second, router reductions that use the same nominalλR\\lambda\_\{R\}can assign different absolute coefficients to each routing decision and different scales relative to the token\-layer mean\. Finally, when the output head and input embeddings are tied, an auxiliary loss applied at the output logits is transported through both output and input pathways, turning a seemingly output\-side regularizer into a coupled update to shared lexical memory\.

For an adaptive optimizer, an isolated additive Z\-loss step need not be well\-defined because the optimizer transforms the total gradient together with the accumulated optimizer state; this optimizer coupling is why Eq\.[4](https://arxiv.org/html/2609.16179#Sx2.E4)placesgθCE\+gθZg\_\{\\theta\}^\{\\mathrm\{CE\}\}\+g\_\{\\theta\}^\{Z\}at the optimizer boundary\. Taking norms in Eq\.[3](https://arxiv.org/html/2609.16179#Sx2.E3)gives

‖∇θℒZ‖F≤2​λ​\|log⁡Z−c\|​‖Jθ‖op​‖p‖2\.\\left\\lVert\\nabla\_\{\\theta\}\\mathcal\{L\}\_\{Z\}\\right\\rVert\_\{F\}\\leq 2\\lambda\|\\log Z\-c\|\\left\\lVert J\_\{\\theta\}\\right\\rVert\_\{\\mathrm\{op\}\}\\left\\lVert p\\right\\rVert\_\{2\}\.\(5\)Eq\.[5](https://arxiv.org/html/2609.16179#Sx2.E5)is the central diagnostic lens\. The magnitude of the transported Z\-loss gradient is not determined by the forward scalar alone; the transported\-gradient magnitude also depends on source shape and the relevant Jacobian gain\.

In matched low\-coefficient regimes, the Z\-loss gradient can remain small relative to the CE gradient\. For a fixed coefficient and a single exact softmax instance, a vanishing scalar Z\-loss also implies a vanishing logit\-space source\. The salient issue is the relationship between scalar aggregates and backward transport: a scalar Z\-loss curve may improve while rare instances and architectural or implementation factors produce large pre\-clip gradient tails, change Adam state, reduce mixed\-precision headroom, or change router update scale\. These effects cannot be diagnosed from validation loss and scalar Z\-loss alone, but the effects determine what the optimizer actually receives\. We therefore treat Z\-loss as a quantity to audit and debug at source, transport, and update levels rather than only as a forward regularizer\.

A reported Z\-loss value specifies the forward penalty, but the transported vectorJθ⊤​δzZJ\_\{\\theta\}^\{\\top\}\\delta\_\{z\}^\{Z\}is what enters the optimizer\. Comparable reports therefore need source\-shape, transport\-path, and optimizer\-facing diagnostics rather than scalar curves alone\.

The finite\-precision rationale is two\-sided: log\-normalizer control can preserve forward representability, while the resulting backward source can consume numerical headroom after architectural transport\.

### Common\-Shift Sensitivity

We first isolate how a token\-wise common logit shift enters both the scalar penalty and the injected logit source\. For a one\-hot next\-token targetyy, letℓCE\\ell\_\{\\mathrm\{CE\}\}denote the per\-token CE loss, in contrast to the sequence\-averaged objectiveℒCE\\mathcal\{L\}\_\{\\mathrm\{CE\}\}above:

ℓCE\(z,y\)≔−∑i=1Vyilogsoftmax\(z\)i\.\\displaystyle\\ell\_\{\\mathrm\{CE\}\}\(z,y\)\\coloneqq\-\\sum\_\{i=1\}^\{V\}y\_\{i\}\\log\\operatorname\{softmax\}\(z\)\_\{i\}\.Ifz′=z\+a​𝟏z^\{\\prime\}=z\+a\\mathbf\{1\}, thensoftmax⁡\(z′\)=softmax⁡\(z\)=p\\operatorname\{softmax\}\(z^\{\\prime\}\)=\\operatorname\{softmax\}\(z\)=p, and CE is invariant both in value and along the common\-shift direction:

∂∂a​ℓCE​\(z\+a​𝟏,y\)=𝟏⊤​\(p−y\)=0\.\\frac\{\\partial\}\{\\partial a\}\\ell\_\{\\mathrm\{CE\}\}\(z\+a\\mathbf\{1\},y\)=\\mathbf\{1\}^\{\\top\}\(p\-y\)=0\.\(6\)This contrast leads to the following common\-shift property of Z\-loss\.

###### Proposition 1\(Z\-loss common\-shift sensitivity\)\.

For any fixeda∈ℝa\\in\\mathbb\{R\}independent ofzz,

ℒZ​\(z\+a​𝟏\)=λ​\(log⁡Z​\(z\)\+a−c\)2,\\displaystyle\\mathcal\{L\}\_\{Z\}\(z\+a\\mathbf\{1\}\)=\\lambda\(\\log Z\(z\)\+a\-c\)^\{2\},and

∇z\[ℒZ​\(z\+a​𝟏\)\]=2​λ​\(log⁡Z​\(z\)\+a−c\)​softmax⁡\(z\)\.\\displaystyle\\nabla\_\{z\}\\left\[\\mathcal\{L\}\_\{Z\}\(z\+a\\mathbf\{1\}\)\\right\]=2\\lambda\(\\log Z\(z\)\+a\-c\)\\operatorname\{softmax\}\(z\)\.Equivalently, the gradient with respect to the shifted logit variablexxatx≔z\+a​𝟏x\\coloneqq z\+a\\mathbf\{1\}has the same displayed value\.

By contrast with Eq\.[6](https://arxiv.org/html/2609.16179#Sx2.E6), summing the components of the source in Proposition 1 gives

∂∂a​ℒZ​\(z\+a​𝟏\)=𝟏⊤​δz\+a​𝟏Z=2​λ​\(log⁡Z​\(z\)\+a−c\)\.\\frac\{\\partial\}\{\\partial a\}\\mathcal\{L\}\_\{Z\}\(z\+a\\mathbf\{1\}\)=\\mathbf\{1\}^\{\\top\}\\delta\_\{z\+a\\mathbf\{1\}\}^\{Z\}=2\\lambda\(\\log Z\(z\)\+a\-c\)\.\(7\)The common\-shift coordinate is therefore forward\-invisible and backward\-inactive for CE, but backward\-active once Z\-loss is added\. Thus, monitoring CE alone cannot reveal common\-shift drift, whereas Z\-loss changes directly with that drift\. Although the prediction is unchanged, the altered source is transported to a dense hidden state asδuZ=WU⊤​δzZ\\delta\_\{u\}^\{Z\}=W\_\{U\}^\{\\top\}\\delta\_\{z\}^\{Z\}and can therefore change parameter gradients and optimizer state\. Centering removes this forward\-invisible but Z\-loss\-active coordinate before source construction\.

### Centered Output Heads

LetWU∈ℝV×dW\_\{U\}\\in\\mathbb\{R\}^\{V\\times d\}be a bias\-free output projection, also called the output head or unembedding matrix,u∈ℝdu\\in\\mathbb\{R\}^\{d\}the final hidden state, andw¯≔V−1​𝟏⊤​WU\\bar\{w\}\\coloneqq V^\{\-1\}\\mathbf\{1\}^\{\\top\}W\_\{U\}\. Define

W~U\\displaystyle\\widetilde\{W\}\_\{U\}≔WU−𝟏​w¯,\\displaystyle\\coloneqq W\_\{U\}\-\\mathbf\{1\}\\bar\{w\},\(8\)z~\\displaystyle\\widetilde\{z\}≔W~U​u=WU​u−\(w¯​u\)​𝟏\.\\displaystyle\\coloneqq\\widetilde\{W\}\_\{U\}u=W\_\{U\}u\-\(\\bar\{w\}u\)\\mathbf\{1\}\.
The following proposition formalizes the standard shift invariance of softmax for output\-head centering\(Blanchardet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib32)\)\.

###### Proposition 2\(Centering preserves CE\)\.

For any fixed hidden stateuuand a bias\-free output head, both the token softmax distribution and token CE remain unchanged in real arithmetic when the deployed logitsz=WU​uz=W\_\{U\}uare replaced byz~=W~U​u=z−α​𝟏\\widetilde\{z\}=\\widetilde\{W\}\_\{U\}u=z\-\\alpha\\mathbf\{1\}, whereα≔w¯​u\\alpha\\coloneqq\\bar\{w\}u\.

Centering changes the Z\-loss path because centering removes the common\-shift channelμ≔w¯​u\\mu\\coloneqq\\bar\{w\}ufrom the deployed\-logit source coordinate\. More generally, write

z=z~\+μ​𝟏,𝟏⊤​z~=0\.\\displaystyle z=\\widetilde\{z\}\+\\mu\\mathbf\{1\},\\qquad\\mathbf\{1\}^\{\\top\}\\widetilde\{z\}=0\.Then

log⁡Z=μ\+log⁡Z~,Z~≔∑iez~i\.\\log Z=\\mu\+\\log\\widetilde\{Z\},\\qquad\\widetilde\{Z\}\\coloneqq\\sum\_\{i\}e^\{\\widetilde\{z\}\_\{i\}\}\.\(9\)
Output\-head centering preserves CE but intentionally changes the Z\-loss source coordinate\. Under the shared\-coefficient, shared\-target convention, output\-head centering removes the deployed source component2​λ​μ​p2\\lambda\\mu p\. The effect of centering should be judged by transport diagnostics such asPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}, both defined later in the subsection on diagnostic metrics, because removing the common\-shift channel can be beneficial or introduce different trade\-offs across model geometries\.

### Tied\-Embedding Coupling

With tied embeddingsWU=EW\_\{U\}=Eand logitsz=E​uz=Eu, the Z\-loss gradient with respect toEEhas two conceptually separate paths\.

###### Proposition 3\(Tied embedding gradient decomposition\)\.

Consider untied copiesEinE\_\{\\mathrm\{in\}\}andEoutE\_\{\\mathrm\{out\}\}and impose the tying constraintEin=Eout=EE\_\{\\mathrm\{in\}\}=E\_\{\\mathrm\{out\}\}=Eafter differentiation\. Then

∇EℒZ=∇EoutℒZ\+∇EinℒZ\.\\nabla\_\{E\}\\mathcal\{L\}\_\{Z\}=\\nabla\_\{E\_\{\\mathrm\{out\}\}\}\\mathcal\{L\}\_\{Z\}\+\\nabla\_\{E\_\{\\mathrm\{in\}\}\}\\mathcal\{L\}\_\{Z\}\.\(10\)

Under this decomposition, the output\-side term is the local head update induced by the logit adjoint\. The input\-side partial derivative follows the transported path\. The sourceδzZ\\delta\_\{z\}^\{Z\}first induces the hidden\-state adjointE⊤​δzZE^\{\\top\}\\delta\_\{z\}^\{Z\}, which is then backpropagated through the Transformer body to the input embedding occurrences\. The input\-side term is transported through

δzZ→E⊤​δzZ→Transformer body→∇EinℒZ\.\\displaystyle\\delta\_\{z\}^\{Z\}\\rightarrow E^\{\\top\}\\delta\_\{z\}^\{Z\}\\rightarrow\\text\{Transformer body\}\\rightarrow\\nabla\_\{E\_\{\\mathrm\{in\}\}\}\\mathcal\{L\}\_\{Z\}\.Thus, an output\-side log\-normalizer regularizer becomes a multi\-path update to a shared lexical table\.

### Output\-to\-Hidden Gain

The hidden\-state injection from Z\-loss is

δuZ≔WU⊤​δzZ=2​λ​\(log⁡Z−c\)​WU⊤​p,\\delta\_\{u\}^\{Z\}\\coloneqq W\_\{U\}^\{\\top\}\\delta\_\{z\}^\{Z\}=2\\lambda\(\\log Z\-c\)W\_\{U\}^\{\\top\}p,\(11\)and hence

‖δuZ‖2≤2​λ​\|log⁡Z−c\|​‖WU‖op​‖p‖2\.\\left\\lVert\\delta\_\{u\}^\{Z\}\\right\\rVert\_\{2\}\\leq 2\\lambda\|\\log Z\-c\|\\left\\lVert W\_\{U\}\\right\\rVert\_\{\\mathrm\{op\}\}\\left\\lVert p\\right\\rVert\_\{2\}\.\(12\)We measure the token\-level shape\-only local gain

at≔‖WU⊤​pt‖2‖pt‖2\+ε\.a\_\{t\}\\coloneqq\\frac\{\\left\\lVert W\_\{U\}^\{\\top\}p\_\{t\}\\right\\rVert\_\{2\}\}\{\\left\\lVert p\_\{t\}\\right\\rVert\_\{2\}\+\\varepsilon\}\.\(13\)We useA¯p≔𝔼t​at\\bar\{A\}\_\{p\}\\coloneqq\\mathbb\{E\}\_\{t\}a\_\{t\}for the mean gain andAp99≔Q0\.99,t​\(at\)A\_\{p\}^\{99\}\\coloneqq Q\_\{0\.99,t\}\(a\_\{t\}\)for the 99th\-percentile \(p99\) gain\. Throughout the paper, unembedding, output head, and output projection refer to the same output\-side matrixWUW\_\{U\}unless a tied input\-embedding path is explicitly being discussed\.

###### Proposition 4\(Gain is anisotropy\-weighted\)\.

LetWU=U​Σ​V⊤W\_\{U\}=U\\Sigma V^\{\\top\},q≔p/‖p‖2q\\coloneqq p/\\left\\lVert p\\right\\rVert\_\{2\}, anda​\(p\)≔‖WU⊤​p‖2/‖p‖2a\(p\)\\coloneqq\\left\\lVert W\_\{U\}^\{\\top\}p\\right\\rVert\_\{2\}/\\left\\lVert p\\right\\rVert\_\{2\}, theε\\varepsilon\-free counterpart of Eq\.[13](https://arxiv.org/html/2609.16179#Sx2.E13)\. Then

a​\(p\)2=∑jσj2​⟨q,uj⟩2≤σmax​\(WU\)2\.a\(p\)^\{2\}=\\sum\_\{j\}\\sigma\_\{j\}^\{2\}\\left\\langle q,u\_\{j\}\\right\\rangle^\{2\}\\leq\\sigma\_\{\\max\}\(W\_\{U\}\)^\{2\}\.\(14\)

The expansion also shows the alignment claim\. For fixed‖p‖2\\left\\lVert p\\right\\rVert\_\{2\}, the contribution from directionuju\_\{j\}is weighted byσj2\\sigma\_\{j\}^\{2\}\. Among feasible source directions, the hidden injection is therefore largest when source mass lies in left singular directions associated with high singular values\. As a linear algebra bound over unit vectors, the upper bound is saturated exactly whenqqlies entirely in the left singular subspace associated withσmax​\(WU\)\\sigma\_\{\\max\}\(W\_\{U\}\)\. For softmax\-derivedqq, this equality condition may be infeasible\. This alignment property explains the operational sense in which unembedding geometry, not onlylog⁡Z\\log Z, predicts transported update size\.

Raw\-logit Z\-loss can reducelog⁡Z\\log Ztails while leavingWU⊤​pW\_\{U\}^\{\\top\}pand rare optimizer\-facing update tails exposed\. Diagnostics such asAp99A\_\{p\}^\{99\}and gradient\-event endpoints are needed to distinguish log\-normalizer\-tail control from transported\-update\-tail control\.

### Router Effective Scale

An MoE router produces, for each token\-layer decision, a vector of expert scoresrt,ℓ∈ℝNexpr\_\{t,\\ell\}\\in\\mathbb\{R\}^\{N\_\{\\mathrm\{exp\}\}\}and applies softmax or top\-kkselection over the expert axis\. Router Z\-loss is therefore the same log\-normalizer penalty as Eq\.[1](https://arxiv.org/html/2609.16179#Sx2.E1), instantiated on the expert\-score axis rather than the vocabulary axis\. For decision\(t,ℓ\)\(t,\\ell\), the router source is the expert\-score gradientδr,t,ℓR≔∂ℒR/∂rt,ℓ∈ℝNexp\\delta\_\{r,t,\\ell\}^\{R\}\\coloneqq\\partial\\mathcal\{L\}\_\{R\}/\\partial r\_\{t,\\ell\}\\in\\mathbb\{R\}^\{N\_\{\\mathrm\{exp\}\}\}, with one component per expert\. The router source has the same softmax\-shaped form as the dense\-head source, while top\-kkrouting and reduction denominators determine the effective scale transported through the router\. LetNtlN\_\{\\mathrm\{tl\}\}be the number of token\-layer router decisions andαt,ℓ\\alpha\_\{t,\\ell\}the reduction weight for decision\(t,ℓ\)\(t,\\ell\)\. Relative to the token\-layer mean, the effective scale isNtl​αt,ℓN\_\{\\mathrm\{tl\}\}\\alpha\_\{t,\\ell\}\. Under the conventions used in the experiments, active\-route mean has a scale of1/k1/k, while active\-route sum has a scale ofkk\.

The sameλR\\lambda\_\{R\}can imply different per\-decision gradients once top\-kkand reduction denominators change\. Router reports should therefore include the reduction convention, the absolute per\-decision coefficient, and the scale relative to the token\-layer mean\.

## Experiments

### Setup

This section empirically evaluates the preceding analysis and diagnostic framework in representative modern Transformer settings where Z\-loss is used\.

We evaluate raw and centered heads, standard and centered Z\-loss, factorized and gain\-aware objectives, and router\-scale variants using text from WikiText\-103 and FineWeb\-Edu\. The comparison spans pretrained\-model diagnostics, matched continued pretraining, implementation audits, stress endpoints, and controlled MoE training\. Unless otherwise stated, the FineWeb\-Edu experiments reported here use streamed blocks from the FineWeb\-Edusample\-10BTtraining split\. Dense runs on models in the GPT\-2 family use matched replications, stated token budgets, 32\-bit floating\-point \(fp32\) master weights with bfloat16 \(bf16\) autocast, and log\-normalizer, gradient\-tail, optimizer\-state, and mixed\-precision endpoints; the implementation and overflow audits also evaluate 16\-bit floating\-point \(fp16\) storage\. The Appendix provides supporting derivations, diagnostic and method conventions, further experimental details, and additional audits and results\.

The experiments follow the diagnostic hierarchy rather than rank the methods as universal alternatives\. A large p99\|μ\|\|\\mu\|motivates a centered\-deployment audit; a large residualPZ99\.9P\_\{Z\}^\{99\.9\}after centering motivates centered Z\-loss; a largeAp99A\_\{p\}^\{99\}motivates gain and spectral audits; forward–backward mismatch motivates shared\-statistics implementation checks; and router\-scale mismatch motivates coefficient matching\. The Centered head row is a gauge\-only control withλaux=0\\lambda\_\{\\rm aux\}=0\. In exact arithmetic, the Centered head row’s CE value and CE gradient equal those of the raw CE baseline, so the row isolates the deployed coordinate rather than adding an auxiliary update\. Because the diagnostic Z\-loss andPZ99\.9P\_\{Z\}^\{99\.9\}are evaluated in each method’s deployed\-logit coordinate, differences between raw and centered coordinates are not coordinate\-free performance gains; in pretrained audits the differences are deterministic consequences to be judged together withAp99A\_\{p\}^\{99\}and PPL\. Validation comparisons therefore report exact matched mean PPL differences within each run regime, making quality changes explicit alongside transport effects\.

### Pretrained Output\-Head Geometry

This comparison tests the common\-shift diagnosis and the geometry\-dependent consequences of centering\. Table[1](https://arxiv.org/html/2609.16179#Sx3.T1)shows that row\-centered output\-head deployment preserves PPL while sharply reducing common\-shift and Z\-loss transport diagnostics for several models\. Raw heads in the GPT\-2 family have very large common\-shift tails\. Centered heads keep CE unchanged while reducingPZ99\.9P\_\{Z\}^\{99\.9\}by factors of4\.94\.9,12\.812\.8, and12\.712\.7for DistilGPT\-2, GPT\-2, and GPT\-2 Medium, respectively\. The larger Pythia\-1B diagnostic provides an important contrast\. The Pythia\-1B raw common\-shift tail is already small, and centering increasesAp99A\_\{p\}^\{99\}\. This increase is predicted by the coordinate analysis rather than treated as an anomalous exception\. In the scalar coordinate, a rawμ\\mucan partially cancellog⁡Z~−c\\log\\widetilde\{Z\}\-c, so removing the raw common\-shift term can increase the centered violation\. In the gain coordinate,W~U⊤​p=WU⊤​p−w¯⊤\\widetilde\{W\}\_\{U\}^\{\\top\}p=W\_\{U\}^\{\\top\}p\-\\bar\{w\}^\{\\top\}, and subtracting the fixed output\-row mean need not reduce the norm when the mean’s alignment withWU⊤​pW\_\{U\}^\{\\top\}pis unfavorable\. The contrast therefore reinforces treating centering as a gauge intervention whose effect is evaluated with transport diagnostics, not as a uniformly beneficial normalization\.

Figure[2](https://arxiv.org/html/2609.16179#Sx3.F2)\(a\) makes the common\-shift geometry visible by showing that centering translates every GPT\-2 hidden\-source shape by the same output\-row mean while leaving the softmax distribution unchanged\. Figure[2](https://arxiv.org/html/2609.16179#Sx3.F2)\(b\) shows that the transported gain is concentrated in a low\-rank, anisotropic subspace rather than determined by scalar source amplitude alone\. Figure[2](https://arxiv.org/html/2609.16179#Sx3.F2)\(c\) directly realizes the identity‖δuZ‖2=Ap​‖δzZ‖2\\left\\lVert\\delta\_\{u\}^\{Z\}\\right\\rVert\_\{2\}=A\_\{p\}\\left\\lVert\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}token by token\. The spread across fixed\-gain rays is the empirical geometry that a scalar Z\-loss curve omits\.

![Refer to caption](https://arxiv.org/html/2609.16179v1/x1.png)Figure 2:Empirical Z\-loss gradient geometry on fp32 GPT\-2 predictions from WikiText\-103\. The left view shows that centering translates every hidden\-source shape by the fixed vector−w¯⊤\-\\bar\{w\}^\{\\top\}while preservingptp\_\{t\}\. The middle view shows that singular contributions closely track exactat2a\_\{t\}^\{2\}\. The raw and centered correlations are0\.9970\.997and0\.9770\.977, respectively\. The right view shows that‖δuZ‖2=Ap​‖δzZ‖2\\left\\lVert\\delta\_\{u\}^\{Z\}\\right\\rVert\_\{2\}=A\_\{p\}\\left\\lVert\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}exposes token\-dependent amplification, and centering lowers p99ApA\_\{p\}from 31\.0 to 18\.5\. Here,λdiag=10−4\\lambda\_\{\\rm diag\}=10^\{\-4\}andc=log⁡50257c=\\log 50257\.Table 1:Pretrained\-model diagnostics on WikiText\-103 validation\. Rows labeled centered use the row\-centered deployed\-logit convention defined above\. Tail metrics are reported on the deployed logits for each row; centering preserves CE but fixes the common\-shift gauge, so centering changes the Z\-loss transport coordinate\. Lower PPL indicates better validation quality, while lower diagnostic Z\-loss, abbreviated as Diag\. Z, and lowerPZ99\.9P\_\{Z\}^\{99\.9\},Ap99A\_\{p\}^\{99\}, and\|μ\|\|\\mu\|indicate smaller penalty or transport tails\. Gradient ratios and cosines characterize auxiliary\-gradient strength and CE alignment; the gradient ratios and cosines have no universally preferred direction unless stated otherwise\.To test whether the WikiText comparison of raw and centered output heads is corpus\-specific, Appendix Table[6](https://arxiv.org/html/2609.16179#A1.T6)repeats the pretrained evaluation on streamed FineWeb\-Edu blocks\. GPT\-2 and GPT\-2 Medium reproduce the large common\-shift pattern, whereas GPT\-2 Large and Pythia\-410M remain in a geometry\-dependent regime\.

Additional tied\-embedding pathway measurements, including the untied Pythia\-160M contrast, confirm that tying changes the transport path, while output\-head common shift remains a separate source factor whose effect depends on model geometry\.

Appendix Table[7](https://arxiv.org/html/2609.16179#A1.T7)reports a low\-rank unembedding spectral audit: dominant singular directions explain much of the token\-level gain in highly anisotropic heads, while the centered\-head results show that reducingσ1\\sigma\_\{1\}does not by itself guarantee a smallerAp99A\_\{p\}^\{99\}\.

### Source\-Consistency and Implementation Audits

These audits test whether forward scalar agreement predicts backward\-source and optimizer\-facing agreement\. Appendix Tables[8](https://arxiv.org/html/2609.16179#A1.T8),[9](https://arxiv.org/html/2609.16179#A1.T9), and[10](https://arxiv.org/html/2609.16179#A1.T10)report source\-space,Triton\-based fused\-kernel, and optimizer\-facing update audits that support the implementation claim used throughout the paper\.111Tritonis the implementation framework for a controlled CE\+Z\-loss GPU kernel\.Together, these audits test whether fused forward and backward computations share the softmax statistics that define the source\. Matching forward Z\-loss scalar values is not sufficient, because the backward source and the transported optimizer\-facing update can differ on corpus\-derived model logits even when the forward scalar error is small\. The source\-space, fused\-kernel, and optimizer\-facing update audits distinguish two failure modes\. Reusing the forward log\-sum\-exp statistics removes inconsistency caused by backward reconstruction, but a source\-consistent formula cannot recover precision already lost through quantized logit storage\.

### Dense Continued\-Pretraining Interventions

This experiment contrasts scalar\-tail control with output\-to\-hidden gain control\. Table[2](https://arxiv.org/html/2609.16179#Sx3.T2)reports the matched\-run GPT\-2 continued pretraining experiment on WikiText\-103\. The CE\-only setting is a strong baseline\. After 4800 optimizer steps, the CE\-only setting reaches a PPL of 21\.046±\\pm0\.012\. The matched mean PPL differences relative to CE are−0\.011\-0\.011for Centered head,\+0\.026\+0\.026for Standard Z\-loss, and−0\.007\-0\.007for Centered Z\-loss\. We report these differences directly rather than infer formal equivalence from three\-run standard deviations\. The larger and more seed\-consistent effect is the geometry change\. Centered head deployment lowersAp99A\_\{p\}^\{99\}from 26\.4 to 17\.0 and removes the common\-shift coordinate; because centered head deployment has no auxiliary loss, the Centered head row isolates the coordinate change\. Standard Z\-loss reducesPZ99\.9P\_\{Z\}^\{99\.9\}from 246\.8 to 21\.8 but leaves theAp99A\_\{p\}^\{99\}gain essentially unchanged\. Centered Z\-loss combines residual log\-normalizer control with the centered coordinate, reducingPZ99\.9P\_\{Z\}^\{99\.9\}to 20\.3 andAp99A\_\{p\}^\{99\}to 16\.8 at the PPL difference\.

Table 2:GPT\-2 continued pretraining on WikiText\-103\. All methods start from the same pretrained checkpoint\.λaux\\lambda\_\{\\rm aux\}is the applied auxiliary coefficient, Diag\. Z is evaluated on the deployed logits for each method, andΔ\\DeltaPPL is the matched mean difference from CE\.Appendix Table[11](https://arxiv.org/html/2609.16179#A1.T11)reports a disjoint\-block GPT\-2 replication on FineWeb\-Edu, and Appendix Table[12](https://arxiv.org/html/2609.16179#A1.T12)extends the comparison to GPT\-2 Medium with fresh tokens\. Both reproduce the same separation\. Centered variants reduceAp99A\_\{p\}^\{99\}, whereas standard raw\-logit Z\-loss does not\.

Appendix Table[13](https://arxiv.org/html/2609.16179#A1.T13)reports a separate 1200\-step GPT\-2 diagnostic with post\-training gradient instrumentation\. The diagnostic confirms that the tied\-table auxiliary gradient is nonzero and path\-coupled and that the alignment of the gradient with CE should be measured rather than assumed\.

### WikiText\-103 High\-Coefficient Stress Tests

With the coefficient increased tenfold toλaux=1×10−3\\lambda\_\{\\rm aux\}=1\\times 10^\{\-3\}, Table[3](https://arxiv.org/html/2609.16179#Sx3.T3)tests optimizer\-facing consequences\. Standard and gain\-aware Z\-loss reduce log\-normalizer tails but leaveAp99A\_\{p\}^\{99\}near CE\. Their p99 gradient norms are 47\.81 and 30\.25, respectively, compared with 5\.58 for CE\. Centered Z\-loss lowers bothPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}, keeps the p99 norm at 5\.69, and improves mean PPL\. The roughly7%7\\%clip rates of the raw\-logit variants reflect mostly sub\-threshold steps punctuated by rare large events rather than safer tails\.

Table 3:High\-coefficient GPT\-2 gradient\-tail stress sweep on WikiText\-103 withλaux=1×10−3\\lambda\_\{\\rm aux\}=1\\times 10^\{\-3\}for auxiliary\-method rows\. CE and centered head rows have no applied auxiliary loss, and Diag\. Z uses the same coefficient scale\. Clip rate is the fraction of optimizer steps whose pre\-clip full\-model gradient norm exceeds the clipping threshold of 1\.0\.
### Router Reduction and MoE Training

These experiments test the predicted reduction scale and whether coefficient matching restores token\-layer\-mean behavior\. The router\-projection audit in Appendix Table[22](https://arxiv.org/html/2609.16179#A1.T22)compares token\-layer mean, active\-route mean, and active\-route sum reductions under a fixed nominal router coefficient\.

End\-to\-end MoE results in Appendix Table[23](https://arxiv.org/html/2609.16179#A1.T23)provide an implementation\-level equivalence check of the same scale accounting\. Matched active\-route variants recover token\-layer\-mean behavior\. Because coefficient matching makes the effective objectives equivalent under matched initialization, this recovery confirms that the predicted reduction\-scale accounting carries through to end\-to\-end MoE training\.

## Conclusion

Z\-loss should be understood not only as a scalar log\-normalizer penalty but as a logit\-gradient source transported by architecture and implementation\. Our backward\-transport formulation explains why forward\-equivalent losses can produce different updates and why reducing raw\-logit tails need not reduce optimizer\-facing tails\. Across dense and MoE training, implementation audits, and optimizer stress tests, we observed that the common\-shift gauge, unembedding anisotropy, tied pathways, and router reductions affect distinct transport coordinates\. Intervention choice should therefore follow the relevant diagnostic, with validation quality and transport metrics reported together\.

## References

- L\. J\. Ba, J\. R\. Kiros, and G\. E\. Hinton \(2016\)Layer Normalization\.CoRRabs/1607\.06450\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p2.1)\.
- S\. Biderman, H\. Schoelkopf, Q\. G\. Anthony, H\. Bradley, K\. O’Brien, E\. Hallahan, M\. A\. Khan, S\. Purohit, U\. S\. Prashanth, E\. Raff, A\. Skowron, L\. Sutawika, and O\. van der Wal \(2023\)Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling\.InICML,Vol\.202,pp\. 2397–2430\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px1.p1.1)\.
- P\. Blanchard, D\. J\. Higham, and N\. J\. Higham \(2019\)Accurate Computation of the Log\-Sum\-Exp and Softmax Functions\.CoRRabs/1909\.03469\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx10.p1.4),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx11.SSSx1.p1.6),[Centered Output Heads](https://arxiv.org/html/2609.16179#Sx2.SSx4.p2.1)\.
- A\. Chowdhery, S\. Narang, J\. Devlin, M\. Bosma, G\. Mishra, A\. Roberts, P\. Barham, H\. W\. Chung, C\. Sutton, S\. Gehrmann, P\. Schuh, K\. Shi, S\. Tsvyashchenko, J\. Maynez, A\. Rao, P\. Barnes, Y\. Tay, N\. Shazeer, V\. Prabhakaran, E\. Reif, N\. Du, B\. Hutchinson, R\. Pope, J\. Bradbury, J\. Austin, M\. Isard, G\. Gur\-Ari, P\. Yin, T\. Duke, A\. Levskaya, S\. Ghemawat, S\. Dev, H\. Michalewski, X\. Garcia, V\. Misra, K\. Robinson, L\. Fedus, D\. Zhou, D\. Ippolito, D\. Luan, H\. Lim, B\. Zoph, A\. Spiridonov, R\. Sepassi, D\. Dohan, S\. Agrawal, M\. Omernick, A\. M\. Dai, T\. S\. Pillai, M\. Pellat, A\. Lewkowycz, E\. Moreira, R\. Child, O\. Polozov, K\. Lee, Z\. Zhou, X\. Wang, B\. Saeta, M\. Diaz, O\. Firat, M\. Catasta, J\. Wei, K\. Meier\-Hellstern, D\. Eck, J\. Dean, S\. Petrov, and N\. Fiedel \(2023\)PaLM: Scaling Language Modeling with Pathways\.J\. Mach\. Learn\. Res\.24,pp\. 240:1–240:113\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.p1.1),[Introduction](https://arxiv.org/html/2609.16179#Sx1.p1.1),[Language\-Modeling Objective and Z\-Loss Setup](https://arxiv.org/html/2609.16179#Sx2.SSx1.p2.14)\.
- A\. Clark, D\. de Las Casas, A\. Guy, A\. Mensch, M\. Paganini, J\. Hoffmann, B\. Damoc, B\. A\. Hechtman, T\. Cai, S\. Borgeaud, G\. van den Driessche, E\. Rutherford, T\. Hennigan, M\. J\. Johnson, A\. Cassirer, C\. Jones, E\. Buchatskaya, D\. Budden, L\. Sifre, S\. Osindero, O\. Vinyals, M\. Ranzato, J\. W\. Rae, E\. Elsen, K\. Kavukcuoglu, and K\. Simonyan \(2022\)Unified Scaling Laws for Routed Language Models\.InICML,Vol\.162,pp\. 4057–4086\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. Ré \(2022\)FlashAttention: Fast and Memory\-Efficient Exact Attention with IO\-Awareness\.InNeurIPS,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1)\.
- A\. de Brébisson and P\. Vincent \(2016\)The Z\-loss: a shift and scale invariant classification loss belonging to the Spherical Family\.CoRRabs/1604\.08859\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.p1.1)\.
- K\. Ethayarajh \(2019\)How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT\-2 Embeddings\.InEMNLP/IJCNLP \(1\),pp\. 55–65\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px3.p1.1)\.
- W\. Fedus, B\. Zoph, and N\. Shazeer \(2022\)Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity\.J\. Mach\. Learn\. Res\.23,pp\. 120:1–120:39\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- T\. Gale, D\. Narayanan, C\. Young, and M\. Zaharia \(2023\)MegaBlocks: Efficient Sparse Training with Mixture\-of\-Experts\.InMLSys,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx16.SSSx1.p2.2)\.
- J\. Gao, D\. He, X\. Tan, T\. Qin, L\. Wang, and T\. Liu \(2019\)Representation Degeneration Problem in Training Natural Language Generation Models\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px3.p1.1)\.
- D\. Goldberg \(1991\)What Every Computer Scientist Should Know About Floating\-Point Arithmetic\.ACM Comput\. Surv\.23\(1\),pp\. 5–48\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx10.p1.4),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx11.SSSx1.p1.6)\.
- E\. Grave, A\. Joulin, M\. Cissé, D\. Grangier, and H\. Jégou \(2017\)Efficient softmax approximation for GPUs\.InICML,Vol\.70,pp\. 1302–1310\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px4.p1.1)\.
- C\. Guo, G\. Pleiss, Y\. Sun, and K\. Q\. Weinberger \(2017\)On Calibration of Modern Neural Networks\.InICML,Vol\.70,pp\. 1321–1330\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p2.1)\.
- N\. Halko, P\. Martinsson, and J\. A\. Tropp \(2011\)Finding Structure with Randomness: Probabilistic Algorithms for Constructing Approximate Matrix Decompositions\.SIAM Rev\.53\(2\),pp\. 217–288\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px11.p1.3)\.
- J\. Hoffmann, S\. Borgeaud, A\. Mensch, E\. Buchatskaya, T\. Cai, E\. Rutherford, D\. de Las Casas, L\. A\. Hendricks, J\. Welbl, A\. Clark, T\. Hennigan, E\. Noland, K\. Millican, G\. van den Driessche, B\. Damoc, A\. Guy, S\. Osindero, K\. Simonyan, E\. Elsen, J\. W\. Rae, O\. Vinyals, and L\. Sifre \(2022\)Training Compute\-Optimal Large Language Models\.CoRRabs/2203\.15556\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px6.p1.1)\.
- H\. Inan, K\. Khosravi, and R\. Socher \(2017\)Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px3.p1.1)\.
- R\. A\. Jacobs, M\. I\. Jordan, S\. J\. Nowlan, and G\. E\. Hinton \(1991\)Adaptive Mixtures of Local Experts\.Neural Comput\.3\(1\),pp\. 79–87\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- J\. Kaplan, S\. McCandlish, T\. Henighan, T\. B\. Brown, B\. Chess, R\. Child, S\. Gray, A\. Radford, J\. Wu, and D\. Amodei \(2020\)Scaling Laws for Neural Language Models\.CoRRabs/2001\.08361\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px6.p1.1)\.
- D\. P\. Kingma and J\. Ba \(2015\)Adam: A Method for Stochastic Optimization\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx16.SSSx1.p2.2),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx2.p1.5)\.
- D\. Lepikhin, H\. Lee, Y\. Xu, D\. Chen, O\. Firat, Y\. Huang, M\. Krikun, N\. Shazeer, and Z\. Chen \(2021\)GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- M\. Lewis, S\. Bhosale, T\. Dettmers, N\. Goyal, and L\. Zettlemoyer \(2021\)BASE Layers: Simplifying Training of Large, Sparse Models\.InICML,Vol\.139,pp\. 6265–6274\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled Weight Decay Regularization\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx16.SSSx1.p2.2),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx2.p1.5)\.
- S\. Merity, C\. Xiong, J\. Bradbury, and R\. Socher \(2017\)Pointer Sentinel Mixture Models\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px1.p1.1)\.
- P\. Micikevicius, S\. Narang, J\. Alben, G\. F\. Diamos, E\. Elsen, D\. García, B\. Ginsburg, M\. Houston, O\. Kuchaiev, G\. Venkatesh, and H\. Wu \(2018\)Mixed Precision Training\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx2.p2.2)\.
- M\. Milakov and N\. Gimelshein \(2018\)Online normalizer calculation for softmax\.CoRRabs/1805\.02867\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx7.SSSx4.p1.9)\.
- J\. Mu and P\. Viswanath \(2018\)All\-but\-the\-Top: Simple and Effective Postprocessing for Word Representations\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px3.p1.1)\.
- R\. Müller, S\. Kornblith, and G\. E\. Hinton \(2019\)When does label smoothing help?\.InNeurIPS,pp\. 4696–4705\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3)\.
- T\. Q\. Nguyen and J\. Salazar \(2019\)Transformers without Tears: Improving the Normalization of Self\-Attention\.InIWSLT,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p2.1)\.
- T\. OLMo, P\. Walsh, L\. Soldaini, D\. Groeneveld, K\. Lo, S\. Arora, A\. Bhagia, Y\. Gu, S\. Huang, M\. Jordan, N\. Lambert, D\. Schwenk, O\. Tafjord, T\. Anderson, D\. Atkinson, F\. Brahman, C\. Clark, P\. Dasigi, N\. Dziri, M\. Guerquin, H\. Ivison, P\. W\. Koh, J\. Liu, S\. Malik, W\. Merrill, L\. J\. V\. Miranda, J\. Morrison, T\. Murray, C\. Nam, V\. Pyatkin, A\. Rangapur, M\. Schmitz, S\. Skjonsberg, D\. Wadden, C\. Wilhelm, M\. Wilson, L\. Zettlemoyer, A\. Farhadi, N\. A\. Smith, and H\. Hajishirzi \(2025\)2 OLMo 2 Furious\.CoRRabs/2501\.00656\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- R\. Pascanu, T\. Mikolov, and Y\. Bengio \(2013\)On the difficulty of training recurrent neural networks\.InICML \(3\),Vol\.28,pp\. 1310–1318\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px2.p1.6)\.
- A\. Paszke, S\. Gross, F\. Massa, A\. Lerer, J\. Bradbury, G\. Chanan, T\. Killeen, Z\. Lin, N\. Gimelshein, L\. Antiga, A\. Desmaison, A\. Köpf, E\. Z\. Yang, Z\. DeVito, M\. Raison, A\. Tejani, S\. Chilamkurthy, B\. Steiner, L\. Fang, J\. Bai, and S\. Chintala \(2019\)PyTorch: An Imperative Style, High\-Performance Deep Learning Library\.InNeurIPS,pp\. 8024–8035\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx11.SSSx1.p1.7),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px4.p1.1)\.
- G\. Penedo, H\. Kydlícek, L\. B\. Allal, A\. Lozhkov, M\. Mitchell, C\. A\. Raffel, L\. von Werra, and T\. Wolf \(2024\)The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale\.InNeurIPS,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px1.p1.1)\.
- G\. Pereyra, G\. Tucker, J\. Chorowski, L\. Kaiser, and G\. E\. Hinton \(2017\)Regularizing Neural Networks by Penalizing Confident Output Distributions\.InICLR \(Workshop\),Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3)\.
- O\. Press and L\. Wolf \(2017\)Using the Output Embedding to Improve Language Models\.InEACL \(2\),pp\. 157–163\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px3.p1.1)\.
- A\. Radford, J\. Wu, R\. Child, D\. Luan, D\. Amodei, I\. Sutskever,et al\.\(2019\)Language models are unsupervised multitask learners\.OpenAI blog1\(8\),pp\. 9\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px1.p1.1)\.
- C\. Riquelme, J\. Puigcerver, B\. Mustafa, M\. Neumann, R\. Jenatton, A\. S\. Pinto, D\. Keysers, and N\. Houlsby \(2021\)Scaling Vision with Sparse Mixture of Experts\.InNeurIPS,pp\. 8583–8595\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- N\. Shazeer, Y\. Cheng, N\. Parmar, D\. Tran, A\. Vaswani, P\. Koanantakool, P\. Hawkins, H\. Lee, M\. Hong, C\. Young, R\. Sepassi, and B\. A\. Hechtman \(2018\)Mesh\-TensorFlow: Deep Learning for Supercomputers\.InNeurIPS,pp\. 10435–10444\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.p1.1),[Introduction](https://arxiv.org/html/2609.16179#Sx1.p1.1),[Language\-Modeling Objective and Z\-Loss Setup](https://arxiv.org/html/2609.16179#Sx2.SSx1.p2.14)\.
- N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. V\. Le, G\. E\. Hinton, and J\. Dean \(2017\)Outrageously Large Neural Networks: The Sparsely\-Gated Mixture\-of\-Experts Layer\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- S\. Shleifer, J\. Weston, and M\. Ott \(2021\)NormFormer: Improved Transformer Pretraining with Extra Normalization\.CoRRabs/2110\.09456\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p2.1)\.
- M\. Shoeybi, M\. Patwary, R\. Puri, P\. LeGresley, J\. Casper, and B\. Catanzaro \(2019\)Megatron\-LM: Training Multi\-Billion Parameter Language Models Using Model Parallelism\.CoRRabs/1909\.08053\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1)\.
- F\. Stollenwerk, A\. Lokrantz, and N\. Hertzberg \(2026\)Output Embedding Centering for Stable LLM Pretraining\.CoRRabs/2601\.02031\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px3.p1.1),[Introduction](https://arxiv.org/html/2609.16179#Sx1.p1.1)\.
- C\. Szegedy, V\. Vanhoucke, S\. Ioffe, J\. Shlens, and Z\. Wojna \(2016\)Rethinking the Inception Architecture for Computer Vision\.InCVPR,pp\. 2818–2826\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3)\.
- G\. Team \(2024\)Gemma 2: Improving Open Language Models at a Practical Size\.CoRRabs/2408\.00118\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3)\.
- P\. Tillet, H\. Kung, and D\. D\. Cox \(2019\)Triton: an intermediate language and compiler for tiled neural network computations\.InMAPL@PLDI,pp\. 10–19\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px4.p1.1),[Introduction](https://arxiv.org/html/2609.16179#Sx1.p2.2)\.
- A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. Polosukhin \(2017\)Attention is All you Need\.InNIPS,pp\. 5998–6008\.Cited by:[Introduction](https://arxiv.org/html/2609.16179#Sx1.p2.2)\.
- L\. Wang, H\. Gao, C\. Zhao, X\. Sun, and D\. Dai \(2024\)Auxiliary\-Loss\-Free Load Balancing Strategy for Mixture\-of\-Experts\.CoRRabs/2408\.15664\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- H\. Wei, R\. Xie, H\. Cheng, L\. Feng, B\. An, and Y\. Li \(2022\)Mitigating Neural Network Overconfidence with Logit Normalization\.InICML,Vol\.162,pp\. 23631–23644\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p2.1)\.
- E\. Wijmans, B\. Huval, A\. Hertzberg, V\. Koltun, and P\. Krähenbühl \(2025\)Cut Your Losses in Large\-Vocabulary Language Models\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px4.p1.1)\.
- T\. Wolf, L\. Debut, V\. Sanh, J\. Chaumond, C\. Delangue, A\. Moi, P\. Cistac, T\. Rault, R\. Louf, M\. Funtowicz, J\. Davison, S\. Shleifer, P\. von Platen, C\. Ma, Y\. Jernite, J\. Plu, C\. Xu, T\. L\. Scao, S\. Gugger, M\. Drame, Q\. Lhoest, and A\. M\. Rush \(2020\)Transformers: State\-of\-the\-Art Natural Language Processing\.InEMNLP \(Demos\),pp\. 38–45\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx8.SSSx1.Px1.p1.1)\.
- R\. Xiong, Y\. Yang, D\. He, K\. Zheng, S\. Zheng, C\. Xing, H\. Zhang, Y\. Lan, L\. Wang, and T\. Liu \(2020\)On Layer Normalization in the Transformer Architecture\.InICML,Vol\.119,pp\. 10524–10533\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p2.1)\.
- Z\. Yang, Z\. Dai, R\. Salakhutdinov, and W\. W\. Cohen \(2018\)Breaking the Softmax Bottleneck: A High\-Rank RNN Language Model\.InICLR,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3)\.
- B\. Zhang and R\. Sennrich \(2019\)Root Mean Square Layer Normalization\.InNeurIPS,pp\. 12360–12371\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px4.p2.1)\.
- Y\. Zhou, T\. Lei, H\. Liu, N\. Du, Y\. Huang, V\. Y\. Zhao, A\. M\. Dai, Z\. Chen, Q\. V\. Le, and J\. Laudon \(2022\)Mixture\-of\-Experts with Expert Choice Routing\.InNeurIPS,Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3)\.
- B\. Zoph, I\. Bello, S\. Kumar, N\. Du, Y\. Huang, J\. Dean, N\. Shazeer, and W\. Fedus \(2022\)Designing Effective Sparse Expert Models\.CoRRabs/2202\.08906\.Cited by:[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.SSS0.Px1.p1.3),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx1.p1.1),[Appendix A](https://arxiv.org/html/2609.16179#A1.SSx10.p1.4),[Introduction](https://arxiv.org/html/2609.16179#Sx1.p1.1),[Language\-Modeling Objective and Z\-Loss Setup](https://arxiv.org/html/2609.16179#Sx2.SSx1.p2.14)\.

## Appendix AAppendix

### Related Work

The closest uses of the log\-normalizer penalty are final\-softmax Z\-loss in Mesh TensorFlow and PaLM and router Z\-loss in Stable and Transferable Mixture\-of\-Experts \(ST\-MoE\)\(Shazeeret al\.[2018](https://arxiv.org/html/2609.16179#bib.bib3); Chowdheryet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib7); Zophet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib6)\)\. To avoid a naming ambiguity, the squared log\-normalizer penalty in Eq\.[1](https://arxiv.org/html/2609.16179#Sx2.E1)is mathematically distinct from an earlier loss with the same name\.de Brébisson and Vincent \([2016](https://arxiv.org/html/2609.16179#bib.bib10)\)proposed a shift\- and scale\-invariant spherical surrogate based on the standardized target\-class score as an alternative to log\-softmax, rather than an auxiliary penalty onlog⁡Z\\log Z\. Relative to output\-centering and logit\-geometry work, our focus is the architecture\- and implementation\-dependent transport of the resulting adjoint\. We develop the broader context on routers, output embeddings, common shifts, logit geometry, systems, and the positioning of this study below\.

##### Z\-loss, router\-logit control, and MoE reporting\.

Conditional computation and sparse experts date back to adaptive mixtures of local experts\(Jacobset al\.[1991](https://arxiv.org/html/2609.16179#bib.bib29)\)and the sparsely\-gated MoE layer\(Shazeeret al\.[2017](https://arxiv.org/html/2609.16179#bib.bib2)\)\. GShard, Switch Transformer, and ST\-MoE then established large sparse Transformer routing conventions and training\-loss\-divergence issues\(Lepikhinet al\.[2021](https://arxiv.org/html/2609.16179#bib.bib5); Feduset al\.[2022](https://arxiv.org/html/2609.16179#bib.bib4); Zophet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib6)\)\. Subsequent sparse language models, routing methods, and systems have continued to emphasize the same design variables\. Examples include BASE Layers, Expert Choice routing, V\-MoE, dropless MoE systems, and routed scaling\-law studies\(Lewiset al\.[2021](https://arxiv.org/html/2609.16179#bib.bib21); Zhouet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib22); Galeet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib47); Riquelmeet al\.[2021](https://arxiv.org/html/2609.16179#bib.bib30); Clarket al\.[2022](https://arxiv.org/html/2609.16179#bib.bib31)\)\. This line of work shows that top\-kkrouting, balancing, capacity, dynamic expert workload, and compute\-normalized scaling remain central design variables\. The log\-normalizer penalty we study follows the Mesh TensorFlow and PaLM final\-softmax convention and the ST\-MoE router adaptation\(Shazeeret al\.[2018](https://arxiv.org/html/2609.16179#bib.bib3); Chowdheryet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib7); Zophet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib6)\)\. Our MoE contribution extends this reporting line\. NominalλR\\lambda\_\{R\}is not sufficient to compare experiments unless the token\-layer denominator, top\-kk, active\-route convention, capacity policy, absolute per\-decision coefficient, and scale relative to the token\-layer mean are reported\. Recent auxiliary\-loss\-free load balancing further shows that router\-logit control and load balancing can be implemented without directly adding auxiliary gradients to the language\-model objective\(Wanget al\.[2024](https://arxiv.org/html/2609.16179#bib.bib55)\)\.OLMoet al\.\([2025](https://arxiv.org/html/2609.16179#bib.bib8)\)document modern dense large language model \(LLM\) training recipes that include logit\-growth\-control terms and transparent recipe reporting\. Our reporting recommendations turn that practice into Z\-loss\-specific backward\-source and transport diagnostics\.

##### Numerical losses and fused implementations\.

Implementation audits in this work build on several lines of research in numerical analysis and systems\. The relevant numerical and systems literature includes mixed\-precision training, floating\-point rounding behavior, large\-scale model\-parallel training stacks, accurate log\-sum\-exp and softmax computation, online softmax normalization, input\- and output\-aware fused kernels, programmable GPU\-kernel languages and compilers such asTriton, adaptive large\-vocabulary softmax, and memory\-efficient large\-vocabulary CE\(Micikeviciuset al\.[2018](https://arxiv.org/html/2609.16179#bib.bib19); Goldberg[1991](https://arxiv.org/html/2609.16179#bib.bib48); Shoeybiet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib54); Blanchardet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib32); Milakov and Gimelshein[2018](https://arxiv.org/html/2609.16179#bib.bib42); Daoet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib20); Tilletet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib26); Graveet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib43); Wijmanset al\.[2025](https://arxiv.org/html/2609.16179#bib.bib33)\)\. The PyTorch CE audit is close to the fp32 reference in our tests, consistent with the stable softmax and CE formulations used by mature frameworks\(Paszkeet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib27)\)\. Adding Z\-loss creates an additional implementation obligation\. When the auxiliary path is fused, quantized, or reconstructed separately, the backward source of the auxiliary path must remain consistent with the forward objective\. Forward scalar agreement alone does not establish this consistency\. Because the Z\-loss source is later transported through tied heads, unembedding geometry, and adaptive optimizer normalization, a fused CE\+Z\-loss path should be validated with source consistency, selected\-parameter gradients, and optimizer directions on target\-model logits from the relevant corpus distribution\.

##### Output embeddings, weight tying, and common\-shift gauges\.

Weight tying and output embeddings affect language\-model parameterization and generalization\(Press and Wolf[2017](https://arxiv.org/html/2609.16179#bib.bib14); Inanet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib15)\)\. Work on embedding degeneration and representation geometry further shows that lexical vectors can concentrate in narrow cones or dominant mean and top directions\. This pattern appears in maximum\-likelihood training, contextual Transformer states, and static embedding spaces\(Gaoet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib38); Ethayarajh[2019](https://arxiv.org/html/2609.16179#bib.bib45); Mu and Viswanath[2018](https://arxiv.org/html/2609.16179#bib.bib46)\)\. That literature characterizes representational sharing, anisotropy, and predictive quality\. Our tied\-embedding diagnostic instead traces how an output\-side auxiliary source returns through both the output head and the input lexical path\.Stollenwerket al\.\([2026](https://arxiv.org/html/2609.16179#bib.bib9)\)directly target common\-shift instability in LLM pretraining through output embedding centering\. We use centering both as a mitigation and as an intervention that exposes the gauge freedom of the output head\. The reason is that CE is invariant to a common shift while Z\-loss is not\. The removed\-source proposition identifies the exact softmax\-shaped component removed by centering and explains why centering can help or introduce different trade\-offs across models and training regimes\.

##### Logit geometry and normalized heads\.

A broad literature reduces or analyzes logit geometry through temperature scaling, label smoothing, confidence penalties, normalized heads, output\-centering methods, final\-logit soft\-capping, activation normalizers, and studies of softmax output\-layer rank\(Guoet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib23); Szegedyet al\.[2016](https://arxiv.org/html/2609.16179#bib.bib49); Mülleret al\.[2019](https://arxiv.org/html/2609.16179#bib.bib50); Pereyraet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib35); Nguyen and Salazar[2019](https://arxiv.org/html/2609.16179#bib.bib25); Weiet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib24); Team[2024](https://arxiv.org/html/2609.16179#bib.bib34); Baet al\.[2016](https://arxiv.org/html/2609.16179#bib.bib36); Zhang and Sennrich[2019](https://arxiv.org/html/2609.16179#bib.bib37); Yanget al\.[2018](https://arxiv.org/html/2609.16179#bib.bib16)\)\. The shared objects are logit scale, shift, confidence, and expressivity\. Relative to this literature, the distinction we draw is the coordinate of intervention\. A method can reducelog⁡Z\\log Zor the centered relative log\-normalizer while leavingWU⊤​pW\_\{U\}^\{\\top\}pand optimizer\-facing update tails nearly unchanged\. This coordinate distinction motivates the experimental diagnostics\. We includeAp99A\_\{p\}^\{99\}, gradient\-event rates, Adam\-state endpoints, and randomized low\-rank unembedding audits\(Halkoet al\.[2011](https://arxiv.org/html/2609.16179#bib.bib28)\)rather than comparing scalar Z\-loss values alone\.

The common\-shift coordinate isolated here is related to temperature calibration and logit\-norm regularization\(Guoet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib23); Weiet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib24)\)\. The common\-shift coordinate is also related to embedding and activation normalization choices in Transformer training\(Nguyen and Salazar[2019](https://arxiv.org/html/2609.16179#bib.bib25); Xionget al\.[2020](https://arxiv.org/html/2609.16179#bib.bib40); Shleiferet al\.[2021](https://arxiv.org/html/2609.16179#bib.bib51)\)\. The coordinate is nevertheless distinct from layer normalization and root\-mean\-square layer normalization\. These methods normalize or recenter hidden states rather than removing the output\-head common\-shift gauge\(Baet al\.[2016](https://arxiv.org/html/2609.16179#bib.bib36); Zhang and Sennrich[2019](https://arxiv.org/html/2609.16179#bib.bib37)\)\.

##### Positioning of this study\.

Taken together, the surrounding literature has treated Z\-loss as a recipe\-level mitigation for training\-loss divergence, a router auxiliary term, a logit\-geometry regularizer, or a numerical\-loss implementation detail\. This paper occupies the intersection of those views but focuses on the logit\-space adjoint created after evaluation of the scalar log\-normalizer penalty and on the transport of the logit\-space adjoint through architecture, precision, routing reductions, and optimizer state\. The contribution is therefore not a new sparse architecture, a new normalization layer, or a compute\-scaling claim\. Instead, the contribution is a transport\-aware measurement, reporting, and intervention framework for existing dense and sparse Transformer recipes\. The source identities are softmax\-axis identities, but the evidence and recommendations intentionally target the Transformer regimes where Z\-loss is deployed to constrain output\-head or router log\-normalizers and where the relevant transport paths arise together\. By deriving source, common\-shift, tied\-path, gain, and router\-scale identities and pairing those identities with transport diagnostics and validation PPL measured under matched run configurations, we make otherwise similar Z\-loss reports comparable in the coordinates that affect updates\. This coordinate\-level comparison marks the point at which the results differ from the common scalar\-control narrative\. A method can improve the reported Z\-loss value while failing to reduce the update coordinates that matter for mixed\-precision headroom and adaptive\-optimizer state; the relevant update coordinates can even increase\. Accordingly, the proposed interventions have distinct roles\. Centering, factorization, gain\-aware weighting, and scale\-matched router reductions act on particular transport coordinates and are not interchangeable ways to lower a scalar Z\-loss curve\.

### Raw\-Coordinate Source and Adjoint

Having fixed the CE\-invariant gauge, we now compare the Z\-loss sources produced by raw and centered deployed coordinates\.

###### Proposition 5\(Raw and centered Z\-loss sources\)\.

Letz=z~\+μ​𝟏z=\\widetilde\{z\}\+\\mu\\mathbf\{1\}and let the raw and centered Z\-losses use coefficients and targets\(λraw,craw\)\(\\lambda\_\{\\mathrm\{raw\}\},c\_\{\\mathrm\{raw\}\}\)and\(λctr,cctr\)\(\\lambda\_\{\\mathrm\{ctr\}\},c\_\{\\mathrm\{ctr\}\}\)\. Sincesoftmax⁡\(z\)=softmax⁡\(z~\)=p\\operatorname\{softmax\}\(z\)=\\operatorname\{softmax\}\(\\widetilde\{z\}\)=p, the raw source and the centered source differ in the respective deployed\-logit coordinates by

δz,rawZ−δz~,centeredZ\\displaystyle\\delta\_\{z,\\mathrm\{raw\}\}^\{Z\}\-\\delta\_\{\\widetilde\{z\},\\mathrm\{centered\}\}^\{Z\}=2​λraw​\(μ\+log⁡Z~−craw\)​p\\displaystyle=2\\lambda\_\{\\mathrm\{raw\}\}\(\\mu\+\\log\\widetilde\{Z\}\-c\_\{\\mathrm\{raw\}\}\)p\(15\)−2​λctr​\(log⁡Z~−cctr\)​p\.\\displaystyle\-2\\lambda\_\{\\mathrm\{ctr\}\}\(\\log\\widetilde\{Z\}\-c\_\{\\mathrm\{ctr\}\}\)p\.

###### Corollary 1\(Removed deployed source component\)\.

When both losses use the same coefficientλ\\lambdaand the same targetcc,

δz,rawZ−δz~,centeredZ=2​λ​μ​p\.\\delta\_\{z,\\mathrm\{raw\}\}^\{Z\}\-\\delta\_\{\\widetilde\{z\},\\mathrm\{centered\}\}^\{Z\}=2\\lambda\\mu p\.\(16\)

###### Proposition 6\(Projected raw\-logit adjoint\)\.

Under the same single\-coefficient, single\-target convention as Eq\.[16](https://arxiv.org/html/2609.16179#A1.E16), if centered logits are implemented asz~=P​z\\widetilde\{z\}=PzwithP≔I−\(1/V\)​𝟏𝟏⊤P\\coloneqq I\-\(1/V\)\\mathbf\{1\}\\mathbf\{1\}^\{\\top\},p=softmax⁡\(P​z\)=softmax⁡\(z\)p=\\operatorname\{softmax\}\(Pz\)=\\operatorname\{softmax\}\(z\), andZ~≔∑iexp⁡\(\(P​z\)i\)\\widetilde\{Z\}\\coloneqq\\sum\_\{i\}\\exp\(\(Pz\)\_\{i\}\), then the adjoint transported back to the raw\-logit coordinate is

∇zℒZ​\(P​z\)=2​λ​\(log⁡Z~−c\)​\(p−1V​𝟏\)\.\\nabla\_\{z\}\\mathcal\{L\}\_\{Z\}\(Pz\)=2\\lambda\(\\log\\widetilde\{Z\}\-c\)\\left\(p\-\\frac\{1\}\{V\}\\mathbf\{1\}\\right\)\.\(17\)

Output\-head centering preserves CE but intentionally changes the Z\-loss source coordinate\. The effect of centering should be judged by transport diagnostics such asPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}, with the complete diagnostic suite specified later, because removing the common\-shift channel can be beneficial or introduce different trade\-offs across model geometries\.

### Factorized Objective

Writingν≔log⁡Z~\\nu\\coloneqq\\log\\widetilde\{Z\}and choosing targets withc=cμ\+crelc=c\_\{\\mu\}\+c\_\{\\mathrm\{rel\}\}, the original Z\-loss expands as

λ​\(μ\+ν−c\)2\\displaystyle\\lambda\(\\mu\+\\nu\-c\)^\{2\}=λ​\(μ−cμ\)2\+2​λ​\(μ−cμ\)​\(ν−crel\)\\displaystyle=\\lambda\(\\mu\-c\_\{\\mu\}\)^\{2\}\+2\\lambda\(\\mu\-c\_\{\\mu\}\)\(\\nu\-c\_\{\\mathrm\{rel\}\}\)\(18\)\+λ​\(ν−crel\)2\.\\displaystyle\+\\lambda\(\\nu\-c\_\{\\mathrm\{rel\}\}\)^\{2\}\.Eq\.[18](https://arxiv.org/html/2609.16179#A1.E18)motivates objective variants that target different transport coordinates\. The factorized objective is a deliberately factorized inductive bias that drops the cross term and allows separate coefficients and targets for shift and relative\-logit normalization\.

ℒfact≔λμ​\(μ−cμ\)2\+λrel​\(log⁡Z~−crel\)2\.\\mathcal\{L\}\_\{\\mathrm\{fact\}\}\\coloneqq\\lambda\_\{\\mu\}\(\\mu\-c\_\{\\mu\}\)^\{2\}\+\\lambda\_\{\\mathrm\{rel\}\}\(\\log\\widetilde\{Z\}\-c\_\{\\mathrm\{rel\}\}\)^\{2\}\.\(19\)

### Architecture\-Aware Objectives

For the gain\-aware Z\-loss runs, we use a detached output\-to\-hidden\-gain weight withβ≥0\\beta\\geq 0,

Ap​\(z\)\\displaystyle A\_\{p\}\(z\)≔‖WU⊤​p‖2‖p‖2\+ε,\\displaystyle\\coloneqq\\frac\{\\left\\lVert W\_\{U\}^\{\\top\}p\\right\\rVert\_\{2\}\}\{\\left\\lVert p\\right\\rVert\_\{2\}\+\\varepsilon\},\(20\)ℒgain\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{gain\}\}≔λ​\(log⁡Z−c\)21\+β​sg⁡\(Ap​\(z\)2\)\.\\displaystyle\\coloneqq\\lambda\\frac\{\(\\log Z\-c\)^\{2\}\}\{1\+\\beta\\operatorname\{sg\}\(A\_\{p\}\(z\)^\{2\}\)\}\.The stop\-gradient is part of the objective used in the experiments; without the stop\-gradient, differentiatingAp​\(z\)A\_\{p\}\(z\)would add extra source terms\. With the detached gain, the logit\-coordinate source is

∇zℒgain=2​λ​\(log⁡Z−c\)1\+β​sg⁡\(Ap​\(z\)2\)​p\.\\nabla\_\{z\}\\mathcal\{L\}\_\{\\mathrm\{gain\}\}=\\frac\{2\\lambda\(\\log Z\-c\)\}\{1\+\\beta\\operatorname\{sg\}\(A\_\{p\}\(z\)^\{2\}\)\}p\.\(21\)

### Router Effective\-Scale Derivations

This subsection gives the details of the router objective and reduction\-scale statements summarized in the main text\. Define

ZR,t,ℓ\\displaystyle Z\_\{R,t,\\ell\}≔∑e=1Nexpexp⁡rt,ℓ,e,\\displaystyle\\coloneqq\\sum\_\{e=1\}^\{N\_\{\\mathrm\{exp\}\}\}\\exp r\_\{t,\\ell,e\},\(22\)ℒR\\displaystyle\\mathcal\{L\}\_\{R\}≔λR​∑t,ℓαt,ℓ​\(log⁡ZR,t,ℓ−cR\)2\.\\displaystyle\\coloneqq\\lambda\_\{R\}\\sum\_\{t,\\ell\}\\alpha\_\{t,\\ell\}\\left\(\\log Z\_\{R,t,\\ell\}\-c\_\{R\}\\right\)^\{2\}\.
###### Proposition 7\(Router coefficient and relative scale\)\.

LetNtlN\_\{\\mathrm\{tl\}\}be the number of token\-layer router decisions, and suppose token\-layer\-mean router Z\-loss usesα=1/Ntl\\alpha=1/N\_\{\\mathrm\{tl\}\}\. If an implementation uses a uniform per\-decision coefficientλR​α\\lambda\_\{R\}\\alpha, the effective gradient scale relative to the token\-layer mean isNtl​αN\_\{\\mathrm\{tl\}\}\\alpha\. More generally, decision\(t,ℓ\)\(t,\\ell\)has a relative scale ofNtl​αt,ℓN\_\{\\mathrm\{tl\}\}\\alpha\_\{t,\\ell\}\.

For the top\-kkconvention statements below, write

ϕt,ℓ≔\(log⁡ZR,t,ℓ−cR\)2\.\\displaystyle\\phi\_\{t,\\ell\}\\coloneqq\\left\(\\log Z\_\{R,t,\\ell\}\-c\_\{R\}\\right\)^\{2\}\.
###### Corollary 2\(Top\-kkactive\-route mean scale\)\.

Under the same token\-layer\-mean reference, assuming exactlykkactive routes per token\-layer decision, if one computes a single violation per token\-layer decision, does not replicate the violation in the numerator, and nevertheless normalizes by the active\-route countNtl​kN\_\{\\mathrm\{tl\}\}k,

ℒRactiveMean≔λR​1Ntl​k​∑t,ℓϕt,ℓ,\\displaystyle\\mathcal\{L\}\_\{R\}^\{\\mathrm\{activeMean\}\}\\coloneqq\\lambda\_\{R\}\\frac\{1\}\{N\_\{\\mathrm\{tl\}\}k\}\\sum\_\{t,\\ell\}\\phi\_\{t,\\ell\},then the effective scale is1/k1/k\.

###### Corollary 3\(Top\-kkactive\-route sum scale\)\.

Under the same token\-layer\-mean reference, assuming exactlykkactive routes per token\-layer decision, if an implementation sumskkactive\-route replicas of the violation before applying the token\-layer mean,

ℒRactiveSum≔λR​1Ntl​∑t,ℓ∑j=1kϕt,ℓ=λR​kNtl​∑t,ℓϕt,ℓ,\\displaystyle\\mathcal\{L\}\_\{R\}^\{\\mathrm\{activeSum\}\}\\coloneqq\\lambda\_\{R\}\\frac\{1\}\{N\_\{\\mathrm\{tl\}\}\}\\sum\_\{t,\\ell\}\\sum\_\{j=1\}^\{k\}\\phi\_\{t,\\ell\}=\\lambda\_\{R\}\\frac\{k\}\{N\_\{\\mathrm\{tl\}\}\}\\sum\_\{t,\\ell\}\\phi\_\{t,\\ell\},then the effective scale iskk\.

##### Top\-kkconventions\.

In the experiments and tables, active\-route mean denotes the single\-violation convention that normalizes by the active\-route count\. Active\-route sum denotes the convention that sums active\-route replicas before taking the token\-layer mean in Corollary[3](https://arxiv.org/html/2609.16179#Thmcorollary3)\. Thus, the reported effective scale is1/k1/kfor active\-route mean by Corollary[2](https://arxiv.org/html/2609.16179#Thmcorollary2)\. The reported effective scale iskkfor active\-route sum by Corollary[3](https://arxiv.org/html/2609.16179#Thmcorollary3)\.

### Formal Details and Reference Implementation

#### Notation

Table 4:Notation summary

#### Proofs

##### Z\-loss common\-shift sensitivity\.

###### Proof\.

A direct computation gives

Z​\(z\+a​𝟏\)\\displaystyle Z\(z\+a\\mathbf\{1\}\)=ea​Z​\(z\),log⁡Z​\(z\+a​𝟏\)=log⁡Z​\(z\)\+a,\\displaystyle=e^\{a\}Z\(z\),\\qquad\\log Z\(z\+a\\mathbf\{1\}\)=\\log Z\(z\)\+a,and therefore

ℒZ​\(z\+a​𝟏\)=λ​\(log⁡Z​\(z\)\+a−c\)2\.\\displaystyle\\mathcal\{L\}\_\{Z\}\(z\+a\\mathbf\{1\}\)=\\lambda\(\\log Z\(z\)\+a\-c\)^\{2\}\.Differentiating this shifted objective with respect toziz\_\{i\}while holdingaafixed yields

∂∂zi​ℒZ​\(z\+a​𝟏\)=2​λ​\(log⁡Z​\(z\)\+a−c\)​ezi∑jezj,\\displaystyle\\frac\{\\partial\}\{\\partial z\_\{i\}\}\\mathcal\{L\}\_\{Z\}\(z\+a\\mathbf\{1\}\)=2\\lambda\(\\log Z\(z\)\+a\-c\)\\frac\{e^\{z\_\{i\}\}\}\{\\sum\_\{j\}e^\{z\_\{j\}\}\},because∂logZ\(z\)/∂zi=softmax\(z\)i\\partial\\log Z\(z\)/\\partial z\_\{i\}=\\operatorname\{softmax\}\(z\)\_\{i\}\. Collecting coordinates gives the displayed gradient\. ∎

##### Centering preserves CE\.

###### Proof\.

By the definition ofw¯\\bar\{w\},

z~\\displaystyle\\widetilde\{z\}=W~U​u\\displaystyle=\\widetilde\{W\}\_\{U\}u=WU​u−𝟏​\(w¯​u\)\\displaystyle=W\_\{U\}u\-\\mathbf\{1\}\(\\bar\{w\}u\)=z−α​𝟏,α≔w¯​u∈ℝ\.\\displaystyle=z\-\\alpha\\mathbf\{1\},\\qquad\\alpha\\coloneqq\\bar\{w\}u\\in\\mathbb\{R\}\.Thus, centering subtracts the same scalar from every logit coordinate of the token\. For each vocabulary indexii,

softmax\(z~\)i=ezi−α∑jezj−α=ezi∑jezj=softmax\(z\)i\.\\displaystyle\\operatorname\{softmax\}\(\\widetilde\{z\}\)\_\{i\}=\\frac\{e^\{z\_\{i\}\-\\alpha\}\}\{\\sum\_\{j\}e^\{z\_\{j\}\-\\alpha\}\}=\\frac\{e^\{z\_\{i\}\}\}\{\\sum\_\{j\}e^\{z\_\{j\}\}\}=\\operatorname\{softmax\}\(z\)\_\{i\}\.Since token CE with a one\-hot target is the negative logarithm of the target softmax probability, the CE loss is identical forzzandz~\\widetilde\{z\}\. The argument relies only on exact algebra, hence the stated real\-arithmetic qualification\. ∎

##### Raw and centered Z\-loss sources\.

###### Proof\.

Forz=z~\+μ​𝟏z=\\widetilde\{z\}\+\\mu\\mathbf\{1\}, the two deployed\-logit coordinates have the same softmaxpp, andZ​\(z\)=eμ​Z~Z\(z\)=e^\{\\mu\}\\widetilde\{Z\}, solog⁡Z​\(z\)=μ\+log⁡Z~\\log Z\(z\)=\\mu\+\\log\\widetilde\{Z\}\. Using Eq\.[2](https://arxiv.org/html/2609.16179#Sx2.E2), the raw and centered Z\-loss sources are therefore

δz,rawZ\\displaystyle\\delta\_\{z,\\mathrm\{raw\}\}^\{Z\}=2​λraw​\(μ\+log⁡Z~−craw\)​p,\\displaystyle=2\\lambda\_\{\\mathrm\{raw\}\}\(\\mu\+\\log\\widetilde\{Z\}\-c\_\{\\mathrm\{raw\}\}\)p,δz~,centeredZ\\displaystyle\\delta\_\{\\widetilde\{z\},\\mathrm\{centered\}\}^\{Z\}=2​λctr​\(log⁡Z~−cctr\)​p\.\\displaystyle=2\\lambda\_\{\\mathrm\{ctr\}\}\(\\log\\widetilde\{Z\}\-c\_\{\\mathrm\{ctr\}\}\)p\.Subtracting the two expressions gives Eq\.[15](https://arxiv.org/html/2609.16179#A1.E15)\. ∎

##### Removed deployed source component\.

###### Proof\.

Setλraw=λctr=λ\\lambda\_\{\\mathrm\{raw\}\}=\\lambda\_\{\\mathrm\{ctr\}\}=\\lambdaandcraw=cctr=cc\_\{\\mathrm\{raw\}\}=c\_\{\\mathrm\{ctr\}\}=cin Eq\.[15](https://arxiv.org/html/2609.16179#A1.E15)\. ∎

##### Projected raw\-logit adjoint\.

###### Proof\.

If the centered loss is viewed as a function of the raw logits throughz~=P​z\\widetilde\{z\}=Pz, then the chain rule gives

∇zℒZ​\(P​z\)=P⊤​δz~,centeredZ\.\\displaystyle\\nabla\_\{z\}\\mathcal\{L\}\_\{Z\}\(Pz\)=P^\{\\top\}\\delta\_\{\\widetilde\{z\},\\mathrm\{centered\}\}^\{Z\}\.SincePPis symmetric andP​p=p−\(1/V\)​𝟏Pp=p\-\(1/V\)\\mathbf\{1\}, Eq\.[17](https://arxiv.org/html/2609.16179#A1.E17)follows\. ∎

##### Tied embedding gradient decomposition\.

###### Proof\.

LetΦ​\(Ein,Eout\)\\Phi\(E\_\{\\mathrm\{in\}\},E\_\{\\mathrm\{out\}\}\)denote the scalar Z\-loss of the untied network and define the tied objective byϕ​\(E\)≔Φ​\(E,E\)\\phi\(E\)\\coloneqq\\Phi\(E,E\)\. For any perturbationHHof the tied embedding table,

D​ϕ​\(E\)​\[H\]=DEin​Φ​\(E,E\)​\[H\]\+DEout​Φ​\(E,E\)​\[H\],\\displaystyle D\\phi\(E\)\[H\]=D\_\{E\_\{\\mathrm\{in\}\}\}\\Phi\(E,E\)\[H\]\+D\_\{E\_\{\\mathrm\{out\}\}\}\\Phi\(E,E\)\[H\],which gives the displayed decomposition after identifying differentials with Frobenius inner products\. ∎

##### Gain is anisotropy\-weighted\.

###### Proof\.

LetWU=U​Σ​V⊤W\_\{U\}=U\\Sigma V^\{\\top\}be a singular value decomposition, with left singular vectorsuju\_\{j\}and singular valuesσj≥0\\sigma\_\{j\}\\geq 0\. Sinceq=p/‖p‖2q=p/\\left\\lVert p\\right\\rVert\_\{2\}andε\\varepsilonis ignored,

a​\(p\)2=‖WU⊤​p‖22‖p‖22=‖WU⊤​q‖22\.\\displaystyle a\(p\)^\{2\}=\\frac\{\\left\\lVert W\_\{U\}^\{\\top\}p\\right\\rVert\_\{2\}^\{2\}\}\{\\left\\lVert p\\right\\rVert\_\{2\}^\{2\}\}=\\left\\lVert W\_\{U\}^\{\\top\}q\\right\\rVert\_\{2\}^\{2\}\.UsingWU⊤=V​Σ⊤​U⊤W\_\{U\}^\{\\top\}=V\\Sigma^\{\\top\}U^\{\\top\}and the orthonormality of the columns ofVV,

‖WU⊤​q‖22=‖Σ⊤​U⊤​q‖22=∑jσj2​⟨q,uj⟩2\.\\displaystyle\\left\\lVert W\_\{U\}^\{\\top\}q\\right\\rVert\_\{2\}^\{2\}=\\left\\lVert\\Sigma^\{\\top\}U^\{\\top\}q\\right\\rVert\_\{2\}^\{2\}=\\sum\_\{j\}\\sigma\_\{j\}^\{2\}\\left\\langle q,u\_\{j\}\\right\\rangle^\{2\}\.Because theuju\_\{j\}form an orthonormal set,∑j⟨q,uj⟩2≤‖q‖22=1\\sum\_\{j\}\\left\\langle q,u\_\{j\}\\right\\rangle^\{2\}\\leq\\left\\lVert q\\right\\rVert\_\{2\}^\{2\}=1\. Hence,

a​\(p\)2≤σmax​\(WU\)2​∑j⟨q,uj⟩2≤σmax​\(WU\)2\.\\displaystyle a\(p\)^\{2\}\\leq\\sigma\_\{\\max\}\(W\_\{U\}\)^\{2\}\\sum\_\{j\}\\left\\langle q,u\_\{j\}\\right\\rangle^\{2\}\\leq\\sigma\_\{\\max\}\(W\_\{U\}\)^\{2\}\.∎

##### Router coefficient and relative scale\.

###### Proof\.

For a fixed token\-layer router decision, define

ϕt,ℓ≔\(log⁡ZR,t,ℓ−cR\)2\.\\displaystyle\\phi\_\{t,\\ell\}\\coloneqq\\left\(\\log Z\_\{R,t,\\ell\}\-c\_\{R\}\\right\)^\{2\}\.In the token\-layer\-mean convention,

ℒRtok≔λR​1Ntl​∑t,ℓϕt,ℓ,\\displaystyle\\mathcal\{L\}\_\{R\}^\{\\mathrm\{tok\}\}\\coloneqq\\lambda\_\{R\}\\frac\{1\}\{N\_\{\\mathrm\{tl\}\}\}\\sum\_\{t,\\ell\}\\phi\_\{t,\\ell\},so the coefficient multiplying each decision isλR/Ntl\\lambda\_\{R\}/N\_\{\\mathrm\{tl\}\}\. Differentiating Eq\.[22](https://arxiv.org/html/2609.16179#A1.E22)for a fixed decision shows that the router Z\-loss gradient is linear in the absolute per\-decision coefficientλR​αt,ℓ\\lambda\_\{R\}\\alpha\_\{t,\\ell\}\. Thus, the scale relative to the token\-layer mean is

λR​αλR/Ntl=Ntl​α\.\\displaystyle\\frac\{\\lambda\_\{R\}\\alpha\}\{\\lambda\_\{R\}/N\_\{\\mathrm\{tl\}\}\}=N\_\{\\mathrm\{tl\}\}\\alpha\.∎

##### Top\-kkactive\-route mean scale\.

###### Proof\.

For the convention that divides a single violation by the active\-route count,α=1/\(Ntl​k\)\\alpha=1/\(N\_\{\\mathrm\{tl\}\}k\), so Proposition[7](https://arxiv.org/html/2609.16179#Thmproposition7)givesNtl​α=1/kN\_\{\\mathrm\{tl\}\}\\alpha=1/k\. ∎

##### Top\-kkactive\-route sum scale\.

###### Proof\.

For active\-route sum before the token\-layer mean, the per\-decision coefficient isk​λR/Ntlk\\lambda\_\{R\}/N\_\{\\mathrm\{tl\}\}, equivalentlyα=k/Ntl\\alpha=k/N\_\{\\mathrm\{tl\}\}, so Proposition[7](https://arxiv.org/html/2609.16179#Thmproposition7)givesNtl​α=kN\_\{\\mathrm\{tl\}\}\\alpha=k\. ∎

#### Reference Python Implementation

Listing[1](https://arxiv.org/html/2609.16179#LST1)gives an example implementation of the dense output\-head operations used in the experiments\. The first function deploys centered logits by subtracting only the fp32 output\-head common shift from already computed raw logits, matching Eq\.[38](https://arxiv.org/html/2609.16179#A1.E38)\. The second function returns the Z\-loss logit adjoint from Eq\.[2](https://arxiv.org/html/2609.16179#Sx2.E2), including the mean\-reduction scaling used by the training and audit code\.

Listing 1:Python implementation example of centered deployed logits and Z\-loss logit\-source construction\.1importtorch

2importtorch\.nn\.functionalasF

5defcentered\_logits\_from\_raw\(raw\_logits,hidden,weight,bias=None\):

6

7mean\_w=weight\.float\(\)\.mean\(dim=0,keepdim=True\)

8mean\_b=NoneifbiasisNoneelsebias\.float\(\)\.mean\(\)\.view\(1\)

9common\_shift=F\.linear\(hidden\.float\(\),mean\_w,mean\_b\)

10returnraw\_logits\.float\(\)\-common\_shift

13defz\_loss\_source\(logits,coef,target,reduction="mean"\):

14

15ifreductionnotin\{"mean","sum","none"\}:

16raiseValueError\(f"unknownreduction:\{reduction\}"\)

17logits\_f=logits\.float\(\)

18log\_z=torch\.logsumexp\(logits\_f,dim=\-1,keepdim=True\)

19probs=torch\.softmax\(logits\_f,dim=\-1\)

20scale=2\.0\*coef\*\(log\_z\-target\)

21ifreduction=="mean":

22scale=scale/logits\_f\[\.\.\.,0\]\.numel\(\)

23returnscale\*probs

### Diagnostics, Interventions, and Coefficient Regimes

#### Diagnostic Decision Framework

We use diagnostics as a decision hierarchy rather than as a leaderboard\. A diagnostic first identifies an active source or transport coordinate; the corresponding intervention targets that coordinate; and optimizer\-facing endpoints, together with validation PPL, determine whether the intervention is acceptable\. Appendix Table[5](https://arxiv.org/html/2609.16179#A1.T5)gives the complete mapping, including the effects that each intervention does and does not guarantee\.

At the source\-coordinate level, p99\|μ\|\|\\mu\|diagnoses common\-shift excursions\. Row\-centered deployment sets this coordinate to zero while preserving the softmax distribution and CE, but row\-centered deployment does not guarantee reductions inPZ99\.9P\_\{Z\}^\{99\.9\}orAp99A\_\{p\}^\{99\}; those remain geometry\-dependent outcomes\. Centered Z\-loss additionally controls the log\-normalizer in the remaining relative\-logit coordinate\. The factorized objective in Eq\.[19](https://arxiv.org/html/2609.16179#A1.E19)is used as a controlled decomposition of the shift and relative\-logit coordinates rather than as a universally preferred replacement for raw or centered Z\-loss\.

At the architectural\-transport level,Ap99A\_\{p\}^\{99\}diagnoses output\-to\-hidden amplification\. Gain\-aware weighting in Eq\.[20](https://arxiv.org/html/2609.16179#A1.E20)attenuates the source conditional on that gain and is evaluated by hidden\-source and optimizer\-facing tails; gain\-aware weighting is not expected to reduce theAp99A\_\{p\}^\{99\}metric\. For tied embeddings, output\- and input\-path gradient norms and cosines provide a pathway audit rather than a separate intervention\. At the implementation level,Δsrc\\Delta\_\{\\mathrm\{src\}\}andΔθ\\Delta\_\{\\theta\}diagnose forward–backward source inconsistency\. Reusing forward log\-sum\-exp statistics removes reconstruction inconsistency, although reuse cannot recover precision already lost through quantized logit storage\. For MoE routers,sRrel=Ntl​αs\_\{R\}^\{\\mathrm\{rel\}\}=N\_\{\\mathrm\{tl\}\}\\alphais a deterministic scale audit, and coefficients are matched wheneversRrel≠1s\_\{R\}^\{\\mathrm\{rel\}\}\\neq 1\.

We distinguish these root\-cause diagnostics from downstream endpoints\. Here, Diag\. Z is computed on the deployed logits, andPZ99\.9≔Q0\.999,t​\(\|log⁡Zt−c\|\)P\_\{Z\}^\{99\.9\}\\coloneqq Q\_\{0\.999,t\}\(\|\\log Z\_\{t\}\-c\|\)measures the corresponding log\-normalizer tail; neither is a coordinate\-free score across raw and centered deployments\. The gain tailAp99A\_\{p\}^\{99\}is defined in Eq\.[13](https://arxiv.org/html/2609.16179#Sx2.E13)\.RZR\_\{Z\}andRauxR\_\{\\rm aux\}describe auxiliary\-gradient strength relative to CE, whileCZC\_\{Z\}andCauxC\_\{\\rm aux\}describe CE alignment; the gradient ratios and cosines are descriptive rather than quantities every intervention should minimize\. Pre\-clip gradient events, Adam\-state tails, and static fp16 overflow margins measure optimizer\-facing consequences\. No tail reduction alone is sufficient\. An intervention is accepted only when the diagnosed coordinate improves and validation quality remains acceptable\.

The low\-coefficient continued\-pretraining regime usesλaux=1×10−4\\lambda\_\{\\rm aux\}=1\\times 10^\{\-4\}to study geometry with limited over\-regularization\. High\-coefficient stress settings useλaux∈\{1×10−3,3×10−3\}\\lambda\_\{\\rm aux\}\\in\\\{1\\times 10^\{\-3\},3\\times 10^\{\-3\}\\\}together with stronger optimization pressure\.

The same intervention can be quality\-neutral at low auxiliary coefficients and reveal trade\-offs between validation quality and both update\-tail behavior and mixed\-precision headroom under stress\. The experiments therefore pair validation quality with log\-normalizer, gain, gradient\-event, optimizer\-state, and overflow\-headroom endpoints\.

#### Diagnostic Metrics

We report a common diagnostic suite\. Here and below,Qα,tQ\_\{\\alpha,t\}denotes the empiricalα\\alpha\-quantile over tokens; the reported tail summaries use p99 and p99\.9\.

RZ\\displaystyle R\_\{Z\}≔‖gZ‖F‖gCE‖F\+ε,\\displaystyle\\coloneqq\\frac\{\\left\\lVert g\_\{Z\}\\right\\rVert\_\{F\}\}\{\\left\\lVert g\_\{\\mathrm\{CE\}\}\\right\\rVert\_\{F\}\+\\varepsilon\},PZ99\.9\\displaystyle P\_\{Z\}^\{99\.9\}≔Q0\.999,t​\(\|log⁡Zt−c\|\),\\displaystyle\\coloneqq Q\_\{0\.999,t\}\\left\(\|\\log Z\_\{t\}\-c\|\\right\),\(23\)CZ\\displaystyle C\_\{Z\}≔⟨gZ,gCE⟩‖gZ‖F​‖gCE‖F\+ε,\\displaystyle\\coloneqq\\frac\{\\left\\langle g\_\{Z\},g\_\{\\mathrm\{CE\}\}\\right\\rangle\}\{\\left\\lVert g\_\{Z\}\\right\\rVert\_\{F\}\\left\\lVert g\_\{\\mathrm\{CE\}\}\\right\\rVert\_\{F\}\+\\varepsilon\},\(24\)A¯p\\displaystyle\\bar\{A\}\_\{p\}≔𝔼t​at,Ap99≔Q0\.99,t​\(at\),\\displaystyle\\coloneqq\\mathbb\{E\}\_\{t\}a\_\{t\},\\qquad A\_\{p\}^\{99\}\\coloneqq Q\_\{0\.99,t\}\(a\_\{t\}\),\(25\)Δsrc\\displaystyle\\Delta\_\{\\mathrm\{src\}\}≔‖δ^zZ−δzZ‖2‖δzZ‖2\+ε,\\displaystyle\\coloneqq\\frac\{\\left\\lVert\\widehat\{\\delta\}\_\{z\}^\{Z\}\-\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}\}\{\\left\\lVert\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}\+\\varepsilon\},\(26\)Δθ\\displaystyle\\Delta\_\{\\theta\}≔‖gZ,θapprox−gZ,θref‖F‖gZ,θref‖F\+ε,\\displaystyle\\coloneqq\\frac\{\\left\\lVert g\_\{Z,\\theta\}^\{\\mathrm\{approx\}\}\-g\_\{Z,\\theta\}^\{\\mathrm\{ref\}\}\\right\\rVert\_\{F\}\}\{\\left\\lVert g\_\{Z,\\theta\}^\{\\mathrm\{ref\}\}\\right\\rVert\_\{F\}\+\\varepsilon\},\(27\)λR,t,ℓabs\\displaystyle\\lambda\_\{R,t,\\ell\}^\{\\mathrm\{abs\}\}≔λR​αt,ℓ,sR,t,ℓrel≔Ntl​αt,ℓ\.\\displaystyle\\coloneqq\\lambda\_\{R\}\\alpha\_\{t,\\ell\},\\qquad s\_\{R,t,\\ell\}^\{\\mathrm\{rel\}\}\\coloneqq N\_\{\\mathrm\{tl\}\}\\alpha\_\{t,\\ell\}\.\(28\)
Table 5:Diagnostic decision map\. A trigger identifies a source, transport, implementation, or reduction coordinate; the action targets that coordinate; and consequence endpoints plus validation quality determine whether the action is acceptable\. The Not guaranteed column records quantities that should not be used as automatic success criteria for that action\.
#### Transport and Numerical Endpoints

For transport and numerical endpoints, we report optimizer\-facing gradient tails, thresholded gradient\-event rates, Adam\-state endpoints, and static fp16 loss\-scale headroom\. We do not rely only on final validation quality\. Here, pre\-clip means measured before any gradient clipping is applied\. We useΦτ\\Phi\_\{\\tau\}, with table columns such asΦ10\\Phi\_\{10\},Φ50\\Phi\_\{50\}, andΦ100\\Phi\_\{100\}, to denote the fraction of optimizer steps whose pre\-clip full\-model gradient norm exceeds thresholdτ\\tau\. These endpoints, including the Adam\-state and static fp16 overflow proxies, are formally defined in the Transport and Numerical Endpoint Definitions subsection\.

Gradient clipping is complementary to these endpoint measurements\. Gradient clipping can cap the post\-backward update norm and is an appropriate training guardrail, but gradient clipping does not change the logit\-space Z\-loss source, the architecture\-dependent transport factor in Eq\.[5](https://arxiv.org/html/2609.16179#Sx2.E5), or source\-consistency errors in fused and low\-precision paths\. We therefore report pre\-clip tails together with clipping or threshold\-event frequencies\. A method may follow a different effective optimization trajectory when the apparent gradient\-tail control of the method depends on frequent clipping\. Adam\-state endpoints should be read under the stated clipping convention\. Clip\-before\-Adam stacks can reduce moment\-state pollution, whereas effectively unclipped stress settings expose the unmitigated transported source\.

#### Interventions and Coefficient Regimes

The corresponding practical interventions follow directly\. Implementations compute the forward log\-sum\-exp sufficient statisticsm≔maxi⁡zim\\coloneqq\\max\_\{i\}z\_\{i\}andSlse≔∑iexp⁡\(zi−m\)S\_\{\\mathrm\{lse\}\}\\coloneqq\\sum\_\{i\}\\exp\(z\_\{i\}\-m\)once and reuse the statistics during the backward pass, mirroring the sufficient\-statistic discipline behind online softmax kernels\(Milakov and Gimelshein[2018](https://arxiv.org/html/2609.16179#bib.bib42)\)\. This shared\-statistics discipline extends standard numerical softmax stabilization to the Z\-loss branch and ensures that fused or low\-precision paths transport the backward source implied by the forward loss\. Row\-centered output\-head deployment fixes the CE\-invariant output\-head common\-shift gauge and removes that coordinate from Z\-loss transport\. The row\-centered deployment is a model\-coordinate intervention, not the internal row\-maximum subtraction used to evaluate softmax stably\. In experimental tables, Centered head denotes row\-centered deployed logits obtained by subtracting the output\-head common shift from raw logits, with no applied auxiliary loss\. The Centered head configuration is a gauge\-only control whose exact\-arithmetic CE gradient equals that of the raw CE baseline; the configuration’s Z\-loss metrics characterize the deployed coordinate rather than an applied auxiliary update\. Centered Z\-loss denotes the same deployed\-logit convention with an applied centered Z\-loss\. Factorized Z\-loss denotes Eq\.[19](https://arxiv.org/html/2609.16179#A1.E19)applied to raw logits withλμ=λrel=λaux\\lambda\_\{\\mu\}=\\lambda\_\{\\mathrm\{rel\}\}=\\lambda\_\{\\rm aux\},cμ=0c\_\{\\mu\}=0, andcrel=log⁡Vc\_\{\\mathrm\{rel\}\}=\\log Vunless stated otherwise\. Gain\-aware Z\-loss denotes Eq\.[20](https://arxiv.org/html/2609.16179#A1.E20)applied to raw logits with the detached gain weight\. Stress and overflow rows useβ=1×10−2\\beta=1\\times 10^\{\-2\}unless stated otherwise\. Comparable router reports include both the absolute coefficientλR​α\\lambda\_\{R\}\\alphaand the scaleNtl​αN\_\{\\mathrm\{tl\}\}\\alpharelative to the token\-layer mean rather than only nominalλR\\lambda\_\{R\}\.

The dense output\-head experiments use two coefficient regimes to separate geometry from stress behavior\. The low\-coefficient continued pretraining regime uses an applied auxiliary coefficient ofλaux=1×10−4\\lambda\_\{\\rm aux\}=1\\times 10^\{\-4\}\. Several matched settings have small paired PPL differences from CE, allowing differences inPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}to track geometry with limited over\-regularization\. We report the paired differences rather than treat overlap of three\-run standard deviations as a formal equivalence test\. The high\-coefficient stress regime usesλaux∈\{1×10−3,3×10−3\}\\lambda\_\{\\rm aux\}\\in\\\{1\\times 10^\{\-3\},3\\times 10^\{\-3\}\\\}and higher learning rates\. This regime isolates scalar\-tail reduction from transported update tails under stronger auxiliary pressure\. These settings are interpreted through transport and numerical endpoints together with paired PPL differences\. In tables, Standard Z\-loss denotes Eq\.[1](https://arxiv.org/html/2609.16179#Sx2.E1)withc=log⁡Vc=\\log Vapplied to raw logits; conventional formulations commonly usec=0c=0\. Diag\. Z uses the reporting coefficientλdiag\\lambda\_\{\\rm diag\}specified for that table or run family\. Unless otherwise specified in a caption, dense Diag\. Z tables useλdiag=1×10−4\\lambda\_\{\\rm diag\}=1\\times 10^\{\-4\}\. For CE\-only or Centered head rows,λaux=0\\lambda\_\{\\rm aux\}=0and Diag\. Z is not an applied auxiliary loss\. In table columns,RauxR\_\{\\rm aux\}denotes the auxiliary\-gradient norm divided by the CE\-gradient norm\.CauxC\_\{\\rm aux\}denotes the auxiliary\-to\-CE gradient cosine\.

The Z\-loss target isc=log⁡Vc=\\log Vunless stated otherwise\. The zero logit vector, equivalently a zero\-mean uniform logit vector, therefore has zero violation\. For centered relative\-logit diagnostics, the relative target uses the samelog⁡V\\log Vreference after removing the token\-wise mean\. We report high\-percentile statistics, especially p99 and the 99\.9th percentile \(p99\.9\), because the hypothesized failure mode is rare\-token or rare\-step numerical pressure\. Mean Z\-loss can improve while the tail ofJ⊤​δzZJ^\{\\top\}\\delta\_\{z\}^\{Z\}worsens\. We therefore treatPZ99\.9P\_\{Z\}^\{99\.9\},Ap99A\_\{p\}^\{99\}, thresholdedΦτ\\Phi\_\{\\tau\}, maximum\-coordinate tails, Adam\-state tails, and static loss\-scale endpoints as primary transport and numerical measurements\. PPL is used to reject interventions that buy tail reduction by degrading the language model\.

### Experimental Protocol

#### Experimental Setup Details

##### Pretrained diagnostics\.

Main diagnostics use corpus text and pretrained language models\. The evaluated models are DistilGPT\-2, GPT\-2, GPT\-2 Medium, and GPT\-2 Large\(Radfordet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib56); Wolfet al\.[2020](https://arxiv.org/html/2609.16179#bib.bib53)\), together with Pythia\-160M and Pythia\-410M\(Bidermanet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib11)\)\. For reproducibility, the latter are theEleutherAI/pythia\-160m\-dedupedandEleutherAI/pythia\-410m\-dedupedcheckpoints, both trained on the deduplicated Pile\. The primary evaluation uses non\-overlapping WikiText\-103 validation blocks of length 256\(Merityet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib12)\)\. DistilGPT\-2 and GPT\-2 use 512 validation sequences, yielding 131,072 next\-token predictions\. GPT\-2 Medium and Pythia\-160M use 256 validation sequences, yielding 65,536 predictions\. GPT\-2 Large and Pythia\-410M use 128 validation sequences, yielding 32,768 predictions\. These smaller evaluations target scale\-sensitive cases where centering can move tail metrics in either direction\. The diagnostics also include a 64\-block WikiText evaluation ofEleutherAI/pythia\-1b\-deduped\. For compact table labels, Pythia\-160M, Pythia\-410M, and Pythia\-1B denote these checkpoints throughout\. To test whether the diagnostic depends on WikiText, we also evaluate GPT\-2 and GPT\-2 Medium on FineWeb\-Edu\. This cross\-corpus check uses 128 streamed blocks from the FineWeb\-Edusample\-10BTtraining sample\(Penedoet al\.[2024](https://arxiv.org/html/2609.16179#bib.bib13)\)\. The cross\-corpus check is paired with 64\-block FineWeb\-Edu evaluations for GPT\-2 Large and Pythia\-410M\.

##### Dense training protocols\.

Dense training interventions continue pretraining GPT\-2 from the same initialization on WikiText\-103 training blocks\. The main run uses 4800 optimizer steps, a micro\-batch size of 16, a gradient\-accumulation factor of 2, a sequence length of 256, and bf16 autocast\. Optimization uses Adam with decoupled weight decay \(AdamW\),β1=0\.9\\beta\_\{1\}=0\.9,β2=0\.95\\beta\_\{2\}=0\.95, a weight decay of 0\.1, a peak learning rate of5×10−55\\times 10^\{\-5\}, 240 warmup steps, cosine decay, and gradient clipping at 1\.0\(Pascanuet al\.[2013](https://arxiv.org/html/2609.16179#bib.bib41)\)\. Here and below, the micro\-batch size denotes the reported number of sequences in each forward and backward accumulation pass before gradient accumulation\. For example, a micro\-batch size of 16 with a gradient\-accumulation factor of 2 gives an optimizer\-step batch size of 32 sequences\. Each 4800\-step run consumes 39\.3M training tokens and is evaluated on 131,072 validation tokens\. Run counts and matched\-run rules appear in the subsection on seed and matched\-run details\. A 1200\-step GPT\-2 diagnostic comparison uses the same three\-run design, 9\.83M tokens per run, and 65,536 validation tokens\. This diagnostic comparison covers CE, standard Z\-loss, and centered Z\-loss\. For this comparison, we measure p99 pre\-clip full\-model gradient norms and the tied\-embedding auxiliary\-to\-CE gradient ratio, and we decompose the tied\-embedding auxiliary gradient into output and input paths for the auxiliary gradient\. A high\-coefficient gradient\-tail stress experiment usesλaux=1×10−3\\lambda\_\{\\rm aux\}=1\\times 10^\{\-3\}for auxiliary\-method rows\. The stress experiment uses 1200 steps, 9\.83M tokens per run, and the same matched\-run design\. In the stress experiment, we compare CE, centered head, standard Z\-loss, centered Z\-loss, and gain\-aware Z\-loss under the same optimizer\. For CE and centered head rows,λaux=0\\lambda\_\{\\rm aux\}=0\. The diagnostic coefficientλdiag\\lambda\_\{\\rm diag\}for the CE and centered head rows is used only for Diag\. Z reporting\.

##### Stress and precision endpoints\.

For all continued pretraining runs, model parameters are kept as fp32 master weights\. bf16 is used only through autocast\. This precision convention prevents optimizer failures caused by fp16 storage from being conflated with the Z\-loss geometry under study\. The FineWeb\-Edu GPT\-2 Medium evaluation includes a 900\-step effectively unclipped web\-corpus update\-tail setting\. The 900\-step setting uses the same train and validation offsets, a learning rate of2×10−42\\times 10^\{\-4\},λaux=1×10−3\\lambda\_\{\\rm aux\}=1\\times 10^\{\-3\}, 3\.69M consumed tokens per run, and a clipping threshold of1×1081\\times 10^\{8\}\. This threshold effectively disables clipping\. The setting probes update\-tail behavior under aggressive optimization\. An update\-outlier setting uses the same model, dataset, and offsets\. The update\-outlier setting uses a learning rate of8×10−48\\times 10^\{\-4\},λaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}, 600 optimizer steps, 2\.46M consumed tokens per run, and the same clipping threshold\. During this experiment, we separately log the fractions of optimizer steps for which the pre\-clip full\-model gradient norm exceeds each of the thresholds 10, 50, and 100, together with the nonfinite\-microbatch and nonfinite\-gradient fractions\. A centered auxiliary\-coefficient sweep usesλaux∈\{1×10−3,3×10−4,1×10−4\}\\lambda\_\{\\rm aux\}\\in\\\{1\\times 10^\{\-3\},3\\times 10^\{\-4\},1\\times 10^\{\-4\}\\\}under the same matched\-run setup, data offsets, optimizer, and clipping threshold\. An Adam\-state endpoint in the same high\-coefficient FineWeb\-Edu setting samples AdamW state every 50 optimizer steps\. This endpoint tests whether raw\-logit transported spikes enter the second\-moment memory of Adam even when no run reaches a nonfinite\-gradient event\.

##### Implementation audits\.

For fused\-kernel sensitivity, we implement a row\-wise GPU kernel inTriton\(Tilletet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib26)\)that fuses CE and Z\-loss forward and backward computations\. The fused kernel reports both total CE\+Z logit\-gradient consistency and Z\-loss source consistency with a PyTorch fp32 reference\(Paszkeet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib27)\)\. Large\-vocabulary CE kernels increasingly fuse, approximate, or avoid materializing the full logit matrix\. This audit therefore targets the same log\-sum\-exp source statistics that such implementations must preserve\(Graveet al\.[2017](https://arxiv.org/html/2609.16179#bib.bib43); Wijmanset al\.[2025](https://arxiv.org/html/2609.16179#bib.bib33)\)\. We run the fused\-kernel audit on eight non\-overlapping GPT\-2 blocks from WikiText\-103 and eight streamed GPT\-2 blocks from FineWeb\-Edu\. The audit uses logit scales 1 and 4 with fp32, bf16, and fp16 logit storage\. To test whether source mismatch reaches actual optimizer inputs, an update audit backpropagates total CE\+Z logit sources computed from fp32 reference logits and storage\-quantized logits through GPT\-2 and GPT\-2 Medium model graphs on corpus batches\. In the update audit, we compare output\-head gradients, whole\-model gradients, and first\-step Adam directions\. The main output\-head audit uses WikiText\-103 and streamed FineWeb\-Edu blocks, logit scales 1 and 4, and the same fp32 model parameters as the training runs\. The whole\-model audit materializes all transported gradients on smaller batches with the same corpus\-derived logits\. This whole\-model audit confirms that the mismatch reaches beyond the output\-head block\.

##### Overflow\-margin audit\.

The overflow audit is a targeted mixed\-precision overflow\-margin endpoint\. The overflow audit separates loss\-scale headroom from optimizer\-state dynamics by computing full\-model gradients on held\-out corpus text under bf16 autocast\. In the overflow audit, we measure the maximum\-coordinate gradient, the maximum safe static fp16 loss scaleκmax,b\\kappa\_\{\\max,b\}from Eq\.[34](https://arxiv.org/html/2609.16179#A1.E34), and overflow batch fractions for static scales\{256,1024,4096,16384,65536\}\\\{256,1024,4096,16384,65536\\\}\. We run this audit on 16 WikiText\-103 GPT\-2 validation blocks for CE, standard Z\-loss, and centered Z\-loss\. We also run the overflow audit on 16 streamed FineWeb\-Edu GPT\-2 Medium blocks for CE, centered head, standard Z\-loss, factorized Z\-loss, gain\-aware Z\-loss, and centered Z\-loss\. The audit usesλaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}for auxiliary\-method rows to match the update\-outlier setting\.

##### Run accounting\.

We separate validation\-quality comparisons from paired diagnostic comparisons\. The experiment descriptions distinguish available from consumed training tokens and report token reuse when it is relevant to the comparison\. This distinction keeps fixed\-subset evaluations separate from fresh\-corpus continued pretraining and compute\- and data\-scale comparisons\(Kaplanet al\.[2020](https://arxiv.org/html/2609.16179#bib.bib52); Hoffmannet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib44)\)\. Here, the consumed\-token count means training\-token instances processed by optimizer steps, including repeated exposures to the same available tokens\. For each matched run pair, we report the target\-minus\-baseline difference\. A negative paired difference indicates improvement in PPL, CE, Diag\. Z, tail, gain, and gradient\-norm metrics\.

##### Run\-level instability analysis\.

Because the failure modes under study are rare and optimizer\-facing, we characterize training behavior using more than the mean loss\. The reported endpoints include training\-loss tails, the step\-to\-step loss\-spike rate, thresholded pre\-clip gradient\-event rates, pre\- and post\-clip maximum\-coordinate and gradient\-to\-parameter tails, finite\-only gradient\-tail summaries, and nonfinite\-microbatch and nonfinite\-gradient fractions\. Paired comparisons hold the seed, precision mode, optimizer hyperparameters, gradient\-clipping threshold, effective batch size, data offsets, and token accounting fixed\. Related deterministic audits use analogous experiment\-specific controls\. This design helps separate differences in data exposure, precision, or clipping from the transport, optimizer\-state, and numerical effects of an intervention\.

##### FineWeb\-Edu runs\.

FineWeb\-Edu continued pretraining experiments extend the dense training evaluation beyond WikiText\. The GPT\-2 experiment uses the same optimizer hyperparameters, 2400 steps, an optimizer\-step batch size of 32 sequences, a sequence length of 256, 120 warmup steps, and bf16 autocast\. Each run consumes 19\.7M training tokens from streamedsample\-10BTblocks\. Each run is evaluated on 131,072 disjoint next\-token predictions from the same split with a 70,000\-sequence offset\. No validation block overlaps the training block set\. The GPT\-2 Medium experiments include a 2400\-step disjoint\-block run with a micro\-batch size of 8, a gradient\-accumulation factor of 2, and 120 warmup steps\. Each run in that experiment uses 9\.83M training tokens and 65,536 validation next\-token predictions\. A fresh\-token experiment uses 9600 steps, a micro\-batch size of 8, a gradient\-accumulation factor of 2, and 480 warmup steps\. The fresh\-token experiment uses 160,000 training sequences, 41\.0M available tokens, 39\.3M consumed tokens per run, and 1024 validation sequences offset by 220,000 sequences\. This configuration yields 262,144 next\-token predictions\. The token\-reuse factor is 0\.96, so this experiment remains fresh\-token within the available training blocks\.

##### Transfer runs\.

To test model\-scale transfer, we continue pretraining GPT\-2 Medium for 2400 optimizer steps\. The run uses a micro\-batch size of 8, a gradient\-accumulation factor of 2, 120 warmup steps, and the same optimizer recipe\. Each GPT\-2 Medium run consumes 9\.83M training tokens and is evaluated on 65,536 validation tokens\. A 4800\-step GPT\-2 Medium fixed\-subset evaluation uses the same micro\-batch size, gradient\-accumulation factor, and optimizer with 240 warmup steps\. Each fixed\-subset run consumes 19\.7M tokens while repeatedly revisiting a fixed 8192\-sequence training subset\. This fixed\-subset setting measures matched\-mechanism behavior under controlled token reuse\.

A Pythia\-160M WikiText\-103 continued pretraining grid tests whether the conclusions for the GPT\-2 family transfer to a different model family with a padded output vocabulary and an untied output projection\. We use 1200 optimizer steps, a micro\-batch size of 8, a gradient\-accumulation factor of 2, and 120 warmup steps\. Each run consumes 4\.9M tokens from 4\.2M available training tokens and is evaluated on 65,536 validation tokens\. For this grid, the Z\-loss target is computed from the actual output\-head vocabulary size rather than the tokenizer length\.

##### Gradient\-geometry visualization\.

Figure[2](https://arxiv.org/html/2609.16179#Sx3.F2)is generated by a deterministic fp32 audit on the first 32 non\-overlapping WikiText\-103 validation blocks, using pretrained GPT\-2, a block length of 256, an evaluation batch size of 2, a random seed of 67,λdiag=10−4\\lambda\_\{\\rm diag\}=10^\{\-4\}, andc=log⁡Vc=\\log V\. The audit computes the exact token\-level logit\-source and hidden\-injection norms, the fixed translation induced by the output\-row mean, and randomized rank\-16 decompositions of both the raw and row\-centered unembedding matrices with four power iterations\. The paired plane uses the normalized output\-row mean as the first axis and the leading principal direction of the centered hidden\-source cloud after projecting out that mean direction as the second axis\.

##### Spectral and router audits\.

We evaluate the output\-to\-hidden gain identity with a spectral diagnostic\. The spectral diagnostic computes randomized rank\-16 singular value decompositions\(Halkoet al\.[2011](https://arxiv.org/html/2609.16179#bib.bib28)\)of the raw and row\-centered unembedding matrices\. The spectral diagnostic evaluates corpus\-derived text logits and decomposes token\-levelat2=‖WU⊤​pt‖22/‖pt‖22a\_\{t\}^\{2\}=\\left\\lVert W\_\{U\}^\{\\top\}p\_\{t\}\\right\\rVert\_\{2\}^\{2\}/\\left\\lVert p\_\{t\}\\right\\rVert\_\{2\}^\{2\}into the contribution of the top singular directions\. For this analysis, we compute the top\-16 unembedding energy fraction, the mean token\-level fraction ofat2a\_\{t\}^\{2\}explained by the top singular directions, and the Pearson correlation between the low\-rank contribution and exactat2a\_\{t\}^\{2\}\. We run the spectral diagnostic on GPT\-2, GPT\-2 Medium, Pythia\-160M, and Pythia\-1B WikiText\-103 validation blocks\. We also run the spectral diagnostic on GPT\-2 Medium, GPT\-2 Large, and Pythia\-410M FineWeb\-Edu blocks\. These measurements estimate dominant\-direction transport without modifying training\.

For sparse routing, we run two experiments\. We apply a randomly initialized 16\-expert router to GPT\-2 hidden states obtained from corpus blocks\. This router projection isolates the reduction convention without changing the hidden\-state distribution\. We also train controlled end\-to\-end top\-22and top\-44MoE language models on WikiText\-103\. The top\-22model runs for 800 steps, and the top\-44model runs for 1200 steps\. Both use three matched random initializations\. The MoE has four Transformer blocks, a width of 256, four attention heads, eight feed\-forward experts, tied input and output embeddings, and a standard load\-balancing loss with a coefficient of1×10−21\\times 10^\{\-2\}\. Router Z\-loss usesλR=1×10−3\\lambda\_\{R\}=1\\times 10^\{\-3\}\.

#### Transport and Numerical Endpoint Definitions

In addition to final validation quality, we report optimizer\-facing gradient tails for the transport and numerical endpoints, as follows\.

G99\\displaystyle G\_\{99\}≔Q0\.99​\(‖gtrain,s‖2\),\\displaystyle\\coloneqq Q\_\{0\.99\}\\left\(\\left\\lVert g\_\{\\mathrm\{train\},s\}\\right\\rVert\_\{2\}\\right\),H99\\displaystyle H\_\{99\}≔Q0\.99​\(‖gtrain,s‖∞\),\\displaystyle\\coloneqq Q\_\{0\.99\}\\left\(\\left\\lVert g\_\{\\mathrm\{train\},s\}\\right\\rVert\_\{\\infty\}\\right\),\(29\)ρ99\\displaystyle\\rho\_\{99\}≔Q0\.99​\(‖gtrain,s‖2‖θs‖2\+ε\),\\displaystyle\\coloneqq Q\_\{0\.99\}\\left\(\\frac\{\\left\\lVert g\_\{\\mathrm\{train\},s\}\\right\\rVert\_\{2\}\}\{\\left\\lVert\\theta\_\{s\}\\right\\rVert\_\{2\}\+\\varepsilon\}\\right\),\(30\)Φτ\\displaystyle\\Phi\_\{\\tau\}≔1Tstep​∑s=1Tstep𝟏​\{‖gtrain,s‖2\>τ\},\\displaystyle\\coloneqq\\frac\{1\}\{T\_\{\\mathrm\{step\}\}\}\\sum\_\{s=1\}^\{T\_\{\\mathrm\{step\}\}\}\\mathbf\{1\}\\\{\\left\\lVert g\_\{\\mathrm\{train\},s\}\\right\\rVert\_\{2\}\>\\tau\\\},\(31\)Φnf\\displaystyle\\Phi\_\{\\mathrm\{nf\}\}≔1Tstep​∑s=1Tstep𝟏​\{nonfinite​\(gtrain,s\)\}\.\\displaystyle\\coloneqq\\frac\{1\}\{T\_\{\\mathrm\{step\}\}\}\\sum\_\{s=1\}^\{T\_\{\\mathrm\{step\}\}\}\\mathbf\{1\}\\\{\\mathrm\{nonfinite\}\(g\_\{\\mathrm\{train\},s\}\)\\\}\.\(32\)Here,ssindexes optimizer steps\.TstepT\_\{\\mathrm\{step\}\}is the number of audited optimizer steps\. In our stress runs,τ∈\{10,50,100\}\\tau\\in\\\{10,50,100\\\}\. When nonfinite gradients occur, we report both nonfinite\-gradient fractions and finite\-only gradient\-tail summaries\. This dual reporting prevents a single overflow event from obscuring the finite update\-pressure distribution\. For AdamW stress runs\(Kingma and Ba[2015](https://arxiv.org/html/2609.16179#bib.bib17); Loshchilov and Hutter[2019](https://arxiv.org/html/2609.16179#bib.bib18)\), we additionally sample the optimizer memory\. Letmsm\_\{s\}andvsv\_\{s\}be the first and second moment tensors of Adam and let

dsproxy≔ms/\(vs\+ϵAdam\)\\displaystyle d\_\{s\}^\{\\mathrm\{proxy\}\}\\coloneqq m\_\{s\}/\(\\sqrt\{v\_\{s\}\}\+\\epsilon\_\{\\rm Adam\}\)denote the uncorrected Adam preconditioner proxy before learning\-rate and weight\-decay scaling\. The corresponding bias\-corrected Adam preconditioner direction is

dsAdam≔ms/\(1−β1ts\)vs/\(1−β2ts\)\+ϵAdam,\\displaystyle d\_\{s\}^\{\\mathrm\{Adam\}\}\\coloneqq\\frac\{m\_\{s\}/\(1\-\\beta\_\{1\}^\{t\_\{s\}\}\)\}\{\\sqrt\{v\_\{s\}/\(1\-\\beta\_\{2\}^\{t\_\{s\}\}\)\}\+\\epsilon\_\{\\rm Adam\}\},wheretst\_\{s\}is the internal Adam step count\. We report the proxy because the uncorrected preconditioner tracks the same second\-moment memory and coordinate\-wise normalization while omitting bias\-correction factors\. The AdamϵAdam\\epsilon\_\{\\rm Adam\}term is not bias\-corrected in the same way asvsv\_\{s\}\. The proxy is therefore a diagnostic normalization rather than a constant multiple of the exact Adam direction\. Specifically, we reportmaxs⁡‖vs‖∞\\max\_\{s\}\\left\\lVert v\_\{s\}\\right\\rVert\_\{\\infty\},Q0\.99,s​‖vs‖∞Q\_\{0\.99,s\}\\left\\lVert v\_\{s\}\\right\\rVert\_\{\\infty\},Q0\.99,s​‖dsproxy‖2/\(‖θs‖2\+ε\)Q\_\{0\.99,s\}\\left\\lVert d\_\{s\}^\{\\mathrm\{proxy\}\}\\right\\rVert\_\{2\}/\(\\left\\lVert\\theta\_\{s\}\\right\\rVert\_\{2\}\+\\varepsilon\), andQ0\.99,s​‖dsproxy‖∞Q\_\{0\.99,s\}\\left\\lVert d\_\{s\}^\{\\mathrm\{proxy\}\}\\right\\rVert\_\{\\infty\}\. These endpoints distinguish smooth scalar\-tail reduction from rare update events and optimizer\-state memory that can dominate mixed\-precision training\.

Following standard mixed\-precision concerns about fp16 dynamic range and loss scaling\(Micikeviciuset al\.[2018](https://arxiv.org/html/2609.16179#bib.bib19)\), we also audit static loss\-scale headroom\. The audited quantity is the largest static fp16 loss scale that would keep already\-computed optimizer\-facing gradients representable if a stack stored scaled gradients in fp16\. For an audited batch gradient vectorgbg\_\{b\}and fp16 maximumM16≔65504M\_\{16\}\\coloneqq 65504, define

κmax,b\\displaystyle\\kappa\_\{\\max,b\}≔M16‖gb‖∞\+ε,\\displaystyle\\coloneqq\\frac\{M\_\{16\}\}\{\\left\\lVert g\_\{b\}\\right\\rVert\_\{\\infty\}\+\\varepsilon\},\(33\)Ψκ\\displaystyle\\Psi\_\{\\kappa\}≔1Nbatch​∑b=1Nbatch𝟏​\{κ‖gb∥∞\>M16\}\.\\displaystyle\\coloneqq\\frac\{1\}\{N\_\{\\mathrm\{batch\}\}\}\\sum\_\{b=1\}^\{N\_\{\\mathrm\{batch\}\}\}\\mathbf\{1\}\\\{\\kappa\\left\\lVert g\_\{b\}\\right\\rVert\_\{\\infty\}\>M\_\{16\}\\\}\.\(34\)This overflow endpoint targets static loss\-scale headroom for stacks that store scaled gradients in fp16\. As a stress endpoint for maximum\-coordinate transport, lower maximum safe scalesκmax,b\\kappa\_\{\\max,b\}and higherΨκ\\Psi\_\{\\kappa\}indicate that rare coordinates leave less numerical headroom\.

#### Seed and Matched\-Run Details

The matched comparisons use explicit run\-count, seed, and token\-accounting conventions\. This matching design ensures that validation quality, transport diagnostics, and numerical endpoints are compared across the same randomized configurations\. Unless otherwise stated, we use the same seed set for every method in grouped comparisons reported as mean±\\pmstandard deviation over three matched runs\. For tables that summarize replicated training endpoints without displaying standard deviations, we report the mean over the same matched seed set\. Each paired difference is computed as target minus baseline after rows have been matched by seed and the rest of the run configuration\. The matching fields include model, corpus, optimizer settings, data offsets, precision, clipping threshold, step count, and coefficient\. Because lower values are preferred for PPL, CE, Diag\. Z, tail, gain, and gradient\-norm metrics, negative paired differences indicate improvement\.

For dense WikiText\-103 continued training, the GPT\-2 comparison in Table[2](https://arxiv.org/html/2609.16179#Sx3.T2)uses three matched seeds\. Rows report mean±\\pmstandard deviation over the three matched runs\. Each run starts from the same pretrained checkpoint, uses 4800 optimizer steps, consumes 39\.3M training tokens, and is evaluated on 131,072 validation tokens\. The GPT\-2 diagnostic run in Appendix Table[13](https://arxiv.org/html/2609.16179#A1.T13)uses a separate three\-seed matched set\. The diagnostic run reports the three\-run mean after 1200 steps, 9\.83M training tokens per run, and 65,536 validation tokens\. Table[3](https://arxiv.org/html/2609.16179#Sx3.T3)reports a high\-coefficient GPT\-2 gradient\-tail stress sweep with three matched seeds\. This sweep design gives 15 total runs across the five reported methods\. Each run uses 1200 steps and 9\.83M training tokens\. Appendix Table[14](https://arxiv.org/html/2609.16179#A1.T14)reports the GPT\-2 Medium scaling experiment with three matched seeds and 9\.83M consumed tokens per run\. Appendix Table[15](https://arxiv.org/html/2609.16179#A1.T15)reports the GPT\-2 Medium fixed\-subset evaluation with three matched seeds and 19\.7M consumed tokens per run\. Appendix Table[16](https://arxiv.org/html/2609.16179#A1.T16)reports the Pythia\-160M continued pretraining experiment with three matched seeds and 4\.9M consumed tokens per run\. These three experiments report mean±\\pmstandard deviation\.

For FineWeb\-Edu continued training, the GPT\-2 experiment in Appendix Table[11](https://arxiv.org/html/2609.16179#A1.T11)uses three matched seeds\. The GPT\-2 experiment reports mean±\\pmstandard deviation over three matched runs with 19\.7M training tokens per run\. Validation uses 131,072 disjoint next\-token predictions from streamed blocks offset by 70,000 sequences from the training blocks\. The GPT\-2 Medium disjoint\-block experiment uses a separate three\-seed matched set\. The fresh\-token GPT\-2 Medium experiment in Appendix Table[12](https://arxiv.org/html/2609.16179#A1.T12)uses another three\-seed matched set\. The fresh\-token experiment reports mean±\\pmstandard deviation over three matched runs with 39\.3M consumed tokens per run\. The fresh\-token training block set contains 41\.0M available tokens, giving a token\-reuse factor of 0\.96\. Validation uses a disjoint 220,000\-sequence offset yielding 262,144 next\-token predictions\. Appendix Table[17](https://arxiv.org/html/2609.16179#A1.T17)reports the effectively unclipped FineWeb\-Edu setting with three matched seeds\. The effectively unclipped setting reports mean±\\pmstandard deviation over three seeds\. Each run uses 900 steps, a learning rate of2×10−42\\times 10^\{\-4\},λaux=1×10−3\\lambda\_\{\\rm aux\}=1\\times 10^\{\-3\}for auxiliary\-method rows, a clipping threshold of1×1081\\times 10^\{8\}, and 3\.69M consumed tokens from the 41\.0M\-token training block set\. Each run uses the same disjoint 262,144\-token validation blocks as Appendix Table[12](https://arxiv.org/html/2609.16179#A1.T12)\. The update\-outlier experiment in Appendix Table[18](https://arxiv.org/html/2609.16179#A1.T18)uses three matched seeds\. The update\-outlier experiment reports the three\-seed mean under 600 steps, a learning rate of8×10−48\\times 10^\{\-4\},λaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}for auxiliary\-method rows, and a clipping threshold of1×1081\\times 10^\{8\}\. Appendix Table[19](https://arxiv.org/html/2609.16179#A1.T19)reports both the centered auxiliary\-coefficient sweep and the matched raw\-logit coefficient sweep under the same three\-seed matched set, offsets, optimizer, and effectively unclipped setup\. The Adam\-state endpoint in Appendix Table[20](https://arxiv.org/html/2609.16179#A1.T20)uses a separate three\-seed matched set\. The Adam\-state endpoint reports the three\-seed mean over 300 steps, a learning rate of8×10−48\\times 10^\{\-4\},λaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}for auxiliary\-method rows, and AdamW state samples every 50 optimizer steps\.

For sparse routing, the random router\-projection audit uses one fixed seed\. Router weights for each top\-kksetting are generated by a deterministic initialization convention\. Appendix Table[23](https://arxiv.org/html/2609.16179#A1.T23)includes end\-to\-end top\-22and top\-44MoE experiments, each with three matched seeds\. The top\-44rows report mean±\\pmstandard deviation over three matched runs\. Scale\-matched active\-route variants reuse the same seed and initialization as the corresponding token\-layer\-mean variants\. This seed reuse is why the matched rows can recover the token\-layer\-mean result to numerical precision\.

For deterministic audit rows, fixed seeds select evaluation blocks or randomized audit components before deterministic computation\. We use one fixed seed per deterministic audit type, covering pretrained evaluation, source\-space,Triton\-based fused CE\+Z\-loss, optimizer\-facing update, spectral, and static fp16 overflow audits\.

#### Additional Configuration Details

The following details clarify the settings used for several Appendix tables and stress endpoints\. Thresholded gradient\-event rates are computed over optimizer steps, and paired comparisons are formed within matched experimental settings before target\-minus\-baseline differences are averaged across seeds\. This matching avoids combining auxiliary configurations that share a model or method label but differ in an experiment\-defining setting\.

The spectral audit in Appendix Table[7](https://arxiv.org/html/2609.16179#A1.T7)uses GPT\-2 Medium on WikiText\-103 with 128 sequences, an audit batch size of 2, and a rank of 16\. The Pythia\-1B diagnostic rows use the same audit settings with 64 sequences and an audit batch size of 2\. We run the fresh\-token GPT\-2 Medium FineWeb\-Edu experiment in Appendix Table[12](https://arxiv.org/html/2609.16179#A1.T12)on four NVIDIA H100 accelerators using the FineWeb\-Edusample\-10BTtraining split\. The fresh\-token experiment uses streamed text, the CE baseline, centered head, standard Z\-loss, and centered Z\-loss\. The fresh\-token experiment uses three matched seeds, 9600 steps, a micro\-batch size of 8, an evaluation batch size of 8, a block size of 256, 160,000 training sequences, 1024 validation sequences, a training offset of 0, and a validation offset of 220,000\. Optimization uses a learning rate of5×10−55\\times 10^\{\-5\}, a weight decay of 0\.1, 480 warmup steps, gradient clipping at 1\.0, a gradient\-accumulation factor of 2, bf16 precision, an applied auxiliary coefficient of1×10−41\\times 10^\{\-4\}, and a gain coefficient of1×10−21\\times 10^\{\-2\}\.

The stress experiment in Appendix Table[17](https://arxiv.org/html/2609.16179#A1.T17)uses the same training and validation offsets as Appendix Table[12](https://arxiv.org/html/2609.16179#A1.T12)\. The stress experiment evaluates GPT\-2 Medium on the FineWeb\-Edusample\-10BTtraining split with streamed text\. In this experiment, we compare the CE baseline, centered head, standard Z\-loss, factorized Z\-loss, gain\-aware Z\-loss, and centered Z\-loss\. The experiment uses three matched seeds, 900 steps, a micro\-batch size of 8, an evaluation batch size of 8, a block size of 256, 160,000 training sequences, and 1024 validation sequences\. Optimization uses a learning rate of2×10−42\\times 10^\{\-4\}, a weight decay of 0\.1, 90 warmup steps, a clipping threshold of1×1081\\times 10^\{8\}, a gradient\-accumulation factor of 2, bf16 precision, an applied auxiliary coefficient of1×10−31\\times 10^\{\-3\}, and a gain coefficient of1×10−21\\times 10^\{\-2\}\.

The update\-outlier experiment in Appendix Table[18](https://arxiv.org/html/2609.16179#A1.T18)uses the same model, dataset, streaming setup, offsets, and method family\. The update\-outlier experiment uses three matched seeds, 600 steps, a micro\-batch size of 8, an evaluation batch size of 8, a block size of 256, 160,000 training sequences, and 512 validation sequences\. Optimization uses a learning rate of8×10−48\\times 10^\{\-4\}, a weight decay of 0\.1, 60 warmup steps, a clipping threshold of1×1081\\times 10^\{8\}, gradient\-alert thresholds of 10, 50, and 100, a gradient\-accumulation factor of 2, bf16 precision, an applied auxiliary coefficient of3×10−33\\times 10^\{\-3\}, and a gain coefficient of1×10−21\\times 10^\{\-2\}\. The Adam\-state endpoint in Appendix Table[20](https://arxiv.org/html/2609.16179#A1.T20)uses the same data offsets, optimizer, methods, and coefficient\. The Adam\-state endpoint uses a separate three\-seed matched set, 300 steps, and optimizer state sampled every 50 optimizer steps\.

For the static fp16 overflow audit in Appendix Table[21](https://arxiv.org/html/2609.16179#A1.T21), we use GPT\-2 Medium on the FineWeb\-Edusample\-10BTtraining split with streamed text\. The audit uses 16 sequences, an audit batch size of 1, all\-parameter gradients, and bf16 precision\. The static fp16 audit compares the CE baseline, centered head, standard Z\-loss, factorized Z\-loss, gain\-aware Z\-loss, and centered Z\-loss\. Auxiliary\-method rows use an applied coefficient of3×10−33\\times 10^\{\-3\}\. The WikiText GPT\-2 rows use the same audit design on WikiText\-103 validation blocks with CE, standard Z\-loss, and centered Z\-loss\.

For the centered auxiliary\-coefficient sweep in Appendix Table[19](https://arxiv.org/html/2609.16179#A1.T19), we use the same three\-seed matched set, offsets, optimizer, and stress setting while applying centered Z\-loss with coefficients 0\.001, 0\.0003, and 0\.0001\. The matched raw\-logit coefficient sweep uses the same stress setting for CE baseline, standard Z\-loss, factorized Z\-loss, and gain\-aware Z\-loss with coefficients 0\.001, 0\.0003, and 0\.0001\. CE baseline rows are grouped by coefficient for matched comparisons\. The coefficient is not applied to CE\-only training\.

Interpretation of the reported tables therefore follows the experiment\-specific settings above: paired results use matched seeds and consistent batch construction, precision, optimizer and clipping conventions, data offsets, and token budgets\. The stress tables additionally summarize loss and gradient tails, while the Adam\-state endpoint reports the corresponding sampled optimizer\-state quantities\.

### Additional Output\-Head Geometry

#### Tied\-Embedding Pathway Diagnostics

The tied\-embedding decomposition also shows pathway coupling\. On GPT\-2, the norm of the Z\-loss gradient with respect to the tied embedding table is3\.5%3\.5\\%of the corresponding CE gradient norm for the diagnostic batch\. The auxiliary\-to\-CE cosine is 0\.022\. Thus, the auxiliary update is nearly orthogonal to the CE update rather than simply reinforcing the likelihood gradient\. The full Z\-loss gradient decomposes into an output\-path norm of 0\.493 and an input\-path norm of 0\.563\. The cosine between the input and output paths is−0\.0098\-0\.0098\. On GPT\-2 Medium, the same diagnostic givesRZ=3\.5%R\_\{Z\}=3\.5\\%, an auxiliary\-to\-CE cosine of 0\.014, an output\-path norm of 0\.698, and an input\-path norm of 0\.375\. Thus, the output auxiliary loss is not a single local regularizer on a tied head\. The auxiliary loss creates a coupled multi\-path update\.

The scale diagnostics sharpen the claim\. Pythia\-160M provides a useful contrast because the Pythia\-160M output projection is untied in this model family\. The Pythia\-160M diagnostic batch has a much larger Z\-loss gradient budget,RZ=17\.5%R\_\{Z\}=17\.5\\%\. The Pythia\-160M auxiliary\-to\-CE cosine is also negative at−0\.149\-0\.149\. There is nevertheless no tied\-table input\-path contribution to the output\-projection parameter\. Z\-loss can still propagate to the untied input embeddings and body throughWU⊤​δzZW\_\{U\}^\{\\top\}\\delta\_\{z\}^\{Z\}\. The large raw common\-shift tail in Table[1](https://arxiv.org/html/2609.16179#Sx3.T1)persists beyond tied embeddings\. Tying changes the transport path, while output\-head common shift can be a separate and very large source factor\. The GPT\-2 Large rows in Table[1](https://arxiv.org/html/2609.16179#Sx3.T1)show how the effect of centering depends on output\-head geometry\. The raw common shift already placeslog⁡Z\\log Zcloser to the default target\. Removing the raw common shift therefore preserves PPL but can increase Diag\. Z andAp99A\_\{p\}^\{99\}\. Pythia\-410M and Pythia\-1B provide the same cautionary contrast\. These rows support the main position of this paper\. Common\-shift removal is a gauge intervention whose effect is evaluated with transport diagnostics\.

#### FineWeb\-Edu Pretrained Diagnostics

To test whether the WikiText results are corpus\-specific, Appendix Table[6](https://arxiv.org/html/2609.16179#A1.T6)repeats the pretrained comparison of raw and centered output heads on streamed FineWeb\-Edu blocks without continued training\. Centering preserves PPL and removes the common\-shift tail in all four model comparisons\. Whether Diag\. Z and the gain metrics increase or decrease still depends on the model geometry\. GPT\-2 and GPT\-2 Medium repeat the large common\-shift pattern from WikiText\. GPT\-2 Large and Pythia\-410M show a geometry\-dependent regime\. Centering removesμ\\mubut increasesAp99A\_\{p\}^\{99\}\. For GPT\-2 Large, centering also increases thePZ99\.9P\_\{Z\}^\{99\.9\}tail\. This cross\-corpus diagnostic replicates the common\-shift pattern beyond WikiText\. The cross\-corpus diagnostic also shows that centering is best evaluated with transport diagnostics\.

Table 6:Pretrained diagnostics on FineWeb\-Edusample\-10BT\. GPT\-2 and GPT\-2 Medium mirror the WikiText common\-shift pattern, while GPT\-2 Large and Pythia\-410M show a geometry\-dependent regime where row\-centered deployed logits increase gain\-tail metrics\.
#### Unembedding Spectral Audit

The gain identity and alignment remark yield a specific prediction\. Large token\-level gainsata\_\{t\}should arise when the normalized softmax source aligns with unembedding directions associated with high singular values\. The measurements in Appendix Table[7](https://arxiv.org/html/2609.16179#A1.T7)numerically confirm the gain identity on corpus\-derived logits\. For GPT\-2 and GPT\-2 Medium, the top 16 singular directions contain between 33% and 36% of the raw unembedding Frobenius energy\. The top 16 singular directions explain between 62% and 64% of token\-levelat2a\_\{t\}^\{2\}on average\. The correlation between the rank\-16 contribution and exactat2a\_\{t\}^\{2\}is at least 0\.997\. Centering removes a large mean\-direction component\. For GPT\-2 Medium, the spectral norm drops from 462\.9 to 87\.9\.Ap99A\_\{p\}^\{99\}drops from 30\.5 to 14\.5 on WikiText and from 28\.0 to 13\.1 on FineWeb\-Edu\. Pythia\-160M is an extreme anisotropy case\. The first raw singular direction alone contains 91\.5% of the unembedding energy\. The top 16 directions explain 97\.8% of token\-levelat2a\_\{t\}^\{2\}\. After centering, the spectral norm drops from 776\.4 to 106\.1 andAp99A\_\{p\}^\{99\}drops from 61\.0 to 6\.39\. The complementary case appears in the FineWeb\-Edu rows for GPT\-2 Large and Pythia\-410M and the WikiText row for Pythia\-1B\. Centering reduces the spectral norm but can increaseAp99A\_\{p\}^\{99\}when the softmax source aligns more closely with the remaining centered directions\. This experiment gives a direct empirical bridge between Eq\.[14](https://arxiv.org/html/2609.16179#Sx2.E14)and the observed gradient\-tail behavior\.

Table 7:Low\-rank unembedding spectral audit on corpus\-derived text logits\.E16E\_\{16\}is the fraction of unembedding Frobenius energy in the top 16 singular directions\.F16F\_\{16\}is the mean token\-level fraction of exactat2a\_\{t\}^\{2\}explained by those directions\.ρ16\\rho\_\{16\}is the Pearson correlation between exactat2a\_\{t\}^\{2\}and the rank\-16 contribution\.

### Why Precision Enters

The connection between Z\-loss and finite precision is two\-sided\. First, production, storage, or communication may round a raw logit before softmax evaluation with a numerically stable implementation\. For a floating\-point value in the normal range under round\-to\-nearest, the storage error obeys the standard model

\|fl⁡\(zi\)−zi\|\\displaystyle\\left\|\\operatorname\{fl\}\(z\_\{i\}\)\-z\_\{i\}\\right\|≤umach​\|zi\|,\\displaystyle\\leq u\_\{\\mathrm\{mach\}\}\|z\_\{i\}\|,\(35\)maxi⁡zi\\displaystyle\\max\_\{i\}z\_\{i\}≤log⁡Z≤maxi⁡zi\+log⁡V,\\displaystyle\\leq\\log Z\\leq\\max\_\{i\}z\_\{i\}\+\\log V,whereumachu\_\{\\mathrm\{mach\}\}is the unit roundoff\(Goldberg[1991](https://arxiv.org/html/2609.16179#bib.bib48)\)\. A large common offset therefore coarsens the absolute spacing of representable logits and can erase small logit differences before the softmax kernel operates\. Subsequent row\-maximum subtraction prevents exponential overflow but cannot reconstruct differences already lost at an earlier low\-precision boundary\. The second inequality in Eq\.[35](https://arxiv.org/html/2609.16179#A1.E35)shows why log\-normalizer control is relevant\. If\|log⁡Z−c\|≤r\|\\log Z\-c\|\\leq r, thenc−r−log⁡V≤maxi⁡zi≤c\+rc\-r\-\\log V\\leq\\max\_\{i\}z\_\{i\}\\leq c\+r\. Keepinglog⁡Z\\log Znear a finite target thus constrains the row maximum and common\-shift excursions\. Such log\-normalizer control can preserve dense\-head probabilities, router gate weights, and threshold\-sensitive routing decisions when low\-precision rounding precedes row\-maximum subtraction\(Zophet al\.[2022](https://arxiv.org/html/2609.16179#bib.bib6); Blanchardet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib32)\)\. The benefit is implementation\-dependent\. Z\-loss does not bound every logit or the full logit range, and Z\-loss does not replace a stable log\-sum\-exp or softmax implementation\.

Second, the Z\-loss branch creates a separate backward precision obligation\. A fused or low\-precision path can match the scalar CE\+Z value while reconstructing differentlog⁡Z\\log Zandppvalues and therefore transporting a differentδzZ\\delta\_\{z\}^\{Z\}\. Once transported, a large source amplitude or architectural gain can create maximum\-coordinate gradients that consume loss\-scale headroom in stacks that store scaled gradients in fp16 or enter adaptive\-optimizer moment state\. Z\-loss can therefore improve forward logit representability while worsening backward numerical headroom\. Our audits distinguish these mechanisms\. Source consistency, selected\-gradient consistency, and optimizer\-direction consistency are compared with an fp32 reference, while gradient tails and static fp16 margins measure the transported cost\. The source\-distortion identity below gives the exact source\-mismatch decomposition and distinguishes Z\-loss source consistency from standard row\-maximum subtraction\.

### Source Distortion and Finite\-Precision Deployment

#### Source\-Distortion Identity and Diagnostics

If a fused or low\-precision implementation produceslog⁡Z^\\widehat\{\\log Z\}andp^\\widehat\{p\}during the backward pass, defineϵZ≔log⁡Z^−log⁡Z\\epsilon\_\{Z\}\\coloneqq\\widehat\{\\log Z\}\-\\log Z,ϵp≔p^−p∈ℝV\\epsilon\_\{p\}\\coloneqq\\widehat\{p\}\-p\\in\\mathbb\{R\}^\{V\}, andδ^zZ≔2​λ​\(log⁡Z^−c\)​p^\\widehat\{\\delta\}\_\{z\}^\{Z\}\\coloneqq 2\\lambda\(\\widehat\{\\log Z\}\-c\)\\widehat\{p\}\. Then

δ^zZ−δzZ=2​λ​\(ϵZ​p\+\(log⁡Z−c\)​ϵp\+ϵZ​ϵp\)\.\\widehat\{\\delta\}\_\{z\}^\{Z\}\-\\delta\_\{z\}^\{Z\}=2\\lambda\\left\(\\epsilon\_\{Z\}p\+\(\\log Z\-c\)\\epsilon\_\{p\}\+\\epsilon\_\{Z\}\\epsilon\_\{p\}\\right\)\.\(36\)Agreement between forward scalar loss values therefore does not imply agreement between the corresponding transported backward vectors\. This distinction is especially important in low precision\. Mathematically equivalent log\-sum\-exp and softmax formulas need not have identical floating\-point behavior\(Goldberg[1991](https://arxiv.org/html/2609.16179#bib.bib48); Blanchardet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib32)\)\. This source\-distortion observation motivates separate source\-space and transported\-gradient consistency diagnostics,

Δsrc\\displaystyle\\Delta\_\{\\mathrm\{src\}\}≔‖δ^zZ−δzZ‖2‖δzZ‖2\+ε,\\displaystyle\\coloneqq\\frac\{\\left\\lVert\\widehat\{\\delta\}\_\{z\}^\{Z\}\-\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}\}\{\\left\\lVert\\delta\_\{z\}^\{Z\}\\right\\rVert\_\{2\}\+\\varepsilon\},\(37\)Δθ\\displaystyle\\Delta\_\{\\theta\}≔‖gZ,θapprox−gZ,θref‖F‖gZ,θref‖F\+ε\.\\displaystyle\\coloneqq\\frac\{\\left\\lVert g\_\{Z,\\theta\}^\{\\mathrm\{approx\}\}\-g\_\{Z,\\theta\}^\{\\mathrm\{ref\}\}\\right\\rVert\_\{F\}\}\{\\left\\lVert g\_\{Z,\\theta\}^\{\\mathrm\{ref\}\}\\right\\rVert\_\{F\}\+\\varepsilon\}\.Row\-maximum subtraction and Z\-loss source consistency address complementary parts of a reliable CE\+Z\-loss implementation\. Subtracting a row maximum is an algebraically invariant way to avoid overflow and reduce harmful underflow, and mature CE paths such as the built\-in PyTorch kernels already use numerically stable softmax or log\-softmax formulations\(Paszkeet al\.[2019](https://arxiv.org/html/2609.16179#bib.bib27)\)\. When a training stack adds Z\-loss, especially in fused or low\-precision CE\+Z\-loss code, the forward scalar and backward source for Z\-loss must be derived from the same sufficient statistics rather than separately rounded or reconstructed quantities\. A source\-consistency audit therefore verifies that the auxiliary signal injected into backpropagation matches the implemented forward objective, complementing standard softmax stabilization in the kernel\.

#### Deployment of Centered Heads in Finite Precision

Throughout, we use fp32, fp16, and bf16 as lowercase data type labels\. In the experiments, centered logits are deployed by subtracting the fp32 output\-head common shift from the already\-computed raw logits, preserving the realized softmax distribution without a second low\-precision output projection\. Eq\.[38](https://arxiv.org/html/2609.16179#A1.E38)gives the bias\-aware form and distinguishes this deployment from internal softmax row\-maximum subtraction\.

With low\-precision weights, token\-dependent residuals can arise whenW~U​u\\widetilde\{W\}\_\{U\}uis formed as a fresh fp16 or bf16 matrix product\. Those residuals do not consist solely of a common shift\. Our implementation therefore deploys centered logits by subtracting the output\-head common shift from the already\-computed raw logits:

zraw\\displaystyle z\_\{\\mathrm\{raw\}\}≔WU​u\+b,z~≔zraw−μhead​𝟏,\\displaystyle\\coloneqq W\_\{U\}u\+b,\\qquad\\widetilde\{z\}\\coloneqq z\_\{\\mathrm\{raw\}\}\-\\mu\_\{\\mathrm\{head\}\}\\mathbf\{1\},\(38\)μhead\\displaystyle\\mu\_\{\\mathrm\{head\}\}≔1V​𝟏⊤​WU​u\+b¯,b¯≔1V​∑vbv\.\\displaystyle\\coloneqq\\frac\{1\}\{V\}\\mathbf\{1\}^\{\\top\}W\_\{U\}u\+\\bar\{b\},\\qquad\\bar\{b\}\\coloneqq\\frac\{1\}\{V\}\\sum\_\{v\}b\_\{v\}\.The formal propositions in the main text use the bias\-free notationμ=w¯​u\\mu=\\bar\{w\}u\. Eq\.[38](https://arxiv.org/html/2609.16179#A1.E38)is the corresponding deployment formula with an optional output bias\. For a bias\-free head,b¯=0\\bar\{b\}=0andμhead=μ\\mu\_\{\\mathrm\{head\}\}=\\mu\. In exact arithmetic, the raw\-logit definition in Eq\.[38](https://arxiv.org/html/2609.16179#A1.E38)gives

1V​𝟏⊤​zraw=μhead,\\displaystyle\\frac\{1\}\{V\}\\mathbf\{1\}^\{\\top\}z\_\{\\mathrm\{raw\}\}=\\mu\_\{\\mathrm\{head\}\},so Eq\.[38](https://arxiv.org/html/2609.16179#A1.E38)is equivalent to subtracting the realized raw\-logit mean in exact arithmetic\. In lower precision, the implementation preserves the realized softmax distribution by subtracting a token\-wise common shift\. The centered\-logit deployment centers with respect to the fp32 output\-head common shift, not necessarily the rounded raw\-logit mean\. This subtraction is also distinct from subtracting the per\-row maximum inside a softmax kernel\. Row\-maximum subtraction is an internal numerical device\. The per\-row maximum is added back in a stable log\-sum\-exp computation, so row\-maximum subtraction does not change the deployed objective coordinate\. Row\-centered deployment subtracts the output\-head common\-shift coordinate before the auxiliary loss is evaluated\. Row\-centered deployment preserves CE because CE is shift\-invariant but intentionally changes the Z\-loss source coordinate\.

### Source\-Consistency and Implementation Audits

For implementation sensitivity, we use a source\-space audit with GPT\-2 logits from WikiText\-103 validation blocks\. We quantize the logits to bf16 or fp16\. We optionally accumulate the softmax statistics in fp32\. We then compare the reconstructed Z\-loss source to the fp32 reference\. Appendix Table[8](https://arxiv.org/html/2609.16179#A1.T8)shows that forward scalar agreement can substantially understate Z\-loss source error\. With bf16 logits at the natural scale, the mean relative forward Z\-loss error is only1\.9×10−41\.9\\times 10^\{\-4\}\. The mean Z\-loss source relative error is 0\.065\. At a logit scale of 4, the source error rises to 0\.183 while the forward relative error remains below5×10−45\\times 10^\{\-4\}\. On GPT\-2 logits from FineWeb\-Edu, the same bf16 reconstruction gives forward relative errors below5×10−45\\times 10^\{\-4\}\. The Z\-loss source errors are 0\.082, 0\.143, and 0\.207 at logit scales 1, 2, and 4\. The fp16 source errors are smaller on both corpora, ranging from 0\.008 to 0\.027\. This source\-audit result illustrates that the relevant issue is not the data type alone but whether the backward path reconstructs the same sufficient statistics\.

We observe similar behavior when applying the same source\-space audit to GPT\-2 Medium logits from WikiText and FineWeb\-Edu\. On WikiText, bf16 source reconstruction has a relative forward Z\-loss error of2\.0×10−42\.0\\times 10^\{\-4\}at the natural scale and2\.6×10−42\.6\\times 10^\{\-4\}at a scale of 4\. The corresponding Z\-loss source relative errors are 0\.057 and 0\.171\. On FineWeb\-Edu, the corresponding forward errors are1\.7×10−41\.7\\times 10^\{\-4\}and2\.3×10−42\.3\\times 10^\{\-4\}\. The Z\-loss source errors are 0\.0687 and 0\.193\. The fp16 source errors are smaller, ranging from 0\.007 to 0\.026 across these GPT\-2 Medium checks\. The direct PyTorch CE backward audit also remains close to the reference for GPT\-2 Medium\. On FineWeb\-Edu, the bf16 total\-source relative differences are 0\.00150, 0\.00138, and 0\.00108 at scales 1, 2, and 4\.

We also run a direct PyTorch backward audit on the same corpus\-derived logits\. We compare built\-incross\_entropyplus manual Z\-loss with a manuallog\_softmaxCE plus the same manual Z\-loss in this audit\. The discrepancies are much smaller than in the Z\-loss source reconstruction stress test\. In bf16, the relative total\-loss backward difference is1\.40×10−31\.40\\times 10^\{\-3\}at the natural scale and1\.04×10−31\.04\\times 10^\{\-3\}at a scale of 4 on WikiText\. On FineWeb\-Edu, the corresponding differences are1\.44×10−31\.44\\times 10^\{\-3\}and1\.05×10−31\.05\\times 10^\{\-3\}\. In fp16, the difference is below2\.8×10−42\.8\\times 10^\{\-4\}\. In fp32, the difference is zero within reporting precision\. These results show that the built\-in PyTorch CE path remains close to the fp32 reference in this setting\. The same results also show the value of source\-consistency checks for the actual fused CE\+Z\-loss implementation used by a training stack\.

As a check of the fused implementation, we also implement a row\-wise fused CE\+Z\-loss GPU kernel inTriton\. The kernel computes the forward scalar, the total CE\+Z logit gradient, and the Z\-loss source from shared softmax statistics\. Appendix Table[9](https://arxiv.org/html/2609.16179#A1.T9)reports this audit on eight GPT\-2 corpus blocks from WikiText\-103 and FineWeb\-Edu\. The fp32 kernel matches the PyTorch fp32 reference with a relative gradient error ranging from1×10−71\\times 10^\{\-7\}to1×10−61\\times 10^\{\-6\}\. This fp32 agreement confirms that the fused formula is source\-consistent\. When the same fused kernel reads bf16 logits, however, the forward Z\-loss relative error stays below3×10−43\\times 10^\{\-4\}\. The Z\-loss source error is 0\.056 on WikiText at the natural scale and 0\.204 on FineWeb\-Edu at a scale of 4\. Thus, even a source\-consistent fused kernel can transport a meaningfully different Z\-loss signal if the sufficient statistics are reconstructed from quantized logits\.

The fused audit with GPT\-2 Medium extends the implementation observation beyond GPT\-2 logits\. The fp32 path again matches the PyTorch reference with a relative gradient error below1×10−61\\times 10^\{\-6\}\. With bf16 logit storage, GPT\-2 Medium has a Z\-loss source relative error of 0\.050 at the natural WikiText scale and 0\.164 at a scale of 4\. On FineWeb\-Edu, the corresponding values are 0\.0696 and 0\.186\. The total CE\+Z logit\-gradient error is smaller but still material, ranging from 0\.033 to 0\.134\. The forward Z\-loss relative error ranges only from1\.3×10−41\.3\\times 10^\{\-4\}to2\.4×10−42\.4\\times 10^\{\-4\}\. This GPT\-2 Medium audit strengthens the conclusion that forward scalar agreement is an inadequate audit criterion even for a source\-consistent fused formula\.

Table 8:Source precision audit on GPT\-2 WikiText\-103 logits\. Forward scalar error can be tiny while the reconstructed backward source differs materially from the fp32 reference\. Logit×\\timesis the multiplier applied to logits before the audit\.Table 9:Audit of aTriton\-implemented fused CE\+Z\-loss GPU kernel on GPT\-2 corpus logits\. The kernel uses shared softmax statistics for forward and backward\.Δtot\\Delta\_\{\\mathrm\{tot\}\}is the relative error in the total CE\+Z logit gradient;ΔZ\\Delta\_\{Z\}is the relative error in the Z\-loss source\. Logit×\\timesis the pre\-audit logit multiplier\.ModelScopeCorpusLogit×\\timesGrad\. rel\. errorGrad\. cosineAdam dir\. rel\. errorGPT\-2headWikiText10\.03070\.99950\.13940\.11940\.99290\.247FineWeb\-Edu10\.03680\.99930\.16140\.12360\.99220\.259GPT\-2 MediumheadWikiText10\.04080\.99920\.12840\.16370\.98660\.228FineWeb\-Edu10\.03970\.99920\.16340\.14530\.98970\.248GPT\-2allWikiText10\.03080\.99950\.14440\.07870\.99680\.192GPT\-2 MediumallWikiText10\.03620\.99930\.14640\.12520\.99210\.181Table 10:Optimizer\-facing update audit\. We backpropagate total CE\+Z logit sources from the fp32 reference and the bf16 reconstruction through model graphs on corpus batches and compare output\-head or whole\-model gradients and first\-step Adam directions\. Whole\-model rows use smaller WikiText batches because all transported gradients are materialized\. Logit×\\timesis the pre\-audit logit multiplier\.Finally, we test whether the same source reconstruction error reaches optimizer\-facing updates in addition to the logit\-space source\. We backpropagate either the fp32 reference total CE\+Z logit source or a storage\-quantized total CE\+Z logit source through the same model graph on corpus batches\. We then compare transported parameter gradients\. The output\-head block includes the direct output path\. For tied heads in the GPT\-2 family, the output\-head block also includes the input\-embedding path that returns through the Transformer body\. Appendix Table[10](https://arxiv.org/html/2609.16179#A1.T10)shows that the mismatch remains visible after transport\. For bf16 storage, the relative difference in the total CE\+Z output\-head gradient ranges from 3\.1% to 4\.1% at the natural scale\. At a logit scale of 4, the relative difference ranges from 11\.9% to 16\.4% across GPT\-2 and GPT\-2 Medium on WikiText and FineWeb\-Edu\. The corresponding first\-step Adam direction mismatch is larger\. The Adam direction mismatch ranges from 12\.8% to 16\.3% at the natural scale and from 22\.8% to 25\.9% at a scale of 4\. Coordinate\-wise adaptive normalization magnifies small sign and relative\-magnitude changes\. A smaller whole\-model audit on WikiText gives comparable errors\. At scales 1 and 4, the total\-gradient relative errors are 3\.1% and 7\.9%, respectively, for GPT\-2 and 3\.6% and 12\.5%, respectively, for GPT\-2 Medium\. The errors in the fp16 rows are much smaller\. The fp16 total\-gradient errors range from 0\.4% to 0\.5% at the natural scale and from 1\.3% to 2\.0% at a scale of 4 in the total\-gradient audits\. This update\-level mismatch supports the practical audit recommendation\. The backward source and update direction used by the actual training stack should be checked alongside the forward fused loss\.

### Additional Dense Continued\-Pretraining Results

#### FineWeb\-Edu GPT\-2 Continued Pretraining

Appendix Table[11](https://arxiv.org/html/2609.16179#A1.T11)reports a replication of the core GPT\-2 comparison on FineWeb\-Edu with disjoint training and validation blocks\. This FineWeb\-Edu comparison evaluates continued pretraining on web\-corpus text rather than adaptation to WikiText validation\. The CE baseline is again strong and low\-variance\. Centered head deployment preserves the mean PPL to within1\.6×10−41\.6\\times 10^\{\-4\}\. Centered head deployment reduces the log\-normalizer tail and the p99 output\-to\-hidden gain by factors of9\.69\.6and1\.71\.7, respectively\. Centered Z\-loss has a mean PPL difference of\+0\.0027\+0\.0027relative to CE in this short comparison\. Centered Z\-loss gives the strongest joint tail reduction\. The value ofPZ99\.9P\_\{Z\}^\{99\.9\}falls from 246\.5 to 19\.5, andAp99A\_\{p\}^\{99\}falls from 26\.6 to 15\.8\. Standard Z\-loss reduces the raw log\-normalizer tail but does not reduceAp99A\_\{p\}^\{99\}\. Standard Z\-loss is also worse in validation PPL at this coefficient\. The FineWeb\-Edu result supports the main claim on a second corpus\. Architecture\-aware methods can reduce backward\-geometry tails while maintaining comparable validation PPL\.

Table 11:GPT\-2 continued pretraining on FineWeb\-Edusample\-10BTwith disjoint streamed validation blocks\.
#### FineWeb\-Edu GPT\-2 Medium Fresh\-Token Continued Pretraining

Appendix Table[12](https://arxiv.org/html/2609.16179#A1.T12)reports the GPT\-2 Medium web\-corpus evaluation with 39\.3M consumed tokens and disjoint validation blocks\. This experiment extends the FineWeb\-Edu evaluation to GPT\-2 Medium while avoiding repeated\-subset ambiguity\. As with the GPT\-2 FineWeb\-Edu experiment, validation\-quality differences are small compared with the geometric effect\. The mean paired PPL difference for centered head relative to CE is−6\.0×10−5\-6\.0\\times 10^\{\-5\}, indicating essentially unchanged validation quality\. Centered head deployment reducesAp99A\_\{p\}^\{99\}by a factor of2\.192\.19\. Standard Z\-loss gives the largest reductions in scalar Z\-loss and the p99 pre\-clip full\-model gradient norm\. Standard Z\-loss leavesAp99A\_\{p\}^\{99\}at the CE level\. Centered Z\-loss gives the strongest joint scalar\-tail and output\-to\-hidden gain reduction\. Centered Z\-loss reducesPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}by factors of7\.57\.5and2\.212\.21, respectively, with a measured PPL difference of\+0\.025\+0\.025relative to the CE mean\.

Table 12:GPT\-2 Medium fresh\-token continued pretraining on FineWeb\-Edusample\-10BT\.A paired analysis of Appendix Tables[11](https://arxiv.org/html/2609.16179#A1.T11)and[12](https://arxiv.org/html/2609.16179#A1.T12)calibrates the same point\. For GPT\-2 on FineWeb\-Edu, the mean paired PPL differences relative to CE are−1\.6×10−4\-1\.6\\times 10^\{\-4\}for centered head deployment and\+2\.7×10−3\+2\.7\\times 10^\{\-3\}for centered Z\-loss\. Both reduceAp99A\_\{p\}^\{99\}by about 10\.8\. In the fresh\-token GPT\-2 Medium grid on FineWeb\-Edu, the corresponding mean paired PPL differences are−6\.0×10−5\-6\.0\\times 10^\{\-5\}for centered head deployment,\+2\.52×10−2\+2\.52\\times 10^\{\-2\}for centered Z\-loss, and\+8\.64×10−3\+8\.64\\times 10^\{\-3\}for standard Z\-loss\. The mean paired differences inAp99A\_\{p\}^\{99\}are−12\.80\-12\.80for centered head,−12\.94\-12\.94for centered Z\-loss, and−0\.046\-0\.046for standard Z\-loss\. These differences confirm that raw\-logit Z\-loss and centered transport target different failure modes\.

#### 1200\-Step GPT\-2 Diagnostic

Appendix Table[13](https://arxiv.org/html/2609.16179#A1.T13)reports a 1200\-step GPT\-2 diagnostic run with post\-training gradient instrumentation\. This diagnostic run emphasizes gradient pathways and alignment rather than validation\-quality rankings\. The norm of the auxiliary gradient with respect to the tied embedding table ranges from0\.4%0\.4\\%to0\.5%0\.5\\%of the corresponding CE gradient norm\. The auxiliary gradient is not zero and splits into output and input paths\. Standard Z\-loss leavesAp99A\_\{p\}^\{99\}close to the CE baseline while reducing the common\-shift tail\. Centered Z\-loss reduces both the log\-normalizer tail andAp99A\_\{p\}^\{99\}\. The auxiliary\-to\-CE cosine is slightly negative in this diagnostic batch\. This negative cosine supports reporting alignment rather than assuming that the log\-normalizer\-regularization update is always CE\-aligned\.

Table 13:GPT\-2 1200\-step diagnostic run on WikiText\-103 with post\-training gradient instrumentation\.
#### Model\-Scale and Model\-Family Transfer

Appendix Table[14](https://arxiv.org/html/2609.16179#A1.T14)reports an extension of the comparison to GPT\-2 Medium across three seeds\. The CE baseline is strong and low\-variance in this run\. Standard Z\-loss improves PPL and reducesPZ99\.9P\_\{Z\}^\{99\.9\}\. Standard Z\-loss leaves the output\-to\-hidden gain tail near the CE baseline\. Centered head deployment removes the common\-shift coordinate and halvesAp99A\_\{p\}^\{99\}\. Centered head deployment targets the log\-normalizer tail less directly\. Centered Z\-loss gives the best mean PPL in this grid\. Centered Z\-loss reducesPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}by factors of10\.510\.5and1\.91\.9, respectively, relative to the CE baseline\. These results reinforce the interpretation that common\-shift removal and output\-to\-hidden gain reduction are distinct interventions that combine constructively in this setting\.

Table 14:GPT\-2 Medium scaling experiment on WikiText\-103\.Appendix Table[15](https://arxiv.org/html/2609.16179#A1.T15)reports a 4800\-step GPT\-2 Medium evaluation on the same fixed training subset\. The fixed\-subset setup yields higher absolute PPL than the 2400\-step scaling experiment\. The higher absolute PPL is consistent with repeated exposure to a small subset rather than fresh large\-corpus pretraining\. Within this fixed\-subset comparison, centered Z\-loss is better than the CE baseline for every matched run\. The mean paired PPL difference for centered Z\-loss relative to CE is−0\.268\-0\.268\. Centered Z\-loss reducesPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}by factors of9\.79\.7and1\.91\.9, respectively\. Standard Z\-loss also improves PPL and strongly reducesPZ99\.9P\_\{Z\}^\{99\.9\}\. Standard Z\-loss leavesAp99A\_\{p\}^\{99\}near the CE baseline\. Centered head deployment reducesAp99A\_\{p\}^\{99\}without the same targeted log\-normalizer\-tail reduction\.

Table 15:GPT\-2 Medium fixed\-subset evaluation on WikiText\-103\. The fixed training subset separates this comparison from large\-corpus pretraining results\.The experiment in Appendix Table[16](https://arxiv.org/html/2609.16179#A1.T16)tests the same interventions on Pythia\-160M\. This model has a much larger raw common\-shift channel than models in the GPT\-2 family\. The experiment is therefore a stress test for architecture\-dependent transport\. Centered head deployment reduces the CE\-baselinePZ99\.9P\_\{Z\}^\{99\.9\}tail andAp99A\_\{p\}^\{99\}by factors of15\.115\.1and8\.78\.7, respectively\. Centered head deployment also gives a mean PPL below that of CE, with larger run\-to\-run variation\. Centered Z\-loss gives the strongest joint reduction\. Centered Z\-loss reducesPZ99\.9P\_\{Z\}^\{99\.9\},Ap99A\_\{p\}^\{99\}, and the p99 pre\-clip full\-model gradient norm by factors of20\.820\.8,8\.98\.9, and approximately1\.51\.5, respectively\. Standard raw\-logit Z\-loss makes the scalar Z\-loss almost vanish and reducesPZ99\.9P\_\{Z\}^\{99\.9\}\. Standard raw\-logit Z\-loss also worsens PPL by 9\.2 points, leavesAp99A\_\{p\}^\{99\}above the CE baseline, and raises the p99 pre\-clip full\-model gradient norm from 522 to 3612\. Thus, nominally successful scalar Z\-loss reduction can correspond to a worse transported update\.

Table 16:Pythia\-160M WikiText\-103 continued pretraining\. The Z\-loss target in this grid is based on the actual output\-head vocabulary size\.The continued pretraining result separates four effects\. Standard Z\-loss is a strong practical baseline for reducing log\-normalizer tails\. The PPL effects of standard Z\-loss depend on the raw common\-shift channel and training regime\. Centered head deployment most directly removes the common\-shift coordinate and lowers output\-to\-hidden gain tails\. In the 4800\-step GPT\-2 run, centered head deployment matches CE within run\-to\-run variation\. Centered Z\-loss is the most balanced intervention in the high\-coefficient GPT\-2 setting and in the high\-common\-shift Pythia\-160M setting\. Centered Z\-loss reduces bothPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}without the rare raw\-logit gradient spikes seen in standard and gain\-aware Z\-loss\. Factorized and gain\-aware alternatives provide additional tuning knobs whose benefits depend on the target diagnostic\. Together, these interventions form architecture\-aware variants of Z\-loss transport\.

### High\-Coefficient Stress and Precision Endpoints

#### Additional WikiText\-103 Stress\-Test Details

The main\-text stress test evaluates optimizer\-facing consequences rather than introducing a new transport mechanism\. The matched GPT\-2 continued pretraining sweep in Table[3](https://arxiv.org/html/2609.16179#Sx3.T3)increases the applied auxiliary coefficient tenfold toλaux=1×10−3\\lambda\_\{\\rm aux\}=1\\times 10^\{\-3\}\. Centered Z\-loss is the only auxiliary method that improves the 1200\-step mean PPL relative to CE while also substantially reducing bothPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}and keeping the p99 pre\-clip full\-model gradient norm near the CE level\. Gain\-aware Z\-loss also improves mean PPL and reduces the log\-normalizer tail\. Because gain\-aware Z\-loss operates on raw logits, the gain\-aware objective leavesAp99A\_\{p\}^\{99\}near the CE baseline and shows rare large gradient\-norm spikes\. Standard Z\-loss has the same raw\-logit geometry issue and a small PPL increase relative to CE in this high\-coefficient setting\. The Centered head row is the gauge\-only control\. The lowerAp99A\_\{p\}^\{99\}of the Centered head row exposes the centered transport geometry, while the row has no auxiliary branch, so the reportedPZ99\.9P\_\{Z\}^\{99\.9\}is a deployed\-coordinate diagnostic rather than the outcome of Z\-loss training\.

Clip rate and p99 gradient norm summarize different parts of the training\-step distribution\. CE, Centered head, and Centered Z\-loss exceed the 1\.0 clipping threshold on nearly every step but remain near the CE\-scale p99 norm\. Standard and Gain\-aware Z\-loss exceed the threshold on only about7%7\\%of steps, yet the p99 norms of Standard and Gain\-aware Z\-loss are much larger\. The lower clip rate therefore does not indicate a safer tail\. The lower rate reflects mostly sub\-threshold steps punctuated by rare, very large events\. This distributional distinction is the reason we report both threshold frequencies and tail magnitudes\.

Appendix Tables[17](https://arxiv.org/html/2609.16179#A1.T17),[18](https://arxiv.org/html/2609.16179#A1.T18),[19](https://arxiv.org/html/2609.16179#A1.T19),[20](https://arxiv.org/html/2609.16179#A1.T20), and[21](https://arxiv.org/html/2609.16179#A1.T21)report further stress and precision endpoints, including effectively unclipped stress tests, update\-outlier and coefficient sweeps, Adam\-state audits, and static fp16 overflow audits\. These endpoints extend Table[3](https://arxiv.org/html/2609.16179#Sx3.T3)\. Raw\-logit scalar\-tail reduction can amplify transported update tails\. Centered transport reduces output\-to\-hidden gain and update\-tail metrics with regime\-dependent quality trade\-offs\.

These endpoints measure update pressure and numerical headroom rather than observed training divergence\. No nonfinite gradients occurred in the reported FineWeb\-Edu stress runs\. The static fp16 endpoint is a held\-out proxy\. The endpoint applies candidate static loss scales to gradients computed in the stated audit under bf16 autocast and asks whether any scaled coordinate would exceed the fp16 range\. The static fp16 endpoint therefore quantifies margin under a specified deployment scenario rather than claiming that the corresponding training run actually overflowed\.

Appendix Tables[14](https://arxiv.org/html/2609.16179#A1.T14),[15](https://arxiv.org/html/2609.16179#A1.T15), and[16](https://arxiv.org/html/2609.16179#A1.T16)report model\-scale and model\-family transfer experiments with GPT\-2 Medium and Pythia\-160M\. These results do not introduce a separate mechanism\. The transfer experiments test the same dense\-head interpretation across scale and model family\. Common\-shift removal and output\-to\-hidden gain reduction are distinct interventions\. Scalar Z\-loss tail reduction can differ from transported\-update tail reduction\.

#### Effectively Unclipped and Update\-Outlier Stress

Appendix Table[17](https://arxiv.org/html/2609.16179#A1.T17)reports results from a 900\-step effectively unclipped update\-tail evaluation using GPT\-2 Medium on FineWeb\-Edu\. The 900\-step setting uses the same disjoint validation offset as Appendix Table[12](https://arxiv.org/html/2609.16179#A1.T12)\. The setting is more aggressive than the low\-coefficient FineWeb\-Edu comparison\. The 900\-step setting uses a learning rate of2×10−42\\times 10^\{\-4\},λaux=1×10−3\\lambda\_\{\\rm aux\}=1\\times 10^\{\-3\}, a clipping threshold of1×1081\\times 10^\{8\}, and 900 optimizer steps\. The result again separates scalar\-tail reduction from transported\-update\-tail reduction\. Standard raw\-logit Z\-loss gives the smallest scalar tail, reducingPZ99\.9P\_\{Z\}^\{99\.9\}by a factor of16\.516\.5\. Under standard Z\-loss, the p99 pre\-clip full\-model gradient norm rises by a factor of4\.14\.1\. Standard Z\-loss also raises the maximum observed pre\-clip full\-model gradient norm by a factor of7\.87\.8\. Factorized and gain\-aware variants confirm that the issue is not unique to the standard scalar objective\. Both reduce the scalar tail\. Because the factorized and gain\-aware variants still leave the raw\-logit output\-to\-hidden gain near CE, the p99 pre\-clip full\-model gradient norms of these two variants rise to 23\.0 and 13\.0\. Centered Z\-loss prioritizes tail reduction at a measured PPL cost in this 900\-step setting\. Centered Z\-loss lowers bothPZ99\.9P\_\{Z\}^\{99\.9\}andAp99A\_\{p\}^\{99\}while leaving p99 and maximum pre\-clip full\-model gradient norms at the CE scale\. Centered head deployment is quality\-neutral in this setting and halvesAp99A\_\{p\}^\{99\}without introducing update outliers\. No run produced gradients containing not\-a\-number \(NaN\) or infinity \(Inf\) values\.

Table 17:Effectively unclipped high\-coefficient GPT\-2 Medium stress on FineWeb\-Edu\.The experiment in Appendix Table[18](https://arxiv.org/html/2609.16179#A1.T18)increases both the learning rate and the auxiliary coefficient to measure optimizer\-facing outliers\. PPL and transport diagnostics rank the methods differently in this setting\. Atλaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}, standard, factorized, and gain\-aware raw\-logit variants all improve 600\-step PPL relative to CE while lowering the scalar log\-normalizer tail\. The backward transport diagnostics reveal the cost\. Standard Z\-loss raises the p99 and maximum pre\-clip full\-model gradient norms by factors of5\.35\.3and22\.322\.3, respectively\. Factorized Z\-loss raises the p99 and maximum norms by factors of5\.55\.5and29\.829\.8, respectively\. Gain\-aware Z\-loss is milder but still raises the p99 and maximum pre\-clip full\-model gradient norms by factors of3\.33\.3and17\.317\.3, respectively\. These high\-coefficient raw\-logit variants have‖g‖\>10\\\|g\\\|\>10events in 1\.56% to 1\.72% of optimizer steps\. The high\-coefficient raw\-logit variants have‖g‖\>100\\\|g\\\|\>100events in 0\.11% to 0\.61% of optimizer steps\.

The coefficient sweeps in Appendix Table[19](https://arxiv.org/html/2609.16179#A1.T19)map this trade\-off\. Raw\-logit variants reduce thresholded\-update rates as the coefficient decreases\. Atλaux=1×10−4\\lambda\_\{\\rm aux\}=1\\times 10^\{\-4\}, standard and gain\-aware Z\-loss are near CE in PPL and have zero‖g‖\>10\\\|g\\\|\>10,‖g‖\>50\\\|g\\\|\>50, and‖g‖\>100\\\|g\\\|\>100events\. Standard and gain\-aware Z\-loss still leaveAp99A\_\{p\}^\{99\}at the CE scale, about 25\.0\. The scalar\-tail reduction from these raw\-logit variants is therefore not output\-to\-hidden gain reduction\. Centered Z\-loss shows the complementary trade\-off\. Atλaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}, centered Z\-loss prioritizes update\-event suppression\. Centered Z\-loss keeps all update\-event rates at zero, with a PPL 0\.60 higher than the CE baseline\. Lowering the coefficient to1×10−41\\times 10^\{\-4\}recovers near\-CE PPL, with 27\.774 for centered Z\-loss and 27\.766 for CE\. Centered Z\-loss still reducesPZ99\.9P\_\{Z\}^\{99\.9\}from 47\.3 to 15\.8 andAp99A\_\{p\}^\{99\}from 25\.0 to 11\.7\. Centered Z\-loss also keeps all‖g‖\>10\\\|g\\\|\>10,‖g‖\>50\\\|g\\\|\>50, and‖g‖\>100\\\|g\\\|\>100event rates at zero\.

Table 18:Stronger GPT\-2 Medium update\-outlier stress on FineWeb\-Edu atλaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}for auxiliary\-method rows\. ColumnsΦ10\\Phi\_\{10\},Φ50\\Phi\_\{50\}, andΦ100\\Phi\_\{100\}are pre\-clip full\-model gradient\-norm event fractions, withΦτ≔Prstep⁡\[‖g‖2\>τ\]\\Phi\_\{\\tau\}\\coloneqq\\Pr\_\{\\rm step\}\[\\\|g\\\|\_\{2\}\>\\tau\]\. CE and centered head rows have no applied auxiliary loss\. No microbatches or gradients contained NaN or Inf values in any run\.Table 19:FineWeb\-Edu update\-outlier coefficient sweeps\. ColumnsΦ10\\Phi\_\{10\},Φ50\\Phi\_\{50\}, andΦ100\\Phi\_\{100\}are pre\-clip full\-model gradient\-norm event fractions, withΦτ≔Prstep⁡\[‖g‖2\>τ\]\\Phi\_\{\\tau\}\\coloneqq\\Pr\_\{\\rm step\}\[\\\|g\\\|\_\{2\}\>\\tau\]\.
#### Optimizer\-State and Overflow Endpoints

Appendix Table[20](https://arxiv.org/html/2609.16179#A1.T20)reports the optimizer\-memory endpoint under the same effectively unclipped FineWeb\-Edu stress setting\. The Adam\-state endpoint samples AdamW state every 50 steps over 300 optimizer steps\. Under this high\-threshold clipping convention, these measurements show how raw\-logit gradient spikes affect an adaptive optimizer\. Standard, factorized, and gain\-aware Z\-loss improve 300\-step PPL and reducePZ99\.9P\_\{Z\}^\{99\.9\}\. The associated raw\-logit gradient spikes are nevertheless written into the second\-moment accumulator\. Relative to CE, standard, factorized, and gain\-aware Z\-loss increase the maximum sampled Adam second\-moment accumulatorvvby factors of36\.736\.7,42\.142\.1, and20\.320\.3, respectively\. The full preconditioned Adam direction norm is partly damped by this largervv, as expected from coordinate\-wise normalization in Adam byv\\sqrt\{v\}\. The maximum\-coordinate Adam direction still increases slightly\. Centered head deployment and centered Z\-loss keep the Adamvvtail close to CE\. Centered head deployment and centered Z\-loss also preserve a zero pre\-clip gradient\-event rate\. Centered Z\-loss shows the corresponding PPL trade\-off at this aggressive coefficient\. This Adam\-state endpoint links the tail diagnostic to optimizer memory and update coordinates, while clip\-before\-Adam training stacks can reduce this particular memory effect by clipping spikes before the Adam update\.

Table 20:Adam\-state endpoint for effectively unclipped GPT\-2 Medium stress on FineWeb\-Edu\.vmaxv\_\{\\max\}is the maximum sampled Adam second moment; thevmaxv\_\{\\max\}ratio to CE is computed relative to the matched CE baseline\. All recorded gradients and sampled optimizer\-state values remained finite\.Appendix Table[21](https://arxiv.org/html/2609.16179#A1.T21)gives a second numerical endpoint that is independent of the Adam state\. We backpropagate the same high\-coefficient objectives on held\-out corpus text\. We inspect full\-model gradients before any optimizer step and measure static fp16 loss\-scale headroom\. The endpoint complements the Adam\-state analysis by measuring maximum\-coordinate headroom\. The ordering is consistent with the effectively unclipped and Adam\-state results\. For GPT\-2 on WikiText\-103, standard raw\-logit Z\-loss increases the mean full\-model gradient norm from 17\.0 to 191\.2\. Standard raw\-logit Z\-loss increases the mean maximum coordinate from 5\.77 to 70\.0 and reduces the worst\-case maximum safe static scale from 6718 to 731\. At a static scale of 4096, standard raw\-logit Z\-loss gradients overflow in every audited batch, while CE and centered Z\-loss do not\. For GPT\-2 Medium on FineWeb\-Edu, standard, factorized, and gain\-aware raw\-logit variants reduce the worst\-case maximum safe scale to 1483, 1437, and 2248, respectively\. The raw\-logit variants produce overflow at a static scale of 4096 in 87\.5% to 100% of audited batches\. Gradients from the centered head and centered Z\-loss stay at the CE scale\. Centered\-head and centered\-Z\-loss gradients have worst\-case maximum safe scales around 8900, and neither centered method produces overflow at a static scale of 4096\. This overflow endpoint links the transported maximum\-coordinate tail to a concrete mixed\-precision numerical margin\.

Table 21:Static fp16 loss\-scale overflow proxy on full\-model gradients from held\-out corpus blocks\. All rows use bf16 autocast and 16 held\-out sequence blocks; auxiliary\-method rows useλaux=3×10−3\\lambda\_\{\\rm aux\}=3\\times 10^\{\-3\}\. Standard, Factorized, and Gain\-aware Z\-loss use raw logits\. Centered head has no auxiliary loss, while Centered Z\-loss uses row\-centered deployed logits\. For Factorized Z\-loss,λμ=λrel=λaux\\lambda\_\{\\mu\}=\\lambda\_\{\\mathrm\{rel\}\}=\\lambda\_\{\\rm aux\},cμ=0c\_\{\\mu\}=0, andcrel=log⁡Vc\_\{\\mathrm\{rel\}\}=\\log V\. Gain\-aware Z\-loss usesβ=1×10−2\\beta=1\\times 10^\{\-2\}\.κmin\\kappa\_\{\\min\}is the minimum over audited batches of the per\-batch maximum safe scale65504/‖g‖∞65504/\\left\\lVert g\\right\\rVert\_\{\\infty\}\. Each ofΨ4096\\Psi\_\{4096\}andΨ16384\\Psi\_\{16384\}denotes the fraction of audited batches whose scaled gradient would exceed the fp16 range at the corresponding static scale, so each lies in\[0,1\]\[0,1\]\.

### Additional Router and MoE Training Results

#### Router Projection Audit

We compare token\-layer mean, active\-route mean, and active\-route sum reductions for a 16\-expert linear router applied to GPT\-2 hidden states from WikiText\-103;δhZ\\delta\_\{h\}^\{Z\}denotes the router Z\-loss correction transported to the router input\.

#### End\-to\-End MoE Training

We then train controlled MoE language models end\-to\-end on WikiText\-103\. Appendix Table[23](https://arxiv.org/html/2609.16179#A1.T23)reports the top\-44setting, where the scale gap between active\-route mean and active\-route sum is largest\. The matched variants use adjusted coefficients so that the scale of each variant relative to the token\-layer mean equals 1\. Consistent with the router effective\-scale accounting above, the matched variants reproduce the token\-layer\-mean result to numerical precision within each matched run\. This matched recovery is an end\-to\-end equivalence check\. Under the reused initialization, coefficient matching makes the effective router objectives equivalent, and their matched recovery confirms the predicted reduction\-scale accounting at the end\-to\-end training level\. Unmatched active\-route mean uses one quarter of the coefficient required to match the token\-layer mean and has lower router entropy\. Active\-route sum uses four times the coefficient required to match the token\-layer mean and has higher router entropy\. The 1200\-step setting yields similar validation CE across rows\. The router metrics follow the predicted scale accounting\. This controlled MoE training experiment shows that absolute per\-decision coefficients and scales relative to the token\-layer mean make reduction conventions comparable\. The top\-22setting shows the same matched\-variant recovery with relative scales 0\.5, 1\.0, and 2\.0\.

Table 22:Router Z\-loss reduction conventions for a 16\-expert linear router applied to GPT\-2 hidden states obtained from WikiText\-103 validation blocks\. Effective scale and hidden correction norm change withkkdespite the same nominalλR=1×10−3\\lambda\_\{R\}=1\\times 10^\{\-3\}\.
Table 23:End\-to\-end top\-44MoE language\-model training on WikiText\-103\. Scale\-matched active\-route variants recover the token\-layer\-mean baseline; unmatched variants move router entropy according to each variant’s effective scale\. Load coefficient of variation \(CV\) reports expert\-load dispersion\.

### Reporting Recommendations and Limitations

#### Z\-Loss Reporting Recommendations

The experiments yield a practical reporting standard for Z\-loss\. For dense output heads, comparable reports should include the applied auxiliary coefficientλaux\\lambda\_\{\\rm aux\}and the diagnostic coefficientλdiag\\lambda\_\{\\rm diag\}when different\. Comparable reports should also include the targetcc, the logit convention, whether embeddings are tied, and whether the auxiliary is applied to raw or deployed logits\. Reports should specify whether the logits are raw, row\-centered deployed logits, or otherwise normalized\. Scalar validation quality is most interpretable alongside transport diagnostics\. These transport diagnostics includePZ99\.9P\_\{Z\}^\{99\.9\},Ap99A\_\{p\}^\{99\}, p99 or p99\.9 gradient norm, maximum\-coordinate gradient tail, gradient\-to\-parameter norm ratio, gradient clipping frequency, auxiliary\-to\-CE gradient ratio, and auxiliary\-to\-CE cosine\. Gradient\-norm reports should also specify the logit convention and the pre\-clip or post\-clip convention\. For tied embeddings, output\-path and input\-path auxiliary gradient norms make the transported update directly visible in diagnostic batches\.

For fused or low\-precision losses, forward loss equality is most informative when paired with backward\-source audits\. We recommend auditing the backward source against a high\-precision reference\. Reports should include relative gradient errors and cosines for both the total CE\+Z logit gradient and the Z\-loss source\. The audit is strongest when the evaluation uses target\-model logits from the relevant corpus distribution as well as random tensors, because the error depends on logit scale and distribution\. When adaptive optimizers such as Adam or AdamW are used\(Kingma and Ba[2015](https://arxiv.org/html/2609.16179#bib.bib17); Loshchilov and Hutter[2019](https://arxiv.org/html/2609.16179#bib.bib18)\), at least one optimizer\-facing check makes the transported effect visible\. Examples include selected\-parameter or whole\-model gradient error and Adam\-direction error after backpropagating the approximate source through the model graph on corpus batches\. Informative transport and numerical reports should first identify nonfinite events through nonfinite\-microbatch rate, nonfinite\-gradient rate, and separate counts for NaN and Inf when such events occur\. The reports should then describe finite\-gradient behavior through finite\-only gradient tails, thresholded gradient\-event rates, maximum\-coordinate gradient tails, and static loss\-scale overflow margins when gradients may be stored in fp16\. Finally, the reports should include optimizer\-memory and loss\-tail endpoints, such as Adam second\-moment tails, Adam update\-state tails, and loss\-spike or training\-loss\-tail endpoints\. For MoE routers, nominalλR\\lambda\_\{R\}is most interpretable with the reduction convention, top\-kk, token\-layer denominator, absolute per\-decision coefficient, and scale relative to the token\-layer mean\. When capacity is enforced, reports should include the capacity policy, dropped\-route fraction, expert load CV, router entropy, and router correction norm\. These details make capacity\-limited routing and dropless sparse kernels easier to compare\(Galeet al\.[2023](https://arxiv.org/html/2609.16179#bib.bib47)\)\.

#### Limitations

The experiments validate the proposed backward\-transport framework across two text corpora, several scales in the GPT\-2 and Pythia families, matched continued pretraining, implementation and precision audits, optimizer endpoints, and controlled top\-kkMoE training\. This controlled design isolates source and transport effects rather than reproducing every production fused kernel, expert\-parallel configuration, or frontier\-scale pretraining setup\. These choices define the scope of the empirical validation; the underlying softmax\-source identities apply more broadly\. As the results show, intervention choice remains model\- and regime\-dependent and should be guided jointly by validation quality and transport diagnostics\. Extending the evaluation to additional production stacks, model families, and training scales is a natural next step\.