Papers
Topics
Authors
Recent
Search
2000 character limit reached

Block AttnRes: Scalable Residual Routing

Updated 15 July 2026
  • Block AttnRes is a scalable residual routing mechanism for deep PreNorm Transformers that replaces fixed additive aggregations with learned, content-dependent mixtures over block summaries.
  • It compresses multiple layer outputs into block-level summaries, reducing memory and computational costs while preserving essential depth-wise information.
  • The WAV extension further augments routing by incorporating phase and split detail bases to capture intra-block dynamics and improve scaling in deeper models.

Block Attention Residuals, commonly abbreviated Block AttnRes, are a family of residual-routing mechanisms for deep Transformer architectures in which the fixed additive residual stream of PreNorm models is replaced by a learned, content-dependent mixture over earlier depth sources that have been compressed at the level of residual blocks rather than individual layers. In the canonical formulation, Block AttnRes inherits the depth-wise selection principle of Attention Residuals while reducing its memory and communication cost by routing over block summaries and the current partial block. Later extensions, notably WAV, augment each block summary with directional detail bases so that routing can exploit not only a block’s total residual displacement but also coarse intra-block structure such as attention-versus-MLP imbalance and early-versus-late dynamics (Team et al., 16 Mar 2026, Wang, 4 Jun 2026).

1. Historical motivation and problem setting

The immediate background to Block AttnRes is the standard PreNorm Transformer residual update. For a decoder-only Transformer with attention sublayer output ala_l and MLP output mlm_l, the conventional update is

al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,

ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.

Under this scheme, every sublayer contributes to the residual stream with fixed coefficient $1$. The core critique advanced by the AttnRes literature is that such uniform aggregation causes hidden-state magnitude growth with depth, dilutes the relative contribution of each layer, and prevents later layers from selectively retrieving specific earlier representations rather than only the most recent recurrent state (Team et al., 16 Mar 2026).

Attention Residuals address this by replacing fixed accumulation with learned, content-dependent routing across depth. In the full formulation, a layer does not merely add its predecessor; instead it forms a convex recombination over earlier depth sources. The practical difficulty is that routing over all preceding layer outputs scales poorly in large models, particularly under activation recomputation and pipeline parallelism. Block AttnRes is the scalable answer: layers are partitioned into blocks, each block is summarized, and later sublayers attend over those summaries plus the current partial block rather than all prior layer states (Team et al., 16 Mar 2026).

This design goal is therefore dual. It is intended to preserve the optimization and representational benefits of content-dependent depth-wise selection, while making the mechanism deployable at LLM scale.

2. Formal definition and routing mechanics

In the conceptual AttnRes formulation, the residual update at layer ll is replaced by softmax attention over earlier sources: hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)}, with

ϕ(q,k)=exp(qRMSNorm(k)).\phi(q,k)=\exp(q^\top \mathrm{RMSNorm}(k)).

In the version reported for large-scale language modeling, the query is a learned pseudo-query ql=wlq_l=w_l, while keys and values are earlier residual sources, with b0=h1b_0=h_1 for the token embedding and subsequent sources derived from layer outputs (Team et al., 16 Mar 2026).

Block AttnRes compresses this source axis. A block is a group of consecutive sublayer updates. In the WAV formulation, a block mlm_l0 contains updates mlm_l1 and is summarized by

mlm_l2

At a later sublayer inside block mlm_l3, the routing source set becomes

mlm_l4

where mlm_l5 is the token embedding source and mlm_l6 is the current partial block sum. The mixer scores each source mlm_l7 using

mlm_l8

Here mlm_l9 is the routed readout supplied as context to the current attention or MLP sublayer, al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,0 is a learned function of the current normalized hidden state, and al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,1 are learnable per-source biases (Wang, 4 Jun 2026).

Two points define the mechanism precisely. First, routing is depth-wise rather than token-neighborhood attention in the usual sequence dimension. Second, Block AttnRes is not a mere scalar gate on the latest residual branch; it is a learned mixture over a structured pool of earlier block-level sources.

3. Computational structure and systems considerations

The principal systems advantage of Block AttnRes is that it reduces the source set from layer-wise scale to block-wise scale. Relative to layer-wise Attention Residuals, the per-sublayer routing cost drops from al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,2 prior states to al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,3 block-level summaries, where al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,4 is depth and al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,5 is the number of residual blocks. In the large-scale AttnRes analysis, Full AttnRes has memory footprint al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,6, whereas Block AttnRes reduces this to al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,7 by storing block representatives rather than all prior layer outputs (Team et al., 16 Mar 2026).

