Cross-entropy loss architecture — log-sum-exp, fused gradients, and vocab-parallel reductions

The logits tensor is the problem

Everything hard about cross-entropy is downstream of one shape. The final hidden state h has shape [tokens, d_model]; the LM head W projects it to logits z of shape [tokens, vocab]. For a modern model that vocabulary is 32k–256k, so the logits tensor is the largest single array the forward pass produces — wider than any attention or MLP activation. On a batch of 65k tokens with a 128k vocabulary in fp32, the logits alone are ~33 GB. The LM-head matmul plus a naive softmax-then-CE can be 20–30% of step memory. So the loss is not a scalar you tack on at the end; it is a memory and communication problem that dictates how the last layer is written. Every technique below exists to avoid holding that tensor, in fp32, all at once, on one device.

Token-level NLL and the reduction that actually matters

At each position the loss is −log p[target] — one gathered entry, never a dot product against a one-hot vector. The subtlety at scale is the reduction. Real batches are packed: many documents concatenated, prompt tokens in SFT masked out, padding at the tail, cross-document boundaries excluded. Positions labelled ignore_index contribute zero and — critically — are removed from the denominator. The loss is the sum of unmasked NLLs divided by the count of unmasked tokens, not by batch size or sequence length. Divide by the wrong denominator and examples with different mask ratios get systematically different gradient weight; the model quietly over-trains on the densely-labelled samples. Token-mean reduction is what makes one gradient step treat every supervised token equally, which is the property the whole training recipe assumes.

log-sum-exp: staying finite in bf16

Numerical stability here is not a textbook nicety — it is a hard constraint of the arithmetic. In bf16 the logits routinely reach ±40, and e^40 ≈ 2×10^17 already strains fp32 accumulation while larger values overflow. The fix is the log-sum-exp identity: subtract the row max m = max_j z_j before exponentiating, and compute the log-probability directly in one expression rather than as log(softmax()).

log p_i = z_i − m − log Σ_j e^(z_j − m)

After the shift every exponent is ≤ 0, every term lands in (0, 1], the largest is exactly e^0 = 1, and the sum can neither overflow nor underflow to zero — so the outer log stays finite, with no tiny intermediate probability to round to log 0 = −∞. Softmax is invariant to a constant shift of the logits, so this changes nothing mathematically and everything numerically. That is why libraries expose a fused cross_entropy(logits, targets) and do the reduction in fp32 even when logits arrive in bf16: the max and sum are computed in high precision precisely because that is where a silent inf would be born. On a CPU SLM, feeding raw logits to the fused loss is the only version that is reliably correct.

The fused gradient: forward and backward share one pass

The reason cross-entropy is written as a single fused kernel — not two composable ops — is the gradient. The derivative of the loss with respect to the logits is ∂z = p − y: the predicted distribution minus the one-hot target. That structure is a systems gift. It means the backward pass needs only p and the target, both of which the forward pass already touched, so a fused kernel can emit the loss and the gradient in one sweep over the logits. The full [tokens, vocab] log-probability tensor never has to be stored for backward — it is recomputed or streamed. Removing that saved tensor is often the difference between a batch that fits and an out-of-memory crash. The clean math from the derivation is precisely what buys the memory win in the implementation.

Chunked cross-entropy: slice the token axis

Fusing removes the stored log-probabilities, but the logits themselves are still huge. Chunked CE attacks that directly: rather than compute the LM-head matmul for all tokens at once, it slices the token dimension into blocks, and for each block computes logits, loss, and the p − y gradient, then accumulates the scalar loss and scatters the gradient into the LM-head backward buffer. Peak memory is one chunk of logits — say 4k tokens × vocab — instead of the whole 65k×vocab array. The arithmetic is identical; only the schedule changes. This is the single most effective trick for training with a large vocabulary on constrained memory, and it composes cleanly with gradient checkpointing: the LM head is recomputed per chunk in backward anyway, so nothing extra needs to be stashed.

Vocabulary-parallel cross-entropy under tensor parallelism

