What a Transformer Loses Without Feed-Forward Layers

A transformer block has two parts. Attention lets a token read from other tokens. The feed-forward network, two matrices with a nonlinearity between them, transforms each token on its own. The feed-forward network holds two thirds of a transformer's non-embedding parameters, and interpretability work describes it as the model's memory, where facts from training are stored as key-value pairs that can be located and edited.

The opposite experiment, deleting it and measuring what is lost with the confounds controlled, had not been run. We ran it, and Needle 2 and Needle 3 are built on the result. The full study is A Controlled Study of Attention-Only Transformers.

The block

A Simple Attention Network, SAN, is a decoder transformer with the feed-forward sublayer removed from every block. Everything else is standard: pre-norm, grouped-query attention, rotary positions, a scalar gate on the residual write, tied embeddings. The control is the same code with a SwiGLU feed-forward network after attention.

What attention can do

For a token with query $q$, and context tokens with keys $k_i$ and values $v_i$, one attention head computes

$w_i = \mathrm{softmax}_i\big(q \cdot k_i\big), \qquad \text{output} = \sum_i w_i\, v_i.$

The softmax makes every weight positive and makes the weights sum to one, so the output is a weighted average of the value vectors. A weighted average of points lies inside the shape the points span, their convex hull, and never outside it.

For a fixed attention pattern the layer is linear in its inputs. The feed-forward layer was the only place where a token could be transformed by a nonlinear function of itself into something new. Without it, a layer selects and transports content that is already in the context. It cannot build a representation the context does not support. That is what context-grounded means here, and every number below is the price of it.

Three ways to ask the question

Removing the feed-forward layer changes three things at once: the parameter count, the compute per token, and the number of nonlinear steps between input and output. Fixing one axis leaves the comparison confounded on the other two, so we ran it three times.

Every arm got its own learning-rate sweep, extended whenever the edge value won, so no architecture ran at another's preferred rate. The best rate differs by 2x between the arms, and a shared rate would have biased the whole comparison. A pair of models that differ only by a mathematically equivalent parameterisation gave the noise floor of the pipeline: 0.0015 nats. All models were pretrained on SYNTH, a reasoning-dense corpus where every document has a query, a reasoning trace and an answer, for up to 105B tokens. Eight predictions were registered before the measurements.

The loss is measured in nats: the average surprise at the true next token. A gap of 0.01 nats means one model gives the true token about 1% more probability than the other.

What the feed-forward layer costs

$\Delta_{\text{depth}} = 0.47, \qquad \Delta_{\text{FLOPs}} = 0.26, \qquad \Delta_{\text{params}} = 0.0055 \text{ nats.}$

The three gaps order by how much of the parameter budget the control lets attention take back. Deleting the layers in place removes 72% of the model and costs 0.47 nats. At matched compute the standard transformer still leads by 0.26, because attention spends part of its budget on the parameter-free quadratic term and so carries fewer weights at equal cost. At matched parameters, with the freed budget put back into attention depth, the gap is 0.0055 and 0.0054 nats on two seed pairs, 0.27% of the loss.

The feed-forward layer's parameters matter. Its functional form, on this distribution, largely does not.

Two more axes agree. The matched-parameter gap shrinks with training: 0.046 nats at 5B tokens, 0.019 at 30B, 0.0055 at 105B, each budget trained and tuned on its own. It does not grow with scale: at a fixed 31.5B-token budget across five matched size pairs, the attention-only model wins at the smallest size, and the gap then stays near 0.02 nats from 16M to 57M non-embedding parameters. Limiting both arms to 2M, 8M or 32M unique documents, up to 18 epochs, costs at most 0.010 nats and shows no interaction with architecture.

Where the gap lives

Every SYNTH document has three regions, a query, a reasoning trace and an answer, separated by marker tokens, so the loss splits exactly by region. The token-weighted sum of the region gaps reproduces the total within 2%.

The deficit sits on query tokens. A query token is at the start of a document, with almost no context to route from, so only knowledge stored in weights can help. At 31B tokens the per-token gap there is five times the overall gap, on tokens that carry 8% of the loss. By 105B the SAN leads on traces and on every answer region, including memorisation exercises, and trails only on queries. Across the size ladder the query deficit is the one invariant, positive at every size and budget, while every other region changes sign.

Benchmarks split the same way. Lambada asks for one specific content word, and for a model trained on SYNTH that is a recall task from outside its distribution; it favours the FFN model at every budget. Sciq puts the answer in a support passage; it favours the SAN, and the margin grows with training in the direction registered in advance, 0.725 to 0.742 for the SAN while the FFN model slides from 0.702 to 0.661.

