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 <!-- TODO: add arXiv/OpenReview link --> 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.

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.

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

Training

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 — [email protected]