When the model is tensor-parallel, each rank owns a column shard of the LM head — a slice of the vocabulary — and therefore only ever holds partial logits. No device should assemble the full-vocab logits, so the softmax normalizer must be reconstructed with collectives. Each rank computes a local row max; an all-reduce(max) yields the true row max m. Each rank computes Σ e^(z − m) over its own columns; an all-reduce(sum) yields the global normalizer Z. The rank whose shard contains a position’s target token contributes the gather term −(z_target − m − log Z); all others contribute zero for that position. Two small scalar collectives replace ever materializing a [tokens, vocab] tensor anywhere. Get the reduction wrong — a missing rank, a stale max — and the model trains on a silently incorrect normalizer, the nastiest class of distributed bug because the loss still looks plausible.

z-loss: penalizing the normalizer itself

At scale, the softmax normalizer log Z can drift upward as the network inflates its logits, and in bf16 that drift shows up as instability — loss spikes, occasional inf. The z-loss regularizer, popularized by PaLM, adds a cheap penalty on the log-normalizer:

L_total = CE + λ · (log Z)^2     (λ ≈ 1e−4)

It pulls log Z toward zero, which keeps the raw logits in a sane range and the exponentials well-conditioned, without meaningfully changing the probabilities the model expresses. Its gradient, 2λ · log Z · ∂log Z, folds into the same fused backward pass for a negligible cost. Think of it as insurance: a small, constant tax on numerical recklessness that buys a much smoother loss curve in low precision. It is the one guardrail the derivation-focused treatment tends to skip, because it only matters once you are training for real.

Label smoothing seen from the logits

Label smoothing replaces the one-hot target with (1 − ε) on the true token and ε spread across the vocabulary. The derivation view is ‘softer targets’; the systems view is more concrete. A one-hot target rewards ever-larger correct logits — the loss keeps dropping as p_c → 1, so logits grow without bound and the model becomes miscalibrated and numerically edgier. Smoothing installs a fixed point: the gradient p − y_smooth reaches zero at a finite logit gap rather than at infinity, so logit norms stop inflating and confidence calibrates. The costs are equally concrete — a constant floor is added to the reported loss (it no longer trends to zero), and the target is no longer a single index, so the gather-only fast path gives way to a slightly heavier reduction. Pretraining often skips it; SFT commonly uses ε = 0.1.

Perplexity and bits-per-byte: the number teams actually watch

The loss has a physical reading, and it is the metric on the dashboard. Cross-entropy in nats is the average code length the model needs per token, and perplexity = e^loss is its intuitive twin: a loss of 2.0 is perplexity ~7.4, meaning the model is on average as uncertain as a fair choice among 7–8 tokens. Because perplexity depends on the tokenizer, cross-model comparisons often prefer bits-per-byte — cross-entropy in base-2 normalized by UTF-8 bytes rather than tokens — which is tokenizer-agnostic. Every scaling law and pretraining leaderboard claim is a statement about this one scalar, so reporting it means being pedantic about the reduction: a perplexity computed over a different mask or denominator is not comparable, and most spurious ‘wins’ trace to exactly that mismatch.

The loss curve is an instrument, not a scalar

Because cross-entropy is built as an observable, a loss curve carries diagnostic signal a bare scalar does not. Instrument the CE path to log, alongside the token-mean loss and its perplexity, the gradient norm and the maximum absolute logit max|z|. When the loss spikes 0.3 nats for a few steps, max|z| jumping to 60 on a pathological document names the cause immediately — a normalizer blow-up the z-loss term is already absorbing — so the on-call watches it recover instead of restarting from a checkpoint. A flat loss whose perplexity disagrees with a sibling run usually means a reduction or masking mismatch, not a modelling difference. Emit these signals from inside the loss kernel, where the logits still exist — after the fused backward they are gone.

Putting the pipeline together

Trace one step of an 8B model, tensor-parallel over four devices, sequence length 8,192, vocab 128k. The last layer emits h replicated on every rank. Each rank multiplies h by its 32k-column head shard — four partial logits tensors, never the full [tokens, 128k] monster. The vocab-parallel kernel runs its two collectives for the exact global normalizer; the ranks owning each target contribute the gather; ~52k unmasked tokens set the denominator while document boundaries and prompt tokens contribute zero. Backward starts from that scalar with no stored softmax: each rank emits p − y for its own vocabulary shard — exactly what that shard of W needs — and z-loss folds in for pennies. The loss lands on the curve the team watches, beside perplexity, grad norm, and max|z|. That is cross-entropy as a system: a formula the derivation earns, engineered into a kernel you can fit in memory, shard across devices, and read like a gauge.