Compute-Optimal Allocation Between Attention and Feed-Forward Width

A transformer block has two sublayers, and their costs scale differently with sequence length. Attention’s flops grow with $n^2 d$ and its parameters with $d^2$; the feed-forward network’s flops grow with $n d^2$ and its parameters with $d^2$ times the expansion ratio. At short sequence length the MLP dominates the flops; at long sequence length attention catches up and then overtakes. Any budget split between them is therefore correct only for the sequence length it was chosen for, and the parameter split is a separate decision from the flop split. Working through both is what makes a width allocation argument concrete rather than a matter of taste.

Counting the block

Take $d$ as the model width, $n$ the sequence length, $L$ the layer count, $h$ the head count with $d_k = d/h$, and $r$ the MLP expansion ratio with hidden width $d_{ff} = r d$. Per token, the flops are:

Component	Multiply-accumulates per token
Q, K, V, O projections	$4 d^2$
Score matrix and value product	$2 n d$
MLP (two matrices)	$2 r d^2$

Per layer, per token. The attention projections and the MLP are independent of $n$; only the score computation carries $n$. At $n = 512$ and $r = 4$, the projections contribute $4d^2$ and the MLP $8d^2$, so the MLP is two-thirds of the per-token work; attention’s score term is only $1024d$, and it stays below the projections and MLP once $d$ exceeds about 85. At $n = 32{,}768$, the score term is $65{,}536 d$, which for $d = 4096$ is $2.7 \times 10^8$ against the MLP’s $1.3 \times 10^8$; attention is now twice the MLP’s cost.

The crossover is where $2 n d = 2 r d^2$, that is $n = r d$. For $r = 4$ and $d = 4096$, the crossover is at $n = 16{,}384$. Below that, widening the MLP is the cheaper way to add capacity; above it, attention dominates and widening the MLP no longer controls the cost.

The parameter counts tell a different story. Attention holds $4 d^2$ parameters across the four projections, output projection included; the MLP holds $2 r d^2$. At $r = 4$ the MLP has $8d^2$ against attention’s $4d^2$, so the MLP is two-thirds of block parameters regardless of $n$. Parameters do not depend on sequence length at all, which is why a model can be widened for quality and then used at any context, while its flop profile changes underneath.

Two budgets, not one

The distinction between parameter-optimal and compute-optimal is the practical core of the scaling literature. A parameter budget fixes $L \cdot d^2$ roughly, and the question is how to split it across depth and width. A compute budget fixes $L \cdot n \cdot d^2$, and the question is how to split it across $L$, $d$, and the amount of data. These give different answers, and conflating them produces architectures that are efficient at one scale and wasteful at another.

For a fixed parameter count, depth is cheaper in parameters than width: doubling $L$ doubles parameters linearly, while doubling $d$ quadruples them in the projections. That makes depth attractive per parameter, but depth costs latency, since each layer is a sequential dependency, and very deep models have gradient propagation problems that residual scaling and normalization must handle. Width is more parallel and more forgiving to optimize, at a quadratic parameter cost.

For a fixed compute budget, the reported fits for language models trained on large corpora put the split near equal scaling of parameters and data: the compute-optimal allocation increases parameters and training tokens in roughly the same proportion as compute grows, rather than holding data fixed and growing the model. The exact exponents depend on the fit and the corpus, and the specific constants are contested between published analyses. What holds across them is the direction: models at a given compute budget should be smaller and trained longer than the earlier practice of scaling parameters alone implied, and the parameter count alone is a poor predictor of quality without the token count beside it.

The architectural split is a separate axis from that trade. Given a total parameter budget, how much goes to attention and how much to the MLP is a within-block decision, and the flop analysis above says it should depend on $n$. A model trained at 2,048 tokens and deployed at 128,000 has an allocation chosen for the wrong regime.

Why the MLP holds most of the capacity

The MLP’s two matrices with a nonlinearity between them form a position-wise function from $d$ dimensions to $d$ dimensions through a $r d$-dimensional hidden layer. Its parameter count is what allows it to store associations: the first matrix selects directions in the input, the nonlinearity gates them, and the second matrix writes combinations back. Analyses of trained models consistently find that the MLP layers hold a large share of the model’s factual associations, in the sense that intervening on specific MLP activations changes specific factual outputs while intervening on attention activations changes more distributed behaviour.

