Sum Head Outputs Before Residual Add

Pre-projection head output

In Appendix A.1 of Identifying and Evaluating Inactive Heads in Pretrained LLMs (Sandoval-Segura et al., arXiv:2504.03889v4, PDF header date March 1, 2026; published as a conference paper at ICLR 2026), the authors separate two definitions of an attention-head output. They report that Vaswani (2017) defines the head output as the convex combination of value states, a dv-dimensional vector with dv = dmodel/h, where dmodel is the hidden size and h is the number of heads. That relation assumes equal head width and that h divides dmodel. It is a definition they attribute to Vaswani, not a measured width of any named checkpoint in the paper. Do not fill it in with a hidden size from a different model.

The same paper’s main text states the operational form. For an input of N tokens, each head produces attention weights A with rows in the probability simplex, then applies those weights to value vectors. The i-th head output headi ∈ ℝN×dv is, at every position, a convex combination of that head’s value vectors. Appendix A.1 writes the same product as Zi = AiVi, with Ai ∈ ℝN×N and Vi ∈ ℝN×dv. This Zi, before the output projection, is their pre-projection head output. They use it because WO mixes every head into the shared model space, and they did not want that mix inside the measurement.

The other definition, which they call circuit-based and attribute to Elhage et al. (2021), multiplies the dv-vector by a dv × dmodel slice of the output projection. The quotation and the partition below are taken from Sandoval-Segura et al.’s appendix, not from an independent reading of Elhage et al. outside that PDF. The appendix says Elhage et al. note that the module output can be written exactly as a sum of per-head terms. It also quotes them as saying this is “equivalent to running heads independently, multiplying each by its own output matrix, and adding them into the residual stream.” That independence claim, as quoted, is not the same statement as “sum the contributions, then add the sum to the residual.” The appendix’s own emphasis is the second statement.

One extracted line in that appendix writes Zi = AiVi ∈ ℝN×dk. The same appendix concatenates H = [Z1; …; Zh] ∈ ℝN×(h dv), and the main text sets headi ∈ ℝN×dv. The product AiVi has the column dimension of Vi, and the concatenation width h dv agrees with each block having width dv. Those two statements agree with each other. They agree with the extracted codomain ℝN×dk only if dk = dv. The paper elsewhere distinguishes dk-dimensional queries and keys from dv-dimensional values. Confirm the appendix line against the PDF before coding a key-width head output. This article does not resolve that line by re-typesetting the source.

Sum head contributions before the residual add

Sandoval-Segura et al. reconstruct the multi-head output as a concatenation followed by one projection, then as an equivalent sum. With Zi the pre-projection head output, they form

H = [Z1; …; Zh] ∈ ℝN×(h dv), then Y = H WO with WO ∈ ℝ(h dv)×dmodel.

They report the row-wise partition used by Elhage et al.:

WO(i) ∈ ℝdv×dmodel, and Y = Σi=1h Zi WO(i).

Under those shapes, each product ZiWO(i) is N × dmodel, and the sum has the same shape as Y. That identity is the appendix’s account of why concatenation-then-project and per-head project-then-sum are the same module output. It is not a license to treat Zi and Y as the same tensor.

The residual boundary is the next sentence in the appendix, and it is the authors’ reading, not the Elhage quotation itself: before the attention-module output is added to the residual stream, the head contributions are added up. That sum is the summation in their Equation (3). Only after that sum do they describe the residual add. They then point out a mechanical consequence: WO can implement cancellations, for example if two head contributions sum to zero, or other linear combinations of heads. A coordinate of the residual stream, after Y has been added, is therefore not a labeled head output. Cross-head cancellation can already have happened inside Y.

Their intervention choice follows from that split. Because they did not want WO inside the score, they define the head output as pre-projection Zi. They also state the algebraic link between the two definitions: whether the i-th head output is defined as Zi or as ZiWO(i), zeroing Zi makes both zero. You do not need a second write to clear the projected term of a head whose Zi is already the zero matrix. That fact does not make the score used to choose the head the same under both definitions. Appendix Table 3 compares selection scores, not the algebra of a zero that has already been applied.

