nthngdy/grad_bottleneck_2B

Gradient Bottleneck — 2B Controlled Pretraining Suite

Research checkpoints for the COLM 2026 paper Lost in Backpropagation: The LM Head is a Gradient Bottleneck by Nathan Godey and Yoav Artzi (Cornell University).

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 8 D 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:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "nthngdy/grad_bottleneck_2B"

model = AutoModelForCausalLM.from_pretrained(model_id, revision="d4096_170k")
tokenizer = AutoTokenizer.from_pretrained(model_id, revision="d4096_170k")

To compare the bottleneck across output ranks at a matched token budget, load the same _170k (or _85k / _130k) revision for different D:

for d in [32, 64, 128, 256, 512, 1024, 2048, 4096]:
    m = AutoModelForCausalLM.from_pretrained(model_id, revision=f"d{d}_170k")
    ...

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

@misc{godey2026lostbackpropagationlmhead,
      title={Lost in Backpropagation: The LM Head is a Gradient Bottleneck}, 
      author={Nathan Godey and Yoav Artzi},
      year={2026},
      eprint={2603.10145},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2603.10145}, 
}

Contact

Nathan Godey — godeynathan@gmail.com