Parallel Attention and FFN

The default transformer block computes attention, adds it to the residual stream, then computes the FFN on the result of that addition (anatomy-of-a-transformer-block). This makes the FFN's input strictly dependent on attention's output: nothing in the FFN sublayer can start until attention has finished and been added back. A handful of large models, GPT-J, GPT-NeoX, and PaLM among them, made a different choice: run both sublayers from the same normalised input, in parallel, and sum both results into the residual stream in one step.

The reformulation

The standard, sequential block is two updates:

x = x + Attention(LN1(x))
x = x + FFN(LN2(x))

The parallel formulation collapses this to one update, and often one shared normalisation:

x = x + Attention(LN(x)) + FFN(LN(x))

Both sublayers read the identical normalised tensor. Neither sees the other's output. Both are added to the residual stream in the same step rather than sequentially. Wang and Komatsuzaki, 2021 used this formulation for GPT-J; Chowdhery et al., 2022 adopted it at far larger scale for PaLM (540B parameters), and it also appears in GPT-NeoX-20B (Black et al., 2022).

Why it helps on real hardware

The sequential form has a hard synchronisation point: the accelerator cannot begin the FFN's matmuls until attention's output has been fully computed and added back. The parallel form removes that dependency, attention and FFN read the same tensor and can be scheduled to overlap, and their linear projections can even be fused into larger combined matmuls, which use accelerator hardware more efficiently than two separate smaller ones. PaLM's authors report this yields a meaningful training throughput improvement, on the order of 15% at their scale, purely from removing that one dependency and halving the number of LayerNorm reads per block (one shared normalisation instead of two).

What it gives up

The restriction is real, not just cosmetic. In the sequential form, the FFN can build on what attention just wrote to the residual stream in the same block, refining a just-mixed representation before the next block sees it. In the parallel form, both sublayers see the same, unmixed input and cannot see each other's contribution until the next block. The model loses one specific kind of within-layer composition. PaLM's own ablations found this cost measurable at smaller model scales, a small quality gap versus the sequential form, that shrank to roughly negligible as scale increased, which is presumably why the parallel formulation surfaces mainly in large models rather than small ones: extra depth elsewhere in a big stack apparently compensates for what any single parallel block gives up.

When it falls down

The saving is about scheduling, not FLOPs. The attention and FFN matrices are typically kept the same size as in the sequential form, so total compute per layer is similar; the win comes from reduced synchronisation and better hardware utilisation, not from doing less work.
Quality cost is scale-dependent. The formulation that PaLM validated at 8B+ parameters is not automatically lossless at smaller scale; ablations at smaller sizes showed a real, if modest, quality gap against the sequential baseline.
A single shared normalisation raises the stakes on that normalisation. Because both sublayers inherit the exact same normalised input, any miscalibration in that one LayerNorm propagates to both branches at once, rather than being isolated to one.
It composes with other design choices, not in isolation. PaLM paired parallel formulation with a specific depth-width aspect ratio (see depth-width-aspect-ratio) chosen partly for its TPU pod topology; attributing PaLM's throughput or quality to the parallel block alone, separate from those other choices, overstates what any single ablation actually isolated.

The reformulation

Why it helps on real hardware

What it gives up

When it falls down

Further reading