Absence of a head label on one residual axis is not the same claim as “residual coordinates are an arbitrary basis.” Privileged bases in the transformer residual stream (Anthropic Interpretability, published March 16, 2023) states that mathematical theories of the Transformer suggest individual residual-stream coordinates should have no special significance, the basis directions should be in some sense arbitrary, and no more likely to encode information than random directions, but that this observation is false in practice. The authors provisionally conclude that the per-dimension normalizers in the Adam optimizer are to blame, and they report confidently ruling out layer normalization and finite-precision floating-point calculations as the source of the observed basis alignment. That basis result does not identify which head entered Y.

A module-level residual add is also not the pre-projection sum. Three-Phase Transformer (Abu Ayyash, arXiv:2604.14430v1, April 15, 2026) writes a decoder-only block that still uses the standard pre-norm residual for attention, h′ℓ = hℓ + Attnℓ(N1,ℓ(hℓ)), where N1,ℓ is that model’s PhaseAwareRMSNorm. A PhaseRotationLayer then overwrites the result, h″ℓ = PRℓ(h′ℓ), rather than adding to it. The paper states that this rotation is an orthogonal map, a direct sum of 2D Givens rotations within each phase block, so the pre-rotation state is recoverable by the inverse rotation. Attnℓ in that equation is the attention-module output added as one residual term. It is not Zi = AiVi, and it is not the scaled Bernoulli sum.

Prerequisites

Keep the two published procedures apart before choosing tensors, gates, or numbers. They are not one shared checkpoint.

For the pre-projection definition you need, per head, attention weights and value states. Their Appendix A.12 listings take attention weights of size (B, Nhead, S, S) and value states of size (B, Nhead, S, dv), and they return attention outputs of size (B, Nhead, S, dv). They state that grouped-query and multi-query attention still construct tensors of those sizes in their implementation, so the listing does not branch on that architecture choice. That is an implementation claim about those listings, not a claim that grouped-query attention has the same head algebra as full multi-head attention in every other respect.

For the Bernoulli procedure you need a per-layer set Π = {π1, …, πH} with each πi ∈ ℝ, a BinConcrete map to gi ∈ {0,1}, and the L0 sparsity loss they add to next-token loss. Shim et al. do not reproduce the BinConcrete density, temperature, or stretch. They point to Louizos, Welling, and Kingma, “Learning sparse neural networks through L0 regularization,” ICLR 2018. Those sampling parameters are unknown from the retrieved paper. Their head output also requires the persistent sets PKi and PVi. In the reported configuration, dh = d/H = 512/8 = 64, and each head has N = 2048 persistent vectors of that width. That 64 is their All-attention head dimension, not a dv for any model in the inactive-head study.

A scalar gi ∈ {0,1} commutes with the head’s linear map: multiplying hi by zero and multiplying hiWOi by zero remove the same term. That algebraic overlap does not merge the procedures. Shim et al. also multiply the whole sum by sg. Their hi includes persistent vectors. Sandoval-Segura et al. zero Zi with a score threshold, not with a BinConcrete sample. Copying either mask onto the other tensor invents a third intervention that neither paper reports.

Steps

The ordered steps below keep definition, residual-bound sum, and Bernoulli drop as separate operations. Steps 2–4 are the Sandoval-Segura et al. account. Steps 5–7 are the Shim et al. forward rule. Do not run step 7 on the tensor from step 2 and treat the result as either published method.

