These models are not intended as general-purpose LMs. They are a controlled
suite built to isolate the optimization effect of the softmax bottleneck — the
compression of gradients as they backpropagate through a low-rank LM head — from
its classical expressivity effect and from any confound of backbone size.
TL;DR
The LM head projects D-dimensional hidden states to V-dimensional logits
(D ≪ V). During backpropagation, the V-dimensional logit gradient is
compressed back through the rank-D head, destroying 95–99% of its norm. This
suite holds the Transformer backbone fixed and varies only the effective
output rank D, so any performance gap is attributable to the gradient
bottleneck alone. Convergence slows by up to ×16 between D=4096 and D=32.
Model Suite
All models share an identical Llama-3-style Transformer backbone (2B total
params, ~1.8B non-embedding). The LM head is factorized as a low-rank product
Wθ = Aθ Bθ with Aθ ∈ ℝ^{V×D} and Bθ ∈ ℝ^{D×dₘ}, so the backbone is equally
expressive across all variants while D controls the strength of the gradient
bottleneck.
Effective dim D
i (D = 2^i)
Notes
32
5
strongest bottleneck (~×16 slower convergence)
64
6
128
7
256
8
512
9
1024
10
2048
11
4096
12
full backbone width — no bottleneck
Checkpoints & Revisions
Each checkpoint lives on its own revision. Revision names follow the pattern:
d{D}_{step}k e.g. d512_170k, d32_85k
{D} — the effective output dimension (one of the 8 values above).
{step}k — the training-step count in thousands. At 512 tokens ×
128 sequences = 65,536 tokens/step, so step count maps to token budget as
below.
Revision suffix
~Tokens
Phase
_20k
1.3B
stable
_40k
2.6B
stable
_60k
3.9B
stable
_80k
5.2B
stable
_85k
5.6B
cooled (≈5B eval point)
_100k
6.6B
stable
_120k
7.9B
stable
_130k
8.5B
cooled (8.5B eval point)
_140k
9.2B
stable
_160k
10.5B
stable
_170k
11.1B
cooled (final / 11B eval point)
The three cooled revisions (_85k, _130k, _170k) exist for all 8D
values and correspond to the LR-cooldown checkpoints used for the downstream
evaluations in the paper (Table 2, Figure 2b). Stable-phase revisions trace the
loss curves in Figure 2a and are available for most — but not all — (D, step)
combinations. See the repo's branch list for exactly what is published.
Architecture & Config
Hyperparameter
Value
Backbone
Llama-3 architecture
Num. layers
6
Hidden size (dₘ)
4096
Intermediate size
16,384
Attention heads
32
Head dimension
128
Effective output dim (D)
32 … 4096 (see table)
LM head
low-rank factorized Aθ Bθ
Weight tying
No
Tokenizer
SmolLM2 (V = 49,152)
Total params
~2.0B (~1.8B non-embedding)
Training
Hyperparameter
Value
Data
FineWeb-Edu (~11B tokens)
Sequence length
512
Batch size
128
Total steps
170,000
Optimizer
AdamW (β₁=0.9, β₂=0.95, ε=1e-8)
Learning rate
3e-4
LR schedule
Warmup-Stable-Decay (WSD)
Warmup steps
2,000
Cooldown
10,000 steps (cosine), applied at ~5B, ~8.5B, ~11B tokens
Weight decay
0.1
Gradient clipping
1.0
Precision
mixed bf16
Hardware
B200 GPUs, ~760 GPU-hours total
Checkpoints are provided at the three WSD cooldown points (≈5B, 8.5B, 11B tokens)
so training dynamics can be compared across D at matched token budgets.
Key Results (from the paper)
95–99% of the logit-gradient norm is projected into the null space of Wθᵀ
and destroyed during backpropagation, across GPT-2, Pythia, Llama 3, OLMo 2 and
Qwen 3; the surviving signal has cosine similarity of only ~0.1–0.3 with the
full gradient.
×16 convergence slowdown between D=32 and D=4096 for the same backbone,
with consistent gaps in validation loss and zero-shot downstream scores
(ARC, HellaSwag, PIQA, SciQ, OpenBookQA, Lambada).
The learned head behaves like a random rank-D projection, not the optimal
(top-D singular) one — i.e. it does not learn to preserve the dominant
gradient directions.
Intended Use
Studying optimization dynamics of the softmax / gradient bottleneck.
Reproducing the paper's analyses (gradient-norm projection, update-direction
efficiency, downstream evaluation across D).
Ablations on LM-head design.
Out of scope: deployment, chat/instruction following, or benchmarking as a
capable general LM. Small-D variants are deliberately handicapped and the whole
suite is undertrained (~11B tokens) by design.
Usage
Pick a variant with revision="d{D}_{step}k". For the final cooled checkpoint of
the full-width (D=4096) model:
To compare the bottleneck across output ranks at a matched token budget, load the
same _170k (or _85k / _130k) revision for different D:
python
1for d in[32,64,128,256,512,1024,2048,4096]:2 m = AutoModelForCausalLM.from_pretrained(model_id, revision=f"d{d}_170k")3...
Limitations & Biases
Trained only on English FineWeb-Edu for a small token budget; not filtered or
aligned for safety. Outputs may be low-quality, repetitive, or biased. Reduced-D
variants are intentionally suboptimal and should never be read as evidence about
model quality per se — only about the bottleneck effect.
Citation
bibtex
1@misc{godey2026lostbackpropagationlmhead,
2 title={Lost in Backpropagation: The LM Head is a Gradient Bottleneck},
3 author={Nathan Godey and Yoav Artzi},
4 year={2026},
5 eprint={2603.10145},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2603.10145},
9}