Attention, by contrast, moves information between positions. Its parameters are projection matrices that decide what to read and what to write, and its function is routing rather than storage. That division of labour is why the parameter split is not symmetric: a model needs enough attention width to route well, and additional attention capacity beyond that competes with the MLP’s storage.

The division also explains the effect of sequence length on the split. Longer context increases the routing problem’s difficulty, with more positions to choose among and more precise addressing required, without increasing the storage problem’s. A model that must handle long context needs relatively more attention capacity, not because attention is more expressive but because the routing task grew.

def block_flops(d, n, r=4.0):
    attn_proj = 4 * d * d
    attn_scores = 2 * n * d
    mlp = 2 * r * d * d
    return {"attn_proj": attn_proj, "attn_scores": attn_scores, "mlp": mlp}

for n in (1024, 8192, 65536):
    f = block_flops(4096, n)
    total = sum(f.values())
    print(n, {k: round(v / total, 3) for k, v in f.items()})
# attn_scores share: 0.040 at 1k, 0.250 at 8k, 0.727 at 65k for d=4096, r=4

Depth, width, and the ratio between them

The layer count and the width are not interchangeable even at equal parameter count, and the reason is in the residual stream’s arithmetic. Each block adds its output to the stream, so the stream’s variance grows with depth unless normalization or residual scaling compensates. The standard treatment divides the sublayer output by $\sqrt{2L}$ before adding, which keeps the stream’s variance roughly constant at initialization when each sublayer’s output has unit variance. Without that scaling, a 100-layer model’s stream variance grows linearly with depth and the final layer’s inputs are far from the scale the readout expects.

Width affects a different quantity: the number of directions the stream can carry. A width of 512 has 512 dimensions of representational space, and every sublayer reads and writes within it. Attention heads partition that space; the MLP expands it transiently and projects back. A model that is deep but narrow has many sequential transformations of a small space, which is a good fit for compositional tasks and a poor fit for storing many independent associations. A wide model has fewer sequential steps and more parallel capacity per step, which fits recall-heavy tasks.

The aspect ratio $d/L$ in trained models spans a wide range across published configurations, and no single value is right. The constraint that does bind is that $d$ must be divisible by $h$, and that $d_k = d/h$ must stay in a range where the attention softmax is well conditioned: the variance argument in why attention divides by the square root of the key dimension ties the temperature to $d_k$, so a very small per-head width softens every head and a very large one sharpens it. That constraint propagates upward: choosing $h$ and $d_k$ fixes a minimum $d$.

Allocation decisions that survive a context change

The flop profile shifting with $n$ means an architecture tuned for pretraining length can be badly proportioned for its deployment length. Three strategies address this, with different costs.

The first is to choose the allocation for the longest context the model will serve. This is safe and wasteful at short context: a model with enough attention capacity for 128,000 tokens spends flops on score computation that a 2,000-token workload does not need, and the extra width in the attention projections is paid on every token regardless of length.

The second is to choose for the training length and accept degradation at longer context. This is what happens implicitly when a model is trained at one length and evaluated at another without adaptation. The degradation is not a cliff; it is a gradual loss of the model’s ability to route at distances it did not see, as discussed in sliding-window attention and effective context reach.

The third is to make the allocation context-dependent within one model, with sparse or windowed attention for most layers and full attention for a few. The per-layer cost is then a mixture, and the average cost tracks the workload rather than the worst case. The design question becomes how many full-attention layers are needed to preserve quality, which is an empirical question with no closed form, and the answer appears to be a small fraction of the total rather than a large one.

A related allocation question concerns the KV cache rather than the flops, and it can invert the flop-based conclusion. At long context the cache is the binding memory constraint, so an architecture that reduces cached heads, as described in grouped-query attention and KV cache sharing, buys batch capacity that no flop argument would suggest. The flop-optimal and memory-optimal allocations diverge at long context, and the deployment constraint decides which one binds.

The last consideration is that all of these counts are per-token multiply-accumulates under the assumption that the hardware achieves peak utilization. It does not: a matrix-vector product at decode time achieves a small fraction of peak because it is bandwidth-bound, and a matrix-matrix product at prefill achieves more. The flop counts predict relative cost between configurations reasonably well and absolute latency poorly, which is why the useful application of this arithmetic is comparing allocations rather than predicting throughput. The bandwidth argument for why the two phases differ is in KV cache size arithmetic, and the kernel-level reason a tiled attention implementation narrows the gap is in IO-aware attention tiling.