1. Name the tensor before you mask it. If the object is a convex combination of value vectors and has not been multiplied by WO, it is the pre-projection head output Zi. If it already has width dmodel, it is at least a projected contribution, or a sum of them. A residual-stream vector after the attention add is neither, until you have undone that add and the cross-head sum.
2. Compute the pre-projection head output. Form attention weights Ai and value states Vi, then Zi = AiVi. Use value width dv. In the usual fused implementation this is the matrix product of attention weights and value states, matching the product in their listings, before the output projection.
3. Add projected contributions before the residual write. Split WO into the head blocks WO(i) given in the appendix, compute ZiWO(i) for each head, and sum those terms to obtain Y. Equivalently, concatenate the Zi and apply the full WO. Add Y, not the individual Zi, to the residual stream. The retrieved pages do not write a fuller residual equation than that prose boundary.
4. Hard-zero a head only if that is the intended intervention. Sandoval-Segura et al. set the selected pre-projection head output to zero before concatenation and output projection. Because the later product is linear, that zero also removes ZiWO(i) from Y. Their published mask is not Bernoulli. For the layer-normalized average head-output norm, they average the row ℓ2 norms of the head output, divide by the average of those scores among heads in the layer, and mark the head inactive when the ratio is below a threshold τ. They give τ = 0.1 only as an example of “less than 10% of the layer average,” not as the experimental threshold. Experimental thresholds come from score quantiles. The paper points to implementation code at github.com/psandovalsegura/inactive-heads; this article does not review that repository.
5. Attach a Bernoulli gate only on the All-attention projected sum. For each layer, take trainable πi and map them through BinConcrete to gi ∈ {0,1}. The retrieved text does not specify BinConcrete’s numeric sampler. Their ungated sum is O = Σi=1H WOi hi. The gated equation is written with the factors in the other order: O = sg Σi=1H gi hi WOi. Treat that order flip as a notation inconsistency to check in the PDF, not as two different sums. hi here already includes persistent keys and values. One gate multiplies the whole head, not one token.
6. Rescale the gated sum, and clip the all-closed case. The paper sets sg = H / Σi gi and clips sg to the maximum value H so that an all-zero sample does not divide by zero. If every gi is 0, the sum of gated terms is still 0, so the clipped scale does not restore a head. They also say the scale compensates for the masked portion and is inversely proportional to (1 − sparsity). Reading sparsity as the fraction of closed gates in that sample gives 1 − sparsity = (Σi gi)/H, hence 1/(1 − sparsity) = H/(Σi gi), which matches sg when the denominator is positive. That reading is an algebraic gloss of their two statements, not a dump of their L0 code. At inference they say sg can be absorbed into WO, so the scale need not be a separate inference kernel. During training it remains explicit.
7. If you need their trained prune rather than one forward sample, add the schedule they say stabilizes it. Start from a converged All-attention model trained without gates. Initialize each πi to 2, which they say leaves the gate open with about 88% probability; zero initialization, they say, opens it with 50% probability. Linearly warm up the sparsity-loss weight λ from 0. In their pruning run that warm-up is the first 4,000 iterations. After 16,000 iterations they stop training the gates so later steps are not still driven by L0 randomness. When a head is decided pruned (gi = 0), they remove {WQi, WKi, WVi, PKi, PVi, WOi}. Those training choices are part of Shim et al.’s prune. They are not part of the definition of Zi or of Y.

Algebraic example, not a reported training step: take their configuration H = 8, set exactly one gate to 0 and the other seven to 1, and assume the clip is not the all-zero branch. Then Σi gi = 7 and sg = 8/7. Rounded half-up to three decimal places, 8/7 = 1.1428… becomes 1.143. This uses only Equation (3) and their stated head count. It is not a measured gate configuration from a WikiText-103 step.

Checks and failure cases

Shape checks

After step 2, Zi should match the value-state width. The appendix concatenation ℝN×(h dv) is the internal check that each block is dv wide. After step 3, each ZiWO(i) and the sum Y should be N × dmodel. If a hook already returns dmodel vectors per head, you are looking at projected contributions, not at the pre-projection definition used for their scores.

For the gate, check that gi is one scalar per head per layer, that it multiplies the All-attention head output which includes persistent vectors, and that sg is applied to the sum. If Σi gi = 0, confirm the implementation clips instead of dividing. Shim et al. write the ungated sum as Σi WOi hi and the gated sum as Σi gi hi WOi. A reproducer should resolve that factor order against the PDF rather than pick one silently.

Merging the checkpoints

Zeroing a residual coordinate does not implement either procedure. By the appendix’s account, cross-head cancellation is possible in Y before the residual add, so one residual axis need not correspond to one head. The Elhage sentence they quote, about adding heads independently into the residual stream, and their own sentence, about summing contributions before that add, should stay quoted as two claims. The retrieved bundle does not contain the Elhage et al. article beyond that citation and quotation.