A substantial part of the method’s practicality comes from the associated execution strategy. The reported implementation uses a two-phase computation strategy. Phase 1 batches all queries in a block against cached inter-block sources. Phase 2 then proceeds sequentially through the current block, computes intra-block attention over the evolving partial sum, and merges inter-block and intra-block contributions with an online softmax update. The merge is written as

al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,8

al=Attn(LN(xl)),xl=xl+al,a_l = \mathrm{Attn}\big(\mathrm{LN}(x_l)\big), \quad x'_l = x_l + a_l,9

This permits efficient reuse of inter-block computations while preserving exact normalization (Team et al., 16 Mar 2026).

The same paper also introduces cache-based pipeline communication for pipeline-parallel training. In the stated formulation, naïve communication cost is

ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.0

whereas cross-stage caching reduces it to

ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.1

This reduces the peak per-transition cost from ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.2 to ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.3, described as a ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.4 improvement. Reported inference overhead is less than ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.5, and training overhead with pipeline parallelism is less than ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.6 (Team et al., 16 Mar 2026).

The architecture is therefore best understood as a joint algorithmic-and-systems proposal: the routing rule alone is insufficient without block compression, online merging, and communication caching.

4. Multi-resolution routing and the WAV extension

A major limitation of basic Block AttnRes is that a single block summary ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.7 records only the block’s total residual displacement. The WAV paper characterizes this as retaining only the low-frequency or DC component of the block trajectory, while discarding directional structure internal to the block. Two specific omissions are emphasized: attention-versus-MLP imbalance and early-versus-late block dynamics (Wang, 4 Jun 2026).

WAV v1 addresses this by augmenting each block with two zero-sum detail bases. The phase basis contrasts attention and MLP updates: ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.8 where ml=MLP(LN(xl)),xl+1=xl+ml.m_l = \mathrm{MLP}\big(\mathrm{LN}(x'_l)\big), \quad x_{l+1} = x'_l + m_l.9 for attention updates and $1$0 for MLP updates. The split basis contrasts the first and second halves of the block: $1$1 where $1$2 for $1$3 and $1$4 for $1$5.

The routed source pool is then expanded to include completed and partial versions of both detail types together with the standard block summaries. Crucially, the same depth-wise softmax mixer is reused; WAV does not alter the attention or MLP modules themselves. Training stability is supported by two mechanisms. First, detail sources are initialized with a negative bias,

$1$6

while embedding and $1$7 sources use zero bias, making the initial model close to ordinary Block AttnRes. Second, each detail tensor is RMS-matched to its associated block summary using a detached scaling factor,

$1$8

The final output head uses only the embedding and $1$9 sources by default, not detail sources (Wang, 4 Jun 2026).

This extension increases the number of routed sources by roughly a factor of three but adds only four scalar biases per Transformer layer, since the attention, MLP, embedding, and output head functions remain unchanged.

5. Empirical behavior and scaling

The most direct small-scale evaluation of Block AttnRes and its WAV extension is reported for character-level GPT language modeling on TinyStories and Text8 with PreNorm RMSNorm, causal self-attention, SwiGLU MLPs, vocabulary size ll0, ll1, ll2 attention heads, MLP dimension ll3, sequence length ll4, depths ll5, and ll6 fixed residual blocks. Under this setup, WAV is not consistently beneficial at ll7 layers, becomes competitive at ll8 layers, and is best at ll9 layers (Wang, 4 Jun 2026).

Depth Text8 validation loss TinyStories validation loss
12 Block AttnRes 0.9801; WAV 1.0085 Block AttnRes 0.5261; WAV 0.5325
24 Block AttnRes 0.9592; WAV 0.9610 Block AttnRes 0.5036; WAV 0.5012
48 Block AttnRes 0.9363; WAV 0.9305 Block AttnRes 0.4960; WAV 0.4738

At hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},0 layers, the reported reduction relative to Block AttnRes is from hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},1 to hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},2 on TinyStories and from hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},3 to hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},4 on Text8, with negligible additional parameters. The draft notes that training curves at hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},5 layers show consistent separation for TinyStories and smaller but consistent improvement by end of training for Text8, while also stating that standard deviations across seeds and error bars are not yet included (Wang, 4 Jun 2026).

