🤖 machine learning

A Controlled Study of Attention-Only Transformers

This study demonstrates that feed-forward networks are not strictly necessary for transformer performance, as attention-only models (SANs) matched for parameters, compute, and depth achieve nearly identical loss to standard transformers, with the remaining minor gap attributable solely to differences in parametric recall rather than architectural limitations.

Original authors: Henry Ndubuaku, Karen Mosoyan, Jakub Mroz, Noah Cylich, Satyajit Kumar, Parkirat Sandhu, Roman Shemet, Justin H Lee

Published 2026-07-22

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Henry Ndubuaku, Karen Mosoyan, Jakub Mroz, Noah Cylich, Satyajit Kumar, Parkirat Sandhu, Roman Shemet, Justin H Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build the ultimate brain for a robot. For years, scientists have been stacking two types of "neurons" on top of each other to make these robots smarter. The first type is like a messenger that looks at what's happening right now and decides where to pay attention. The second type is like a notebook that stores facts and rules learned from the past. The robot's brain is built with a repeating pattern: a messenger, then a notebook, then a messenger, then a notebook.

In the world of artificial intelligence, these "messengers" are called attention mechanisms, and the "notebooks" are called feed-forward networks (FFNs). For a long time, everyone assumed the notebook was essential. In fact, in modern robot brains, these notebooks take up about two-thirds of the total space (parameters). The big question was: Is the notebook actually necessary? Or could the messenger do all the work if we just gave it enough space? This paper is a giant, controlled experiment to find out if we can build a super-smart robot using only messengers, with no notebooks at all.

The Great Brain Surgery Experiment

The researchers from Cactus Compute, Inc. decided to perform a very precise "brain surgery" on a standard AI model. They took a standard model and, layer by layer, ripped out the "notebook" (the feed-forward network) entirely, leaving only the "messenger" (the attention mechanism). They called this new, stripped-down model a Simple Attention Network (SAN).

But here's the tricky part: you can't just delete a huge chunk of a brain and expect it to work the same way. If you remove the notebook, you suddenly have a lot of empty space. So, the researchers ran three different versions of the experiment to see what really matters:

Same Depth: They kept the number of layers the same but deleted the notebooks. (This made the SAN much smaller and weaker).
Same Compute: They adjusted the size so both models used the exact same amount of electricity (FLOPs) to think.
Same Size: They took the empty space left by deleting the notebooks and used it to add more layers of messengers. This is the fairest test: "If we give the messenger the same total brainpower as the notebook, can it do the job?"

The Shocking Result: The Notebook Isn't Magic

The results were surprising. When they just deleted the notebooks without giving the messengers more space, the robot got significantly dumber. But when they took that "deleted" brainpower and used it to build deeper messenger stacks, the gap almost completely vanished.

In the fairest test (matching the total size), the standard model with notebooks and the new "messenger-only" model were nearly identical. The difference was a tiny 0.006 nats (a unit of measurement for how wrong a prediction is). To put that in perspective, that's a difference of just 0.27% in their performance. It's so small that if you ran the experiment with different random seeds, the results were reproducible to one part in ten thousand.

The paper concludes that for the type of data they used (a reasoning-heavy synthetic corpus), the "notebook" isn't a magical, irreplaceable component. It's just a way to store information. If you give the "messenger" enough depth, it can store that information itself.

Where the Tiny Gap Hides

So, if they are almost the same, why isn't the difference exactly zero? The researchers dug deep to find out where the "messenger-only" model struggled. They found the problem wasn't everywhere; it was very specific.

The messenger-only model was slightly worse at answering questions where the context (the story or text provided) gave it absolutely nothing to work with.

The "Context-Grounded" Win: When the answer was hidden somewhere in the text provided (like finding a fact in a Wikipedia article), the messenger-only model was actually better or equal.
The "Memory" Loss: The tiny gap only appeared when the model had to pull a fact out of thin air—pure memory recall where the text offered no clues.

Think of it like this: The messenger is great at saying, "Hey, I see a clue in the text, let's look there!" But the notebook was slightly better at saying, "I remember this fact from my training, even though there's no clue here." The researchers found that the messenger-only model could still memorize things, but it had to work a bit harder to store those facts in its "attention" layers instead of a dedicated "notebook" layer.

How the Brain Actually Works

The paper also looked under the hood to see how these models learned. They discovered a fascinating rhythm in how the brain grows:

The Routing Crystallizes Early: The part of the brain that decides where to look (the routing) gets set very quickly, in the first quarter of training, and then stays frozen.
The Content Grows Slowly: The part that actually stores the information keeps growing and getting more complex throughout the entire training process.

When they removed the notebook, the "storage" job didn't disappear; it just moved. The messenger's output layer took over the job of storing facts. It's like if you removed a filing cabinet from an office; the employees (the messengers) would just start carrying the files in their pockets instead. It works, but it's a slightly different way of organizing.

The Secret Sauce: Normalization

One of the most practical findings was about keeping these deep "messenger-only" brains from falling apart. The researchers found that a specific trick called QK-normalization was absolutely critical. Without it, the deep messenger-only models would just crash and stop learning. It turns out that the "notebook" wasn't the thing keeping the brain stable; it was this specific normalization technique.

The Bottom Line

This paper doesn't say that notebooks are useless. It says that for the specific type of reasoning-heavy data they tested, the "notebook" isn't a fundamental requirement for intelligence. If you have a limited budget for brain size, you can trade the notebook for more layers of messengers and get almost the same result.

The only time the notebook really shines is when the model needs to recall facts that have nothing to do with the current text. But even then, the difference is tiny. The main takeaway is that the "messenger" is incredibly powerful and flexible, capable of doing the job of the "notebook" if you just give it enough depth. The gap between the two is so small that it's almost negligible, proving that the architecture of the transformer is more flexible than we thought.

