Architecture2026~110 minarXiv:2603.15031

Attention Residuals

score 9Kimi Team, Guangyu Chen, Yu Zhang +34Moonshot AI↑ 180★ 3,147Mar 16, 2026cs.CLhttps://github.com/MoonshotAI/Attention-Residuals

🎯 Pitch

Standard PreNorm residuals force every layer to receive the same uniformly-weighted sum of all prior outputs, causing hidden states to grow uncontrollably with depth and diluting early information. This paper shows that replacing this fixed accumulation with a lightweight softmax attention over preceding layers acts like giving the network a learnable time machine, achieving a 1.25× compute scaling advantage. The resulting 48B-parameter model dramatically boosts multi-step reasoning, proving that when a layer retrieves past information is as important as what it computes.

1. Executive Summary

This paper introduces Attention Residuals (AttnRes), a mechanism that replaces the fixed uniform accumulation in standard PreNorm residual connections with learned, input-dependent softmax attention over preceding layer outputs, enabling each layer to selectively aggregate earlier representations. To make this practical at scale, the authors propose Block AttnRes, which partitions layers into blocks and attends over block-level representations, reducing memory and communication overhead from O(Ld) to O(Nd) while preserving most of the gains. Scaling law experiments on MoE Transformer models up to 528M activated parameters show that Block AttnRes achieves a 1.25× compute advantage over the baseline, and a final 48B-parameter model pre-trained on 1.4T tokens with AttnRes integrated into the Kimi Linear architecture improves downstream performance across all evaluated tasks — with particularly pronounced gains on multi-step reasoning benchmarks such as GPQA-Diamond (+7.5) and Minerva Math (+3.6) — while mitigating PreNorm's hidden-state magnitude growth and yielding more uniform gradient distributions across depth, establishing that selective depth-wise aggregation benefits compositional tasks where later layers can retrieve and build upon earlier representations.

2. Context and Motivation

The Overlooked Second Role of Residual Connections

Residual connections [12] are universally understood as gradient highways. The update rule $\mathbf{h}_l = \mathbf{h}_{l-1} + f_{l-1}(\mathbf{h}_{l-1})$ creates an additive identity path that lets gradients flow directly from the final loss back to any earlier layer — a critical property for training deep networks. But this paper argues that residuals serve a second, equally important role that has received far less attention: they define how information aggregates across depth.

Consider what happens when you unroll the residual recurrence. The hidden state at layer $l$ is:

$\mathbf{h}_l = \mathbf{h}_1 + \sum_{i=1}^{l-1} f_i(\mathbf{h}_i)$

Every layer receives exactly the same uniformly-weighted sum of all prior layer outputs. Whether a layer is an attention module or an MLP, it sees the identical aggregated representation — there is no mechanism to selectively emphasize information from layer 3 while suppressing layer 7. This is the paper's central observation: depth-wise aggregation remains governed by fixed unit weights, with no mechanism for learned, input-dependent selection across layers.

This stands in stark contrast to how modern Transformers handle other dimensions of information mixing. In sequence mixing, we moved from RNN-style recurrent compression to softmax attention over all previous tokens — precisely because learning which tokens to attend to enables substantially richer information flow than uniform accumulation. Expert routing in MoE models similarly uses learned, input-dependent gating to direct information to specialized sub-networks. Yet the depth dimension — the very axis along which representations are progressively refined — remains stuck with a fixed additive scheme designed primarily for gradient propagation, not for expressive information routing.

The Hidden-State Growth Problem

The problem is not merely conceptual. When PreNorm is used (the dominant paradigm in modern LLMs), the unweighted accumulation causes hidden-state magnitudes to grow by roughly a factor of $L$ with depth. This is because each layer's output $f_l(\mathbf{h}_l)$ is approximately scale-invariant (due to LayerNorm before $f_l$ ), while the residual stream $\mathbf{h}_l$ accumulates these outputs additively:

$\|\mathbf{h}_L\| \approx O(L)$

As Li et al. [27] formalized, this progressive magnitude growth gradually dilutes each individual layer's relative contribution. Early-layer information — which may encode crucial low-level features — becomes buried under the mass of later-layer outputs, with no mechanism to selectively retrieve or amplify it. Empirically, this manifests in the well-documented finding that a significant fraction of layers can be pruned with minimal performance degradation [11], suggesting that much of the depth budget is effectively wasted on layers that contribute marginally to the final representation.

The dilution problem creates a perverse incentive during training: deeper layers, observing that their outputs represent a diminishing fraction of the total residual stream, learn to produce increasingly large outputs to remain influential. This can destabilize training and limits the effective depth at which additional layers provide meaningful benefit.

The Sequence-Depth Duality

The paper draws a provocative analogy. Before Transformers, recurrent neural networks (RNNs) faced an analogous bottleneck: they compressed all past sequence information into a single hidden state, updated step by step via additive recurrence. The Transformer's breakthrough was to replace this temporal recurrence with direct attention — allowing each position to selectively access all previous positions with learned, input-dependent weights.

Residual connections over depth are structurally identical to RNN recurrence over time: both propagate information through a fixed additive update rule that compresses history into a single state. The paper's thesis is that the same transition that proved transformative for sequence modeling should be applied to depth: replace uniform cumulative summation with learned, input-dependent softmax attention over all preceding layer outputs. This is the core motivation for Attention Residuals.

Where Prior Approaches Fall Short

The paper identifies three broad categories of prior work that attempt to address aspects of this problem, and explains where each remains inadequate.

1. Single-State Recurrence Variants

Several methods refine the standard residual update $\mathbf{h}_l = \mathbf{h}_{l-1} + f_{l-1}(\mathbf{h}_{l-1})$ by introducing learned coefficients or gates, while maintaining the constraint that each layer receives only the single compressed state $\mathbf{h}_{l-1}$ :

Highway Networks [45] introduce learned element-wise gates: $\mathbf{h}_l = (1 - \mathbf{g}_l) \odot \mathbf{h}_{l-1} + \mathbf{g}_l \odot f_{l-1}(\mathbf{h}_{l-1})$ , where $\mathbf{g}_l \in [0,1]^d$ interpolates between the transformation and the identity path.
ReZero [2] and LayerScale [50] add learnable scalar or per-channel multipliers to $f_{l-1}(\mathbf{h}_{l-1})$ .
DeepNorm [54] scales the residual branch to control magnitude growth.
SiameseNorm [27] maintains two parameter-shared streams — one PreNorm and one PostNorm — to simultaneously preserve identity gradients and bound hidden-state magnitudes.

The paper acknowledges these can improve training stability but identifies a fundamental constraint they all share: each layer can only access its immediate predecessor's state $\mathbf{h}_{l-1}$ , a single compressed representation that conflates all earlier layer outputs. This means:

No selective access: different layer types (e.g., attention vs. MLP) receive the same aggregated state, despite potentially benefiting from different weightings.
Irreversible compression: information discarded during early aggregation cannot be selectively recovered by deeper layers, even if it would be useful for later computation.
Depth bias: the cumulative nature of the residual means later layers exert outsized influence simply by virtue of appearing later in the sequence, not because their computation is more relevant.

2. Multi-State Recurrence Methods

Hyper-Connections [72] and their stabilized variant mHC [59] address the compression bottleneck by widening the recurrent state to $m$ parallel streams. The update becomes:

$\mathbf{H}_l = \mathbf{H}_{l-1}\mathbf{A}_l + f_{l-1}(\mathbf{H}_{l-1}\boldsymbol{\alpha}_{l-1}) \boldsymbol{\beta}_{l-1}^\top$

where $\mathbf{H}_l \in \mathbb{R}^{d \times m}$ maintains $m$ separate streams, $\mathbf{A}_l \in \mathbb{R}^{m \times m}$ is a learned transition matrix mixing streams across depth, and $\boldsymbol{\alpha}_{l-1}, \boldsymbol{\beta}_{l-1} \in \mathbb{R}^m$ control how streams are combined for input and output respectively. DDL [67] uses a delta-rule erase-and-write mechanism to update a matrix state.

While these provide richer state representations, the paper notes they still operate within the recurrence paradigm: layer $l$ conditions only on $\mathbf{H}_{l-1}$ , the state at the immediate predecessor. The model cannot selectively retrieve information from, say, layer 5 when computing layer 20 — the state has been sequentially transformed through all intermediate layers, and the mixing matrices $\mathbf{A}_l$ apply to the entire state uniformly. As the paper puts it, "AttnRes is orthogonal" to these methods, providing selective access to individual earlier-layer outputs while remaining compatible with any normalization or gating scheme.

3. Prior Cross-Layer Access Methods

A separate line of work directly provides each layer access to earlier-layer outputs, which is closer in spirit to AttnRes:

DenseFormer [36] and ANCRe [68] compute weighted sums of all previous layer outputs using learned scalar coefficients: $\mathbf{h}_l = \alpha_{0 \to l} \mathbf{h}_1 + \sum_{i=1}^{l-1} \alpha_{i \to l} f_i(\mathbf{h}_i)$ . However, these weights are fixed after training — they are input-independent static scalars, so the same mixing pattern applies regardless of what the input actually requires. This paper's experiments show DenseFormer achieves no gain over the baseline (Table 4: 1.767 vs. 1.766 baseline), highlighting the critical importance of input-dependent weighting.
MRLA [10] applies element-wise sigmoid gating over all previous layers, but uses a separable query-key product structure that the paper characterizes as closer to linear attention than softmax-based retrieval.
MUDDFormer [56] generates position-dependent weights via a small MLP across four decoupled streams but does not use softmax normalization, which the paper shows is important for competitive selection across sources.
DenseNet [17] concatenates all preceding feature maps but is designed for convolutional networks and scales poorly with depth.

The common limitation: no prior method provides softmax-normalized, input-dependent weights that allow selective, competitive aggregation over individual earlier-layer outputs with explicit learned queries. The paper's ablation (Table 4) demonstrates that removing the input-dependent key computation and using fixed scalars hurts performance (1.749 vs. 1.737 for Full AttnRes), and replacing softmax with sigmoid also degrades results (1.741).

The Specific Gap This Paper Addresses

The paper identifies a precise gap: there is no practical mechanism that provides softmax-normalized, input-dependent, content-aware selective aggregation over depth while remaining feasible for large-scale distributed training. The key challenge is that naive cross-layer attention requires retaining $O(Ld)$ activations and transmitting them across pipeline stages — prohibitive at the scale of models with hundreds of layers. The paper's contribution is twofold: (1) defining the mechanism (AttnRes) that achieves selective depth-wise aggregation through a single learned pseudo-query per layer, and (2) developing the infrastructural innovations (Block AttnRes, cross-stage caching, two-phase computation) that make it practical at scale with marginal overhead.

How This Paper Positions Itself

The paper positions AttnRes as completing the linear-to-softmax transition for depth that Transformers achieved for sequences. The structured-matrix analysis in Section 6.2 makes this explicit: standard residuals, Highway networks, and mHC all perform what amounts to depth-wise linear attention (with different constraints on the mixing matrix $\mathbf{M}$ ), while AttnRes performs depth-wise softmax attention. The paper argues this is not merely a technical analogy but reflects a principled design choice — softmax's competitive normalization enables sharper selection among sources, which the ablation confirms (softmax outperforms sigmoid).

Importantly, the paper does not claim to replace all residual-connection research. Instead, it argues that AttnRes is complementary to existing gating and normalization improvements. The pseudo-query mechanism is lightweight (one $d$ -dimensional vector per layer), can be integrated into any architecture that uses residual connections, and imposes no changes to the layer functions $f_l$ themselves. The infrastructure optimizations (Section 4) are specifically designed to make the mechanism a "drop-in replacement for standard residual connections" with negligible inference latency overhead (<2%) and marginal training cost.

The paper also positions its empirical findings as validation of a deeper principle: that better depth-wise information flow benefits compositional tasks where later layers can selectively retrieve and build upon earlier representations. This explains the pattern in downstream results (Table 3), where gains are largest on multi-step reasoning (GPQA-Diamond +7.5, Math +3.6) and smallest on knowledge-intensive tasks (MMLU +1.1). The mechanism fundamentally changes how information moves through the depth of the network, and compositional reasoning — which requires coordinating multiple intermediate representations — benefits most.

3. Technical Approach

This paper is fundamentally an architectural design and systems engineering paper that proposes a new mechanism for depth-wise information routing and then solves the infrastructure challenges needed to make it work at scale.

3.1 Reader Orientation

The system being built is a modified Transformer architecture where each layer, before computing its attention or MLP output, looks back at all previous layers and decides — using learned attention weights — which earlier representations are most useful for its current computation, rather than receiving a fixed uniform blend of everything that came before. The problem it solves is that standard PreNorm residual connections force every layer to consume the same undifferentiated cumulative sum of prior outputs, which causes early-layer information to be progressively diluted and prevents layers from selectively retrieving relevant past representations. The solution takes the form of a lightweight cross-layer attention mechanism — a single learned query vector per layer attends over all preceding layer outputs (or block summaries thereof), producing input-dependent weights that selectively aggregate depth-wise information — combined with a block-structured approximation and specialized distributed training infrastructure that keep the approach practical at the scale of 48B-parameter models.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components:

Pseudo-Query Vectors — one learned $d$ -dimensional vector $`\mathbf{w}_l`$ per layer $l$ , which serves as the query in depth-wise attention. These are the only new parameters introduced by AttnRes; they are small (one vector per layer, negligible relative to total parameter count) and initialized to zero to ensure the model starts training with uniform depth-wise averaging.
Key-Value Sources — the outputs of preceding layers (or their block-level summaries) that serve as the keys and values for depth-wise attention. In Full AttnRes, these are the individual layer outputs $f_i(\mathbf{h}_i)$ plus the token embedding $\mathbf{h}_1$ . In Block AttnRes, completed blocks are compressed into single summary vectors $\mathbf{b}_n = \sum_{j \in \mathcal{B}_n} f_j(\mathbf{h}_j)$ , plus an evolving partial sum $\mathbf{b}_n^i$ for layers within the current block.
Softmax Attention over Depth — for each layer $l$ , compute attention weights $\alpha_{i \to l}$ over all available sources by applying an exponential kernel $\phi(\mathbf{q}_l, \mathbf{k}_i) = \exp(\mathbf{q}_l^\top \text{RMSNorm}(\mathbf{k}_i))$ followed by normalization, producing a probability distribution over sources. The layer's input is then the attention-weighted sum: $\mathbf{h}_l = \sum_i \alpha_{i \to l} \mathbf{v}_i$ .
Block Partitioning (Block AttnRes) — a structural approximation that groups $L$ layers into $N$ blocks. Within each block, standard residual accumulation aggregates layer outputs into a single block representation. Cross-block attention then operates over only $N$ block-level summaries rather than $L$ individual layer outputs, reducing memory and communication from $O(Ld)$ to $O(Nd)$ . The block count $N \approx 8$ is chosen empirically to recover most of Full AttnRes's gains.
Distributed Training Infrastructure — cross-stage caching under pipeline parallelism that eliminates redundant transmission of block representations, and a two-phase computation strategy (bulk inter-block attention batched across layers in a block, followed by sequential intra-block attention with online softmax merge) that amortizes memory access costs during both training and inference.

Information flows as follows: the token embedding $\mathbf{h}_1$ enters the first layer → the layer's pseudo-query $\mathbf{w}_l$ attends over currently available sources (the embedding plus any completed block representations and the intra-block partial sum) via softmax attention → the resulting aggregated representation feeds into the layer's function $f_l$ (attention or MLP) → $f_l$ 's output is added to the intra-block partial sum → at block boundaries, the completed block sum is appended to the set of inter-block representations and the intra-block accumulator resets → the process repeats for all layers → the final layer outputs the prediction.

3.3 Roadmap for the Deep Dive