At larger scale, the original AttnRes paper reports scaling-law fits

hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},6

At hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},7 PFLOP/s-days, Block AttnRes reaches hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},8 validation loss versus hl=i=0l1αilvi,αil=ϕ(ql,ki)j=0l1ϕ(ql,kj),h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i, \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)},9 for the baseline, described as approximately a ϕ(q,k)=exp(qRMSNorm(k)).\phi(q,k)=\exp(q^\top \mathrm{RMSNorm}(k)).0 compute advantage. In the same study, Block AttnRes is integrated into the Kimi Linear architecture with ϕ(q,k)=exp(qRMSNorm(k)).\phi(q,k)=\exp(q^\top \mathrm{RMSNorm}(k)).1B total and ϕ(q,k)=exp(qRMSNorm(k)).\phi(q,k)=\exp(q^\top \mathrm{RMSNorm}(k)).2B activated parameters and pretrained on ϕ(q,k)=exp(qRMSNorm(k)).\phi(q,k)=\exp(q^\top \mathrm{RMSNorm}(k)).3T tokens, where it improves downstream performance across all evaluated tasks and yields more uniform output magnitudes and gradient distribution across depth (Team et al., 16 Mar 2026).

The consistent empirical pattern is therefore depth sensitivity. Block-level routing by itself is beneficial; directional detail routing becomes increasingly useful as the network deepens and each block accumulates richer internal dynamics.

6. Interpretability, scope, and common confusions

Because Block AttnRes exposes routing weights explicitly in the forward pass, it presents an obvious interpretability target. The strongest cautionary result so far is that architectural exposure is necessary but not sufficient for mechanistic interpretation. In a causal-probing study on two same-scale ϕ(q,k)=exp(qRMSNorm(k)).\phi(q,k)=\exp(q^\top \mathrm{RMSNorm}(k)).4B Qwen3 checkpoints, a wrapped vanilla baseline with zero routing projection and recency bias ϕ(q,k)=exp(qRMSNorm(k)).\phi(q,k)=\exp(q^\top \mathrm{RMSNorm}(k)).5 produced routing weights that were content-independent and exactly matched the analytic schedule. The trained Block AttnRes checkpoint instead showed three localized routing motifs: an embedding-source pathway through early-layer MLP, a current-state pathway through early-layer attention and MLP, and an older-history pathway through late-layer attention. Yet the paper also found a sharp dissociation between average routing mass and causal importance, including a source family with appreciable mass but no detectable causal role under intervention (Javadov, 11 Jun 2026).

This interpretability result helps resolve several common confusions.

First, Block AttnRes is not synonymous with block-sparse attention. MoBA, NABLA, and “Block-attention for Efficient Prefilling” all operate on block-structured attention masks, routing, or KV reuse in the sequence dimension, but each explicitly retains standard Transformer residual connections rather than replacing the residual stream with depth-wise source mixing (Lu et al., 18 Feb 2025, Mikhailov et al., 17 Jul 2025, Ma et al., 2024).

Second, the term has broadened in domain-specific adaptations. DeRes introduces a dual-path CTR architecture with an Identity residual path plus a Block Attention Residual path over earlier block summaries, while BARFI-Q uses adaptive block-attention residual aggregation inside a time-series forecasting backbone. These systems inherit the broad idea of adaptive cross-depth reuse, but their exact semantics differ from the canonical Transformer depth-routing formulation of AttnRes and Block AttnRes (Cheng et al., 6 Jun 2026, Dastagir et al., 6 May 2026).

Third, higher routing mass should not be read as proof of functional importance. The causal-probe evidence suggests that routing visualizations are best treated as hypotheses to be tested by intervention rather than as direct evidence of mechanism (Javadov, 11 Jun 2026).

Taken together, the literature defines Block AttnRes most precisely as a scalable, block-compressed form of Attention Residuals for deep PreNorm Transformers, aimed at mitigating residual dilution through learned depth-wise routing. Its later extensions widen the representational bandwidth of block summaries, and its interpretability value appears real but conditional: the routing tensor is informative only when it has been learned as part of optimization and only when descriptive analyses are checked against causal tests.

Topic to Video (Beta)

No one has generated a video about this topic yet.

Whiteboard

No one has generated a whiteboard explanation for this topic yet.

Follow Topic

Get notified by email when new papers are published related to Block Attention Residuals (Block AttnRes).