Issue 30 · Pick 04 AI / ML ✓ read
A Controlled Study of Attention-Only Transformers
TL;DR: Take a transformer, delete every feed-forward layer, and spend the freed parameters on more attention layers instead. On a reasoning-dense corpus at 24M parameters and 105B tokens, this attention-only model lands within 0.006 nats of validation loss — 0.27% — of a standard transformer with the identical parameter count, reproducibly across seeds. The residual deficit isn't diffuse: three independent measurements localize it to parametric recall — predicting tokens when the context gives you nothing to work with — and a numerically pre-registered prediction on knowledge-dense web text (0.02–0.05 nats) came in at 0.040. The FFN, the paper argues, is not a distinct computational primitive; it's parameter capacity attached to a write path, and attention depth can substitute for almost all of it.
The experiment nobody ran under controls
Feed-forward networks hold roughly two thirds of a transformer's non-embedding parameters, and a whole interpretability literature treats them as the model's parametric memory: key–value stores over training facts (Geva et al.), the locus of editable associations (ROME), the home of "knowledge neurons." But all of that evidence comes from probing and editing — localizing function without deleting it. The converse experiment, remove the FFN and measure what's actually lost, had never been run with proper controls. The closest prior attempt (Sukhbaatar et al. 2019) removed the FFN but snuck equivalent capacity back in as persistent memory vectors, leaving the necessity question open. And the classic theory result — that pure attention stacks suffer rank collapse (Dong et al. 2021) — applies to a setting without residual connections that nobody actually trains.
Why hadn't anyone run it? Because deleting the FFN perturbs three things simultaneously: parameter count, compute per token, and the number of nonlinear composition steps. Any single-control comparison is confounded no matter which axis you fix. If you match depth, the attention-only model has 72% fewer parameters. If you match parameters by making the attention-only model deeper, it now burns roughly 2× the FLOPs per token, because attention pays a parameter-free quadratic cost in sequence length. There is no single "fair" comparison — so the paper runs all three.
The setup, concretely
The attention-only model ("Simple Attention Network," SAN) is a pre-norm decoder block with the FFN simply deleted: zero-centered RMSNorm → GQA attention (8 query / 4 KV heads) with RoPE and per-head QK normalization → output projection scaled by a scalar sigmoid gate → residual add. That's the whole block. The control arm is identical plus a SwiGLU FFN (d_{ff}=4d) after attention, sharing one code path. The only per-position nonlinearities in the SAN are the attention softmax and the norms — there is no learned per-position feature map, which is precisely the manipulation.
The paper formalizes what's removed with a few short results worth internalizing. For any fixed attention pattern, a SAN layer's update is a linear function of the normalized context; per head, the update lies in the convex hull of \{W_v \hat u_j\} over prior positions. So a SAN layer can select and transport content present in context but cannot synthesize representations unsupported by it. Most starkly (Lemma 1): at sequence length 1, the entire depth-L network collapses to a fixed linear map modulated by exactly L scalars (the running RMS values). All per-position nonlinearity flows through that L-scalar bottleneck; cross-token attention is the only escape. Pricing that restriction is what the experiments do.
The fairness engineering is unusually careful and is half the paper's value. Every arm×optimizer cell got its own learning-rate sweep with a boundary-extension rule (17 cells; the tuned Muon rate differs by 2× between arms, so a shared rate silently biases any comparison). The noise floor was calibrated with an optimizer-equivalent control pair: two provably-identical-under-Adam parameterizations trained anyway, yielding a same-seed floor of 0.0015 nats. Eight predictions were registered before their measurements; two were falsified, and one falsification produced the central finding.
The main result
And the iso-param gap isn't static — it shrinks with training. Each budget below was separately trained and separately tuned (different schedule horizons), so this is not one curve read at three points:
On the size axis, at a fixed 31.5B-token budget across five matched pairs spanning a 29× range of non-embedding parameters, the gap starts negative at the tiniest size (the 2-layer FFN partner hits a depth floor), crosses zero, and plateaus near 0.02 nats from 16M to 57M non-embedding parameters. The gap does not grow with scale in the tested range — though 87M parameters is a very small window from which to extrapolate.
The authors are careful to claim a tiny real effect, not a null: 0.006 nats is 3.7× their calibrated noise floor and reproduces to 10^{-4} across clean seed pairs. The FFN's parameters matter enormously; its functional form — the per-position learned nonlinearity — apparently mostly does not, on this distribution.
Where the residual 0.006 nats lives
This is the part that connects to the FFN-as-memory literature, and it's the paper's best contribution. The SYNTH corpus formats every document into query / reasoning-trace / answer regions delimited by marker tokens, so the loss gap can be decomposed exactly (the token-weighted sum of region gaps reproduces the aggregate to within 2%).
At 31B tokens, the per-token deficit on query regions — positions where there's minimal context to route from, so only weights can help — is +0.052 nats, five times the aggregate gap, while those tokens carry only 8% of corpus loss. By 105B tokens the localization is complete: the SAN leads on every answer region (including memorization exercises) and on reasoning traces, and trails only on query tokens (+0.038). Across the entire size ladder, the query deficit is the one invariant — positive at every size and budget while every other region flips sign.
Zero-shot benchmarks tell the same story at task level. Lambada — predicting a specific content word, which for SYNTH-trained models is an out-of-distribution recall task — favors the FFN arm at every budget. Sciq — where the answer sits in a provided passage, a pure routing task — favors the SAN with a margin that grows with training (0.725→0.742 while the FFN model regresses 0.702→0.661, non-overlapping seed ranges) and falls monotonically with FFN capacity. Amusingly, the 87M iso-depth model, the best in-distribution model, is the worst at lambada.
Then the confirmatory shot: if the FFN's advantage is low-context prediction, and natural web text is mostly low-context prediction, the gap should widen on web text. Before launching the run, the authors registered a numerical prediction — 0.02 to 0.05 nats on fineweb-edu, bounded above by the pure query-region deficit. The measured gap: 0.0398. Notably, the original registered account (gap largest on memorization answers) was falsified by the decomposition; only the revised, query-localized account made this successful prediction, and the paper says so plainly.
One subtle and genuinely useful side-finding: the same fineweb-trained pair reverses on lambada (SAN beats FFN). When training data is distribution-matched to the task, the passage suffices to infer the final word and lambada becomes a routing task. Storage-vs-routing is not a fixed property of a benchmark; it's a property of the match between task and training distribution. That reframes a lot of small-model benchmark comparisons.
The mechanism: routing crystallizes, content accumulates, storage relocates
Tracking stable rank (\|W\|_F^2/\|W\|_2^2, a smooth proxy for how many directions a matrix actively uses) of every weight matrix over training checkpoints reveals two regularities across all architectures, sizes, and budgets in the study:
First, routing matrices (Q/K) spectrally crystallize within the first quarter of training and never move again — the optimizer sets the level (Muon holds spectra 2–3× flatter than AdamW, in both architectures alike), but the schedule is universal. Second, the matrices that write content into the residual stream accumulate stable rank for as long as the stable phase lasts: the FFN down-projection in standard models, and in SANs the attention output projection W_o, which simply inherits the job. This is the mechanistic-interpretability QK/OV story (Elhage et al. 2021) — routing vs. content — showing up as a training dynamic at realistic depth. Storage isn't a property of the FFN; it's a property of whatever write path exists, and it relocates when you remove one.
The classical rank-collapse regime is never approached: with residuals and normalization in place, representation effective rank stays high everywhere (minimum layer rank 173 of 512 at 20 layers, 48-layer stacks train fine).
The one load-bearing component
The ablations (Table 6) contain a finding the authors did not predict: QK normalization — RMS-normalizing queries and keys per head before RoPE — is the single component that keeps deep attention-only stacks trainable. Removing it produces the study's only outright divergence (loss 8.28). Meanwhile the scalar residual gates, which the authors expected to matter, are performance-neutral at every depth from 8 to 48 layers and in both architectures; sandwich normalization is a free −0.009 nats; depth at iso-param is U-shaped with a 20-layer optimum. There's also a curious robustness asymmetry: every instability event in the study occurred in the FFN arm — including a wonderful diagnostic where, at 16× the tuned learning rate, the FFN model's FFN gates collapsed to mean 0.07 while its attention gates stayed structured. Under rate stress, the standard transformer self-pruned toward attention-only form.
What changes if this holds, and what to doubt
If the result generalizes, the FFN's two-thirds share of transformer parameters is best understood as fungible storage capacity, not an architecturally necessary nonlinearity — attention depth can absorb it, at the cost of ~2× FLOPs per token at 2048 context. The practical regime the authors claim (and only that regime) is where parameters, not FLOPs, bind: on-device memory-limited inference, plus settings where a single repeated interpretable primitive has analysis value. It also sharpens the interpretability picture: the storage/routing division is real, causal, and quantitatively small on trace-rich data — with the deficit exactly where a "context-transport-only" model should fail, on tokens with nothing to transport.
The reasons for skepticism are stated by the authors and worth repeating. Everything is at ≤87M total parameters and ≤105B tokens — three orders of magnitude below frontier scale. The headline near-parity is on SYNTH, a reasoning-dense synthetic corpus where most tokens are inferable from context; the one knowledge-dense control (fineweb-edu) already shows a 7× larger gap, and the storage account itself predicts the gap should widen on storage-heavy mixtures at scale — the natural next experiment, which the paper explicitly names. Token-axis convergence was measured only at the base size, size-flatness only at 31.5B tokens; their conjunction at larger scale is extrapolation. The region decomposition uses training-exposed data (valid for paired localization, not held-out measurement), reports point estimates without per-block intervals, and the confirmatory out-of-distribution result rests on a single fineweb point. And MMLU-class benchmarks are at chance throughout, so nothing here speaks to broad-knowledge capability. The FLOP overhead also means iso-FLOP — arguably the deployment-relevant matching for cloud serving — decisively favors the standard transformer.
Still, as an exercise in experimental hygiene — three matchings, per-arm LR sweeps, a calibrated noise floor, pre-registered predictions with two honest falsifications — this is a model of how to run an architecture ablation, and the registration record (Appendix B) is refreshingly candid about which predictions failed. If you read two sections, make them "Where the Gap Lives" (the decomposition, benchmark split, and pre-registered fineweb test) and the weight-spectrum mechanism section — together they turn a near-null loss comparison into an actual explanation of what feed-forward layers are for.