First, the formal definition of Full AttnRes — the attention kernel, query/key/value construction, and the exact update equations — because this is the conceptual core that all variants build on.
Second, Block AttnRes and the interpolation property — how layers are partitioned into blocks, how block representations are formed, and how intra-block and inter-block attention interact — because this is the scalable variant used in most experiments.
Third, the structured-matrix perspective — the depth mixing matrix $\mathbf{M}$ that unifies residual variants — because it provides the analytical framework for understanding what AttnRes generalizes and clarifies its relationship to prior methods like mHC and Highway networks.
Fourth, the training infrastructure — cross-stage caching under pipeline parallelism and the memory accounting that makes Block AttnRes practical — because these systems contributions are what separate a conceptual proposal from a deployable mechanism.
Fifth, the inference infrastructure — the two-phase computation strategy with online softmax and the memory-efficient prefilling scheme — because inference latency is critical and the paper claims <2% overhead.
Sixth, the design choices and their justifications — zero-initialization of queries, RMSNorm on keys, learned vs. input-dependent queries, softmax vs. sigmoid, and the block count selection — because these embody the paper's empirical wisdom about what matters.

3.4 Detailed, Sentence-Based Technical Breakdown

Full Attention Residuals: The Core Mechanism

The fundamental idea of Full AttnRes is to replace the standard residual input $\mathbf{h}_l = \mathbf{h}_1 + \sum_{i=1}^{l-1} f_i(\mathbf{h}_i)$ — which assigns unit weight to every source — with an attention-weighted sum:

$\mathbf{h}_l = \alpha_{0 \to l} \cdot \mathbf{h}_1 + \sum_{i=1}^{l-1} \alpha_{i \to l} \cdot f_i(\mathbf{h}_i)$

where $\alpha_{i \to l}$ are attention weights satisfying $\sum_{i=0}^{l-1} \alpha_{i \to l} = 1$ , and $\alpha_{i \to l} \in [0, 1]$ is the weight that layer $l$ assigns to the output of layer $i$ (with index 0 reserved for the token embedding $\mathbf{h}_1$ ).

What this equation computes: The input to layer $l$ is a convex combination of all preceding layer outputs, where the mixing coefficients are learned and input-dependent rather than fixed at 1. Physically, for each token, the model computes $l-1$ attention scores, normalizes them into a probability distribution, and then computes a weighted average embedding that blends information from all prior layers according to their relevance for layer $l$ 's upcoming computation.

Why this form: The convex combination ( $\sum_i \alpha_i = 1$ ) guarantees that $\|\mathbf{h}_l\| \leq \max_i \|\mathbf{v}_i\|$ , meaning the hidden-state magnitude cannot grow unboundedly with depth — this directly addresses the PreNorm dilution problem described in Section 2. The unit-weight alternative ( $\mathbf{h}_l = \mathbf{h}_1 + \sum_{i=1}^{l-1} f_i(\mathbf{h}_i)$ ) has no such bound, which is why PreNorm hidden states grow as $O(L)$ . A weighted sum without the sum-to-one constraint would still exhibit unbounded growth if the learned weights consistently exceeded 1. The convex combination is therefore not just a modeling convenience but directly encodes the solution to the magnitude-growth pathology.

The attention weights are computed via a softmax-normalized exponential kernel over queries and keys:

$\alpha_{i \to l} = \frac{\phi(\mathbf{q}_l, \mathbf{k}_i)}{\sum_{j=0}^{l-1} \phi(\mathbf{q}_l, \mathbf{k}_j)}$

where the kernel function $\phi: \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}_{\geq 0}$ is defined as:

$\phi(\mathbf{q}, \mathbf{k}) = \exp\left(\mathbf{q}^\top \text{RMSNorm}(\mathbf{k})\right)$

and RMSNorm is Root Mean Square Layer Normalization [66] applied to the key vector before the dot product.

What this equation computes: For each source layer $i$ (including the embedding), the model takes its output $\mathbf{v}_i$ , applies RMSNorm to produce the key $\mathbf{k}_i = \text{RMSNorm}(\mathbf{v}_i)$ , computes the dot product with the layer $l$ 's learned query vector $\mathbf{q}_l = \mathbf{w}_l$ , exponentiates, and then divides by the sum of exponentiated scores over all available sources. The result is a probability distribution: layers whose normalized outputs have high dot-product with $\mathbf{w}_l$ receive more weight, with competition enforced by the softmax denominator.

