Why do graph attention networks concatenate multiple heads in hidden layers but average them at the output?

## What a head is One attention head is a complete copy of the layer's machinery: its own linear transform `W_k`, its own scoring vector `a_k`, and therefore its own softmax over each node's neighbourhood. `K` heads run on the same graph and the same inputs, and produce `K` different coefficient distributions and `K` different output vectors per node. Nothing couples them during the forward pass; they interact only through how their outputs are combined and through the gradient that flows back. ## Hidden layers — concatenate ``` h_i' = concat over k=1..K of sigma(sum_j alpha_ij^k W_k z_j) ``` With `K = 8` heads of `F' = 16` features each, a node leaves the layer as a 128-dimensional vector. Two things are gained. **Several weightings survive.** One head may concentrate on the single highest-scoring neighbour, another may spread its mass over a broader set. Concatenation keeps both intact for the next layer to use, so the network is not forced to pick one story about the neighbourhood at this depth. **Training is steadier.** A single head's coefficients are a high-variance quantity early in training — the softmax is sharp, the scorer is barely trained, and one unlucky score can dominate a node's update. With several heads, no single coefficient distribution controls the whole representation, and the empirical effect is a less brittle optimisation. The price is width. Concatenating `K` heads multiplies the next layer's input dimension by `K`, so the next layer's transform matrix grows by the same factor, and so does the activation memory carried into it. Going from 1 head to 8 does not cost 8 times the current layer — it costs 8 times the *next* layer's parameters too. ## The output layer — average ``` h_i' = sigma_out((1/K) * sum over k=1..K of sum_j alpha_ij^k W_k z_j) ``` At the prediction layer the output width is not a free hyperparameter: a 7-class node-classification task needs exactly 7 logits per node. Concatenating 8 heads that each emit 7 numbers would give 56, which is not a class distribution. Three options exist in principle — shrink each head to `classes/K` features, add a projection after concatenation, or average — and the standard convention is to average, because it is parameter-free and keeps every head predicting in the same output space where the average is meaningful. Note the ordering: the heads are averaged **before** the final nonlinearity, not after. Averaging `K` softmax distributions and taking a softmax of the averaged logits are different operations, and the convention is the latter. Averaging is also a variance reduction. Each head's prediction carries noise from its own randomly initialised scorer; a mean over `K` heads damps the part of that noise that is uncorrelated across heads. That is a real but modest benefit — the heads share inputs and gradients, so their errors are far from independent. ## Do the heads actually learn different things? Nothing enforces it. The heads differ only through random initialisation, and it is entirely possible for several of them to converge onto near-identical coefficient patterns, at which point the extra width is buying redundancy rather than diversity. This is worth checking directly: compare the coefficient distributions across heads on the same nodes, or look at the pairwise correlation of their output vectors. If they are near-duplicates, fewer heads with a wider `F'` each is usually the better trade — same width, fewer softmaxes, fewer per-edge coefficients to store. ## Memory, concretely Per layer, per head, the model stores one coefficient per edge. Eight heads on a graph with ten million edges is eighty million coefficients per layer during the backward pass. On a graph with heavy-tailed degrees this is dominated by a handful of hub nodes. Head count is therefore not a free knob: it multiplies both the width flowing forward and the per-edge state held for the gradient. ## The answer an interviewer is listening for The crisp version is that concatenation and averaging are answering two different constraints. In the middle of the network, width is a resource and you want to keep the heads distinct. At the end of the network, width is fixed by the task, so the heads must be collapsed, and averaging is the parameter-free way to collapse them while staying in the output space. A candidate who says only "averaging is more stable" has half the answer; the shape constraint at the output layer is the part that makes the convention non-negotiable.