Residual connections are one of the reasons very deep neural networks train at all: each layer can add a correction to a running hidden state instead of having to rewrite the representation from scratch. In modern PreNorm Transformers, that running state is typically a simple additive accumulation. Every attention block and feed-forward block contributes with fixed unit weight, and the next block receives the sum.
Attention Residuals change that assumption. Instead of treating the residual stream as an unweighted sum, each layer selects from earlier representations with a learned, input-dependent attention distribution over depth. The mechanism asks a depth-wise question: which previous layer outputs are useful for this token at this point in the network?
The practical goal is not to replace token self-attention. Token attention still mixes information across sequence positions. Attention Residuals add a second, orthogonal aggregation axis across layer history. This can reduce PreNorm hidden-state growth, make gradient flow more evenly distributed across depth, and let later blocks retrieve earlier representations rather than inheriting a diluted sum.
A standard residual stream accumulates many transformed versions of the same token representation. As depth increases, each individual layer contribution becomes a smaller part of the total state, while activation magnitude can grow. This is useful for stable optimization, but it gives the model no explicit mechanism for saying that a particular earlier representation is more relevant than the latest accumulated state.
Attention Residuals keep previous layer outputs as values and compute softmax weights over them. The query is not the usual token query matrix used for sequence attention; the source describes a learned pseudo-query per layer. The resulting weights choose how much each earlier representation contributes to the current hidden state.
Full depth-wise attention over every earlier layer has memory and communication costs proportional to the number of stored layer outputs. Block Attention Residuals reduce that cost by grouping layers into blocks. Within a block, the model can accumulate normally; across blocks, it attends over block-level representations plus the current partial block state.
The residual stream stops being a passive sum and becomes a learned router over historical representations. That preserves the optimization benefits of skip connections while giving the model a content-dependent path to recover earlier features.
For a current layer, collect the prior hidden representations, normalize them, score each one with a learned projection, apply a softmax along the depth axis, and form a weighted sum. The attention axis is the list of previous depth states, not the token sequence length.
def depth_attention_residual(previous_states, pseudo_query, norm):
# previous_states: list of [batch, tokens, dim] tensors
values = torch.stack(previous_states, dim=0) # [depth, batch, tokens, dim]
keys = norm(values)
scores = torch.einsum("d,nbtd->nbt", pseudo_query, keys)
weights = torch.softmax(scores, dim=0)
return torch.einsum("nbt,nbtd->btd", weights, values)Block Attention Residuals store completed block summaries rather than every layer output. Before an attention or MLP sublayer, the model attends over completed block representations and the current block’s partial accumulation. At a block boundary, the partial accumulation is appended to the block list and a new block begins.
def block_attn_res(blocks, partial_block, proj, norm):
states = torch.stack(blocks + [partial_block])
keys = norm(states)
logits = torch.einsum("d,nbtd->nbt", proj.weight.squeeze(), keys)
weights = logits.softmax(dim=0)
return torch.einsum("nbt,nbtd->btd", weights, states)The mechanism is inserted where the residual stream would normally be read before a sublayer. A block can apply depth-wise aggregation before self-attention, add the attention output to the partial block state, apply depth-wise aggregation again before the MLP, then add the MLP output. This makes the replacement local to the Transformer block interface.
The most direct use case is training deeper Transformer or linear-attention models where residual accumulation becomes a bottleneck. A model can keep access to earlier states without relying on an ever-growing uniform sum, which is especially relevant when depth increases faster than width.
The reported downstream improvements are strongest on multi-step reasoning and code-generation benchmarks. That suggests a plausible application area: models that benefit from preserving intermediate abstractions across depth, such as theorem-style reasoning, program synthesis, and long dependency chains.
Attention Residuals are a compact architectural change compared with introducing a new attention kernel or a new expert routing stack. That makes them attractive for controlled ablations: train small baselines, compare full depth attention versus block depth attention, and measure whether the gains persist across model size and data scale.
Do not flatten depth into batch or sequence dimensions. Keeping a tensor layout such as [depth, batch, tokens, dim] makes the softmax axis explicit, simplifies debugging, and avoids accidentally mixing token attention with residual attention.
class BlockAttentionResidual(nn.Module):
def __init__(self, dim):
super().__init__()
self.norm = RMSNorm(dim)
self.query = nn.Parameter(torch.empty(dim))
nn.init.normal_(self.query, std=0.02)
def forward(self, completed_blocks, partial):
states = torch.stack([*completed_blocks, partial], dim=0)
keys = self.norm(states)
scores = torch.einsum("d,nbtd->nbt", self.query, keys)
weights = scores.softmax(dim=0)
return torch.einsum("nbt,nbtd->btd", weights, states)A block should be large enough to reduce memory overhead but small enough to preserve useful depth resolution. If each Transformer layer contains both attention and MLP sublayers, decide whether a block boundary counts full layers, sublayers, or pairs of sublayers and keep that convention consistent in checkpoints.
The mechanism is intended to address hidden-state growth and gradient concentration. Track activation norms by layer, attention mass over depth, and gradient norms across blocks. If attention collapses to the newest block every time, the implementation may behave like a more expensive standard residual path.