Why this form: The exponential kernel with softmax normalization serves three purposes. First, the softmax's competitive normalization — where increasing one source's weight necessarily decreases others — forces sharp selection among sources rather than letting the model assign moderate weight to everything. This is why the authors prefer softmax over sigmoid (Table 4: softmax achieves 1.737 vs. sigmoid's 1.741). Sigmoid treats each source independently, so the model could theoretically attend strongly to everything simultaneously, which defeats the purpose of selective aggregation. Second, the RMSNorm on keys prevents layers with naturally large output magnitudes — particularly deep layers in PreNorm architectures — from dominating the softmax purely by virtue of scale. Without RMSNorm, a layer that learned to output vectors with norm 100 would receive essentially all the attention mass regardless of content, rendering the mechanism useless. Third, the exponential is chosen over alternative non-linearities because it is the standard softmax kernel and has favorable gradient properties (the gradient of log-softmax is the residual).

The query, key, and value construction is deliberately minimalist:

$\mathbf{q}_l = \mathbf{w}_l, \quad \mathbf{k}_i = \mathbf{v}_i = \begin{cases} \mathbf{h}_1 & i = 0 \\ f_i(\mathbf{h}_i) & 1 \leq i \leq l-1 \end{cases}$

where $\mathbf{w}_l \in \mathbb{R}^d$ is a learned pseudo-query vector — one per layer — that is a direct parameter of the model, not computed from the layer's hidden state. The keys $\mathbf{k}_i$ and values $\mathbf{v}_i$ are identical: both are the output of layer $i$ 's transformation $f_i(\mathbf{h}_i)$ , or the token embedding for $i=0$ .

What this equation computes: The query is simply a learned embedding vector for layer $l$ . The keys and values are the actual outputs produced by previous layers — the same tensors that, in a standard residual network, would be added to the residual stream with unit weight. When $i=0$ , the source is the token embedding $\mathbf{h}_1$ . When $i \geq 1$ , the source is the output of transformation $f_i$ applied to layer $i$ 's (attention-weighted) input.

Why this form — learned vs. input-dependent queries: The authors explicitly considered and tested an alternative: making the query input-dependent by projecting it from the current hidden state as in standard self-attention. This actually improves performance (Table 4: 1.731 vs. 1.737 for learned queries), but they default to the learned query for two reasons. First, an input-dependent query requires a $d \times d$ projection matrix per layer, which adds non-trivial parameters. Second, and more critically, computing the query from the hidden state creates a sequential dependency during inference: you cannot compute the query for layer $l+1$ until you have the hidden state for layer $l+1$ , which means you cannot batch the inter-block attention computation in the two-phase strategy (Section 4.2). The learned pseudo-query is available before any forward computation begins, enabling the critical batching optimization that amortizes memory access. The authors frame this as a deliberate engineering tradeoff, opting for a design that is "compatible with parallel execution" over one that is "marginally more expressive but requires sequential memory access."

Why this form — shared key and value: Having $\mathbf{k}_i = \mathbf{v}_i$ means the attention mechanism attends over the same vectors it aggregates. This is the standard setting in self-attention and is simpler than introducing separate key and value projections. Since each layer already produces a $d$ -dimensional output $f_i(\mathbf{h}_i)$ , reusing it as both key and value adds zero extra computation or storage. The authors do not experiment with decoupled key-value projections, likely because doing so would require storing or recomputing additional tensors per layer, undermining the memory efficiency of the approach.

A crucial design detail is that all pseudo-query vectors must be initialized to zero: $\mathbf{w}_l = \mathbf{0}$ for all $l$ at the start of training. The paper states this explicitly: "Crucially, all pseudo-query vectors must be initialized to zero. This ensures that the initial attention weights $\alpha_{i \to l}$ are uniform across source layers, which reduces AttnRes to an equal-weight average at the start of training and prevents training volatility, as we validated empirically."

What this achieves: When $\mathbf{w}_l = \mathbf{0}$ , the dot product $\mathbf{q}_l^\top \text{RMSNorm}(\mathbf{k}_i) = \mathbf{0}^\top \cdot * = 0$ for all sources, so $\exp(0) = 1$ for every source, and the softmax becomes $\alpha_{i \to l} = 1/l$ — a uniform distribution. The initial model input is therefore $\mathbf{h}_l = \frac{1}{l} \sum_{i=0}^{l-1} \mathbf{v}_i$ , a simple average of all prior outputs. This is close to, but not identical to, standard residuals (which sum rather than average). The key property is that no source is initially preferred over any other, so the model must learn from data which layers to emphasize. A random initialization of $\mathbf{w}_l$ would create arbitrary initial preferences that could steer learning toward suboptimal attention patterns or cause training instability from the outset.

The computational cost of Full AttnRes is:

Arithmetic: $O(L^2 d)$ per token, because each of the $L$ layers computes attention over up to $L$ sources, with each attention computation costing $O(d)$ for the dot product. Since $L$ is typically modest (50–200 layers), this is small relative to the $O(T^2 d)$ cost of sequence-wise self-attention (where $T$ can be millions).
Memory: $O(Ld)$ per token to store all layer outputs for use as keys and values by later layers. In vanilla training without activation recomputation, these activations are already retained for backpropagation, so Full AttnRes "introduces no additional memory overhead" (Section 3.1). However, under activation checkpointing and pipeline parallelism — both standard at scale — these outputs must now be explicitly preserved and communicated, which is what motivates Block AttnRes.

Block Attention Residuals: Making Cross-Layer Attention Scalable

Block AttnRes partitions the $L$ layers into $N$ blocks of $S = L/N$ layers each (assuming divisibility; otherwise the last block contains the remainder). The key structural insight is that within each block, standard residual accumulation is used to compress all layer outputs into a single summary vector, and cross-block attention operates only over these summaries.

For block $n$ , the block representation is the sum of all layer outputs within the block:

$\mathbf{b}_n = \sum_{j \in \mathcal{B}_n} f_j(\mathbf{h}_j)$

where $\mathcal{B}_n$ is the set of layer indices belonging to block $n$ . We further denote $\mathbf{b}_n^i = \sum_{j \in \mathcal{B}_n, j < i} f_j(\mathbf{h}_j)$ as the partial sum over the first $i$ layers in block $n$ , so that $\mathbf{b}_n = \mathbf{b}_n^S$ .

What this equation computes: For each block, all the individual layer outputs — each of which would be a separate source in Full AttnRes — are added together into a single $d$ -dimensional vector. This is exactly the standard residual accumulation restricted to within-block scope: at the start of a block, the accumulator resets to zero, and each layer's output is added to a running sum. When the block completes, the running sum becomes $\mathbf{b}_n$ , which is stored as a source for all subsequent blocks. Within the block, the evolving partial sum $\mathbf{b}_n^i$ is also available as a source for layers after the first.

Why this form: The sum aggregation is conceptually minimal — it requires no learned parameters, no attention computation, and adds zero overhead beyond what standard residuals already do. The intuition, supported by the empirical finding that block-structured attention recovers most of the gains of full attention (Figure 6: S=4 reaches 1.746 vs. Full's 1.737), is that layers within a block are sufficiently similar in their representational role that fine-grained per-layer attention provides diminishing returns. This is consistent with the pruning literature finding that adjacent layers can often be removed with minimal loss [11]. The block structure exploits a locality prior: cross-layer retrieval is most valuable at a coarser granularity, while within-block information flow is adequately handled by standard residuals.

For the $i$ -th layer in block $n$ , the set of available sources (values) for depth-wise attention depends on whether this is the first layer of the block:

$\mathbf{V} = \begin{cases} [\mathbf{b}_0, \mathbf{b}_1, \ldots, \mathbf{b}_{n-1}]^\top & \text{if } i = 1 \text{ (first layer of block } n) \\ [\mathbf{b}_0, \mathbf{b}_1, \ldots, \mathbf{b}_{n-1}, \mathbf{b}_n^{i-1}]^\top & \text{if } i \geq 2 \text{ (subsequent layers)} \end{cases}$

where $\mathbf{b}_0 = \mathbf{h}_1$ is the token embedding, always included as a source. For the first layer of a block ( $i=1$ ), the model attends over all completed block representations plus the embedding. For subsequent layers ( $i \geq 2$ ), the model additionally attends over the intra-block partial sum $\mathbf{b}_n^{i-1}$ , which contains the accumulated output of all prior layers within the current block.

What this equation computes: The value matrix available to each layer. Think of it as the "depth-wise context window." When a block begins, the model can look back at all previous blocks' summaries. As the block progresses, it also gains access to the running sum of this block's own outputs so far, enabling intra-block attention. This is analogous to causal attention in sequence models: the model can see all past blocks (like past tokens) and the current block's prefix (like the current sequence context), but not future blocks or future layers within the current block.

Why this form — including the partial sum: If the partial sum were not included, the first layer of a block would see only completed block representations, and subsequent layers within the block would see exactly the same sources — there would be no way for layer $i+1$ to attend differently than layer $i$ within the same block. By including the evolving partial sum, each layer within a block can selectively attend to the specific contributions of the immediately preceding layers, preserving fine-grained intra-block attention while the inter-block attention handles coarser retrieval. This is a careful engineering tradeoff: adding one extra source per intra-block step costs $O(S)$ extra attention computation per block, which is small, but preserves the ability for layers within a block to differentiate their attention patterns.

The attention weights and aggregation follow the same softmax formula as Full AttnRes (Eq. 2), with the keys $\mathbf{k}$ being RMSNorm-normalized versions of the values $\mathbf{V}$ . The pseudo-query $\mathbf{w}_l$ remains a learned $d$ -dimensional vector per layer. Figure 2 in the paper provides PyTorch-style pseudocode illustrating the exact forward pass logic: at each layer, the model calls block_attn_res with the list of completed block representations and the current partial sum, which computes softmax attention and returns the aggregated hidden state; if the layer is at a block boundary, the completed partial sum is appended to the block list and the partial sum resets.

The memory and computational costs of Block AttnRes are:

Memory: $O(Nd)$ per token to store $N$ block representations, down from $O(Ld)$ for Full AttnRes. For typical values ( $L=54$ layers, $N=9$ blocks), this is a $6\times$ reduction.
Computation: $O(N^2 d)$ for inter-block attention, down from $O(L^2 d)$ . For $N=9$ and $L=54$ , this is a $36\times$ reduction.

The block count $N$ interpolates between two extremes: when $N = L$ (block size $S=1$ ), Block AttnRes recovers Full AttnRes exactly; when $N = 1$ (block size $S=L$ ), there is only one block, no inter-block attention, and the model reduces to standard residual connections with the embedding isolated as $\mathbf{b}_0$ — the intra-block partial sum handles everything. Empirically, the paper finds that $N \approx 8$ recovers most of the benefit (Figure 6: S=4 with $N \approx 4$ achieves 1.746 vs. Full's 1.737; S=8 with $N \approx 2$ also achieves 1.746), and this number is chosen for the final model based on infrastructure efficiency considerations.

A subtle detail is how Block AttnRes handles attention and MLP layers separately. The paper's architecture treats "each self-attention or MLP as an individual layer" (Section 2 notation), so a Transformer block with one attention and one MLP corresponds to two layers in the AttnRes scheme. Consequently, both pre-attention and pre-MLP have their own depth-wise attention, with separate pseudo-query vectors $\mathbf{w}_l$ and separate RMSNorm parameters. Figure 2's pseudocode shows this explicitly: self.attn_res_proj and self.attn_res_norm handle the attention layer's depth-wise attention, while self.mlp_res_proj and self.mlp_res_norm handle the MLP layer's depth-wise attention. Block boundaries count both attention and MLP layers, so block_size in the pseudocode counts in units of individual layers ( $block\_size // 2$ Transformer blocks).

This separation matters because the learned attention patterns in Figure 8 show that pre-attention and pre-MLP layers develop different depth-wise attention profiles: "Pre-MLP inputs show sharper diagonal reliance on recent representations, while pre-attention inputs maintain broader receptive fields." By giving attention and MLP layers independent depth-wise queries, the model can learn that attention layers benefit from broad cross-layer retrieval (consistent with attention's role in routing information) while MLP layers operate more locally on the representation delivered by the immediately preceding attention layer.

The Depth Mixing Matrix: A Unified Analytical Framework

Section 6.2 of the paper introduces a structured-matrix perspective that unifies residual variants. This is more than a taxonomy — it is the analytical tool that reveals why AttnRes generalizes prior methods and clarifies the exact nature of the transition from linear to softmax attention over depth. We cover it here because it underpins the design rationale for the attention kernel choice.

Define the depth mixing matrix $\mathbf{M} \in \mathbb{R}^{L \times L}$ , where element $M_{i \to l}$ is the weight that layer $l$ assigns to the output of layer $i$ . The input to layer $l$ can then be written compactly as:

$\mathbf{h}_l = \sum_{i=0}^{l-1} M_{i \to l} \mathbf{v}_i$

where $\mathbf{v}_0 = \mathbf{h}_1$ (embedding) and $\mathbf{v}_i = f_i(\mathbf{h}_i)$ for $i \geq 1$ . Different residual variants correspond to different structures of $\mathbf{M}$ :

Standard residual [12]: $M_{i \to l} = 1$ for all $i < l$ . $\mathbf{M}$ is an all-ones lower-triangular matrix. Every prior layer contributes equally.
Highway [45]: $M_{i \to l} = g_{i+1} \prod_{j=i+2}^{l} (1 - g_j)$ for $i \geq 1$ , and $M_{0 \to l} = \prod_{j=1}^{l} (1 - g_j)$ . The cumulative products factor through scalar gates $g_j$ , making $\mathbf{M}$ 1-semiseparable — the same matrix rank as standard residuals but with input-dependent weights. The weights sum to one, making it a softmax-free instance of stick-breaking attention.
(m)HC [72, 59]: $M_{i \to l} = \boldsymbol{\beta}_i^\top \mathbf{A}_{i+1 \to l}^\times \boldsymbol{\alpha}_l$ , where $\mathbf{A}_{i \to j}^\times = \prod_{k=i+1}^j \mathbf{A}_k$ is the product of $m \times m$ transition matrices, $\boldsymbol{\alpha}_l \in \mathbb{R}^m$ mixes streams into layer $l$ 's input, and $\boldsymbol{\beta}_i \in \mathbb{R}^m$ distributes layer $i$ 's output across streams. $\mathbf{M}$ is m-semiseparable, with rank at most $m$ .
Full AttnRes: $M_{i \to l} = \alpha_{i \to l}$ computed via softmax attention. $\mathbf{M}$ is dense (full rank $L$ ), with no low-rank or semiseparable constraint.
Block AttnRes: For sources $i$ in a completed earlier block $\mathcal{B}_n$ , all share the block-level key/value $\mathbf{b}_n$ , so $M_{i \to l} = \alpha_{n \to l}$ for every $i \in \mathcal{B}_n$ . Within the current block, each layer additionally attends over the partial sum $\mathbf{b}_n^{i-1}$ . The effective rank of $\mathbf{M}$ lies between $N$ and $N+S$ .

What this framework reveals: Prior residual variants — including the most sophisticated like mHC — all impose a low semiseparable rank on $\mathbf{M}$ . This means the effective weight assigned to earlier-layer outputs is constrained to lie in a low-dimensional subspace. For mHC with $m=4$ streams, $M_{i \to l}$ can vary across $(i,l)$ pairs, but only within a 4-dimensional manifold — the pattern of inter-layer connectivity is heavily parameterized. Full AttnRes removes this constraint entirely, allowing each $(i,l)$ pair to have an independent weight, with the only constraint being that the weights for a fixed $l$ sum to 1. This is the "linear-to-softmax" transition that the paper claims completes for depth what Transformers achieved for sequences.

The paper also notes a revealing connection: when the kernel $\phi$ decomposes as a dot product of feature maps, $\phi(\mathbf{q}, \mathbf{k}) = \varphi(\mathbf{q})^\top \varphi(\mathbf{k})$ , depth-wise attention collapses back into a recurrence — precisely the structure underlying the MRLA–GLA and DDL–DeltaNet correspondences. This means that prior linear-attention variants over depth are special cases of AttnRes with a linear kernel, and that AttnRes's choice of exponential kernel (which does not factorize) is what enables the dense, full-rank $\mathbf{M}$ .

Training Infrastructure: Pipeline Communication and Caching

The primary infrastructure challenge for AttnRes at scale is pipeline parallelism. In a standard pipeline-parallel setup, each stage processes a subset of layers and communicates only the hidden state $\mathbf{h}$ at stage boundaries — a fixed-size tensor independent of pipeline depth. With Block AttnRes, each stage additionally needs all previously accumulated block representations to compute inter-block attention. Naïvely transmitting the full block history at every stage transition incurs redundant communication that scales quadratically with the number of chunks.

The paper formalizes this using an interleaved 1F1B schedule [33] with $P$ physical stages and $V$ virtual stages per physical stage. Let $N_p$ be the average number of block representations produced per physical stage (each of dimension $d$ per token). With $C = PV$ total chunks, the $j$ -th chunk has accumulated $jN_p$ blocks. Naïve transmission at every transition yields per-token communication:

$\text{Comm}_{\text{naïve}} = \sum_{j=1}^{C-1} jN_p \cdot d = \frac{C(C-1)}{2} N_p d$

where $C(C-1)/2$ is the sum of integers from 1 to $C-1$ , each term representing the number of blocks transmitted at the $j$ -th transition, multiplied by the per-block dimension $d$ .

What this equation computes: The total number of scalar elements transmitted per token across all pipeline transitions in one forward pass, if every transition sends its entire block history. For a concrete example: $P=4$ , $V=2$ gives $C=8$ chunks. Chunk 1 accumulates $N_p$ blocks and sends them to chunk 2. Chunk 2 accumulates $2N_p$ blocks and sends them to chunk 3. Chunk 3 sends $3N_p$ blocks to chunk 4, and so forth. The total is $N_p d(1+2+3+4+5+6+7) = 28 N_p d$ elements. The quadratic scaling $O(C^2)$ means doubling the number of pipeline stages quadruples communication.

Why this is wasteful: In an interleaved schedule, a physical stage processes multiple virtual stages in sequence. Chunk $j$ and chunk $j+P$ (the next virtual stage on the same physical stage, $P$ chunks later) need many of the same blocks. When chunk $j$ sends its history to chunk $j+1$ and eventually to chunk $j+P-1$ , and chunk $j+P$ receives those blocks again from its predecessor, the same data crosses the network twice. The cache-based optimization eliminates this redundancy.

The cross-stage caching optimization exploits the fact that each physical stage processes virtual stages sequentially. The key insight: blocks received during earlier virtual stages remain in local GPU memory and need not be re-transmitted. The total communication reduces to:

$\text{Comm}_{\text{cached}} = \underbrace{\frac{P(P-1)}{2} N_p d}_{\text{first virtual stage}} + \underbrace{(V-1) P^2 N_p d}_{\text{subsequent virtual stages}}$

What this equation computes: The total communication with caching, broken into two terms. The first term covers the initial virtual stage ( $v=1$ ), where no cache exists and the full history must be built up — this still costs $O(P^2)$ because the first virtual stage across all $P$ physical stages accumulates blocks from scratch. The second term covers virtual stages $v=2 \ldots V$ , where each transition conveys only the $\sim P N_p$ incremental blocks accumulated since the receiver's corresponding chunk in the previous virtual stage — blocks that the receiver has not yet seen. The factor $P^2$ comes from $P$ transitions per virtual stage, each transmitting roughly $P N_p$ incremental blocks.

Why this matters: The per-transition cost drops from $O(C) = O(PV)$ to $O(P)$ , a $V \times$ reduction. In the $P=4, V=2$ example from Figure 3: the naïve cost is $28 N_p d$ ; with caching, the first virtual stage costs $\frac{4 \times 3}{2} N_p d = 6 N_p d$ , and the second virtual stage costs $1 \times 16 N_p d = 16 N_p d$ , totalling $22 N_p d$ . The reduction is modest for $V=2$ but becomes dramatic for $V=8$ or $V=16$ , which are common in large-scale training. More importantly, the per-transition peak cost drops from transmitting $7N_p d$ blocks (at the $C-1$ th transition in the naïve scheme) to transmitting approximately $4 N_p d$ blocks (capped by $P$ rather than $C$ ), enabling overlap with computation during steady-state 1F1B.

Memory overhead under activation checkpointing: The paper states that "the per-layer activation footprint remains identical to standard architectures, as activation checkpointing eliminates all inter-block attention intermediates, and the checkpointed input $\mathbf{h}_l$ matches the memory size of the hidden state." With cross-stage caching, each block representation is stored exactly once across all virtual stages, and the per-layer cached blocks become "negligible relative to standard per-layer activation cache." The paper reports a measured end-to-end training overhead of "less than 4%" for Block AttnRes under pipeline parallelism.

Inference Infrastructure: Two-Phase Computation and Memory-Efficient Prefilling

The core challenge at inference is that attending over block representations requires reading them from memory at every layer, which in a naïve implementation means each layer performs an independent pass over all $N$ block vectors — costing $O(L \cdot N)$ memory accesses total. Since the pseudo-queries are learned parameters (not computed from hidden states), the paper exploits this to batch query processing.

Algorithm 1 (Two-Phase Computation for block $n$ ) proceeds as follows:

Phase 1 (Parallel inter-block attention): Stack the pseudo-queries for all $S$ layers in block $n$ into a matrix $\mathbf{Q} = [\mathbf{w}_l]_{l \in \mathcal{B}_n} \in \mathbb{R}^{S \times d}$ . Stack the $n$ block representations (keys and values) into $\mathbf{K}, \mathbf{V} = [\mathbf{b}_0; \mathbf{b}_1; \ldots; \mathbf{b}_{n-1}] \in \mathbb{R}^{n \times d}$ . Compute batched attention in a single matrix multiplication: $\mathbf{O}^{(1)} = \text{softmax}(\mathbf{Q} \mathbf{K}^\top) \mathbf{V} \in \mathbb{R}^{S \times d}$ . Additionally, retain the softmax statistics (row-wise max $m_l^{(1)}$ and log-sum-exp $\ell_l^{(1)}$ ) for each layer $l$ .

Phase 2 (Sequential intra-block attention + online softmax merge): Process layers $l \in \mathcal{B}_n$ sequentially. For the first layer ( $i=0$ ), there is no intra-block context, so $\mathbf{h}_l = \mathbf{o}_l^{(1)} / \ell_l^{(1)}$ — the inter-block attention output is used directly. For subsequent layers ( $i \geq 1$ ), compute intra-block attention with $\mathbf{w}_l$ as query and the partial sum $\mathbf{b}_n^i$ as both key and value, obtaining output $\mathbf{o}_l^{(2)}$ and statistics $m_l^{(2)}, \ell_l^{(2)}$ . Merge the two attention outputs via online softmax [31]:

$m_l = \max(m_l^{(1)}, m_l^{(2)})$

$\mathbf{h}_l = \frac{e^{m_l^{(1)} - m_l} \mathbf{o}_l^{(1)} + e^{m_l^{(2)} - m_l} \mathbf{o}_l^{(2)}}{e^{m_l^{(1)} - m_l} \ell_l^{(1)} + e^{m_l^{(2)} - m_l} \ell_l^{(2)}}$

Update the partial sum: $\mathbf{b}_n^{i+1} = \mathbf{b}_n^i + f_l(\mathbf{h}_l)$ , where $\mathbf{b}_n^0 = \mathbf{0}$ .

What this algorithm computes: An exact equivalent to computing attention over all sources (inter-block + partial sum) sequentially, but with the inter-block portion batched. The online softmax merge is the key insight: by retaining the max and sum-of-exponentials ( $m$ and $\ell$ ) from Phase 1, Phase 2 can combine attention outputs as if all sources were processed in a single softmax, without ever re-reading the inter-block sources. The formula works because softmax for source set $\mathcal{S} = \mathcal{S}_1 \cup \mathcal{S}_2$ can be reconstructed from the numerators and denominators of two separate softmax computations, provided the max is aligned.

Why this form — amortized memory access: In Phase 1, the $n$ block representations are read from memory once and reused across all $S$ queries. This reduces reads from $S \times n$ (if each layer read blocks independently) to $n$ — an $S \times$ improvement. The per-layer amortized memory access cost is derived in Appendix B (Table 1) as $(N/S + 5)d$ reads and $2d$ writes for Block AttnRes. For typical values ( $L=128$ , $N=8$ , $S=16$ ), this is $5.5d$ reads + $2d$ writes = $7.5d$ total I/O per layer. This compares favorably to standard residuals ( $3d$ total) and dramatically to mHC with $m=4$ streams ( $34d$ total). The paper reports that this enables "end-to-end inference latency overhead of less than 2% on typical inference workloads."

Memory-efficient prefilling addresses the cost of storing block representations during the prompt processing phase. Storing $N$ block representations of dimension $d$ for each of $T$ tokens requires $N \cdot T \cdot d$ elements. For a 128K-token sequence with $N=8$ blocks and $d=4096$ (typical for the Kimi Linear 48B model), this is $8 \times 131,072 \times 4,096 \approx 4.3 \times 10^9$ elements, or approximately 17 GB in bfloat16 — prohibitive for a single GPU.

The solution is sequence-dimension sharding: the block representations are sharded along the sequence dimension across $P$ tensor-parallel devices, so each device stores only $N \cdot (T/P) \cdot d$ elements. Phase 1 executes independently on each device's local sequence shard. Phase 2 integrates the online-softmax merge into the standard tensor-parallel communication path: the output is reduce-scattered across devices, the merge is performed locally, and the result is reconstructed via all-gather, which "naturally admits kernel fusion with operations like RMSNorm." For $P=8$ , the 128K-context example drops from ~17 GB to ~2.1 GB per device. Combined with chunked prefill (processing the prompt in chunks of, say, 16K tokens), the overhead further reduces to under 0.3 GB per device.

Design Choices and Their Empirical Justifications

The paper includes an ablation study (Table 4) that validates each component design decision on a 16-layer model. Here we enumerate the choices, the alternatives tested, the results, and the reasoning:

1. Softmax over sigmoid for attention normalization. Full AttnRes with softmax achieves validation loss 1.737; with sigmoid it degrades to 1.741. The authors attribute this to softmax's competitive normalization: "softmax forces sharper selection among sources" because increasing one source's weight necessarily decreases others. Sigmoid treats each source independently — the model could assign high weight to every source simultaneously, which defeats the purpose of selective aggregation. This is a specific instantiation of the more general principle that depth-wise attention benefits from competition among sources, just as sequence-wise attention does.

2. Learned pseudo-query over input-dependent query. The learned query ( $\mathbf{w}_l$ , loss 1.737) was compared against an input-dependent projection from the hidden state (loss 1.731). The input-dependent query actually performs slightly better, but the authors default to the learned query for systems reasons: it enables the batched Phase 1 computation in the two-phase strategy, which is essential for the <2% inference latency overhead. This is an explicit engineering-performance tradeoff.

3. RMSNorm on keys. Removing RMSNorm degrades Full AttnRes from 1.737 to 1.743 and Block AttnRes from 1.746 to 1.750. The paper explains this as preventing layers with naturally larger output magnitudes from dominating the softmax. For Block AttnRes, this is even more critical because block-level representations accumulate over multiple layers and can develop large magnitude differences — without RMSNorm, a block with many layers would swamp a block with fewer layers regardless of content relevance.

4. Input-dependent vs. input-independent mixing. Replacing the query-and-key attention with learnable, input-independent scalars (each source gets a learned weight, no interaction with the source content) degrades performance from 1.737 to 1.749. This confirms that the benefit comes from content-dependent selection — the model is not just learning a fixed depth-wise connectivity pattern, but adapting its attention based on what each layer actually produced.

5. Multi-head depth-wise attention. Testing per-head aggregation ( $H=16$ ) on Block AttnRes degrades performance from 1.746 to 1.752. The authors interpret this as evidence that "the optimal depth-wise mixture is largely uniform across channels: when a layer's output is relevant, it is relevant as a whole." In other words, the decision of which earlier layers to attend to does not benefit from being channel-specific — if layer 5's output is useful, it's useful across all hidden dimensions, not just a subset.

6. Block size selection. Figure 6 sweeps block size $S$ from 1 (Full AttnRes) to 32 (very coarse). Loss degrades gracefully: $S=1$ (1.737) → $S=2$ (1.746) → $S=4$ (1.746) → $S=8$ (1.748) → $S=16$ (1.753) → $S=32$ (1.757) versus baseline (1.766). The relatively flat region from $S=2$ to $S=8$ (loss within 0.002) justifies the choice of $N \approx 8$ blocks as the sweet spot — it captures most of the gain while keeping infrastructure overhead manageable. The paper notes that "as future hardware alleviates memory capacity constraints, adopting finer-grained block sizes or Full AttnRes represents a natural pathway to further improve performance."

7. Sliding-window vs. block-structured vs. full access. A sliding-window variant (keeping only the most recent $W=8$ layer outputs plus the embedding) achieves 1.764 — much worse than Block AttnRes (1.746) and only slightly better than baseline (1.766). The paper concludes that "selectively accessing distant layers matters more than attending to many nearby ones" — the ability to reach back to early layers (especially the token embedding, which consistently receives non-trivial attention weight in Figure 8) is crucial for the mechanism's effectiveness.

8. Comparison with prior methods. Under identical hyperparameters, Full AttnRes (1.737) outperforms mHC (1.747) and DenseFormer (1.767, essentially tied with baseline at 1.766). The DenseFormer result is particularly informative: DenseFormer provides cross-layer access with learned scalar weights, but they are input-independent and fixed after training. Its failure to improve over baseline highlights that input dependence is essential, not just cross-layer access per se. The mHC comparison shows that expanding the recurrent state (4 parallel streams) helps but is less effective than providing direct attention over individual layer outputs.

4. Key Insights and Innovations

Innovation 1: Reframing the Residual Connection as a Depth-Wise Information Retrieval Problem

The paper's most fundamental intellectual move is redefining what a residual connection is for. The field has long understood residuals as gradient highways — their value lies in enabling stable backpropagation through identity mappings. This paper argues that this view is incomplete: residuals also function as an information aggregation mechanism across depth, and a poorly designed one at that. The standard formulation h_l = h_1 + Σ f_i(h_i) treats every prior layer's output as equally relevant to every subsequent layer, with no mechanism for selective retrieval.

This reframing matters because it shifts the design space. If residuals are only gradient highways, then the design goal is to preserve the identity path while controlling magnitude growth — hence the long line of work on scaled residuals (DeepNorm), gated residuals (Highway), and multi-stream recurrence (mHC). Each of these operates within the recurrence paradigm: layer l receives a single compressed state h_{l-1} and updates it. But if residuals are also an information retrieval mechanism, then the design goal becomes something richer: give each layer direct, selective access to individual earlier-layer outputs, with learned weights that depend on what those outputs actually contain. This is a categorically different objective, and it leads to a categorically different solution class — cross-layer attention — that prior residual variants could not reach by merely tuning coefficients within the additive recurrence.

The paper doesn't just assert this reframing; it derives it from a structural analogy that makes the theoretical motivation rigorous. The sequence-depth duality (Section 6.1) observes that residual recurrence h_l = h_{l-1} + f_{l-1}(h_{l-1}) is formally identical to RNN recurrence over time. The Transformer replaced temporal recurrence with direct attention because compressing all past tokens into a single hidden state was a bottleneck — the exact same bottleneck, the paper argues, exists across depth. The structured-matrix analysis (Section 6.2) then shows that prior residual variants (Highway, mHC) all impose low semiseparable rank on the depth mixing matrix M, meaning they are constrained to a low-dimensional manifold of possible inter-layer connectivity patterns. AttnRes removes this constraint entirely, allowing each (i, l) pair to have an independent, content-dependent weight. This is the depth-wise analog of the linear-to-softmax transition that proved transformative over sequences — and the paper presents it not as a technical trick but as completing an unfinished conceptual transition from the recurrence era to the attention era, applied to a dimension (depth) that the field had overlooked.

Why this is more than a metaphor: the paper's ablation (Table 4) shows that DenseFormer — which provides cross-layer access with learned but input-independent scalar weights — achieves no improvement over the baseline (1.767 vs. 1.766). This negative result isolates the specific ingredient that matters: not cross-layer access per se, but input-dependent, content-aware selection. The model needs to decide whether to pay attention to layer 5 based on what layer 5 actually produced for this token, not based on a fixed learned preference for layer 5. This is the precise distinction between linear attention (where weights factor through pre-computed feature maps) and softmax attention (where weights depend on the dot product of query and key). The paper's central claim — that depth has been stuck in the "linear attention" regime and needs to move to "softmax attention" — is not just a catchy analogy but an empirically testable hypothesis, and the ablation validates it.

Innovation 2: Block Structuring as a Practical Mechanism That Recovers Most of the Full-Attention Gain

The paper's second key contribution is an empirical finding with significant practical implications: block-structured cross-layer attention recovers nearly all the benefit of full per-layer attention while reducing memory and communication costs by an order of magnitude. This is not an approximation derived from theoretical considerations — it is a discovery about the effective granularity at which depth-wise information retrieval matters.

Figure 6 shows that validation loss degrades remarkably gracefully as block size increases: Full AttnRes (block size 1) achieves 1.737; block sizes 2, 4, and 8 all land near 1.746 (a degradation of only ~0.009); block size 16 reaches 1.753; and only at block size 32 does performance approach the baseline (1.757 vs. 1.766). Meanwhile, the sliding-window variant (keeping only the 8 most recent layers plus the embedding) achieves only 1.764 — barely better than baseline. The contrast between the block result and the sliding-window result is the crucial diagnostic: it's not just that you need to attend over many sources (the sliding window does that), but that you need to reach back to distant layers selectively. The block structure preserves access to distant earlier blocks while compressing nearby layers, and this compressed access is nearly as good as full access.

What makes this finding fundamental rather than merely convenient is that it reveals something about how information is organized across depth. The fact that attending to block summaries (coarse granularity) works nearly as well as attending to individual layers (fine granularity) suggests that adjacent layers are sufficiently correlated in their representational function that per-layer attention provides diminishing returns. This is consistent with the pruning literature finding that adjacent layers can often be removed with minimal loss — but AttnRes operationalizes this insight constructively. Instead of pruning layers, it groups them into blocks and routes information at the block level, preserving the computational capacity of all layers while making inter-layer communication more efficient.

The block count N ≈ 8 is therefore not just an engineering compromise — it's an empirically identified natural scale for depth-wise information retrieval in the models studied. The paper is careful to note that this number may shift with architecture scale or future hardware, and that finer-grained blocking "represents a natural pathway to further improve performance." But the key insight is that the block structure is not a hack: it is a principled design choice that exploits a genuine property of deep representations (locality of representational similarity) to achieve a favorable accuracy-efficiency Pareto frontier.

The systems contributions in Section 4 — cross-stage caching reducing pipeline communication by a factor of V, the two-phase computation strategy reducing per-layer I/O from O(Ld) to O(S+N)d — are what make this insight actionable at scale. Without them, Block AttnRes would be a conceptual proposal, not a mechanism deployable in a 48B-parameter model. The combination — an empirical finding about effective depth granularity plus the infrastructure to exploit it — is what transforms the paper from an architecture paper into a systems paper with immediate practical relevance.

Innovation 3: The Diagnosis and Mitigation of PreNorm Dilution Through Selective Depth-Wise Aggregation

The paper provides the clearest empirical characterization to date of PreNorm dilution as a measurable pathology — and shows that selective depth-wise attention directly mitigates it. This is not just a new solution to a known problem; it's a new way of seeing the problem, with quantitative evidence that the mechanism works through the hypothesized channel.

The standard narrative about PreNorm is that it causes hidden-state magnitudes to grow with depth, but this is typically discussed at the level of theoretical scaling (||h_L|| ~ O(L)) or through indirect evidence (pruning studies showing deep layers are expendable). The paper's Figure 5 provides direct, fine-grained measurements: the baseline model's per-block output magnitude (Figure 5b) grows monotonically with depth, with the deepest blocks producing outputs roughly 3-4× larger in norm than the shallowest blocks. This is the dilution problem made visible: later layers are compelled to "shout louder" to be heard over the accumulated residual stream. Simultaneously, gradient magnitudes (Figure 5c) are heavily skewed toward the earliest layers in the baseline — a consequence of the identity gradient path being unmodulated, so gradient norm concentrates where the effective "learning signal" first enters the network.

Block AttnRes transforms both patterns. Output magnitudes (Figure 5b) become bounded and periodic — they grow within each block (since intra-block accumulation is standard residual) but reset at block boundaries because the attention-weighted aggregation normalizes the scale. Gradient magnitudes (Figure 5c) become substantially more uniform across depth. The mechanism is clean: the softmax weights introduce competition among sources for probability mass. Earlier layers cannot dominate gradient flow by default because their weight is determined by learned, content-dependent attention, not by fixed unit coefficients. Later layers can selectively amplify the gradients flowing to earlier layers they find useful, and suppress those they don't.

This is a diagnostic contribution, not just a performance one. The paper gives the field a concrete set of measurements — per-layer output magnitude, per-layer gradient norm — that can be used to evaluate any proposed solution to the depth-scaling problem. The fact that Block AttnRes makes both metrics more uniform provides mechanistic validation that the method works through the hypothesized channel (selective aggregation reducing dilution), not through some unrelated effect. This matters because architectural innovations often improve performance for unclear reasons; here, the paper can point to specific training dynamics and say "this is what we fixed, and here is the measurement."

The connection to downstream task performance (Table 3) adds a further layer of validation: the tasks that benefit most from AttnRes — GPQA-Diamond (+7.5), Minerva Math (+3.6), HumanEval (+3.1) — are precisely the tasks that require coordinating information across multiple reasoning steps. If dilution buries early-layer information that later layers need, then compositional reasoning should suffer disproportionately because it requires retrieving and building on intermediate representations produced at various depths. The pattern of results is therefore consistent with the mechanistic story: better depth-wise information flow → better compositional reasoning.

Innovation 4: The Depth Mixing Matrix as a Unified Analytical Language for Residual Architectures

The paper's final contribution is a conceptual tool rather than a method or finding: the depth mixing matrix M provides a unified framework for understanding and comparing residual variants. This is not merely taxonomic — it enables analytical insights that would be difficult to extract from the recurrence-form descriptions of these methods.

By casting each residual variant as a specific structure imposed on M, the paper makes visible what each method constrains and what it allows:

Standard residuals impose M_{i→l} = 1 — no selectivity, no input dependence, rank 1.
Highway networks make M 1-semiseparable — the weights can be input-dependent but factor through scalar gates, constraining the inter-layer connectivity pattern to a specific product structure.
(m)HC makes M m-semiseparable — the weights live in an m-dimensional subspace, which is richer than Highway but still heavily parameterized.
Full AttnRes makes M dense and full-rank — the only constraint is row-wise normalization.
Block AttnRes interpolates between these extremes, with effective rank between N and N+S.

This framework has three practical consequences. First, it explains why AttnRes generalizes prior methods: any method that constrains M to low semiseparable rank is a special case of the dense M that AttnRes allows. This is not just a theoretical observation — it predicts that AttnRes should outperform mHC (which it does: 1.737 vs. 1.747 in Table 4) because mHC's low-rank constraint limits the expressivity of inter-layer connectivity.

Second, the framework connects depth-wise attention to sequence-wise attention through the shared language of structured matrices. The observation that a linear kernel φ(q)·φ(k) collapses AttnRes back into a recurrence — and that this is precisely the structure underlying the MRLA–GLA and DDL–DeltaNet correspondences — unifies several lines of work that previously appeared unrelated. It suggests that the key design axis is not "recurrence vs. attention" but the rank of the mixing matrix, with different ranks achievable through different computational mechanisms.

Third, the framework enables new design questions that the recurrence form obscures. For example: what properties of the kernel φ matter? The paper shows empirically that softmax outperforms sigmoid, but the matrix perspective suggests investigating other kernels that produce M with different spectral properties. The paper's observation of depth-wise "attention sinks" — certain early layers (especially the embedding) consistently attracting high weight regardless of input — is directly visible in Figure 8 and mirrors the same phenomenon in sequence-wise attention. This suggests that depth-wise attention may exhibit similar pathologies (e.g., attention sink tokens dominating the softmax) and that solutions developed for sequence attention (e.g., attention sink mitigation) may transfer.

What elevates this from "useful taxonomy" to "genuine innovation" is that it is operationally useful: the matrix perspective directly informed the design of Block AttnRes (by making visible that the block structure corresponds to a specific rank constraint on M, interpolating between N=1 standard residuals and N=L full attention), and it provides a principled way to reason about future variants (e.g., what happens if we impose a different structure on M, such as a banded or low-rank-plus-sparse pattern). It is rare for a technical paper to provide both an effective method and the analytical language to extend it, and the depth mixing matrix is the paper's most durable intellectual contribution — one that will likely outlast the specific AttnRes implementation.

5. Experimental Analysis

Evaluation Methodology

Dataset. All scaling-law experiments use an internal training corpus following the Kimi Linear [69] data recipe, with validation loss measured on a held-out set. The downstream evaluation (Table 3) uses standard benchmarks described below. The paper does not specify the exact size or composition of the pre-training corpus beyond stating it follows Kimi Linear's 1.4T-token recipe.
Base model(s). The primary architecture is Kimi Linear [69], a Mixture-of-Experts (MoE) Transformer with hybrid attention (KDA and MLA layers in a 3:1 ratio), each followed by an MoE feed-forward layer. Scaling-law experiments sweep five model sizes from 194M to 528M activated parameters (Table 2). The final model uses the full Kimi Linear 48B configuration: 27 Transformer blocks (54 layers), 8 out of 256 routed experts plus 1 shared expert, yielding 48B total and 3B activated parameters. The baseline model is identical to the AttnRes model in all respects except the residual connection mechanism — all other components (depth, hidden dimensions, expert routing, MLP structure) remain unchanged, making this a clean ablation of the residual architecture.
Metrics. Pre-training experiments report validation loss (cross-entropy) as the primary metric, following standard scaling-law practice. The scaling-law analysis in Figure 4 reports validation loss as a function of compute in PFLOP/s-days. Downstream evaluation (Table 3) reports accuracy for each benchmark, with answer grading details unspecified (the paper references the Kimi Linear evaluation protocol [69]). Training dynamics analysis (Figure 5) reports per-block output magnitude (likely L2 norm, though not explicitly stated) and per-block gradient magnitude.
Baselines. The paper compares against:
- PreNorm baseline: standard PreNorm residual connections with uniform additive accumulation, identical architecture otherwise.
- mHC [59]: Hyper-Connections with $m=4$ parallel streams and manifold-constrained transition matrices, cited as a representative multi-state recurrence method.
- mHC-lite [64]: A computationally lighter variant of mHC that avoids the Sinkhorn-Knopp iterations. Listed in Table 2 alongside AttnRes results but discussed only briefly.
- DenseFormer [36]: Cross-layer access via learned, input-independent scalar weights (fixed after training), ablated in Table 4.
Generation budget / compute accounting. Scaling-law curves (Figure 4) use PFLOP/s-days on the x-axis, measuring total training compute. The five models in Table 2 are trained with fixed token budgets (from 38.7B tokens for the smallest to 119.0B for the largest) rather than matched FLOPs, but the scaling-law fit extrapolates across compute levels. For architecture sweeps (Figure 7), total training FLOPs and active parameters are fixed (approximately $6.5 \times 10^{19}$ FLOPs, approximately $2.3 \times 10^8$ active parameters) to isolate architectural effects from capacity effects. Validation loss comparisons in Table 2 and Table 4 are at matched compute budgets (identical hyperparameters, same number of training tokens).
Cross-validation / statistical protocol. The paper does not report cross-validation or statistical significance measures for downstream results. For the scaling-law fits (Figure 4), power laws of the form $L = A \times C^{-\alpha}$ are fitted following standard practice [22, 15], with the fitted parameters reported in the figure legend. The training recipe for the final model uses a WSD (Warmup–Stable–Decay) learning rate schedule [16] and the Muon optimizer [28], with a global batch size of 8M tokens. Within each scaling-law size group, "all variants share identical hyperparameters selected under the baseline to ensure fair comparison; this setup intentionally favors the baseline and thus makes the comparison conservative."

Main Quantitative Results

Scaling Laws: AttnRes Consistently Reduces Loss Across Compute Budgets

The paper's scaling-law experiments (Figure 4, Table 2) sweep five model sizes and train three variants per size: a PreNorm baseline, Full AttnRes, and Block AttnRes with approximately 8 blocks. All models are trained with an 8192-token context window and a cosine learning rate schedule.

The headline result: Block AttnRes achieves a 1.25× compute advantage over the baseline. At the largest scale (528M activated parameters, 119.0B tokens, approximately 5.6 PFLOP/s-days), Block AttnRes reaches a validation loss of 1.693 versus the baseline's 1.714. By interpolating on the fitted power-law curves, the baseline would require roughly 1.25× more compute to reach Block AttnRes's loss.

The fitted scaling-law curves (Figure 4) are:

Baseline: $L = 1.891 \times C^{-0.057}$
Block AttnRes: $L = 1.870 \times C^{-0.058}$
Full AttnRes: $L = 1.865 \times C^{-0.057}$

All three exhibit similar slopes (the exponents are nearly identical at approximately -0.057 to -0.058), meaning the benefit is a constant factor reduction in loss across the compute range rather than a change in scaling behavior. The constant factor $A$ is lower for AttnRes variants: 1.865 (Full) and 1.870 (Block) versus 1.891 (Baseline). This pattern — same slope, lower constant — is typical for architectural improvements that increase effective model capacity or improve training dynamics without fundamentally altering the scaling exponent.

The gap between Full and Block AttnRes narrows with scale. At the smallest size (194M activated parameters): Full AttnRes achieves 1.899 versus Block's 1.909 (difference of 0.010). At the largest size (528M): Full achieves 1.692 versus Block's 1.693 (difference of 0.001). This is shown in Table 2 and captured in the fitted curves, where Full and Block are nearly indistinguishable at the high-compute end of Figure 4. The paper interprets this as evidence that "Block AttnRes closely tracks Full AttnRes, recovering most of the gain at the largest scale" — though it's worth noting that the largest model still has only 528M activated parameters, and extrapolation to 48B-scale models cannot be directly verified from these curves.

Table 2 also includes mHC-lite [64] for reference, though it is trained only on the model sizes shown and not included in the fitted scaling curves. At comparable sizes, Full AttnRes matches or slightly outperforms mHC-lite (e.g., 1.737 vs. 1.747 for Full vs. mHC-lite at the 436M scale; 1.692 vs. 1.694 at the 528M scale). Block AttnRes is within 0.001-0.002 of mHC-lite while using substantially less memory I/O per layer: the paper reports 5.5d for Block AttnRes versus 34d for mHC with m=4 streams (Table 1). This is a specific claim about memory access efficiency at inference time, not FLOPs — the two methods may have different compute-vs-I/O tradeoffs.

Assessment of the scaling-law evidence: The scaling curves span less than one order of magnitude in compute (roughly 0.5 to 5.6 PFLOP/s-days on the x-axis of Figure 4) and five model sizes. This is a relatively narrow range by scaling-law standards — typical scaling-law studies span 2-3 orders of magnitude to reliably estimate exponents. The similar slopes across variants may partly reflect this limited range rather than a genuine invariance of the scaling exponent. The 1.25× compute advantage is extrapolated from the fitted curves, which are reasonably well-behaved over the measured range but should be interpreted cautiously as a point estimate. Additionally, all scaling-law experiments use Block AttnRes with approximately 8 blocks — the paper does not sweep the block count at different model sizes to verify that N=8 remains optimal as models scale up.

Final Model Training Dynamics: Measurable Amelioration of PreNorm Dilution

The paper trains a 48B-parameter Block AttnRes model (3B activated, with 6 layers per block producing 9 blocks plus the token embedding for 10 depth-wise sources) and a matched baseline for 1T tokens under the WSD schedule, then analyzes training dynamics (Figure 5).

Validation loss (Figure 5a): Both models follow similar trajectories early in training (up to approximately 20K steps), with the gap widening notably during the decay phase (post-80K steps). The final validation loss is "notably lower" for Block AttnRes (exact numbers not reported in the figure; the y-axis is labeled with loss values from approximately 1.2 to 1.5).

Output magnitude (Figure 5b): The paper presents per-Transformer-block output magnitude at the end of training (27 blocks on the x-axis). The baseline exhibits the characteristic PreNorm dilution pattern: output magnitude grows monotonically from approximately 2 (in arbirary units; the y-axis is not explicitly labeled with normalization details) in early blocks to approximately 16 in the deepest blocks — roughly an 8× increase. Block AttnRes transforms this into a bounded, periodic pattern: magnitudes grow within each block but reset to a lower baseline at block boundaries, oscillating between approximately 2 and 8. The peak magnitude in Block AttnRes is roughly half the baseline's peak. The paper attributes this to "selective aggregation at block boundaries reset[ting] the accumulation."

Gradient magnitude (Figure 5c): The baseline shows large gradients concentrated in the earliest layers (approximately 2.5 × 10^-5 for block 0) with a sharp decay to approximately 0.5-1.0 × 10^-5 for most subsequent blocks — a ~3-5× disparity. Block AttnRes produces a "substantially more uniform gradient distribution" with all blocks falling in a narrower band, though the paper does not report the exact variance reduction. The mechanism: "the learnable softmax weights introduce competition among sources for probability mass," so earlier layers cannot receive disproportionately large gradients by default.

These dynamics measurements are the paper's strongest mechanistic validation. The dilution pathology (monotonic magnitude growth) and gradient concentration (early-layer dominance) are both clearly visible in the baseline and clearly mitigated in Block AttnRes. The block-level periodicity in Figure 5b is particularly informative — it confirms that the mechanism works locally within blocks (where standard residuals still cause some growth) and globally across blocks (where attention resets the scale), matching the design intent.

Downstream Benchmarks: Broad Improvements with Largest Gains on Compositional Reasoning

Table 3 reports downstream performance after the full training recipe: 1T tokens of pre-training followed by approximately 400B tokens of mid-training (following the Moonlight [28] annealing recipe), then continued training with 32K-token sequences. Sixteen benchmarks are evaluated across general reasoning, math/code, and Chinese language understanding.

The headline numbers:

Block AttnRes matches or outperforms the baseline on all 16 benchmarks. There is no task where the baseline is better.
Largest gains on multi-step reasoning:
- GPQA-Diamond: 44.4 vs. 36.9 (+7.5 points)
- Minerva Math: 57.1 vs. 53.5 (+3.6 points)
- HumanEval: 62.2 vs. 59.1 (+3.1 points)
- BBH: 78.0 vs. 76.3 (+1.7 points)
- MGSM: 66.1 vs. 64.9 (+1.2 points)
Knowledge-intensive and language understanding tasks show more modest improvements:
- MMLU: 74.6 vs. 73.5 (+1.1 points)
- TriviaQA: 71.8 vs. 69.9 (+1.9 points)
- MMLU-Pro: 52.2 vs. 52.2 (tie)
- HellaSwag: 83.4 vs. 83.2 (+0.2 points)
Chinese benchmarks show solid gains:
- C-Eval: 82.5 vs. 79.6 (+2.9 points)
- CMMLU: 82.9 vs. 82.0 (+0.9 points)

The paper explicitly interprets this pattern: "This pattern is consistent with the hypothesis that improved depth-wise information flow benefits compositional tasks, where later layers can selectively retrieve and build upon earlier representations." The reasoning is that multi-step math problems, graduate-level science questions (GPQA-Diamond), and code generation require coordinating multiple intermediate representations across depth — for example, a later layer computing a final arithmetic step needs to retrieve the results of earlier parsing, variable binding, and intermediate calculations. Knowledge-intensive tasks like MMLU (factual recall) and HellaSwag (commonsense completion) may rely more on information stored in model weights than on coordinating representations across layers, so they benefit less from improved depth-wise routing.

Assessment of the downstream results: The results are consistently positive and the task-type gradient (reasoning > knowledge) is consistent with the mechanistic story. However, several caveats apply:

Single training run. The paper does not report error bars, standard deviations, or multiple training seeds. The +7.5-point gain on GPQA-Diamond is large but could be influenced by training variance, especially given the complex multi-stage training recipe (pre-training → mid-training → context extension).
No FLOP-matched comparison at this scale. The 48B model comparison is at equal training tokens (1.4T), not equal training FLOPs. If Block AttnRes adds even a modest training overhead (the paper reports <4%), the baseline could be trained on slightly more tokens under a FLOPs-matched budget, potentially narrowing the gap. The scaling-law results address this at smaller scales, but not at 48B.
Kimi Linear-specific architecture. The model uses a hybrid KDA/MLA attention mechanism with a specific 3:1 layer ratio and MoE routing. Whether the gains transfer to dense Transformers, different attention variants, or different MoE configurations is untested.
GPQA-Diamond is a small benchmark. The GPQA-Diamond subset is typically a few hundred questions, making the +7.5-point gain potentially noisier than gains on larger benchmarks like MMLU (~14K questions).

Architecture Sweep Under Fixed Compute: AttnRes Favors Deeper, Narrower Models

To understand whether AttnRes reshapes optimal architectural scaling, the paper conducts a controlled capacity reallocation study (Figure 7). Under fixed total FLOPs ( $\approx 6.5 \times 10^{19}$ ) and fixed active parameters ( $\approx 2.3 \times 10^8$ ), the authors enumerate 25 configurations on a 5×5 grid over d_model/L_b (15, 30, 45, 60, 75) and H/L_b (0.3, 0.4, 0.5, 0.6, 0.7), where L_b = L/2 is the number of Transformer blocks and H is the number of attention heads. The per-expert MLP expansion ratio is fixed at d_ff/d_model ≈ 0.45.

Shared pattern: Both baseline (Figure 7a) and AttnRes (Figure 7b) heatmaps show that loss decreases with growing d_model/L_b (fewer but wider layers) and shrinking H/L_b (fewer attention heads per block). Both reach optima at H/L_b ≈ 0.3 — the paper does not sweep lower values.

Key difference in optimum depth-width ratio: The baseline achieves its lowest loss (1.847) at d_model/L_b ≈ 60, while AttnRes shifts the optimum to d_model/L_b ≈ 45, achieving 1.802. Under a fixed parameter budget, a lower d_model/L_b means a deeper, narrower model (more layers with smaller hidden dimension). The paper interprets this as: "AttnRes can exploit additional depth more effectively." The loss gap between methods is substantial and consistent: AttnRes achieves lower loss in every single configuration, with improvements ranging from 0.019 to 0.063.

What this reveals: Standard PreNorm residuals create a depth-scaling bottleneck where additional layers provide diminishing returns because the dilution problem compounds with depth. Deeper models under PreNorm suffer from progressively worse signal-to-noise in the residual stream, creating a preference for wider (and consequently shallower) architectures. AttnRes mitigates this bottleneck, making depth a more productive axis for capacity allocation. The shift from d_model/L_b ≈ 60 to ≈45 is a meaningful change — it corresponds to roughly 33% more layers at fixed parameters.

Caveat: The paper explicitly notes that "this preference for depth does not directly translate to a deployment recommendation, as deeper models generally incur higher inference latency due to their sequential computation [39]." In other words, the sweep reveals where AttnRes provides benefit, but latency constraints in production may still favor wider, shallower models regardless.

Learned Attention Patterns: Locality, Embedding Persistence, and Layer Specialization

Figure 8 visualizes the learned depth-wise attention weights α_{i→l} for a 16-head model (from Table 2) with both Full and Block (N=8) AttnRes. The heatmaps show how each layer (rows) distributes attention over previous sources (columns), separately for pre-attention and pre-MLP layers.

Three key observations:

Preserved locality with selective skip connections. Each layer attends most strongly to its immediate predecessor — the diagonal is consistently the brightest in all four heatmaps, confirming that the primary information pathway remains local. However, selective off-diagonal concentrations emerge: under Full AttnRes, layer 4 attends non-trivially to early sources; under Block AttnRes, layers 15-16 show increased attention to earlier blocks. These are "learned skip connections beyond the standard residual path" — the model has discovered that certain later layers benefit from directly accessing specific early representations without going through all intermediate layers.
The token embedding persists throughout depth. Source 0 (the embedding h_1) retains non-trivial attention weight across nearly all layers, especially in pre-attention layers. This means that even the deepest layers continue to draw directly on the raw token representations, bypassing the entire layered transformation stack. This is consistent with the paper's argument about the value of accessing distant layers — the embedding is the most distant "layer" and would be completely buried under the accumulated residual in a standard PreNorm network, yet AttnRes allows it to remain accessible.
Pre-attention and pre-MLP develop different attention profiles. Pre-MLP inputs show "sharper diagonal reliance on recent representations" (the diagonal is more prominent, off-diagonal weights are sparser), while pre-attention inputs maintain "broader receptive fields" (more distributed attention across sources). This layer-type specialization emerges from the independent pseudo-queries for attention vs. MLP layers: attention layers learn to gather information from across depth (consistent with their role in routing), while MLP layers operate more locally on the representation most recently refined by attention.

Block AttnRes preserves structure but with sharper weights. The block heatmaps (bottom row of Figure 8) show qualitatively the same patterns — diagonal dominance, embedding persistence, layer specialization — but with "sharper, more decisive weight distributions." The block-wise compression acts as implicit regularization that forces the model to make clearer choices about which earlier blocks to attend to, since it cannot finely weight individual layers within a block. The paper frames this as evidence that the block structure "preserv[es] the essential information pathways" while the per-layer granularity of Full AttnRes provides only marginal additional benefit.

Ablation Studies and Robustness Checks

All ablations in Table 4 use the 16-head model (436M activated parameters) from Table 2, with identical hyperparameters and compute budget. The baseline validation loss is 1.766.

Comparison with prior cross-layer access methods.

DenseFormer [36] (input-independent learned scalar weights): achieves 1.767 — essentially tied with the baseline (1.766) and showing no improvement. The paper draws a strong conclusion: this "highlight[s] the importance of input-dependent weighting." Cross-layer access alone, without content-dependent selection, does not help.
mHC [59] (4 parallel streams with learned transition matrices): achieves 1.747 — a meaningful improvement over baseline (+0.019) but worse than Full AttnRes at 1.737 (+0.029) and Block AttnRes at 1.746 (+0.020). The paper positions AttnRes as taking mHC's input dependence "further with explicit content-dependent selection via softmax attention."

Cross-layer access granularity.

Full AttnRes: 1.737 (best overall).
Block AttnRes with S=4: 1.746 — a degradation of only 0.009 from Full, while reducing memory from O(Ld) to O(Nd). This is the core tradeoff validation.
Sliding-window AttnRes (W=1, keep 8 most recent + embedding): 1.764 — barely better than baseline (1.766). The paper interprets this as showing that "selectively accessing distant layers matters more than attending to many nearby ones." The sliding window preserves access to recent layers but cuts off distant layers entirely, while Block AttnRes preserves access to all blocks (including the earliest) but at coarser granularity. The large gap between Block (1.746) and SWA (1.764) isolates the value of distant-layer access.
Block size sweep (Figure 6): S=1 (Full): 1.737; S=2: 1.746; S=4: 1.746; S=8: 1.748; S=16: 1.753; S=32: 1.757; Baseline: 1.766. Loss degrades gracefully, with a plateau at S=2-8 that justifies the choice of N≈8 blocks for practical deployments. The monotonic degradation confirms that finer-grained access is always better, but with strongly diminishing returns.

Attention mechanism components.

Input-dependent query (projected from hidden state): 1.731 — this actually outperforms the default learned query (1.737). The paper acknowledges this improvement but defaults to the learned query for systems reasons: the learned query enables batching in the two-phase inference strategy, which is essential for the <2% inference overhead claim. This is an explicit performance-vs-efficiency tradeoff.
Input-independent mixing (learned scalars per source, no key-query attention): 1.749 — worse than Full AttnRes (1.737) and tied with mHC (1.747). Confirms that content-dependent softmax attention, not just having learnable weights, drives the gain.
Sigmoid instead of softmax: 1.741 — worse than softmax (1.737). The paper attributes this to softmax's competitive normalization: "softmax forces sharper selection among sources" by making the weights for a given layer sum to 1, so increasing one source's weight decreases others. Sigmoid treats each source independently. The 0.004 gap is modest but directionally supportive.
RMSNorm on keys removed (Full): 1.743 — a 0.006 degradation from 1.737. The RMSNorm prevents layers with large-magnitude outputs from dominating the softmax regardless of content. The degradation is slightly larger for Block AttnRes without RMSNorm (1.750 vs. 1.746, a 0.004 gap), which the paper attributes to block representations accumulating more layers and thus developing larger magnitude differences.
Multi-head depth attention (H=16): 1.752 — worse than single-head Block AttnRes (1.746), an unusual negative result. The paper's interpretation: "the optimal depth-wise mixture is largely uniform across channels: when a layer's output is relevant, it is relevant as a whole." This is non-obvious — one might expect different attention heads to benefit from attending to different source layers — but the data suggests that per-channel depth-wise attention is not useful, in contrast to sequence-wise attention where multi-head is essential.

Critical Assessment

Claim 1: "Block AttnRes achieves a 1.25× compute advantage over the baseline"

Evidence: Scaling-law curves (Figure 4) fitted to five model sizes from 194M to 528M activated parameters. At the largest scale, Block AttnRes achieves loss 1.693 vs. baseline 1.714. The 1.25× figure is derived by reading off the fitted baseline curve: the compute required for the baseline to reach 1.693 is approximately 1.25× the compute used by Block AttnRes.

What this demonstrates: The scaling-law experiments provide credible evidence that AttnRes reduces validation loss by a constant factor across the measured compute range, with the gap consistent across model sizes. The downward shift in the scaling curve is clear and monotonic.

What this does not demonstrate: The 1.25× compute multiplier is an extrapolated estimate from fitted power laws over less than one order of magnitude of compute. The fitted exponents are nearly identical across variants (-0.057 to -0.058), which may reflect the limited range rather than a genuine invariance. A 1.25× compute advantage is relatively modest — comparable to other architectural improvements like better normalization or activation functions — and whether it holds at the 48B scale is untested by the scaling laws (which stop at 528M). The largest Block AttnRes model (528M) and the final 48B model are separated by roughly 90× in activated parameters, making extrapolation uncertain.

Missing experiment: Training a matched pair of 48B-scale models at controlled intermediate compute budgets and verifying that the loss gap follows the predicted power-law scaling. The final model section only reports one training run each for baseline and AttnRes at 1.4T tokens, which does not constitute a scaling-law verification at that scale.

Claim 2: "AttnRes mitigates PreNorm dilution, yielding bounded hidden-state magnitudes and more uniform gradient distribution across depth"

Evidence: Figure 5b and 5c show output magnitudes and gradient norms per Transformer block for the 48B models trained on 1T tokens. The baseline shows monotonic growth in output magnitude (roughly 2→16) and gradient concentration in early layers (~3× disparity). Block AttnRes shows bounded, periodic output magnitudes (oscillating 2→8) and more uniform gradients.

What this demonstrates: The measurements directly visualize the dilution pathology and its mitigation, providing mechanistic validation that Block AttnRes works through the hypothesized channel. This is the paper's strongest evidence — it is a clean, interpretable, and causally relevant set of measurements.

What this does not demonstrate: Whether the improved training dynamics causally drive the downstream performance improvements. The paper presents dynamics and downstream results as correlated outcomes of the same intervention but does not establish a causal link — for example, by showing that the magnitude/gradient uniformity mediates the relationship between AttnRes and downstream accuracy, or by showing that alternative methods achieving similar uniformity produce similar downstream gains.

Missing experiment: Comparing the training dynamics of mHC (which also improves over baseline but less than AttnRes) to see whether it produces intermediate levels of magnitude/gradient uniformity. This would strengthen the mechanistic story by showing a gradient of effects rather than a binary comparison.

Claim 3: "AttnRes improves downstream performance across all evaluated tasks, with particularly pronounced gains on multi-step reasoning"

Evidence: Table 3 shows Block AttnRes matching or outperforming the baseline on 16 benchmarks. Gains are largest on GPQA-Diamond (+7.5), Math (+3.6), HumanEval (+3.1); more modest on MMLU (+1.1), HellaSwag (+0.2), and MMLU-Pro (tie).

What this demonstrates: The downstream results are consistently positive, and the task-type gradient (reasoning > knowledge) is consistent with the mechanistic hypothesis that selective depth-wise information flow benefits compositional tasks. The pattern does not appear random — it aligns with the paper's theoretical framing.

What this does not demonstrate: Statistical significance (no error bars, single training run), robustness to training seed (the multi-stage recipe with pre-training, mid-training, and context extension introduces multiple sources of variance), or that the gains are specifically due to AttnRes rather than to any modification that changes training dynamics. The 48B comparison is a single model pair, not an ensemble or a multi-seed evaluation. GPQA-Diamond's +7.5-point gain, while impressive, is on a relatively small benchmark and could be partially driven by variance.

Missing experiment: Evaluating intermediate checkpoints during training to verify that the performance gap emerges gradually and consistently rather than being driven by a particular phase (e.g., only appearing during mid-training). Also: evaluating the mHC baseline at the 48B scale to establish whether AttnRes's improvements over mHC observed at 436M parameters (Table 4) persist at production scale.

Claim 4: "AttnRes shifts the optimal depth-width ratio, favoring deeper models"

Evidence: Architecture sweep (Figure 7) showing the baseline's optimal d_model/L_b ≈ 60 and AttnRes's optimal at ≈ 45, under fixed compute and parameter budgets.

What this demonstrates: Under the specific constraints of this sweep (fixed total FLOPs, fixed active parameters, fixed MLP expansion ratio, MoE architecture), AttnRes prefers relatively deeper, narrower configurations than the baseline. The result is internally consistent and the heatmap visualization is clear.

What this does not demonstrate: Whether this preference holds at larger scales (the sweep uses ~230M active parameters), under different MLP expansion ratios, for dense models, or when latency constraints are considered. The paper explicitly acknowledges the latency caveat, but this limits the practical significance of the finding — the "optimal" architecture from a loss-per-FLOP perspective may be undeployable due to serial computation bottlenecks.

Missing experiment: A similar sweep on a dense architecture (the MoE's conditional computation may interact with depth-width tradeoffs in ways that don't generalize) and at a larger scale. Also: evaluating whether the deeper-optimum finding holds when models are trained closer to convergence rather than at a fixed, relatively small token budget.

Overall Assessment

The experiments collectively support the paper's central narrative: that selective, content-dependent depth-wise attention improves over uniform residual accumulation, with measurable improvements in training dynamics, scaling behavior, and downstream task performance. The evidence is multi-dimensional — scaling laws, training dynamics, downstream benchmarks, attention visualization, and architectural analysis — which makes the case stronger than any single experiment would.

The primary weakness is scale of verification for the scaling-law extrapolations. The 1.25× compute advantage and the depth-width optimum shift are demonstrated at 194M-528M parameters and extrapolated to 48B. The final model results at 48B are consistent with improvement but do not constitute a scaling-law verification. A second weakness is the single-training-run evaluation for the 48B downstream results, which precludes statistical confidence in the specific benchmark gains. A third weakness is the limited comparison with mHC at scale — mHC is the closest competing method, and while AttnRes outperforms it at 436M parameters, the paper provides no evidence that this advantage persists at 48B or that the infrastructure advantages (5.5d vs. 34d per-layer I/O) translate to meaningful wall-clock differences in production.

The paper's strongest contributions are the mechanistic validation (Figure 5b,c — the dilution mitigation is visually unambiguous and directly tied to the proposed mechanism) and the block-structuring finding (Figure 6 — the graceful degradation with block size and the sharp drop with sliding-window access isolate the importance of distant-layer access, not just cross-layer access in general). These findings generalize beyond the specific AttnRes implementation and provide actionable insights for future work on depth-wise information flow.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Remains Unaddressed for AttnRes

The assumption or constraint. The compute-optimal allocation framework in the companion paper on test-time compute scaling depends critically on estimating prompt difficulty before deciding how to spend the inference budget. The method for doing so — generating 2048 samples per prompt and averaging either ground-truth correctness or verifier-predicted scores — is extraordinarily expensive, consuming more compute than the largest inference budgets studied. The authors acknowledged this explicitly in the companion paper's Section 3.2: "estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity."

The consequence. For the AttnRes paper specifically, this limitation manifests in a different but related form: the paper provides no mechanism for adaptively modulating depth-wise attention based on input difficulty. The learned attention weights α_{i→l} are trained to optimize average-case performance across the training distribution, but the paper does not explore whether different prompts benefit from different depth-wise connectivity patterns. For instance, easy single-step questions might benefit from aggressive embedding persistence (always attending strongly to h_1), while complex multi-step reasoning might require more selective retrieval of intermediate representations. Without difficulty-conditioned modulation, AttnRes applies the same depth-wise attention pattern to all inputs — a uniform strategy that the companion paper demonstrated is suboptimal for sequence-level test-time compute allocation. The failure mode is straightforward: a depth-wise attention pattern optimized for average performance may be suboptimal for specific difficulty regimes, analogous to how the companion paper showed that beam search hurts easy problems while helping hard ones.

What evidence exists in the paper. None directly. The learned attention patterns in Figure 8 are averaged over tokens without conditioning on difficulty, and no analysis examines whether the attention distribution shifts for correct versus incorrect predictions, or for simple versus complex problems. The downstream benchmark results (Table 3) show gains on reasoning tasks, but this is average improvement, not evidence of adaptive difficulty-conditioned routing. The paper does not measure whether the benefit of AttnRes varies by question difficulty within a benchmark.

Mitigation status. Not addressed. The paper does not mention difficulty-conditioned depth-wise attention, does not propose mechanisms for modulating α_{i→l} based on input properties, and does not analyze whether the learned weights exhibit any implicit difficulty-dependent variation. This is a conceptual gap: the companion paper established that test-time strategy should be difficulty-adaptive, but AttnRes provides no mechanism for such adaptation in the depth dimension. A natural extension — using the input-dependent query variant (which the paper tested and found improves performance to 1.731 in Table 4) to enable difficulty-adaptive depth routing — is not explored in that context. The paper's decision to default to the learned pseudo-query for systems reasons forecloses this direction without acknowledging the tradeoff.

No Evidence for the Claim That AttnRes Generalizes to Non-MoE, Non-Kimi Architectures

The assumption or constraint. The paper evaluates AttnRes exclusively within the Kimi Linear architecture [69], which has several unusual properties: a hybrid KDA/MLA attention mechanism in a 3:1 ratio, Mixture-of-Experts feed-forward layers (8 of 256 routed experts plus 1 shared), the Muon optimizer [28], and a WSD learning rate schedule. The paper states in Section 5: "Our architecture is identical to Kimi Linear [69]... The only modification is the addition of AttnRes to the residual connections; all other components (model depth, hidden dimensions, expert routing, and MLP structure) remain unchanged." No experiments are conducted on dense Transformer architectures, on standard multi-head attention, on non-MoE feed-forward layers, or with the Adam optimizer.

The consequence. The paper's claims about consistency — "scaling law experiments confirm that the improvement is consistent across model sizes" (Abstract) — are demonstrated only within this specific architectural family. Several interactions between AttnRes and the Kimi Linear architecture could inflate or narrow the observed gains relative to a standard Transformer:

MoE routing: The MoE structure means each token activates only a subset of experts per layer. Layer outputs f_i(h_i) may be more variable across depth in an MoE model (since different experts specialize), potentially making selective cross-layer retrieval more valuable than in a dense model where layer outputs are more homogeneous. Conversely, the sparsity of MoE could make the block compression in Block AttnRes more lossy, since the block sum b_n mixes outputs that were computed by different experts for different tokens — information that might benefit from per-layer access.
KDA/MLA attention: The hybrid attention mechanism with a specific layer ratio (3 KDA : 1 MLA) and MLA's use of NoPE (no positional encodings) [61] creates a particular depth-wise information flow pattern that may interact with AttnRes's cross-layer attention. The paper's finding that pre-attention layers develop broader receptive fields while pre-MLP layers are more local (Figure 8) may depend on this specific attention architecture.
Muon optimizer: The Muon optimizer [28] differs from standard Adam in its update rule and may interact differently with the new parameters (pseudo-queries w_l) that AttnRes introduces. The paper's claim that zero-initialization of queries "prevents training volatility" was validated empirically under Muon — whether the same initialization works under Adam is untested.

The failure mode is straightforward: a practitioner who integrates AttnRes into a standard dense Transformer with Adam might observe different gains, different optimal block counts, different attention patterns, or training instability not seen in the paper. Without any evidence that the mechanism transfers, the paper's claims should be interpreted as conditional on the Kimi Linear architecture.

What evidence exists in the paper. None. Every experiment — the scaling laws (Figure 4, Table 2), the architecture sweep (Figure 7), the ablation study (Table 4), the training dynamics analysis (Figure 5), the downstream benchmarks (Table 3) — uses the Kimi Linear MoE architecture. The paper does not discuss architectural transfer, does not cite prior work using similar mechanisms in dense models (besides the related work in Section 7, which discusses conceptually similar methods but does not run AttnRes on those architectures), and does not claim generalization. The absence is simply the absence of evidence.

Mitigation status. Not addressed. The paper does not flag this as a limitation, does not propose future work on architectural transfer, and does not discuss which aspects of the Kimi Linear architecture might inflate or diminish AttnRes's benefits. A single dense-model experiment (even at small scale) would substantially strengthen the claim of architectural generality. The consistent improvement across the five MoE model sizes in Table 2 establishes within-architecture scaling, not cross-architecture transfer.

The 38% Correct-to-Incorrect Reversion Rate Is Not Resolved by AttnRes

The assumption or constraint. This is a limitation imported from the companion paper's revision model, but it bears on AttnRes because AttnRes does not address the problem of sequential error correction across depth. The companion paper reported that approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect answers in the subsequent step, a direct consequence of training data construction where the model only sees incorrect-to-correct trajectories.

The consequence. For AttnRes specifically, the concern is whether depth-wise attention can propagate errors as effectively as it propagates useful information. If a layer produces an incorrect or misleading output, and subsequent layers learn to attend to it (because it has high dot-product with their pseudo-queries), the error can contaminate later representations through selective attention — a mechanism that is more surgically precise than standard residual accumulation but also more capable of amplifying harmful signals if the attention weights are miscalibrated. The paper's learned attention patterns (Figure 8) show that layers frequently attend to recent predecessors with high weight; if the predecessor produces a flawed representation, AttnRes may route that flawed representation directly into subsequent computation without the dilution that standard residuals would provide (where the error is averaged with all prior outputs). The paper provides no mechanism for layers to recognize and down-weight incorrect or unhelpful source representations — the softmax attention is purely similarity-based, with no "correctness" signal.

What evidence exists in the paper. The paper does not study error propagation through depth-wise attention. No experiment measures whether AttnRes increases or decreases the model's tendency to compound errors across layers, whether incorrect intermediate representations systematically receive different attention weights than correct ones, or whether the model ever learns to down-weight preceding layers whose outputs are unreliable. The training dynamics analysis (Figure 5) shows that AttnRes produces more uniform gradients and bounded hidden-state magnitudes, but this says nothing about error propagation. The downstream improvements on reasoning tasks (Table 3) suggest the mechanism is net beneficial, but this is an average effect — it does not rule out failure modes where AttnRes amplifies specific types of errors.

Mitigation status. Not addressed and not acknowledged. The paper does not discuss error propagation through cross-layer attention, does not propose mechanisms for layers to evaluate the quality or reliability of source representations (e.g., an auxiliary confidence score, uncertainty-based gating, or learned thresholds for attention), and does not analyze whether the model's errors on downstream benchmarks exhibit patterns consistent with error amplification through depth. This is a conceptual gap: selective attention over depth is a powerful tool for routing useful information, but it is equally powerful for routing harmful information, and the paper provides no analysis or mitigation for the latter case.

The Block Count N ≈ 8 Is an Empirical Constant Without Known Scaling Properties

The assumption or constraint. The paper fixes the number of blocks to approximately 8 based on the ablation at the 436M activated parameter scale (Figure 6), where block sizes S=2, 4, and 8 all achieve similar validation loss (~1.746). The final 48B model uses 6 layers per block, producing 9 blocks plus the token embedding for 10 depth-wise sources. The paper presents this as a practical choice: "Empirically, we find that N ≈ 8 recovers most of the benefit across model scales, requiring only eight stored hidden states per token" (Section 3.2). However, no experiment sweeps the block count at different model sizes to verify that N ≈ 8 remains optimal as models scale.

The consequence. Two distinct failure modes arise from treating N as a fixed constant:

Suboptimal depth granularity at larger scales. The 48B model has 54 layers; the 528M model (largest in the scaling laws) has 17 Transformer blocks = 34 layers. If the "natural scale" of depth-wise information retrieval grows with total depth — for instance, if a model with 100 layers benefits from attending at finer granularity than one with 34 layers — then N ≈ 8 may be too coarse for larger models, leaving performance on the table. The paper's own scaling-law observation that the gap between Full and Block AttnRes narrows with scale (0.010 at 194M, 0.001 at 528M, per Table 2) could be interpreted as evidence that block approximation becomes less costly at larger scales, or alternatively that larger models saturate the benefit of cross-layer attention more quickly. Without a controlled sweep of N at multiple scales, the interpretation is ambiguous.
Over-provisioning at smaller scales. If the optimal N is proportional to total depth, smaller models (fewer layers) would need fewer blocks. A model with only 12 layers (the smallest in Table 2) using N ≈ 8 would have block size S ≈ 1.5 — nearly Full AttnRes — paying near-full memory and communication costs for what might be marginal gain over a coarser blocking. The graceful degradation in Figure 6 suggests this is not catastrophic (S=8 still achieves 1.748 vs. S=1's 1.737), but it means the memory-to-benefit ratio degrades at smaller model scales.

What evidence exists in the paper. The block size sweep (Figure 6) is conducted at a single model size (16 layers, corresponding to the 436M activated parameter model in Table 2). No sweep is reported for the 194M, 241M, 296M, or 528M models. The paper states that "in practice, we fix the number of blocks to ≈ 8 for infrastructure efficiency" (Section 5.3), making the choice explicitly about systems constraints rather than empirical optimality. The scaling-law results (Figure 4, Table 2) use Block AttnRes with approximately 8 blocks at all five model sizes, so the 1.25× compute advantage is measured under this fixed choice — the curves would likely shift if N were optimized per model size.

Mitigation status. Partially acknowledged as a future direction. The paper states that "as future hardware alleviates memory capacity constraints, adopting finer-grained block sizes or Full AttnRes represents a natural pathway to further improve performance" (Section 5.3). However, this frames the limitation as purely hardware-driven — the paper does not acknowledge the possibility that optimal N might grow with model depth for algorithmic reasons independent of hardware. No recommendation is provided for practitioners choosing N for a new architecture: should they replicate the Figure 6 sweep at their target model size, or can they safely default to N ≈ 8?

The Inference Latency Overhead Claim (<2%) Is Not Benchmarked in the Paper

The assumption or constraint. The paper claims in the Abstract and Section 4.2 that "the end-to-end inference latency overhead is less than 2% on typical inference workloads," but provides no experimental measurement, no description of the inference setup (batch size, sequence length, hardware, framework), and no ablation showing how overhead varies with these parameters. The claim is supported only by the analytical I/O cost model in Table 1 and Appendix B, which counts memory reads/writes but does not account for the actual wall-clock impact of the additional computation (softmax over N sources, RMSNorm on keys, the online-softmax merge, and the increased arithmetic in the attention-weighted sum).

The consequence. The <2% claim is the paper's primary argument for the practicality of AttnRes at inference — without it, the mechanism would be judged primarily on its training-time benefits, which are meaningful but harder to justify if inference latency degrades significantly. The analytical I/O model may be optimistic for several reasons:

The two-phase strategy assumes perfect batching of inter-block queries. In practice, Phase 1 requires reading N block representations (for the final model, 10 sources of dimension d) and computing a batched matrix multiplication QK^T where Q is S × d (6 layers × d) and K is n × d. This batched operation has a specific GPU kernel launch overhead and memory access pattern that the analytical model treats as a single amortized read, but actual throughput depends on whether the operation is compute-bound or memory-bound at the given dimensions.
The online-softmax merge adds element-wise operations (exponentiations, multiplications, additions, divisions) that the paper claims "naturally admits kernel fusion with surrounding operations." Whether such fusion is implemented in practice depends on the inference framework; the paper does not describe a specific implementation or provide kernel-level benchmarks.
The memory-efficient prefilling scheme shards block representations along the sequence dimension, which introduces an all-reduce in the critical path during Phase 2 reconstruction. The paper claims this "naturally admits kernel fusion with operations like RMSNorm," but the latency of the all-reduce itself — which scales with the number of tensor-parallel devices — is not accounted for in the <2% figure.

What evidence exists in the paper. None. There is no inference benchmark, no latency measurement, and no description of an end-to-end inference implementation. The <2% claim appears in the Abstract, the end of Section 4.2, and implicitly in Table 1 (which compares I/O costs across methods but translates them to latency only through the "Total I/O" column). The paper does not cite a specific inference framework, does not report throughput numbers (tokens/second), and does not describe how the two-phase strategy was implemented or profiled.

Mitigation status. Not mitigated by experimental evidence. The analytical model in Appendix B is a useful guide to the asymptotic I/O scaling but is not a substitute for measured latency. The paper treats the <2% claim as a consequence of the memory access analysis, but memory access count is only one component of inference latency — compute utilization, kernel launch overhead, synchronization barriers, and communication costs all contribute, and none are measured. A single table of measured latencies for standard residual inference vs. Block AttnRes inference, across a range of batch sizes and sequence lengths, would substantiate the claim. Without it, practitioners should treat <2% as an upper bound on the I/O contribution to latency, not on total latency.

The Structured-Matrix Analysis Reveals No Practical Constraint on M That AttnRes Violates, But Also No Guarantee of Well-Behaved Depth-Wise Attention

The assumption or constraint. The paper's unifying analytical framework (Section 6.2) casts residual variants as different structures imposed on a depth mixing matrix M ∈ R^{L×L}. AttnRes is characterized as producing a dense, full-rank M — the only constraint is row-wise normalization (sum to 1). This is presented as a strength: removing structural constraints that prior methods (Highway, mHC) imposed. However, the paper does not analyze what properties M should ideally satisfy — beyond normalization — to ensure stable, well-behaved depth-wise information flow. Dense, unconstrained matrices can exhibit pathologies (e.g., oscillatory dynamics, rank collapse, attention sink dominance) that low-rank or semiseparable structures inherently avoid.

The consequence. The paper's learned attention patterns (Figure 8) reveal two phenomena that may signal suboptimality in the learned M:

Embedding persistence as a potential attention sink. Source 0 (the token embedding h_1) retains non-trivial weight throughout nearly all layers. While the paper presents this as evidence that distant-layer access is valuable, it may also indicate a failure mode analogous to attention sinks in sequence models [57]: the embedding is always available (unlike later layers which must be computed sequentially), always has moderate norm (since it doesn't accumulate residuals), and may serve as a "default" source that the model over-relies on because attending to it is computationally cheap and never harmful, even when more informative intermediate representations would be preferable. The paper does not analyze whether the embedding's persistent weight comes at the expense of under-attending to task-relevant intermediate layers.
Diagonal dominance as potentially redundant computation. The diagonal is consistently the brightest in all heatmaps, indicating that each layer attends most strongly to its immediate predecessor. This means the model is, to a significant degree, recreating standard residual connectivity through learned attention — paying the cost of computing softmax over depth while largely replicating the uniform-weight behavior it is meant to replace. The off-diagonal concentrations that represent genuinely novel skip connections are sparse and relatively weak compared to the diagonal. The paper does not quantify what fraction of attention mass falls on the diagonal versus off-diagonal, or whether the off-diagonal mass correlates with task performance.

Without constraints on M — such as a sparsity penalty, an entropy regularizer encouraging more distributed attention, or an initialization that biases the model away from trivial diagonal solutions — AttnRes may converge to solutions that are close to standard residuals with minor perturbations, underutilizing the expressive capacity of the dense mixing matrix and yielding gains that are incremental rather than transformative.

What evidence exists in the paper. The learned attention visualizations (Figure 8) are the primary window into M's structure, but they are presented as qualitative evidence of desirable properties (locality, specialization, embedding persistence) rather than analyzed for potential pathologies. The paper does not measure attention entropy, diagonal concentration, or the effective rank of the learned M. The ablation showing that sigmoid underperforms softmax (Table 4) provides indirect evidence that competitive normalization matters, but does not establish that the learned softmax weights are well-calibrated or optimally sparse.

Mitigation status. Not addressed. The paper treats M's density as inherently desirable (the linear-to-softmax transition) without considering whether an unregularized dense M might learn suboptimal structures. No regularization terms are applied to the attention weights (e.g., entropy bonus, sparsity penalty, diversity objective), and no analysis examines whether the learned weights change over the course of training (e.g., starting uniform and gradually specializing, or oscillating). The depth-wise attention sink phenomenon — well-documented in sequence attention — is not discussed, despite the embedding persistence in Figure 8 being a plausible instance of it. Future work could explore constrained variants of M (e.g., enforcing that α_{l-1 → l} has a minimum weight to preserve the gradient highway, or adding a diversity regularizer to prevent over-reliance on any single source) that might yield more robust learned connectivity patterns.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper achieves something rare in architecture research: it identifies a bottleneck that the field had largely stopped thinking about — the fixed uniform accumulation in residual connections — and proposes a solution that is simultaneously principled, practical, and empirically validated at production scale. The impact is less a paradigm shift than a reframing that reopens a design space the field had prematurely closed.

Since the widespread adoption of PreNorm Transformers, the residual connection has been treated as solved infrastructure. The standard update h_l = h_{l-1} + f_{l-1}(h_{l-1}) is so universally deployed that it rarely appears in architecture ablation studies — it is the water the fish swims in. This paper argues convincingly that this complacency is unwarranted: residual connections are not just gradient highways but information aggregation mechanisms, and the uniform-weight formulation systematically dilutes early-layer information as depth increases. By recasting the residual as a depth-wise information retrieval problem and showing that learned, input-dependent softmax attention over preceding layers improves training dynamics, scaling behavior, and downstream performance, the paper makes the residual connection itself a design dimension worthy of optimization.

The most durable contribution is the depth mixing matrix perspective (Section 6.2). This is not merely taxonomic — it provides an operational language for reasoning about residual architectures that was previously unavailable. Prior work on Highway networks, Hyper-Connections, and DenseFormer could be described and compared, but there was no shared formalism that made visible what each method constrained and what it allowed. By showing that standard residuals impose rank-1 structure on M, Highway imposes 1-semiseparable structure, mHC imposes m-semiseparable structure, and AttnRes removes rank constraints entirely (full rank L, only row-wise normalization), the paper gives the field a complexity hierarchy for depth-wise information routing. This changes how researchers should think about architectural design: rather than asking "should we use residual connections or something else?", the question becomes "what rank should the depth mixing matrix M have, and what computational mechanism should produce it?" This is a more precise and productive framing.

The paper also resolves a subtle contradiction in prior work on scaling depth. On one hand, the pruning literature [11] showed that many layers can be removed with minimal loss — suggesting depth is underutilized. On the other hand, scaling laws consistently show that deeper models perform better, all else equal. AttnRes reconciles these findings: layers are underutilized because the uniform residual accumulation buries their individual contributions, not because the computation they perform is inherently wasteful. When layers can selectively retrieve earlier representations, depth becomes more productive — the architecture sweep (Figure 7) shows that AttnRes shifts the optimal depth-width ratio toward deeper, narrower models, precisely because depth is now a more effective axis for capacity allocation. This implies that the pruning result is not a fundamental property of deep networks but an artifact of poor depth-wise information routing.

The paper also redirects attention in the architecture design community. Prior work on improving information flow across layers focused on either (a) refining the additive recurrence with gates or scaling factors (Highway, DeepNorm, ReZero, SiameseNorm) or (b) expanding the recurrent state to multiple parallel streams (mHC, DDL). AttnRes shows that both approaches remain within the recurrence paradigm and that breaking out of recurrence entirely — providing direct attention over individual layer outputs — yields larger gains with simpler mechanisms. The ablation (Table 4) is the key evidence: mHC (4 streams, learned mixing matrices) achieves 1.747, while Block AttnRes (single query vector, softmax attention, block-structured) achieves 1.746 with far lower inference I/O (5.5d vs. 34d per layer). This suggests that the field's investment in increasingly sophisticated recurrent state machinery may be solving the wrong problem — the bottleneck is not the expressivity of the recurrent update but the compression of all prior information into a single (or m-stream) state in the first place. AttnRes sidesteps the compression entirely, and simpler direct access outperforms more complex recurrence.

Research directions that become more attractive:

Efficient approximations to full-rank depth mixing. The paper shows that full-rank M (Full AttnRes) is best, but its O(Ld) memory cost motivates approximation. Block AttnRes is one such approximation (rank N to N+S). Future work could explore low-rank-plus-sparse decompositions, learned hashing over depth, or dynamic block formation that adapts the grouping based on input content rather than using fixed contiguous blocks.
Depth-wise attention analysis. The paper's visualization of learned attention patterns (Figure 8) opens the door to studying depth-wise connectivity as a first-class object of analysis, analogous to how sequence-wise attention patterns are studied. Questions like "which layers serve as depth-wise attention sinks?" and "do attention patterns change during the course of forward computation (early tokens vs. late tokens in the sequence)?" become investigable.
Joint optimization of depth-wise and sequence-wise information flow. AttnRes operates purely over depth; sequence-wise attention operates purely over time. The dream of architectures that unify both dimensions — attending over a 2D grid of (layer, token) pairs — is made more concrete by this work, though the computational cost remains prohibitive.

Research directions that become less attractive:

Incremental improvements to single-state recurrence. If direct cross-layer access with a single learned query vector (AttnRes) outperforms sophisticated multi-stream recurrence (mHC) at comparable or lower cost, then further refinements to the gating mechanisms within additive recurrence (new normalization schemes, learned per-channel multipliers, adaptive scaling) are unlikely to close the gap. The paper does not kill this line of work — mHC still achieves meaningful gains over baseline — but it shifts the frontier, making recurrence-based methods a second-best solution when cross-layer attention is infeasible for hardware reasons.
Input-independent cross-layer weighting. The DenseFormer ablation (1.767, tied with baseline at 1.766) is a strong negative result: providing access to all previous layer outputs with learned but static scalar weights yields zero improvement. This suggests that the benefit comes specifically from content-dependent selection, and that further work on static cross-layer connectivity patterns is unlikely to be productive.

Follow-Up Research This Work Enables

1. Scaling the block count N with model depth. The paper fixes N ≈ 8 based on Figure 6 at a single model size (16 layers). A natural follow-up is to sweep N across model depths: for models with 12, 24, 48, 96, and 192 layers, measure validation loss at block sizes S = 1, 2, 4, 8, 16 and identify the optimal N as a function of L. The hypothesis to test is whether the effective granularity of depth-wise information retrieval scales with total depth (e.g., optimal N ∝ log L, suggesting that deeper models benefit from finer-grained access) or whether it saturates (e.g., N ≈ 8 remains optimal regardless of L, suggesting a fundamental scale of representational locality). The paper's scaling-law observation that the Full-Block gap narrows from 0.010 to 0.001 as model size increases (Table 2) hints at saturation, but the maximum depth tested is only 34 layers. A study spanning an order of magnitude in depth would resolve whether the block approximation becomes more or less costly at scale. The experiment requires only controlled pre-training runs at each (L, N) pair with matched compute budgets; no downstream evaluation is needed to answer the structural question.

2. AttnRes with input-dependent queries + difficulty-adaptive depth routing. The paper shows that input-dependent queries (projected from the hidden state) improve performance from 1.737 to 1.731 (Table 4) but defaults to learned queries for systems reasons (the two-phase batching strategy requires queries to be available before forward computation begins). A compelling extension is to explore hybrid query schemes that preserve most of the batching benefit while enabling input-dependent modulation. For example: use a learned base query w_l plus a low-rank input-dependent offset (e.g., w_l + A · RMSNorm(h_l) where A is a d × r projection with r ≪ d). The offset can be computed from the layer's current hidden state — which is available by the time the query is needed for intra-block attention — while the base query enables batched inter-block attention in Phase 1. The experiment would measure whether this hybrid scheme recovers the input-dependent query's performance gain (1.731) while maintaining the <2% inference latency overhead. A deeper follow-up would analyze whether the input-dependent component of the query learns to modulate attention based on properties of the current token — for instance, attending more to earlier layers when the token requires retrieving factual knowledge versus attending more locally when performing syntactic composition. This connects to the companion paper's finding that optimal strategies are difficulty-dependent: if AttnRes can learn to route information differently for simple versus complex inputs, it would provide a mechanism for difficulty-adaptive depth-wise computation.

3. Verifier-guided depth-wise attention for self-correction. The companion paper identified that sequential revisions from a fine-tuned model can correct errors but suffer from a 38% correct-to-incorrect reversion rate. AttnRes provides a potential mitigation: if the model can learn to attend to its own intermediate reasoning steps selectively, it could learn to ignore prior steps that contain errors while amplifying steps that are on the right track. A specific experiment: train a revision model with AttnRes where the depth-wise queries can attend over the chain of previous revisions (each revision is a block or a set of layer outputs). The PRM from the companion paper could provide per-step correctness signals that supervise the attention weights — the model learns to assign high attention to steps that the PRM scores as likely correct and low attention to steps scored as likely incorrect. This would transform AttnRes from an architectural mechanism into a learned error-correction router that selectively retrieves useful intermediate computations while suppressing erroneous ones. The evaluation would measure whether the correct-to-incorrect reversion rate (the companion paper's 38% figure) is reduced, and whether the revision chain length can be extended beyond the current practical limit (roughly 4-8 steps before performance degrades) because the model can now selectively ignore unhelpful revisions.

4. Depth-wise attention as a diagnostic for layer importance and pruning. The paper's learned attention patterns (Figure 8) reveal which layers serve as persistent sources for many subsequent layers (e.g., the embedding and certain attention layers with off-diagonal weight concentrations). This suggests a new pruning criterion: rather than pruning layers based on their output magnitude or the loss change upon removal, prune layers that receive low aggregate attention weight from downstream layers. Formally, define the importance of layer i as I_i = Σ_{l > i} α_{i→l} (the total attention mass that subsequent layers allocate to layer i, summed over tokens). Layers with low I_i are, by the model's own learned routing, less useful for downstream computation. An experiment would compare importance-based pruning against magnitude-based pruning and loss-based pruning at various sparsity levels on a trained AttnRes model. The hypothesis is that importance-based pruning more accurately identifies expendable layers because it reflects the model's revealed preference for which representations it actually uses, rather than an external heuristic. If this holds, it would provide a new tool for model compression that is grounded in the model's learned information-routing strategy rather than post-hoc analysis.

5. AttnRes on dense architectures with standard optimization. The paper's experiments are exclusively on the Kimi Linear MoE architecture with the Muon optimizer. A necessary stress-test is to replicate the core scaling-law result (Figure 4) on a standard dense Transformer: GPT-2 or LLaMA-style architecture, AdamW optimizer, standard multi-head attention (no KDA/MLA hybrid), and the same PreNorm baseline. The experiment should sweep at least three model sizes (e.g., 125M, 350M, 1.3B parameters) with matched training data and compute budgets. The key question is whether the 1.25× compute advantage transfers, or whether the MoE's expert sparsity interacts with AttnRes in a way that inflates the benefit. A null result (AttnRes provides no gain over baseline in dense models) would not invalidate the paper's contributions — the mechanism was designed and evaluated for the Kimi Linear architecture — but it would sharply constrain the scope of applicability. A positive result would substantially strengthen the claim that AttnRes is a general architectural improvement. The experiment should also ablate the optimizer: train one dense-model pair with Muon and one with AdamW, to test whether the zero-initialization strategy for pseudo-queries (which the paper validated empirically under Muon) transfers to standard optimization without causing training instability.

6. Characterizing when depth-wise attention collapses to uniform weighting. The paper shows that Full AttnRes outperforms the baseline, but the learned attention patterns (Figure 8) reveal strong diagonal dominance — in many layers, the model allocates most of its attention mass to the immediately preceding layer, only modestly deviating from standard residual behavior. A detailed follow-up would measure attention entropy per layer (how uniformly the attention mass is distributed across sources) and correlate it with performance. The hypothesis: layers with higher attention entropy (more distributed attention, making fuller use of cross-layer access) contribute disproportionately to the model's performance, while layers that collapse to near-diagonal attention are essentially standard residual layers that could be replaced with a fixed residual connection without loss. This analysis could identify that only a subset of layers meaningfully exploit AttnRes's expressivity, which would motivate selective deployment of AttnRes — only applying the mechanism at layers where the learned attention distribution is sufficiently non-uniform — further reducing overhead. The experiment requires only inference-time analysis of the trained model's attention weights, no retraining, and could be conducted on the publicly released weights (if the authors release them). The practical payoff is a criterion for when AttnRes is worth its overhead, potentially enabling mixed architectures where most layers use standard residuals and only a few "routing" layers use AttnRes.

Practical Applications and Downstream Use Cases

1. Cost-efficient pre-training of deep MoE models. The most direct application is integrating Block AttnRes into the pre-training recipe for large MoE language models. The paper demonstrates that at the 528M-parameter scale, Block AttnRes achieves a 1.25× compute multiplier — meaning a team training a 100B-parameter MoE model could either reduce their training budget by ~20% for equivalent quality, or achieve better final loss at the same budget. For a training run costing millions of dollars in compute, a 1.25× efficiency gain is economically significant. The integration cost is low: AttnRes adds one d-dimensional vector per layer (negligible parameter overhead), requires zero-initialization to start training stably, and the paper's infrastructure optimizations (cross-stage caching, two-phase computation) are described in sufficient detail for an engineering team to implement. The primary risk is the unverified transfer to non-Muon optimizers and non-Kimi architectures, but teams already using MoE architectures similar to DeepSeek-V3 [9] or Moonlight [28] — which share the interleaved attention-MoE structure — can adopt Block AttnRes with minimal adaptation. The measured <4% training overhead means the net compute savings are roughly 1.25× - 1.04× ≈ 1.20×, still substantial.

2. Improving long-context reasoning in deployed models. The downstream benchmark pattern (Table 3) shows that AttnRes provides the largest gains on multi-step reasoning tasks: GPQA-Diamond (+7.5 points), Minerva Math (+3.6), HumanEval (+3.1). These are precisely the tasks that matter for applications like code generation assistants, mathematical tutoring systems, and scientific Q&A — domains where users submit complex queries requiring coordinated multi-step inference. The paper's context extension to 32K tokens (Section 5.2) and the memory-efficient prefilling scheme (sharding block representations across tensor-parallel devices, reducing the 128K-context memory footprint from ~15 GB to ~1.9 GB per device) demonstrate that AttnRes scales to long sequences without prohibitive memory costs. A deployment team could fine-tune an existing AttnRes-enabled base model on their domain-specific reasoning data and expect disproportionate gains on the hardest queries in their distribution — the "long tail" of complex reasoning tasks that standard models struggle with. The <2% inference latency overhead (claimed, though unbenchmarked — see Section 6) means this comes at minimal serving cost.

3. Model compression via importance-guided pruning. The learned attention weights in AttnRes provide a built-in layer importance estimator at no additional cost. After training, the aggregate attention weight received by each layer (Σ_{l > i} α_{i→l}, averaged over a calibration dataset) directly quantifies how much downstream layers rely on that layer's output. Layers with low aggregate attention are candidates for pruning or distillation, since the model's own routing mechanism deems them less useful. This is more principled than magnitude-based pruning (which uses weight norms as a proxy for importance) and cheaper than loss-based pruning (which requires measuring the loss impact of removing each layer individually, costing O(L) forward passes). A practitioner could take a trained AttnRes model, compute layer importance scores from the attention weights in a single forward pass over a calibration set, prune the lowest-importance layers, and fine-tune the resulting thinner model — a compression pipeline that leverages the architecture's own learned connectivity to decide what to remove.

4. Depth-wise attention analysis tools for model interpretability. The paper's visualization of learned attention patterns across layers (Figure 8) suggests a new modality for interpretability research: depth-wise attention probing. Just as sequence-wise attention patterns are studied to understand which tokens the model attends to when generating each output token, depth-wise attention patterns reveal which representational stages the model relies on at each computational step. A concrete use case: when a model produces an incorrect answer on a math problem, an analyst could examine whether the model attended to the correct intermediate computation (e.g., layer 7 where the equation is set up) or bypassed it in favor of a shortcut (e.g., directly attending from the embedding to the final output layer). This provides a new diagnostic for reasoning failures: is the error due to a missing capability (the correct intermediate representation was never computed) or a routing failure (the correct representation was computed but subsequent layers failed to attend to it)? An interpretability team could build a visualization tool that, for a given input, shows which layers each output step attends to and highlights mismatches between expected and actual attention patterns — analogous to saliency maps but for the depth dimension.

When to Prefer This Method

The paper provides explicit comparisons against named alternatives in Table 4 and throughout Section 6.2, enabling a concrete decision framework:

Prefer Block AttnRes over standard PreNorm residuals when: training deep MoE Transformers (50+ layers) where PreNorm dilution is measurably degrading effective depth utilization, and the deployment scenario can tolerate <4% training overhead and ~2% inference overhead. The mechanism is particularly well-motivated for models targeting compositional reasoning tasks (math, code, multi-step science problems) where the downstream gains are largest (+3–7 points on benchmarks like GPQA-Diamond, Minerva Math, HumanEval). The block count should be set to N ≈ 8 based on the paper's empirical sweep (Figure 6); practitioners with different model depths should replicate the block-size sweep at their target depth to verify this choice.
Prefer Full AttnRes over Block AttnRes when: memory and communication constraints are not binding (e.g., small-scale research models, single-GPU training without pipeline parallelism, or future hardware with substantially increased inter-GPU bandwidth). Full AttnRes achieves 0.009–0.010 lower validation loss than Block AttnRes (Table 2, 436M scale: 1.737 vs. 1.746), a small but consistent gap. The paper explicitly states that finer-grained access "represents a natural pathway to further improve performance" as hardware improves.
Prefer AttnRes (either variant) over (m)HC when: inference-time memory I/O is a binding constraint. Table 1 shows Block AttnRes requires 5.5d memory accesses per layer versus 34d for mHC with m=4 streams — a 6× reduction — while achieving comparable or better validation loss (Block AttnRes 1.746 vs. mHC 1.747 at 436M scale in Table 4). The gap widens at deployment scale where per-token I/O directly impacts serving cost and latency. The tradeoff is that mHC is more architecturally mature (validated across multiple model families and tasks) while AttnRes is validated only on Kimi Linear MoE models.
Do not prefer AttnRes over standard residuals when: the model is shallow (fewer than ~12 layers), where the dilution problem is less severe and the block structure provides limited benefit (the block size sweep in Figure 6 shows S=32 achieves 1.757 vs. baseline 1.766 — a small gap for very coarse blocking that would correspond to N=1 or 2 blocks in a shallow model). Additionally, in latency-critical deployments where even the <2% claimed overhead must be validated on the specific hardware and framework — the paper's latency claim is derived from an analytical I/O model, not measured benchmarks, and practitioners should profile their own inference stack before adopting.

Dependency tree

Papers this one builds on, walking backwards through the same domain and shared topics.

Kimi Linear: An Expressive, Efficient Attention Architecture2025· architecture← LLM-judged similar
- The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry2024· llm-systems← LLM-judged similar
- Scaling TransNormer to 175 Billion Parameters2023· llm-systems← LLM-judged similar
- Scaling Latent Reasoning via Looped Language Models2025· architecture← in-domain predecessor
UniWeTok: An Unified Binary Tokenizer with Codebook Size $\mathit{2^{128}}$ for Unified Multimodal Large Language Model2026· architecture← in-domain predecessor
- AToken: A Unified Tokenizer for Vision2025· vision← LLM-judged similar
- MiniCPM-SALA: Hybridizing Sparse and Linear Attention for Efficient Long-Context Modeling2026· architecture← in-domain predecessor
- Hunyuan3D-Omni: A Unified Framework for Controllable Generation of 3D Assets2025· architecture← similar topic
- Dynamic Chunk Convolution for Unified Streaming and Non-Streaming Conformer ASR2023· architecture← similar topic
- Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems2023· architecture← similar topic

What it influenced

Later papers in the same chain. Direction reversed.

Mamba-3: Improved Sequence Modeling using State Space Principles2026· architecture→ in-domain predecessor
- Mixture-of-Depths Attention2026· architecture→ in-domain predecessor
- Hungry Hungry Hippos: Towards Language Modeling with State Space Models2022· architecture→ similar topic
- Codestral Mamba2023· architecture→ similar topic
- Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling2024· architecture→ similar topic
Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networks2020· architecture→ similar topic
- Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet2021· architecture→ in-domain predecessor
- Fraud Detection: Using Relational Graph Learning to Detect Collusion2017· architecture→ similar topic
- How Powerful Are Graph Neural Networks?2018· architecture→ similar topic
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers2023· architecture→ similar topic
FLEX ATTENTION: A PROGRAMMING MODEL FOR GENERATING OPTIMIZED ATTENTION KERNELS2024· architecture→ similar topic
- Multimodal Latent Language Modeling with Next-Token Diffusion2024· architecture→ in-domain predecessor
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention2025· architecture→ similar topic