pytorch - 💡(How to fix) Fix [RFC] nn.BlockAttentionResidual: depth-wise softmax attention over layer outputs as a drop-in residual replacement [3 comments, 3 participants]

# Current behaviour — every layer in TransformerEncoder:
x = x + self._sa_block(self.norm1(x), ...)   # fixed weight: 1
x = x + self._ff_block(self.norm2(x))        # fixed weight: 1

---

import torch
import torch.nn as nn

encoder = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True),
    num_layers=12,
)

residual = nn.BlockAttentionResidual(
    d_model=512,
    num_layers=12,
    num_blocks=8,     # N; paper finds N≈8 recovers most of Full AttnRes gain
)

# Usage: residual manages layer output history, encoder layers delegate to it
src = torch.randn(2, 64, 512)
out = residual(encoder, src)

---

encoder = nn.TransformerEncoder(
    layer,
    num_layers=12,
    residual="block_attention",   # "standard" (default) | "block_attention"
    attn_residual_num_blocks=8,
)

---

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List


class BlockAttentionResidual(nn.Module):
    """
    Block Attention Residuals (Block AttnRes) as described in:
        "Attention Residuals", Kimi Team, 2025.
        https://github.com/MoonshotAI/Attention-Residuals

    Replaces standard additive residual connections in a Transformer stack
    with learned, input-dependent softmax attention over block-level
    summaries of preceding layer outputs.

    Instead of:
        h_l = h_{l-1} + f_{l-1}(h_{l-1})          # uniform weight = 1

    Each layer computes:
        h_l = Σ_{i} α_{i→l} · v_i                  # softmax weights α

    where v_i are block-level summaries of earlier layer outputs and
    α_{i→l} = softmax( w_l^T · RMSNorm(v_i) ) for a learned pseudo-query w_l.

    Layers are grouped into N blocks of size S = num_layers / num_blocks.
    Within each block, layer outputs are summed into a single representation b_n.
    Cross-block attention is then applied over [b_0, b_1, ..., b_{n-1}] plus
    the evolving intra-block partial sum, reducing memory from O(Ld) to O(Nd).

    Args:
        d_model:    Hidden dimension of the transformer.
        num_layers: Total number of transformer sublayers (attn + ffn each count
                    as one sublayer, so a standard TransformerEncoderLayer with
                    one attn + one ffn = 2 sublayers).
        num_blocks: Number of blocks N. Paper finds N≈8 recovers most of the
                    gain of Full AttnRes across scales. Must divide num_layers.

    Complexity:
        Memory:     O(N · T · d) for block representations vs O(L · T · d) for
                    Full AttnRes (T = sequence length).
        Training:   Negligible overhead without pipeline parallelism; <4% with it.
        Inference:  <2% latency overhead on typical workloads.

    Notes:
        - All pseudo-query vectors w_l are zero-initialized, ensuring uniform
          attention weights at the start of training (= standard residual).
          This is critical for training stability (validated in the paper).
        - RMSNorm on keys prevents layers with large-magnitude outputs from
          dominating the softmax. Especially important for block-level
          representations that accumulate over multiple layers.
        - The token embedding is always included as source b_0, allowing
          any layer to attend back to the original input.
    """

    def __init__(self, d_model: int, num_layers: int, num_blocks: int = 8) -> None:
        super().__init__()

        if num_layers % num_blocks != 0:
            raise ValueError(
                f"num_layers ({num_layers}) must be divisible by num_blocks ({num_blocks})"
            )

        self.d_model = d_model
        self.num_layers = num_layers
        self.num_blocks = num_blocks
        self.block_size = num_layers // num_blocks  # S

        # One learned pseudo-query per sublayer, zero-initialized (critical).
        # Shape: (num_layers, d_model)
        self.pseudo_queries = nn.Parameter(torch.zeros(num_layers, d_model))

        # One RMSNorm per sublayer for key normalization.
        self.key_norms = nn.ModuleList([
            nn.RMSNorm(d_model) for _ in range(num_layers)
        ])

    def _attn_residual_op(
        self,
        layer_idx: int,
        block_reps: List[torch.Tensor],
        partial_block: torch.Tensor,
    ) -> torch.Tensor:
        """
        Compute the attention-weighted mixture of block representations
        for a single sublayer.

        Args:
            layer_idx:    Global sublayer index (0-based).
            block_reps:   List of completed block representations [b_0, ..., b_{n-1}].
                          b_0 is always the token embedding.
            partial_block: Intra-block partial sum b_n^i (current block so far).

        Returns:
            h_l: Input to sublayer `layer_idx`, shape (B, T, d_model).
        """
        # Stack all sources: completed blocks + current partial sum
        # V: (num_sources, B, T, d_model)
        sources = block_reps + [partial_block]
        V = torch.stack(sources, dim=0)

        # Normalize keys (prevents magnitude-dominant blocks from winning)
        norm = self.key_norms[layer_idx]
        K = norm(V)  # (num_sources, B, T, d_model)

        # Compute attention logits with the learned pseudo-query w_l
        # w_l: (d_model,)  →  logits: (num_sources, B, T)
        w_l = self.pseudo_queries[layer_idx]
        logits = torch.einsum("d, n b t d -> n b t", w_l, K)

        # Softmax over the source (depth) dimension
        weights = torch.softmax(logits, dim=0)  # (num_sources, B, T)

        # Weighted sum
        h = torch.einsum("n b t, n b t d -> b t d", weights, V)
        return h

    def forward(
        self,
        embedding: torch.Tensor,
        sublayers: nn.ModuleList,
        **sublayer_kwargs,
    ) -> torch.Tensor:
        """
        Run all sublayers with Block AttnRes residual connections.

        Args:
            embedding:       Token embeddings, shape (B, T, d_model).
                             Treated as b_0 — always available as a source.
            sublayers:       nn.ModuleList of sublayer callables (each takes a
                             tensor and returns a tensor of the same shape).
                             Typically the attention and FFN blocks of a
                             TransformerEncoderLayer, unrolled into a flat list.
            **sublayer_kwargs: Passed through to each sublayer (e.g., attn_mask).

        Returns:
            Output tensor of shape (B, T, d_model).
        """
        assert len(sublayers) == self.num_layers, (
            f"Expected {self.num_layers} sublayers, got {len(sublayers)}"
        )

        # b_0 is always the token embedding — allows any layer to attend back to input
        block_reps: List[torch.Tensor] = [embedding]
        partial_block: torch.Tensor = torch.zeros_like(embedding)

        for layer_idx, sublayer in enumerate(sublayers):
            # Check if we are at a block boundary (start of new block)
            pos_in_block = layer_idx % self.block_size

            if pos_in_block == 0 and layer_idx > 0:
                # Completed the previous block — save its representation and reset
                block_reps.append(partial_block)
                partial_block = torch.zeros_like(embedding)

            # Compute attention-weighted input for this sublayer
            h = self._attn_residual_op(layer_idx, block_reps, partial_block)

            # Apply the sublayer (e.g., self-attention or FFN)
            sublayer_out = sublayer(h, **sublayer_kwargs)

            # Accumulate into the current block's partial sum
            partial_block = partial_block + sublayer_out

        return partial_block

---

import torch
import torch.nn as nn

d_model, nhead, num_layers, num_blocks = 512, 8, 12, 6
batch_size, seq_len = 2, 64

# Build sublayers as a flat ModuleList (attn + ffn per TransformerEncoderLayer)
# In a real integration this would hook into TransformerEncoderLayer internals.
sublayers = nn.ModuleList([
    nn.Linear(d_model, d_model)   # placeholder; real layers are attn/ffn blocks
    for _ in range(num_layers)
])

attn_res = BlockAttentionResidual(d_model=d_model, num_layers=num_layers, num_blocks=num_blocks)

embedding = torch.randn(batch_size, seq_len, d_model)
output = attn_res(embedding, sublayers)
print(output.shape)  # torch.Size([2, 64, 512])