Delete
Transcript

A Controlled Study of Attention-Only Transformers

Delete 72% of a transformer's parameters — the feed-forward network — and the model loses 0.47 nats. Give those parameters back as attention depth, and the gap shrinks to 0.006 nats, agreeing to 10⁻⁴

Unknown · arxiv.org

Gist

1.

Delete 72% of a transformer's parameters — the feed-forward network — and the model loses 0.47 nats. Give those parameters back as attention depth, and the gap shrinks to 0.006 nats, agreeing to 10⁻⁴ across clean seeds. The FFN is not a computational primitive. It is a write path.

Logic

2.

Three matchings, three answers — and only one is honest

  • Iso-depth (87M vs. 24M parameters) costs 0.470 nats; iso-FLOP (43M vs. 24M) costs 0.263 nats; iso-param (24.12M vs. 24.13M) costs +0.0055 nats — the gap reverses sign
  • No single matching is fair because removing the FFN perturbs parameters, compute, and depth simultaneously; each control fixes one axis while the other two vary
  • The iso-param gap agrees to 10⁻⁴ across clean seed pairs, is 3.7× the same-seed noise floor (0.0015 nats), and shrinks monotonically from 0.046 nats at 5B tokens to 0.0055 at 105B

3.

The gap lives on query tokens — nowhere else

  • At 31B tokens, the per-token deficit on query regions is +0.052 nats — five times the sample aggregate — while carrying only 8% of corpus loss; the token-weighted sum of region gaps reproduces the aggregate within 2%
  • By 105B, the SAN leads on every answer region, including memorization exercises, and on reasoning-trace tokens; the deficit survives only on query tokens (+0.038 nats), the positions with the least context to route from
  • The pre-registered fineweb-edu prediction — gap 0.02–0.05 nats, measured 0.0398 — confirms the query-localized account on a knowledge-dense distribution

4.

QK-normalization is the load-bearing component — not gates, not residuals

  • Removing QK-normalization diverges outright at the tuned learning rate (8.28 nats), the only divergence in the entire study and a finding none of the authors' predictions anticipated
  • Scalar residual gates are performance-neutral at every depth (20–48 layers) and in both architectures; sandwich normalization is the only variant to beat the baseline, at −0.009 nats
  • The depth optimum is U-shaped at 20 layers, but 48-layer attention-only stacks train without incident — the classical rank-collapse regime (minimum layer rank 173 of 512) is never approached with residuals and normalization in place

5.

Routing freezes early; content accumulates indefinitely

  • QK matrices spectrally crystallize within the first quarter of training and do not move thereafter; Muon holds their spectra 2–3× flatter than AdamW, alike in both architectures
  • Content matrices — FFN down-projection in standard transformers, W_o in SANs — accumulate stable rank through the entire stable phase, contracting only in the learning-rate decay tail
  • The FFN's functional form largely does not matter; what matters is that it provides parameter capacity attached to a write path into the residual stream, and removing it relocates that accumulation to W_o

6.

The FFN arm broke in ways the SAN arm never did

  • At 16× the tuned learning rate, all FFN gates collapsed to mean 0.07 while attention gates stayed structured — the standard transformer self-pruned toward attention-only form
  • At the tuned rate under Muon, FFN gates ended low (mean 0.16–0.17) in three separate runs, echoing the same self-pruning behavior
  • One of three FFN seeds at 105B suffered a terminal-phase instability: gradient norm tripled over the final 800 steps at the learning-rate floor; no SAN run showed any of these behaviors

Counter-Argument

7.

The paper's own predictions failed — and the central finding is the consolation prize

  • Of eight pre-registered predictions, two were falsified outright and one was budget-dependent; the query-localized account emerged from the failure of the memorization-peak prediction, not from a confirmed hypothesis
  • The headline 0.006 nats gap is measured on SYNTH, a synthetic reasoning-dense corpus deliberately overtrained relative to compute-optimal; the fineweb-edu gap is 0.040 nats, and the authors' own storage-account theory predicts the gap widens further at larger scale — the paper's own framework undermines its own conclusion
  • The SAN pays ≈2× FLOPs per token at 2048 context, iso-FLOP favors the FFN by 0.263 nats, and MMLU-class benchmarks sit at chance throughout — the "regime, not superiority" framing is doing heavy lifting to contain a result that is small, distribution-dependent, and measured on a corpus no one actually deploys

Steelman

8.

The paper's real contribution is not the gap — it is the first causal map of what a transformer actually loses

  • Both the thesis and the counter-argument share a hidden assumption: that the gap's size is the paper's value. But before this work, no controlled experiment had ever deleted the FFN and measured the loss under matched parameters, FLOPs, and depth — the field had only probing studies, editing experiments, and theoretical analyses of attention without residuals
  • The spectral dynamics — routing matrices freezing early, content matrices accumulating rank through the stable phase, relocation to W_o when the FFN is removed — are a mechanistic discovery about how transformers learn, not just a measurement of a gap; they confirm at realistic depth the QK/OV decomposition that the interpretability literature had only shown in small models
  • The fineweb reversal on lambada (SAN 0.203 vs. FFN 0.181) is not a limitation — it is the paper's most important finding. A task's storage/routing identity is relative to the training distribution, which means every benchmark comparison in the field has been silently confounded by the match between task and corpus. The paper's lasting contribution is not "the FFN barely matters" but "we now have the instruments to ask what matters, where, and why"

Original

Continue Reading