Technical Summary: A Controlled Study of Attention-Only Transformers

Problem Statement
The feed-forward network (FFN) constitutes roughly two-thirds of a transformer's non-embedding parameters and is widely interpreted as the model's parametric memory (Geva et al. 2021, 2022; Meng et al. 2022; Dai et al. 2022). However, the necessity of the FFN has never been rigorously tested under controlled conditions. Previous attempts to remove FFNs either folded their capacity back into attention via persistent memory vectors (Sukhbaatar et al. 2019), leaving the necessity question open, or analyzed theoretical settings without residual connections that are not used in practice (Dong, Cordonnier, and Loukas 2021). The core experimental challenge is that removing the FFN simultaneously perturbs three critical variables: parameter count, compute per token (FLOPs), and the number of nonlinear composition steps (depth). A single-axis control is therefore insufficient to isolate the true cost of the FFN.

Methodology
The authors conduct a "necessity test" by pretraining Simple Attention Networks (SANs)—transformers with FFNs removed—against standard transformers. To control for confounding variables, the study employs three distinct matching strategies (Table 1):

Iso-depth: Matching the number of layers (20 layers). This removes 72% of parameters from the standard transformer.
Iso-param: Matching the total non-embedding parameter count. The freed budget is reallocated into attention depth (20 SAN layers vs. 4 standard layers).
Iso-FLOP: Matching training compute (GFLOPs per token). Since attention incurs a parameter-free quadratic cost in sequence length, the standard transformer carries more parameters in this regime.

Experimental Controls and Fairness:

Hyperparameter Fairness: Each architecture arm undergoes its own learning-rate sweep (0.5× to 16×) at a 5B token checkpoint to ensure no arm competes at another's preferred hyperparameters.
Optimizer Control: Training uses Muon against an AdamW control, with a calibrated noise floor of 0.0015 nats established via optimizer-equivalent controls.
Data Regime: Models are pretrained on SYNTH, a reasoning-dense synthetic corpus, up to 105B tokens. A separate pre-registered test is conducted on Fineweb-edu (knowledge-dense).
Architecture: SANs use pre-norm attention with Grouped Query Attention (GQA), RoPE, and specific normalization strategies (ZCN, QK-normalization).

Key Results
The study decomposes the performance gap between SANs and standard transformers into a monotone sequence based on the matching strategy:

Cost of Removal (Iso-depth): Deleting FFNs in place at matched depth results in a significant loss of 0.47 nats in validation loss.
Compute Efficiency (Iso-FLOP): At matched training FLOPs, the standard transformer still leads by 0.26 nats, as the attention-only model must spend more compute on the parameter-free quadratic term, leaving fewer parameters for capacity.
Parameter Efficiency (Iso-param): When the freed parameter budget is reallocated into attention depth, the gap collapses to a small, highly reproducible +0.0055 to +0.0054 nats (0.27% of total loss) on clean seed pairs. This difference is stable across 5B, 30B, and 105B token budgets and remains near 0.02 nats across a 29× range of non-embedding parameter sizes.

Localization of the Residual Gap
The remaining performance deficit is not diffuse; it is strictly localized:

Token Regions: The deficit concentrates entirely on low-context query tokens (where the context offers little to route). On reasoning traces and answer regions (including memorization exercises), the SAN often outperforms or matches the FFN model. By 105B tokens, the SAN leads on all answer regions, with the deficit surviving only on query tokens.
Task Types: On benchmarks requiring out-of-distribution recall (e.g., Lambada), the FFN model consistently wins. On tasks where the answer is provided in the context (e.g., SciQ), the SAN wins, with the margin growing as training progresses.
Distribution Dependence: A pre-registered test on Fineweb-edu (knowledge-dense text) measured a gap of 0.040 nats, confirming the prediction that the FFN's advantage is largest on low-context, knowledge-dense text. Conversely, on Lambada trained on Fineweb-edu, the preference reversed, indicating that "storage vs. routing" is a property of the task-distribution match, not an intrinsic task property.

Mechanism: Weight-Spectrum Dynamics
The paper identifies the mechanism behind these results through weight spectrum analysis:

Routing vs. Content: Routing matrices (Query/Key) spectrally crystallize (freeze) within the first quarter of training. Content matrices accumulate stable rank throughout the stable training phase.
Relocation of Storage: In standard transformers, content accumulation occurs in the FFN down-projection. In SANs, this accumulation relocates to the attention output projection ( $W_o$ ), which becomes the sole write path into the residual stream.
Stability: The study identifies QK-normalization as the critical component for training deep (48-layer) attention-only stacks. Removing it causes immediate divergence. Scalar residual gates were found to be performance-neutral but useful for diagnostics.

Significance and Claims
The paper claims to provide the first controlled necessity test of the transformer FFN. Its primary conclusions are:

The FFN is not a distinct computational primitive: Its removal does not break the model; rather, it represents a specific allocation of parameter capacity to a write path.
Functional Form vs. Capacity: On reasoning-dense data, the functional form of the FFN largely does not matter; what matters is the parameter capacity. Reallocating this capacity into attention depth recovers nearly all performance.
Specific Deficit: The only remaining cost is a small, predictable deficit on tokens where context provides no information for routing, forcing reliance on parametric memory.
Robustness: SANs exhibit greater training stability than standard transformers under learning-rate stress, with no divergence events observed in SAN runs compared to instability in FFN runs.

The authors conclude that within the tested regime, "attention does the rest," provided the parameter budget is matched and the architecture is stabilized with QK-normalization. The study reframes the FFN not as an essential structural component but as a specific mechanism for storing knowledge that can be effectively replaced by deeper attention layers, with the trade-off being a localized performance hit on low-context, knowledge-heavy prediction tasks.