Looking at one Transformer attention matrix can answer a local question: which positions a token attends to in that layer. It does not directly tell you how much an input token can influence a representation several layers later.

The reason is mixing. After one layer, a token representation already contains information gathered from other positions. The next layer attends to those mixed representations, not to untouched input tokens. Residual connections add another path that carries each representation forward. Reading only the final layer therefore skips the paths through earlier layers.

Attention rollout is a simple post-hoc method for following those paths. It combines attention matrices across layers, with an adjustment for residual connections, to estimate how representations at a later layer connect back to input positions. This article builds the method from a small example, shows how to implement it correctly, and explains why rollout is useful for inspection but should not be treated as proof of causal importance.

Start with a path, not a heatmap

Imagine a three-token sequence:

[A, B, C]

Suppose a token at layer 1 attends mostly to B. At layer 2, another token attends strongly to that layer-1 representation. The layer-2 token can therefore receive information originating from B even if its layer-2 attention matrix does not point directly to the original B position in the way a reader might expect.

The useful mental model is a layered graph:

layer 2:       x2_A   x2_B   x2_C
                 \     |     /
layer 1:       x1_A   x1_B   x1_C
                 \     |     /
input:            A     B     C

Each attention edge connects a representation to representations in the previous layer. To estimate a later representation’s connection to the input, we need to compose paths through the graph.

Matrix multiplication gives us exactly that composition.

Compose two attention layers

For the smallest useful example, ignore residual connections for a moment. Let A1 and A2 be attention matrices for two consecutive layers. Rows represent destination positions and columns represent source positions.

A1 = [[0.8, 0.2],
      [0.1, 0.9]]

A2 = [[0.5, 0.5],
      [0.0, 1.0]]

Every row sums to 1. The first row of A2 says that the first layer-2 position mixes the two layer-1 positions equally.

But those layer-1 positions are already mixtures of the two inputs. Multiplying the matrices traces the composition:

A2 @ A1

= [[0.5, 0.5],
   [0.0, 1.0]]
  @
  [[0.8, 0.2],
   [0.1, 0.9]]

= [[0.45, 0.55],
   [0.10, 0.90]]

The first output row now connects the first layer-2 representation back to the original inputs. Its weight for the second input is 0.55, even though the immediate layer-2 attention assigned only 0.5 to the second layer-1 position. Earlier mixing changed the effective path weights.

For layers 1 through L, the basic composition is:

R_L = A_L @ A_(L-1) @ ... @ A_1

This cumulative matrix is the core of attention rollout.

Residual connections add identity paths

A standard Transformer block does not replace a representation with only its attention output. A residual connection also carries the block input forward. A rollout calculation that ignores this path misses an important route through the network.

A common simplified adjustment is to add the identity matrix to an attention matrix and renormalize its rows:

A_tilde = normalize_rows(A + I)

When A is already row-normalized, adding I makes each row sum to 2, so the simplified form becomes:

A_tilde = (A + I) / 2

Then rollout composes the adjusted matrices:

R_L = A_tilde_L @ ... @ A_tilde_1

The identity term represents a path from a position’s previous representation directly into its next representation.

This adjustment is a modeling approximation. It does not establish that the residual branch and attention branch contribute equally to the actual hidden state. Their vector magnitudes, learned projections, normalization, and later computation can change their effects. The equal weighting comes from the simplified rollout construction, not from a universal Transformer guarantee.

Handle multiple attention heads deliberately

Real Transformer layers usually contain multiple attention heads. Rollout therefore needs a rule for turning head-specific matrices into the layer matrix that will be composed.

A simple teaching implementation averages the heads:

A_layer = mean(A_head_1, ..., A_head_H)

This produces one row-normalized attention matrix when each head is row-normalized. You can then add the residual identity path and compose layers.

Head averaging is convenient, but it discards head-specific structure. A head that tracks a particular relation can disappear into the average, while a head with diffuse attention can dilute a sharper pattern. Other aggregation rules are possible, but they encode different assumptions rather than revealing a uniquely correct attribution.

If your investigation concerns a particular head, averaging may answer the wrong question. If you want a compact layer-level view, averaging can be a reasonable exploratory choice as long as you record that choice.

Implement the calculation with explicit invariants

The algorithm itself is small. Pseudocode keeps the important assumptions visible:

rollout = identity(sequence_length)

for layer in layers_from_input_to_output:
    attention = aggregate_heads(layer.attention)
    adjusted = attention + identity(sequence_length)
    adjusted = adjusted / row_sum(adjusted)
    rollout = adjusted @ rollout

After the final iteration, row i of rollout describes the accumulated attention paths from final position i back to input positions under this model of information flow.