Using the inactive-head mask as if it were gi fails the Bernoulli requirement. The mask in their listings is a comparison of a head score with a threshold. Average head-output norm, including the layer-normalized form, is computed from the head output itself. The authors say that this best score therefore cannot save the compute of the head it might later skip. A Bernoulli sample drawn without BinConcrete, without sg, or on Zi from a standard residual block, is not the Shim et al. forward equation either.

Omitting the scale is a reported failure mode, not a cosmetic one. Shim et al. state that without output scaling they observe a consistent increase in training loss. On the WikiText-103 test set at λ = 0.02, their ablation table reports these perplexity increases relative to the run that uses all three techniques: removing λ warm-up, +0.12; removing the large positive gate initialization, +0.61; removing attention-head output scaling, +1.25; removing all three, +1.48. Those deltas are for that All-attention prune. They are not MMLU accuracy changes.

Leaving the gates under L0 pressure is another stated failure. They report that if gating parameters keep training, sparsity continues to increase and fine-tuning becomes harder. Freezing the gates after 16,000 pruning iterations is how they remove that randomness. A single inference-time Bernoulli draw, with no freeze and no decision to delete {WQi, WKi, WVi, PKi, PVi, WOi}, is only the forward multiply. It is not their full prune.

What the published numbers do not transfer

Appendix Table 3 in Sandoval-Segura et al. is a check on score definition, not a Bernoulli result. The caption says they measure the percent of model heads that can be removed while maintaining accuracy within 1% of baseline, higher meaning more heads removed, on MMLU, for three models. The main intervention section specifies MMLU in the 5-shot setting through lm-evaluation-harness. The Table 3 caption does not restate the shot count. Figure 7’s caption says that, for those three models on MMLU, the normalized circuit-based head-output definition is comparable to the original definition, and that dashed lines denote layer-normalized scores. The figure’s curve values are not in the retrieved text. The table is:

The authors interpret the circuit-based and original AHON (LN) columns as comparable. The pairs are not identical: 10.97 versus 9.88, 7.29 versus 5.95, and 7.54 versus 9.04. Comparable removable fraction also does not mean the two scores selected the same heads. Elsewhere they show that layer normalization can reorder heads relative to an unnormalized score. Table 3 is not their full 14-model intervention study, and it is not evidence about a Bernoulli gate.

Shim et al.’s corresponding outcome, under a different intervention, is a language-modeling perplexity. On WikiText-103 test, with embedding and output-projection parameters excluded (they occupy 19.6M parameters), the ungated All-attention baseline has 54.6M parameters and perplexity 23.24. At λ = 0.02, they report 43.8% sparsity, 30.7M parameters, and perplexity 24.85. That run depends on All-attention, their depth and width, the three stabilization techniques, and the rest of the training recipe in the paper. It does not measure how many heads can be zeroed in Llama-3.1-8B-Instruct, Qwen2.5-3B, or Qwen2.5-7B.

Two further limits are stated by Sandoval-Segura et al. and apply only to their hard zero. Their accuracy results are log-likelihood evaluations of multiple-choice benchmarks, especially MMLU. They warn that generation is a different setting: zeroing one incorrectly identified head can ruin the rest of the generated sequence. They also assume, for that study, that inactive heads are heads that can be zeroed before concatenation and output projection. That is an operational assumption of the intervention study, not a theorem about Bernoulli gates.

Unknowns that this article does not fill in: BinConcrete’s temperature and stretch; the residual-add program line for All-attention, which is not in the retrieved equations; whether any later revision of either PDF changes Appendix A.1 or Equation (3); and any reimplementation. The equations here are those printed in the retrieved arXiv:2504.03889v4 and arXiv:2110.03252v1 files. Head-dropout papers cited by Shim et al. use related output scaling for regularization during training. Their citations do not make dropout the same procedure as this L0 gate.

Editorial Team

Editorial Team

Publisher description: The J-Space Editorial Team curates and synthesizes cutting-edge research in mechanistic interpretability, AI alignment, and cognitive science. With a commitment to precision, empirical rigor, and clear exposition, the team distills complex technical developments into accessible yet authoritative analyses for researchers, engineers, and policymakers working at the frontier of artificial intelligence.