#54 · Logit Rescale

running

◷train_gpt.py▸23 × sigmoid((logits+5) / 7.5)◎learner＄no GPU

Blocks

Quick summary

The LM head raw logits can grow unboundedly during training, causing loss spikes and instability. Instead of…

Full explanation below the code →

fba-lab — train_gpt.py · 23 × sigmoid((logits+5) / 7.5)executing

// block: 23 × sigmoid((logits+5) / 7.5) · lines 1164–1167$ study train_gpt.py --block logit-softcapThe LM head raw logits can grow unboundedly during training, causing loss spikes and instability. In… ✓

Explanation

The LM head raw logits can grow unboundedly during training, causing loss spikes and instability. Instead of tanh(x/c)·c (Gemma 2 soft-cap), this uses a shifted sigmoid: 23·σ((x+5)/7.5). The +5 shift biases toward the active range, 23 sets the output scale, 7.5 controls sharpness.

Think about

Why does logit bounding stabilize training compared to letting logits grow freely?

// architecture

Live diagram

Record #54 · PR #181 · Logit rescale · −40 steps, −2.9 s

Training stability

Bounding logits prevents the loss spikes that kill speedruns

The language model head produces raw logits — one score per vocabulary token. During training these scores can grow very large. When a few logits become enormous, the softmax denominator explodes and almost all probability mass collapses onto one token. Gradients for every other token approach zero. The loss spikes and recovery costs dozens of wasted steps.

PR #181 applies a shifted sigmoid to squash logits into a bounded range before computing the loss: 23 · σ((x + 5) / 7.5). The +5 shift centers the active zone where logits naturally live; 23 sets the output ceiling; 7.5 controls how gradually the cap kicks in. Every vocabulary token receives meaningful gradient throughout training — no spikes, no recovery steps lost.

# lines 1164–1167 — bounded logits via shifted sigmoid
logits = 23 * torch.sigmoid((logits + 5) / 7.5)
loss = F.cross_entropy(logits, targets)

Why does a logit spike kill training?

When one logit becomes very large, softmax assigns near-100% probability to that token. All other tokens see near-zero probability — and near-zero gradient. The optimizer effectively stops learning from most of the vocabulary. Recovery requires many steps to redistribute probability mass, wasting the compute budget.

Why sigmoid instead of tanh?

Gemma 2 uses tanh(x/c)·c (symmetric around zero). This PR uses a shifted sigmoid because the logit distribution during training is skewed positive — the +5 shift puts the sigmoid's inflection point in the natural operating range, giving sharper gradients where they matter most.

What changed vs #53

+         # @Grad62304977 added tanh softcapping following Gemma 2 paper, @KoszarskyB reduced it from 30 to 15
+         # @YouJiacheng shifted it by +15 (2*sigmoid(2*x)=tanh(x)+1). @classiclarryd updated to 23*sigmoid((logits+5)/7.5)
+         logits = 23 * torch.sigmoid((logits+5) / 7.5)

Bounding logits prevents the loss spikes that kill speedruns

# lines 1164–1167 — bounded logits via shifted sigmoid
logits = 23 * torch.sigmoid((logits + 5) / 7.5)
loss = F.cross_entropy(logits, targets)

Why does a logit spike kill training?

Why sigmoid instead of tanh?

Lab

#54 · Logit Rescale

Live diagram

Bounding logits prevents the loss spikes that kill speedruns

What changed vs #53

Logit softcap

#54 · Logit Rescale

Live diagram

Bounding logits prevents the loss spikes that kill speedruns

What changed vs #53

Logit softcap