Several details should be checked in a real implementation:

  • Confirm the tensor axes. Libraries differ in whether attention is returned as [batch, heads, query, key] or another layout.
  • Confirm the direction of rows and columns before multiplying matrices. Reversing query and key axes changes the meaning.
  • Apply the model’s attention mask. Padding or forbidden causal positions should not acquire probability mass accidentally.
  • Preserve layer order. Rollout composes transformations from earlier layers toward later ones.
  • Check row sums after every normalization step. Small floating-point deviations are normal, but large deviations usually indicate a bug.

These checks matter more than making the implementation compact.

Read a rollout map as accumulated routing

Suppose you inspect the final token in a classification-style sequence and obtain:

input token:    server   is   not   healthy
rollout weight:  0.18   0.07  0.31   0.44

A useful interpretation is:

Under the chosen attention aggregation and residual approximation, more accumulated attention-path mass from this final representation reaches healthy and not than the other input positions.

That can guide investigation. Perhaps the model routes information about negation and state words toward the inspected representation.

A stronger statement such as “healthy caused the prediction because it has weight 0.44” is not justified by rollout alone.

Attention weights are only part of the computation. Attention also applies value projections; multi-head outputs are transformed; residual streams contain earlier information; feed-forward blocks modify representations; normalization changes their geometry; and the final prediction depends on later computations. A large routing weight does not tell you the magnitude or direction of the vector contribution that traveled along that route.

Validate hypotheses with interventions

Attention rollout is most useful when it generates a hypothesis that you test with a different method.

For example, if rollout consistently assigns high accumulated weight to a product identifier when a classifier predicts a certain category, you can ask whether the identifier is genuinely important. Perturb or remove it in controlled evaluation examples and measure how the model output changes. You can also compare rollout with gradient-based or other attribution methods when those methods fit the model and question.

These methods answer different questions, so agreement is informative but disagreement is not automatically a bug. Rollout summarizes attention paths. Input ablation measures what happens under a particular intervention. Gradients measure local sensitivity around an input representation. None should be silently substituted for another.

The original attention-rollout work reported stronger correlation with input-gradient and input-ablation importance measures than raw attention in its experiments. That is evidence for the method in those evaluated settings, not a guarantee that rollout is faithful for every architecture, task, or checkpoint.

Watch for common failure modes

The most common mistake is treating rollout as a general explanation layer rather than a specific approximation.

Using only the last attention layer. This ignores how earlier layers mixed token information. If the question is about paths back to the input, compose the layers.

Ignoring residual paths. A representation can continue through a block without traveling through the attention mixture. Rollout should account for that path if the architecture has residual connections around attention.

Assuming equal residual and attention influence is physically exact. Adding identity and renormalizing is a convenient approximation. Actual hidden-state contributions depend on vectors and block architecture.

Averaging heads without recording it. Head aggregation changes the result. Keep the aggregation rule with the visualization or analysis output.

Comparing weights across incompatible setups. Different tokenization, sequence lengths, masks, aggregation rules, or model architectures can change rollout distributions. A weight of 0.4 has no universal meaning independent of those choices.

Equating routing with causality. Rollout follows attention connectivity. It does not include every operation that determines the prediction.

Know when rollout is the right tool

Use attention rollout when you have access to internal attention matrices and want a compact view of how attention paths accumulate across layers. It is especially useful for debugging, exploratory model analysis, and comparing routing patterns under a fixed architecture and calculation procedure.

It is less suitable when you only have a hosted model’s text API, because the required internal attention tensors may not be exposed. It is also a poor final answer when you need a strong causal claim about why a model produced a prediction. In that case, intervention-based evaluation or a combination of attribution methods is usually more defensible.

A simpler method may also be enough. If you are debugging whether a causal mask is wired correctly, inspecting a single layer can directly answer the question. Rollout adds value when the question specifically spans multiple layers.

Finally, architecture details matter. The simple formulation assumes that attention matrices can be aligned across layers and that the residual adjustment is meaningful for the blocks being analyzed. Variants with unusual routing, token merging, cross-attention, or other structural changes may require a different graph and a different composition rule.

Conclusion

Attention rollout turns a stack of local attention matrices into an accumulated view of paths back to input tokens. The central idea is simple: earlier layers already mix information, so later attention must be composed with those earlier transformations. Residual connections add identity paths that should also be represented.

The method is useful because it is cheap, transparent, and easy to inspect. Its limitation is equally important: it models attention flow, not the full computation that causes a prediction. Use rollout to find patterns and form hypotheses, then use interventions or complementary attribution methods when the decision requires stronger evidence.

Artificial Intelligence Transformers Attention Interpretability