This account predicts what happens on text we had not trained on. Natural web text is mostly low-context prediction, so we registered, before launching the run, that a matched-parameter pair trained on fineweb-edu would show a gap between 0.02 and 0.05 nats. It measured 0.040. The same pair reverses on lambada, 0.203 against 0.181 in the SAN's favour: trained on matching text, the passage is enough to infer the final word and lambada becomes a routing task. Whether a task is storage or routing depends on the match between the task and the training distribution, not on the task alone.

The mechanism, in the weights

Stable rank measures how many independent directions a matrix uses. Tracked over training, it gives the same picture in every architecture, size and budget. The routing matrices, $Q$ and $K$, settle within the first quarter of training and do not move again; the optimiser sets the level, with Muon holding them two to three times flatter than AdamW, but the schedule is universal. The matrices that write content into the residual stream keep gaining rank for as long as training continues: the down-projection in FFN models, and in attention-only models the output projection, which inherits the role because it is the SAN's only write path.

Selection structure is learned early and frozen. Content capacity accrues for as long as training continues, in whichever matrices can hold it. Representation rank stays high throughout, with a minimum layer rank of 173 of 512 at 20 layers, so the rank collapse known for pure attention never appears once residuals and normalisation are in place.

What keeps a deep attention-only stack trainable

Component ablations at 20 layers give three facts.

1. QK-normalisation is load-bearing. Removing it diverges training at the tuned learning rate, the only divergence in the study and the one result none of the predictions anticipated.
2. Scalar residual gates are performance-neutral at every depth from 20 to 48 layers and in both architectures. Their trajectories were still useful: they showed the FFN model pruning itself toward attention-only form under learning-rate stress.
3. Post-attention sandwich normalisation is the only variant that beats the baseline, by 0.009 nats.

Depth at matched parameters is U-shaped with a 20-layer optimum, and 48-layer attention-only stacks train without incident. Every instability in the study happened in the FFN arm.

When to use one

An attention-only model is the better choice when parameters are the limit and FLOPs are not. At matched parameters a SAN spends about twice the FLOPs per token at 2048 context, and at matched FLOPs the model with feed-forward layers is ahead. On a device, memory decides how large the model can be, and a token spends more time waiting on memory than on arithmetic, so the extra FLOPs are cheap and the saved parameters are not. Attention-only models also do best on data where the answer can be assembled from the context.

Needle's tasks are of that kind. In tool calling and structured extraction the function names, the argument values and the schema are all in the prompt, and the model's job is to select them and place them into a call. Needle 2 was a Simple Attention Network with an engram memory for the facts that are not in the context. Needle 3 keeps that stack, puts a Hadamard MLP in the feed-forward slot at 25K parameters per layer, and adds a ladder so that every depth from 2 to 20 layers is a deployable model. The 0.27% of loss measured above is the price of the design, and the engram is how we pay back the part of it that sits on query tokens.

The limits of the study:

• all models are at or below 87M parameters and 105B tokens;
• the data is one reasoning-dense corpus plus one knowledge-dense control pair;
• MMLU-class benchmarks are at chance at these sizes, so the comparison rests on loss;
• the storage account predicts a wider gap on storage-heavy text at larger scale, and that experiment has not been run.

Code, training curves, evaluation reports and every checkpoint are published with the paper.

References

1. Ndubuaku et al., A Controlled Study of Attention-Only Transformers, 2026.
2. Vaswani et al., Attention Is All You Need, 2017.
3. Geva et al., Transformer Feed-Forward Layers Are Key-Value Memories, 2020.
4. Meng et al., Locating and Editing Factual Associations in GPT, 2022.
5. Dong et al., Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth, 2021.
6. Shazeer, GLU Variants Improve Transformer, 2020.
7. Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023.
8. Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021.
9. Henry et al., Query-Key Normalization for Transformers, 2020.
10. Ding et al., CogView: Mastering Text-to-Image Generation via Transformers, 2021.
11. Jordan et al., Muon: An optimizer for hidden layers in neural networks, 2024.
12. PleIAs, SYNTH, 2025.
13. Penedo et al., The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale, 2024.
14. Paperno et al., The LAMBADA dataset: Word prediction requiring a broad discourse context, 2016.
15. Welbl et al., Crowdsourcing Multiple Choice Science Questions, 2017.