ArXiv: 2603.15619

🎯 Pitch

Deep Transformers suffer a silent failure mode: informative features from early layers are progressively smeared away by residual updates, yet existing fixes explode parameter counts. MoDA solves this by letting each attention head directly retrieve key–value memories from all preceding layers inside a single softmax, recovering prior signals with negligible overhead (3.7% FLOPs, 97.3% of FlashAttention-2 speed) and boosting downstream performance by over 2 percentage points.

1. Executive Summary

This paper introduces mixture-of-depths attention (MoDA), a unified attention mechanism that allows each attention head to jointly attend to standard sequence KV pairs of the current layer and depth KV pairs from all preceding layers (fusing historical cross-layer information into a single softmax operator), addressing the information dilution problem where informative features formed in shallow layers degrade through repeated residual updates in deep Transformers. The authors validate MoDA on decoder-only language models trained with the OLMo2 recipe at 700M and 1.5B scales over 400B tokens, and develop a hardware-efficient fused kernel with chunk-aware depth-KV layout and group-aware indexing that reaches 97.3% of FlashAttention-2 efficiency at 64K sequence length. At 1.5B parameters, MoDA improves average perplexity by 0.2 across 10 validation benchmarks and increases average downstream performance by 2.11% on 10 tasks with only 3.7% FLOPs overhead, while also showing that combining MoDA with post-norm yields better performance than pre-norm in deeper models — establishing that explicit depth-aware retrieval is a practical primitive for scaling Transformer depth, with the strongest gains emerging when FFN-side depth KV projections are added and when the model already has sufficient representational capacity for the task.

2. Context and Motivation

The Core Problem: Deep Transformers Lose Information Through Their Depth

The fundamental problem this paper tackles is that as Transformer models get deeper, the quality of information flowing through the network degrades — a phenomenon the authors call information dilution. While scaling model depth is theoretically appealing because deeper stacks can support richer hierarchical computation and have been a key driver of performance improvements in other domains (computer vision, for example), modern LLMs "often fail to convert additional layers into proportional benefits" (Section 1). The gains from adding layers diminish, and in some cases, making the model deeper can even hurt optimization stability.

Why does this happen? The paper identifies the root cause in the residual connection mechanism that is the default method for stacking Transformer blocks. In the standard ResNet-style residual architecture, each layer reads the output of the previous layer (the "read" step is identity), performs some computation (the "operate" step — attention or FFN), and then writes its result back by addition to the incoming stream (the "write" step is add). This creates a single hidden-state trajectory where all depth history is continuously compressed into one fixed-size tensor X_l ∈ R^{T×D}. The problem, as the paper states in Section 2.2, is that "the depth stream is continuously compressed into a fixed-size tensor via repeated superposition, which dilutes salient features and leads to signal degradation."

Think of this as an information bottleneck along the depth dimension: features that were clearly formed in shallow layers — say, a grammatical structure detected at layer 3, or a named entity identified at layer 5 — get progressively overwritten and muddied by dozens of subsequent additive updates. By the time the model reaches layer 48 or layer 64, those once-sharp features are still technically there in the residual sum, but they're diluted by the accumulation of everything that came after. The model must expend capacity to "recover" them from the noisy superposition, and often fails to do so effectively.

This is not merely an optimization problem (vanishing/exploding gradients, which residual connections already largely solved). It is a representational problem: the architecture's fixed connectivity pattern — identity read, additive write — imposes a bottleneck that limits how effectively deep networks can utilize their early-layer computations. The paper frames this through a specific diagnostic lens in Section 1: "informative features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers."

Why This Matters: Depth Scaling Is the Under-Exploited Dimension

The paper situates this problem within a broader narrative about LLM scaling. Recent progress has been driven by scaling along four major dimensions (Section 1):

  • Context length — making models process longer sequences (via methods like sparse attention, Transformer-XL-style recurrence, etc.)
  • Training data — training on more tokens (the "Chinchilla scaling" paradigm)
  • Model width — increasing the hidden dimension and number of attention heads
  • Model depth — stacking more Transformer layers

The authors argue that scaling is "often realized more through data, context, and especially width, whose optimization behavior and system efficiency are generally easier to realize at scale," while "depth, by contrast, remains comparatively under-exploited despite its strong representational appeal" (Section 1). This is a specific and consequential claim: the community has been systematically under-investing in depth scaling not because depth is inherently less valuable, but because our current architectures can't effectively convert additional depth into proportional performance gains.

The practical implications are significant. If depth scaling were as effective as width scaling, we could build models that are deeper and narrower — potentially achieving better representational capacity at lower total parameter counts, or better throughput characteristics. There are theoretical arguments for depth's appeal: deeper networks can in principle learn more compositional representations, with each layer performing a specific transformation that builds on earlier ones. But these theoretical benefits remain unrealized because of the information dilution bottleneck.

The problem also has a self-improvement dimension. In the current paradigm, when we want to make a model better, we often train a larger one from scratch. If depth could be scaled more efficiently, we might instead add layers to existing models (with proper architectural support) and recover more of the investment in pretraining. The paper's framing — "how can a model scale depth while maintaining optimization stability and preventing information dilution?" (Section 1) — speaks directly to this desire for more efficient depth utilization.

Prior Approaches and Their Limitations

The paper identifies three broad families of prior attempts to address the depth information flow problem, each with specific shortcomings:

1. Standard Residual Connections (ResNet-style). This is the baseline: each layer's output is added to its input, creating shortcut paths for gradient flow. While residual connections solved the vanishing gradient problem that plagued early deep networks (enabling the training of 100+ layer models), the paper argues they leave the information dilution problem largely unresolved. The additive "write" operation continuously compresses all depth history into one tensor. As the paper puts it in Section 2.2: "this formulation alleviates vanishing gradients and enables training deep networks. However, the depth stream is continuously compressed into a fixed-size tensor X_l ∈ R^{T×D} via repeated superposition, which dilutes salient features and leads to signal degradation."

This is a subtle but crucial distinction that the "read, operate, write" framing in Section 2.2 makes clear. The "read" step in depth residual is identity — you always read the most recent state, which is the accumulated sum of all previous operations. The "write" step is addition — you always add to this running sum. There is no mechanism for selectively retrieving information from specific earlier layers; everything is mixed into one uniformly weighted sum (where weights are implicitly 1.0 from the residual path). This fixed connectivity means that even if layer 3 computed something extremely useful, layer 47 has no way to attend to it specifically — it only sees the accumulated superposition.

2. Dense Cross-Layer Connections (DenseNet-style). Inspired by the success of DenseNet in computer vision, some prior work has applied dense connectivity to Transformers, where each layer receives the concatenated outputs of all preceding layers as input. Methods like DenseFormer (Pagliardini et al., 2024) and earlier cross-layer connection schemes (cited in the paper as [20, 28]) take this approach.

Dense connections solve the information dilution problem in principle because they preserve layer-wise history losslessly: concatenation doesn't compress or overwrite earlier states, so every layer has direct access to every previous layer's raw output. As the paper states, "depth-dense connections propagate information through depth losslessly, because concatenation does not compress the historical set."

However, this approach has a fatal computational cost problem that has prevented its adoption at LLM scale. The paper's complexity analysis in Table 1 makes this concrete: depth-dense methods incur O(L²D²) parameter costs (because each layer must linearly project the growing concatenated history back to width D), O(TL²D²) prefilling FLOPs, and O(L²D²) decoding FLOPs. At LLM scale — where L might be 64–128 and D might be 1024–8192 — these quadratic-depth and quadratic-width terms become utterly prohibitive. The paper calls this "prohibitive for large models" (Section 2.2).

Beyond cost, dense connections also enforce a fixed connectivity pattern: every layer reads from every previous layer with the same fixed projection weights (after the linear projection back to width D). There's no data-dependent selectivity — layer 47 can't decide that layer 3's output is especially relevant for this particular token while layer 12's is not. The mixing is static and uniform.

3. Other Residual Connection Upgrades. The paper briefly acknowledges other attempts to improve upon standard residual connections, citing methods like hyper-connections (Zhu et al., 2025) and manifold-constrained hyper-connections (Xie et al., 2025) as well as virtual width networks (Li et al., 2025). These methods modify the residual pathway — for example, by introducing learnable gating, multiple parallel residual streams, or more sophisticated write operations than simple addition. While these approaches can improve optimization stability and gradient flow, the paper argues they still fundamentally rely on the additive superposition paradigm: depth history is compressed into a small number of fixed-size tensors, and the read operation remains largely fixed-pattern rather than data-dependent.

The paper's critique of these methods is implicit but important: they modify the write operation (how layers contribute to the depth stream) but leave the read operation unchanged (each layer still reads the most recent state). Information dilution is partly about how information is written (superposition causes interference), but more fundamentally about how information is retrieved — and fixed read patterns can't adaptively recover what was lost through writing.

The Key Insight: Borrowing from Sequence Attention for Depth

The paper's critical conceptual move is to draw an analogy between the depth dimension and the sequence dimension. In the original Transformer, the sequence dimension was modeled with a fixed connectivity pattern (recurrence or convolution) until attention replaced it with data-dependent dynamic mixing: each token decides, based on its own content, which other tokens to attend to and how strongly. This was transformative because it allowed the model to retrieve information selectively rather than through a fixed aggregation pattern.

The paper asks: can we apply the same principle to the depth dimension? Rather than each layer reading the fixed accumulated sum of all previous layers, what if each layer could attend to specific earlier layers in a data-dependent way — paying more attention to layer 3's output when it's relevant, and less when it's not? This is the core insight that motivates both the intermediate "Depth Attention" formulation and the final MoDA design.

The authors state this explicitly in Section 1: "The success of attention in sequence modeling suggests a broader principle: data-dependent dynamic mixing can preserve and retrieve historical information more effectively than fixed-pattern aggregation. This motivates extending the same principle from sequence modeling to depth modeling, i.e., enabling each layer to adaptively read useful states from earlier layers."

This framing positions depth attention not as a radical departure from the Transformer paradigm, but as a natural extension of its core innovation — dynamic, content-based retrieval — to a dimension where it hasn't been systematically applied before.

Where This Paper Positions Itself

The paper positions MoDA as an intermediate point in the depth-stream design space that combines the benefits of both previous approaches while avoiding their drawbacks:

  • Like depth-dense connections, MoDA provides direct access to historical layer states, mitigating information dilution. But unlike depth-dense, it avoids the quadratic-depth parameter growth by using attention (which doesn't require projecting the entire concatenated history) and achieves data-dependent selectivity (each token-layer pair decides which historical layers are relevant).

  • Like standard residual connections, MoDA has manageable computational cost (the paper reports only 3.7% FLOPs overhead over the baseline). But unlike residuals, it provides explicit, selective retrieval rather than passive accumulation.

  • Unlike prior residual upgrades, MoDA fundamentally changes the read operation (from identity to attention) rather than just the write operation, and does so in a way that fuses sequence and depth retrieval into one unified softmax.

The complexity analysis in Table 1 is central to this positioning. It shows that MoDA reduces parameter complexity from O(L²D²) (depth dense) to O(LD²/G) (where G is the GQA group size), while keeping FLOPs at O(L²D) for both decoding and prefilling — a factor of D reduction from depth dense. This is what the paper means when it says MoDA "occupies an efficient point that preserves data-dependent depth retrieval without dense cross-layer overhead" (Section 1).

The paper also positions itself relative to a practical engineering reality: hardware efficiency matters as much as theoretical expressivity. The authors acknowledge that "adaptive cross-layer retrieval is therefore promising, yet practical designs still require a better balance among expressivity, efficiency, and hardware friendliness" (Section 1). This motivates the substantial engineering effort in Section 3 to make MoDA's non-contiguous depth-KV access patterns compatible with FlashAttention-style tiled computation. Many prior cross-layer methods have been proposed in the literature but failed to see adoption partly because their hardware characteristics were poor. By showing that MoDA can reach 97.3% of FlashAttention-2 efficiency at 64K sequence length, the paper is making an explicit argument about deployability.

Finally, the paper positions its contribution through a specific design philosophy articulated in the abstract and Figure 3: all Transformer stacking methods can be understood through a "read, operate, write" lens. By explicitly defining the design space — what is read (identity, linear projection, attention), what operates on it (attention, FFN), what is written (addition, concatenation) — the paper creates a taxonomy that makes MoDA's position clear and principled rather than ad-hoc. The progression from Depth Residual → Depth Dense → Depth Attention → MoDA in Figure 3 is not just a historical narrative; it's a logical derivation that shows each design choice and its consequences.

3. Technical Approach

3.1 Reader Orientation

This paper proposes a new attention mechanism — not a new model architecture from scratch, but a drop-in replacement for the standard self-attention operation in Transformer decoders — that lets each layer explicitly look back at what earlier layers computed, adaptively selecting which historical information is useful for the current computation, instead of only seeing the accumulated sum of all previous layers through the residual stream. The system solves the information dilution problem: as Transformers get deeper, distinctive features formed in early layers get progressively "smeared out" by dozens of additive residual updates, making them hard for later layers to recover; MoDA solves this by giving each attention head access to a depth memory — the key-value pairs from all preceding layers at the same token position — and fusing retrieval from this depth memory together with standard sequence attention into a single, jointly-normalized softmax operation.

3.2 Big-Picture Architecture (Diagram in Words)

The MoDA system consists of five interrelated components that together modify how information flows through a Transformer decoder's depth dimension:

  1. Depth KV Cache — a persistent store that accumulates key-value pairs from every layer as the forward pass proceeds. After each attention layer and each FFN layer processes its input, the resulting key and value projections (or lightweight FFN-specific KV projections) are appended to this cache. The cache is organized by token position and layer index, so for each position $t$ at layer $l$, all $\{K_{i,t}, V_{i,t}\}_{i=0}^{l-1}$ are available.

  2. Unified Query Projection — the standard query projection from the current hidden state, used for both sequence attention and depth attention. Unlike Depth Attention (which required a separate depth-query projection), MoDA reuses the same query, saving parameters and ensuring sequence and depth retrieval operate in a shared representational space.

  3. Mixture-of-Depths Attention Operator — the core computation: for each attention head, the query attends over a concatenated key-value space consisting of (a) the sequence KV pairs from the current layer (the standard self-attention keys and values, with causal masking) and (b) the depth KV pairs from all preceding layers at the same token position. All attention scores — both sequence-to-sequence and sequence-to-depth — are normalized jointly under a single softmax, producing a unified attention distribution.

  4. Write Operations (Two Types) — after the attention layer computes its output, the layer's keys and values are appended to the depth KV cache so subsequent layers can access them. After each FFN layer, an additional lightweight linear projection produces FFN-specific depth KV pairs that are also appended to the cache. This means the depth memory contains contributions from both attention and FFN sublayers.

  5. Hardware-Efficient Fused Kernel — a custom CUDA/Triton implementation that performs the combined sequence+ depth attention in a single tiled pass, using shared online-softmax states, chunk-aware depth-KV layout to reduce memory traffic, and group-aware indexing that exploits GQA's query-grouping structure to further reduce wasted computation.

The forward pass through one layer proceeds as follows: the hidden state $X_{l-1}$ enters the attention sublayer → the standard Q, K, V projections are computed from $X_{l-1}$ → the MoDA kernel computes attention over current-layer sequence KV plus historical depth KV from the cache → the attention output is produced and added to the residual stream → the current layer's KV are appended to the depth cache → the residual stream enters the FFN sublayer → a lightweight KV projection maps the FFN input to FFN-specific depth KV pairs, which are appended to the cache → the FFN output is added to the residual stream → $X_l$ is passed to the next layer.

3.3 Roadmap for the Deep Dive

  • First, the "read, operate, write" design space for Transformer stacking (Section 2.2 from the paper), because understanding what MoDA changes requires understanding what the baseline mechanisms do. I will walk through Depth Residual, Depth Dense, and the intermediate Depth Attention formulation, establishing the design axes along which MoDA innovates.

  • Second, the MoDA operator itself — how the unified softmax works, what keys and values are included, how masking is applied, and what the "mixture" in mixture-of-depths means concretely. This is the algorithmic core.

  • Third, the write operations — how depth KV pairs are produced, the difference between reusing attention-layer KV versus adding FFN-side KV projections, and why this distinction matters for the performance-parameter tradeoff.

  • Fourth, the complexity analysis (Table 1) that justifies MoDA's efficiency relative to depth-dense alternatives. Understanding the parameter, cache, and FLOPs scaling is essential to understanding why MoDA is practical while depth-dense is not.

  • Fifth, the hardware-efficient implementation (Section 3), because MoDA's non-contiguous depth-memory access patterns are precisely the kind of irregular memory access that kills GPU utilization. The chunk-aware layout, group-aware indexing, and fused online-softmax are what make the difference between a theoretically elegant operator and one that actually runs fast.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architectural innovation paper whose core idea is that the depth dimension of Transformers should be modeled with the same data-dependent dynamic retrieval that made attention successful on the sequence dimension, and that this retrieval should be fused with standard sequence attention into a single unified softmax operator to avoid representational fragmentation and parameter bloat.

The "Read, Operate, Write" Design Space for Transformer Stacking

Before presenting MoDA, the paper establishes a taxonomy of how Transformers can be stacked along the depth dimension. Every stacking method can be decomposed into three operations: read (what information the current layer receives from the depth history), operate (what computation the layer performs on that information), and write (how the layer's output is stored for future layers). This decomposition makes the design choices behind MoDA explicit and principled rather than ad-hoc.

Depth Residual (the baseline). This is the standard ResNet-style residual connection used in virtually all modern Transformers. The formulation, from Equation 3, is:

Xl=X0+∑i=1l−1F(Xi,Wi)X_l = X_0 + \sum_{i=1}^{l-1} F(X_i, W_i)

where $X_0 \in \mathbb{R}^{T \times D}$ is the input embedding (or the output of an initial embedding layer), $X_i \in \mathbb{R}^{T \times D}$ is the hidden state after layer $i$, $F(\cdot, W_i)$ is the token-mixing operator (attention or FFN) at layer $i$ with trainable weights $W_i$, and $X_l \in \mathbb{R}^{T \times D}$ is the hidden state at the output of layer $l-1$ (using 0-indexed layers, so $X_1$ is after layer 0's computation, etc.).

What it computes: the hidden state at any depth $l$ is the sum of the original input $X_0$ plus the accumulated outputs of every preceding token-mixing operation. The read step is identity — you always read $X_{l-1}$, which is this running sum. The operate step is $F(X_{l-1}, W_{l-1})$. The write step is addition — you add the result to the running sum.

Why this form: the additive write creates shortcut paths that allow gradients to flow directly from the loss at the top of the network back to any earlier layer without attenuation (the gradient of the sum with respect to any term is 1). This solved the vanishing gradient problem that had prevented training networks deeper than roughly 20–30 layers prior to ResNet. However, the additive write also means that information from different layers is combined without any selectivity — every layer's contribution gets equal weight in the sum, and there is no mechanism for a later layer to preferentially retrieve information from a specific earlier layer. If layer 3 computed a very useful syntactic feature and layer 27 wants to use it, layer 27 can only access it through the accumulated sum, where it is intermixed with 24 other layers' outputs.

Depth Dense (DenseNet-style). This is the concatenation-based alternative. From Equation 4:

{Xi}i=0l={X0,F({X0},W1),F({X0,X1},W2),…,F({Xi}i=0l−1,Wl)}\{X_i\}_{i=0}^l = \{X_0, F(\{X_0\}, W_1), F(\{X_0, X_1\}, W_2), \ldots, F(\{X_i\}_{i=0}^{l-1}, W_l)\}

where $\{X_i\}_{i=0}^{l-1}$ is the set of all previous layer outputs, and $F$ at each layer takes this entire set as input (after some linear projection to compress it back to width $D$).

What it computes: each layer receives the concatenated outputs of all previous layers, linearly projects this growing history back to a fixed width $D$, applies its token-mixing operator, and then concatenates its own output to the historical set. The read step is a linear projection of the concatenated history. The operate step is $F$ on the projected representation. The write step is concatenation — the new output is appended to the set without modifying existing entries.

Why this form: concatenation preserves historical information losslessly — earlier outputs are stored verbatim and can be read directly by any later layer. This solves information dilution in principle. However, the cost is catastrophic at LLM scale because of the linear projection from the growing concatenated history back to width $D$. At layer $l$, the input dimension is $l \times D$ (since you've concatenated $l$ previous outputs of dimension $D$ each), and you must project this to $D$ with a matrix of size $(lD) \times D$. Summed over all layers, the total parameter cost grows as $O(L^2 D^2)$ — quadratic in both depth and width. For $L=64$ and $D=1024$, this is billions of parameters just for the depth-mixing projections. Additionally, the connectivity is fixed: the linear projection weights are learned but static after training, so the model cannot dynamically decide which previous layers are relevant for a particular token or context.

Depth Attention (the intermediate bridge). The paper introduces this as a conceptual stepping stone. Instead of linearly projecting the concatenated history, depth attention uses the attention mechanism to read from it:

Xlin=Attention(Ql−1,{Ki}i=0l−1,{Vi}i=0l−1)X_l^{\text{in}} = \text{Attention}(Q_{l-1}, \{K_i\}_{i=0}^{l-1}, \{V_i\}_{i=0}^{l-1})

where $Q_{l-1} \in \mathbb{R}^{T \times D/G}$ is a query projection from the previous layer's output (in the GQA-group view, where $D/G$ is the per-group dimension), and $\{K_i\}_{i=0}^{l-1}$ and $\{V_i\}_{i=0}^{l-1}$ are the key-value pairs from all preceding layers at the same dimensions. The attention is performed along the depth dimension only: for each token position $t$, the query $Q_{l-1,t}$ attends to the depth keys $\{K_{i,t}\}_{i=0}^{l-1}$ and depth values $\{V_{i,t}\}_{i=0}^{l-1}$ from the same token position across all previous layers. There is no sequence mixing in this step — tokens do not attend to other sequence positions.

After this depth-attention read, the resulting $X_l^{\text{in}}$ is fed into the standard token-mixing operator $F$ (which performs sequence attention or FFN computation). At the write step, new query, key, and value projections are computed from the layer output $X_l^{\text{out}}$:

Ql=XloutWQ,lW,Kl=XloutWK,lW,Vl=XloutWV,lWQ_l = X_l^{\text{out}} W_{Q,l}^W, \quad K_l = X_l^{\text{out}} W_{K,l}^W, \quad V_l = X_l^{\text{out}} W_{V,l}^W

where $W_{Q,l}^W, W_{K,l}^W, W_{V,l}^W \in \mathbb{R}^{D \times D/G}$ are trainable "write" projection matrices specific to layer $l$. The $K_l$ and $V_l$ are concatenated to the depth store; $Q_l$ is passed forward to the next layer for its depth-attention read.

What this achieves: data-dependent depth retrieval — each token at each layer can decide, based on its current query representation, which previous layers' outputs are most relevant, and attend to them with learned weights. The computational cost drops from $O(L^2 D^2)$ to $O(L D^2)$ in parameters and $O(L^2 D)$ in FLOPs (Table 1), because attention doesn't require projecting the entire concatenated history through a dense matrix — it computes pairwise similarities between the query and each depth key, which is $O(LD/G)$ per token per layer.

Why it's intermediate, not final: Depth Attention separates depth retrieval (done by the dedicated depth-attention read step) from sequence retrieval (done by the standard self-attention in the operate step). This means depth and sequence information are processed in separate softmax operations with separate learned projections. This has two drawbacks: (1) the model must learn two different query representations — one for reading from depth, one for reading from sequence — which may not share information effectively, and (2) the depth-attention step requires its own dedicated query projection $Q_{l-1}$, adding parameters. MoDA addresses both by fusing depth and sequence retrieval into a single softmax with a single query projection.

Mixture-of-Depths Attention (MoDA): The Unified Operator

MoDA's key innovation is combining depth attention and sequence attention into one joint softmax operation. The formulation, which builds on the standard attention equation (Equation 2), is:

For each attention head $h$ with query $Q_h \in \mathbb{R}^{T \times d}$, the head attends over a concatenated key-value space consisting of:

  1. Sequence KV: $K^{\text{seq}} \in \mathbb{R}^{T \times d}$ and $V^{\text{seq}} \in \mathbb{R}^{T \times d}$ — the standard self-attention keys and values from the current layer, containing information from all sequence positions.

  2. Depth KV: $K^{\text{depth}} \in \mathbb{R}^{L \times d}$ and $V^{\text{depth}} \in \mathbb{R}^{L \times d}$ — for each token position $t$, the depth keys and values from layers $0$ through $l-1$ at that same token position. Note that while the sequence KV dimension is $T$ (number of tokens), the depth KV dimension is $L$ (number of layers), because each token looks at its own depth history across layers, not across token positions.

The joint attention computation is:

MoDA(Qh,Kseq,Vseq,{Kidepth}i=0l−1,{Videpth}i=0l−1)=softmax(Qh[Kseq∥Kdepth]Td+M)[Vseq∥Vdepth]\text{MoDA}(Q_h, K^{\text{seq}}, V^{\text{seq}}, \{K_i^{\text{depth}}\}_{i=0}^{l-1}, \{V_i^{\text{depth}}\}_{i=0}^{l-1}) = \text{softmax}\left(\frac{Q_h [K^{\text{seq}} \| K^{\text{depth}}]^T}{\sqrt{d}} + M\right) [V^{\text{seq}} \| V^{\text{depth}}]

where $[A \| B]$ denotes concatenation along the key-value dimension, $M$ is the combined mask matrix, and $d$ is the head dimension. For a query at token position $t$ in layer $l$, the concatenated key space has dimension $T + l$: $T$ sequence keys (from tokens $0$ through $T-1$, with causal masking restricting to $j \leq t$) and $l$ depth keys (from layers $0$ through $l-1$ at token position $t$).

What it computes: a single unified attention distribution over both the sequence and depth dimensions. The query $Q_{h,t}$ produces $T + l$ attention scores: one for each sequence position (how relevant is token $j$'s content to the current computation?) and one for each previous layer (how relevant is the representation that layer $i$ produced at this token position?). All scores are normalized together by the same softmax, meaning that probability mass allocated to depth keys is probability mass that is not allocated to sequence keys, and vice versa. This creates a direct competition: the model must decide, for each attention head at each token at each layer, whether historical depth information or current sequence context is more useful for the task at hand.

The masking scheme has two components applied jointly:

  • Sequence causal mask: standard causal masking where query at position $t$ cannot attend to sequence keys at positions $j > t$. This is $M_{t,j}^{\text{seq}} = 0$ if $j \leq t$, and $-\infty$ otherwise. Under GQA, this becomes grouped causal masking: for query head $h$ mapped to key-value head $\phi(h)$, the mask is $\lfloor i_q / G \rfloor \geq i_k$ where $i_q$ is the query row index, $i_k$ is the key row index, and $G$ is the GQA group size.

  • Depth matching mask: a query at token position $t$ in layer $l$ can only attend to depth keys from the same token position $t$ across previous layers. A query from token $t$ cannot attend to depth keys from token $t' \neq t$. This is because depth keys represent "what earlier layers computed for this specific token," and there is no meaningful cross-token depth relationship. In the flattened index notation used in Algorithm 1, this mask is $\text{mask}(i_q, j_d) = \mathbb{1}[\lfloor i_q / G \rfloor = \lfloor j_d / L \rfloor]$, where $i_q$ is the query row, $j_d$ is the flattened depth-column index (which encodes both token position and layer), and $L$ is the number of layers. Operationally, for query row $i_q$ with base-time index $t_{\text{base}}(i_q) = \lfloor i_q / G \rfloor$, only depth columns with indices in the range $[t_{\text{base}} L, (t_{\text{base}} + 1) L)$ are valid — these correspond to the $L$ depth keys for that token position.

Why this form matters — five key properties:

First, representational unity: by normalizing sequence and depth attention together, the model learns a single attention distribution that jointly optimizes which sequence tokens and which historical layers are relevant. The alternative — separate softmaxes for sequence and depth — would allow the model to independently allocate 100% of depth attention and 100% of sequence attention, which doesn't force the tradeoff decisions that produce sparse, interpretable attention patterns. The joint softmax means that if a head allocates 30% of its probability mass to depth layer 5, it only has 70% left for all sequence positions combined.

Second, parameter efficiency through query reuse: MoDA uses the same query projection for both sequence and depth attention. The standard attention query $Q_h = X W_Q^h$ is computed once and used for both retrieval tasks. This is in contrast to Depth Attention, which required a separate query projection $W_{Q,l}^W$ specifically for the depth read step. As Table 1 shows, this parameter reuse reduces MoDA's parameter complexity from $O(L D^2)$ (for Depth Attention) to $O(L D^2 / G)$, because only the grouped key/value projections (for the depth write step) are needed, not separate query projections. In GQA settings where $G = 8$, this is an 8× reduction in depth-specific parameters.

Third, complementary retrieval channels: sequence attention provides spatial context (what other tokens are saying), while depth attention provides temporal context (what earlier computational stages produced for this token). These are fundamentally different kinds of information. Sequence attention might retrieve a noun phrase from five tokens ago to resolve a pronoun reference; depth attention might retrieve a syntactic parse feature that layer 3 computed and that layer 27 needs for semantic composition. The joint softmax forces the model to balance these complementary information sources.

Fourth, implicit gating against information dilution: because depth attention allows direct retrieval of early-layer features, later layers don't need to "recover" those features from the diluted residual stream. A later layer can simply attend directly to the depth KV from layer 3 if layer 3's computation is relevant, bypassing the accumulated noise from layers 4 through 26. This is fundamentally different from residual connections, where the only way to access layer 3's contribution is through the sum that includes everything else.

Fifth, attention sink redistribution: the paper observes (Section 4.3.2, Figure 5) that MoDA heads allocate substantial probability mass to depth KV entries. In standard Transformers, some heads exhibit "attention sink" behavior — they allocate a large fraction of their probability mass to a few fixed token positions (often the first token or punctuation tokens), which serves as a kind of "no-op" or bias term but doesn't contribute task-relevant information. By providing an alternative productive destination for probability mass (depth keys from informative earlier layers), MoDA may naturally reduce reliance on these uninformative sink positions.

The "mixture" terminology reflects this joint softmax: each head produces a mixture of sequence and depth attention weights, with the mixing ratio determined dynamically per token per layer by the relative magnitudes of the sequence and depth attention scores before softmax normalization.

The depth KV at different layers: At layer $0$, there are no preceding layers, so the depth KV component is empty and MoDA reduces to standard causal self-attention. At layer $1$, the depth KV contains only layer $0$'s key-value pairs. At layer $l$, the depth KV contains $l$ entries per token — one for each preceding layer. This means the relative importance of depth attention grows as the network gets deeper: early layers have little depth history to draw on, while later layers have rich depth histories spanning dozens of previous computational stages. This is exactly where information dilution would otherwise be most severe, so the architecture naturally provides more depth-retrieval capacity where it's most needed.

Write Operations: Populating the Depth KV Cache

MoDA's write operations determine what gets stored in the depth KV cache for future layers to read. There are two types of write operations, corresponding to the two sublayers in each Transformer block.

Attention-layer KV reuse (the default). After the attention sublayer computes its output, the keys and values that were computed for the sequence attention operation — $K^{\text{attn}}_l$ and $V^{\text{attn}}_l$ — are appended to the depth KV cache. This is the "reuse" strategy: the same KV pairs that served as sequence keys/values for layer $l$'s self-attention now also serve as depth keys/values for layers $l+1, l+2, \ldots$. This introduces zero additional parameters — no extra projection matrices are needed — and only the storage cost of retaining the KV tensors.

FFN-layer KV projection (the "Extra FFN KV Proj." in Table 3). The FFN sublayer does not naturally produce key-value pairs, since it operates per-token without any attention mechanism. To incorporate FFN outputs into the depth memory, the paper adds a lightweight linear projection that maps the FFN's input $X$ (or equivalently, the hidden state entering the FFN) to FFN-specific depth keys and values:

Klffn=XWK,ffn,Vlffn=XWV,ffnK^{\text{ffn}}_l = X W_{K,\text{ffn}}, \quad V^{\text{ffn}}_l = X W_{V,\text{ffn}}

where $W_{K,\text{ffn}}, W_{V,\text{ffn}} \in \mathbb{R}^{D \times D/G}$ are trainable projection matrices. These are separate from the attention projections — the FFN gets its own key-value projection weights, which are learned to encode whatever information from the FFN's computation is useful for downstream depth retrieval. These FFN KV pairs are also appended to the depth cache, interleaved with the attention KV pairs in layer order.

Design choice — reuse vs. separate projections for attention KV: The paper explores whether to reuse the attention sublayer's existing KV projections or add separate "Extra Attn KV Proj." weights specifically for depth. Table 3, row 5 shows that adding separate attention-side depth projections (total parameters 742.4M vs. 705.7M for reuse-only) yields only marginal gains: +0.07 train PPL, +0.04 C4 validation PPL, +0.10 downstream average. The authors conclude this is "overly saturated" — the attention KV already encode sufficient information for depth retrieval, and adding a second set of projections is redundant. The default MoDA configuration (row 4) therefore reuses attention KV and adds only the FFN-side projections.

Design choice — FFN KV from input vs. output: The projection maps from the FFN's input $X$, not its output. This is a pragmatic choice: the FFN input is available before the FFN computation, allowing the KV projection to happen in parallel with the FFN forward pass, reducing serialization. The FFN input already contains the residual stream after attention, so it encodes the state just before the FFN transformation — which is the natural "read point" for capturing what the FFN sublayer will operate on.

Why FFN depth KV matters: The FFN sublayers perform non-linear transformations that are complementary to attention — they process each token independently, applying learned weight matrices that can store factual knowledge, perform simple computations, or transform representations. By including FFN outputs in the depth memory, later layers can retrieve not just what earlier attention mechanisms attended to, but also what earlier FFNs transformed. Table 3 shows this matters empirically: row 3 (attention KV only) to row 4 (attention + FFN KV) improves C4 validation PPL by 0.27 and downstream average by 0.77 at 700M scale, with the parameter increase from 669.0M to 705.7M being modest relative to the gains.

Complexity Analysis: Why MoDA Is Practical

Table 1 in the paper provides the asymptotic complexity comparison between Depth Dense, Depth Attention, and MoDA. Understanding this table is essential to understanding why MoDA is deployable while depth-dense is not.

Parameters. Depth Dense requires, at each layer $l$, a linear projection from $l \times D$ (the concatenated history) to $D$ (the operating width). Each such projection has $l D^2$ parameters. Summed over $L$ layers, the total is $\sum_{l=1}^L l D^2 = \frac{1}{2} L^2 D^2$, which is $O(L^2 D^2)$. For $L=64$ and $D=1024$, this is roughly 2 billion parameters just for depth mixing — comparable to the entire parameter budget of a 1.5B model.

Depth Attention removes the dense projection by using attention for the read operation. The parameter cost comes from the per-layer write projections $W_{Q,l}^W, W_{K,l}^W, W_{V,l}^W$, each of size $D \times D/G$. With $H_k = D/(Gd)$ key-value heads, the total is $3 L H_k \times D \times (D/G)$ which simplifies to $O(L D^2)$ — linear in depth, quadratic in width. This is a factor of $L$ smaller than Depth Dense's dominant term: for $L=64$, Depth Attention has roughly 64× fewer depth-specific parameters.

MoDA further reduces parameters by reusing the sequence attention query and only adding group-aware key-value write projections for the FFN side. The parameter cost is dominated by $2 G L H_k^2 d^2$ for the grouped KV projections, which simplifies to $O(L D^2 / G)$. This is a factor of $G$ smaller than Depth Attention's parameter cost: for $G=8$, MoDA uses 8× fewer depth-specific parameters.

Cache (memory for stored KV pairs). During autoregressive decoding, the model must cache all historical KV pairs to avoid recomputation. Depth Dense must cache $L$ layer outputs of size $T \times D$, giving $O(T L D)$ cache. Depth Attention and MoDA must cache the per-layer key-value pairs: $2 L$ tensors of size $T \times D/G$, giving $O(T L D / G)$. For $G=8$ and $D=1024$, MoDA's KV cache is 8× smaller than storing full hidden states.

FLOPs. For decoding (generating one token), Depth Dense requires projecting the growing history at each layer and computing the token-mixing operators. The dominant term is $O(L^2 D^2)$ from the depth-mixing projections. Depth Attention and MoDA both have decoding FLOPs dominated by the attention operations: sequence attention costs $O(L T D)$ (for attention over the growing sequence length $T$) and depth attention costs $O(L^2 D)$ (for attention over $L$ depth keys at each of $L$ layers). The combined dominant term is $O(L^2 D)$ — linear in width, quadratic in depth. This is a factor of $D$ smaller than Depth Dense's quadratic-width term: for $D=1024$, MoDA's decoding FLOPs are roughly 1000× smaller.

For prefilling (processing a full sequence of length $T$ in parallel), the same patterns hold with an additional factor of $T$: Depth Dense has $O(T L^2 D^2)$, while MoDA has $O(T L^2 D)$. The key practical implication is that MoDA's FLOPs grow linearly with width rather than quadratically, which is what enables the paper's reported 3.7% FLOPs overhead over the baseline at 1.5B scale.

The asymptotic takeaway: MoDA keeps the data-dependent, selective retrieval capability of attention (unlike fixed-connectivity residual or dense approaches) while avoiding the quadratic-width parameter and FLOPs growth that makes depth dense impractical. The $O(L D^2/G)$ parameter scaling means that the depth-specific overhead shrinks relative to the total model size as width grows — because the total model parameters grow as $O(L D^2)$ (from the standard Transformer blocks), the depth-specific overhead as a fraction of total parameters is $O(1/G)$, which is constant (e.g., ~12.5% for $G=8$). This is a crucial property for scalability: the relative cost of depth awareness doesn't increase as models get wider.

Hardware-Efficient Implementation

MoDA's algorithmic design introduces a memory access pattern that is fundamentally hostile to GPU architecture if implemented naively. Each query at token position $t$ needs to read depth KV pairs from layers 0 through $l-1$ at that same position $t$. In the natural memory layout — where depth KV are stored as $[L, T, D/G]$ or similar — this requires gathering elements from non-contiguous memory locations (stride $T \times D/G$ between consecutive depth layers), which prevents the coalesced memory access that GPUs depend on for high bandwidth utilization. The paper's hardware-efficient implementation (Algorithm 1) resolves this through three progressively more sophisticated layout and indexing strategies.

Flash-compatible depth-KV layout. The first optimization is to flatten the depth cache into a contiguous layout along a single axis of length $T \times L$. For each sequence position $t$, its $L$ depth states are stored contiguously in the range $[tL, (t+1)L)$ of this flattened tensor. This means a query at token $t$ can read all its depth KV pairs as a single contiguous block, rather than as $L$ separate gather operations.

However, this still leaves an efficiency problem. When processing a query block covering multiple token positions (say, tokens $t$ through $t+B-1$), the valid depth KV region for the entire block spans from $tL$ to $(t+B)L$ in the flattened layout. But within this region, only the block-diagonal entries are valid — query row $i_q$ (corresponding to some token position) can only attend to depth entries $[\lfloor i_q/G \rfloor L, (\lfloor i_q/G \rfloor + 1)L)$. The fraction of valid entries in the full $B \times (B L)$ attention score submatrix is only $1/B$. The paper defines this as depth utilization $\eta_{\text{depth}}$: if computed densely over the full $T \times (TL)$ matrix (the attention scores between all queries and all depth KV), $\eta_{\text{depth}} = \frac{T \cdot L}{T \cdot (T \cdot L)} = 1/T$. At $T=4096$, this means 99.98% of the computed depth attention scores would be masked out — an enormous waste of compute and memory bandwidth.

Chunk-aware depth-KV layout. To address the low depth utilization, the paper reorganizes depth KV into chunks. Queries are divided into chunks of size $C$ (the paper uses $C=64$ in experiments). For a query chunk covering token positions $[t, t+C)$, the kernel only loads the depth KV corresponding to those $C$ positions — a region of size $C \times L$ rather than the global $T \times L$. The depth utilization improves to $\eta_{\text{depth}} = \frac{T \cdot L}{T \cdot (C \cdot L)} = 1/C$, a factor of $T/C$ improvement. For $T=4096$ and $C=64$, this is roughly 64× better utilization — 1.56% valid entries instead of 0.024%.

Figure 4 (right panel) illustrates this chunk-aware layout. Instead of one global depth KV cache of size $T \times L$, the cache is partitioned into $T/C$ chunks, each of size $C \times L$. Each query chunk only accesses its corresponding depth chunk, reducing both HBM traffic (fewer wasted bytes loaded) and compute (fewer wasted FLOPs on masked entries).

Group-aware depth-KV calculation. The third optimization exploits the GQA structure. Under GQA with group size $G$, $G$ adjacent query rows share the same key-value head and, critically, the same base-time index $t_{\text{base}}(i_q) = \lfloor i_q / G \rfloor$. This means that for a query chunk of length $C$, there are only $C/G$ unique base-time indices, not $C$. Consequently, the valid depth-KV region for the chunk is of size $(C/G) \times L$, not $C \times L$. The depth utilization further improves to $\eta_{\text{depth}} = G/C$.

The mechanism works as follows: for a query block $b_q$ aligned to $G$ (so block size is divisible by $G$), the kernel computes $t_{\text{base}}^{\text{start}} = \min_{i_q \in b_q} \lfloor i_q / G \rfloor$ and $t_{\text{base}}^{\text{end}} = \max_{i_q \in b_q} \lfloor i_q / G \rfloor + 1$. The depth loop then only iterates over depth blocks with indices $b_d \in [t_{\text{base}}^{\text{start}} L, t_{\text{base}}^{\text{end}} L)$, which is a factor of $G$ smaller than iterating over the full $C \times L$ region. For $G=8$ and $C=64$, $\eta_{\text{depth}} = 8/64 = 12.5\%$, meaning 12.5% of computed depth attention scores are valid — a dramatic improvement over the naive 0.024%.

Why query-block alignment to $G$ matters: if query blocks weren't aligned to $G$, a single block could contain query rows that span two different GQA groups, requiring the kernel to handle cross-group boundaries within one tile. This would complicate the masking logic and potentially require loading additional depth KV blocks. By enforcing block sizes divisible by $G$, each query block maps cleanly to a contiguous range of base-time indices, simplifying the depth loop bounds.

Fused online-softmax (Algorithm 1). The entire sequence + depth attention is computed in a single pass with shared online-softmax states. The kernel maintains three running statistics for each query row: $m$ (the running maximum logit seen so far), $acc$ (the running sum of exponentiated logits, i.e., the softmax denominator), and $o$ (the running weighted sum of values, i.e., the softmax numerator before final normalization). These statistics are updated incrementally as the kernel processes sequence key blocks and then depth key blocks, using the standard online softmax update:

m′=max⁡(m,rowmax(S)),acc′=acc⋅2m−m′+∑2S−m′,o′=o⋅2m−m′+∑2S−m′Vm' = \max(m, \text{rowmax}(S)), \quad acc' = acc \cdot 2^{m - m'} + \sum 2^{S - m'}, \quad o' = o \cdot 2^{m - m'} + \sum 2^{S - m'} V

where $S$ is the current block's attention score matrix (queries × keys), and the sums are over the key dimension. After all blocks are processed, the final output is $o / acc$.

Why fused online-softmax matters: without fusion, sequence attention and depth attention would require separate softmax operations, each with their own HBM reads/writes of intermediate attention matrices. The fused version processes sequence keys and depth keys as if they were one long concatenated key sequence, building up a single softmax incrementally. This eliminates the memory traffic for storing and reloading intermediate sequence attention outputs, and it ensures that the normalization denominator correctly reflects the sum over both sequence and depth logits — which is what makes the mixture a true joint distribution.

The two sequence attention loops: Algorithm 1 processes sequence keys in two phases. For key blocks where all keys are strictly before the query block's base-time range ($b_s < t_{\text{base}}^{\text{start}}$), no causal masking is needed — the softmax update is applied without modification. For key blocks within the causal boundary ($t_{\text{base}}^{\text{start}} \leq b_s < t_{\text{base}}^{\text{end}}$), grouped causal masking $\lfloor i_q / G \rfloor \geq i_k$ is applied, zeroing out (setting to $-\infty$ before softmax) any attention scores where the key position is ahead of the query position in the GQA-group sense. This two-phase structure is standard in FlashAttention but adapted here for the GQA-grouped indexing.

The depth attention loop: After sequence accumulation completes, the kernel enters the depth loop, iterating over depth blocks $b_d$ in the range $[t_{\text{base}}^{\text{start}} L, t_{\text{base}}^{\text{end}} L)$. For each depth block, the depth matching mask $\mathbb{1}[\lfloor i_q / G \rfloor = \lfloor j_d / L \rfloor]$ is applied, where $j_d$ is the flattened depth-column index. This mask ensures that query row $i_q$ only attends to depth entries whose token position (encoded in $\lfloor j_d / L \rfloor$) matches the query's base-time index $\lfloor i_q / G \rfloor$. All valid depth attention scores are accumulated into the same online-softmax state as the sequence scores. Finally, $o \leftarrow o / acc$ produces the normalized output.

Efficiency results (Table 2). The incremental ablation in Table 7 shows the impact of each optimization at a small scale (T=1024, G=8, L=64, C=64): naive PyTorch takes 2128.9 ms, flash-compatible layout alone reduces this to 13.1 ms (~162× speedup), adding chunk-aware layout brings it to 6.3 ms (another 2.1×), and adding group-aware indexing brings it to 1.46 ms (another 4.3×), for a total 1458× speedup over naive.

Table 2 reports end-to-end forward+backward runtime against FlashAttention-2 Triton under realistic scaling settings. Key patterns: (1) As sequence length $T$ grows from 4096 to 65536, MoDA's extra time over FlashAttention-2 decreases from 25.86% to 2.73%, because the sequence computation dominates and the fixed depth overhead is amortized. (2) As GQA group size $G$ grows from 2 to 32, the extra time decreases from 27.07% to 2.84%, because larger $G$ means higher depth utilization (fewer unique base-time indices per chunk). (3) As depth $L$ grows from 64 to 256, extra time increases from 8.59% to 30.52%, because the depth-KV processing cost scales with $L$ while the sequence cost stays constant. This last result is the fundamental scaling limitation: deeper models spend proportionally more time in the depth attention loop, and the paper acknowledges this as motivation for the bounded depth-KV caching discussed in Section 6.2.

Numerical precision: The paper states that the fused kernel maintains "numerical precision within the allowed range" (contribution summary), implying that the online-softmax accumulation with shared states across sequence and depth blocks doesn't introduce meaningful numerical error compared to computing the two attentions separately and combining them. This is expected because online softmax is an exact algorithm (the running statistics are mathematically equivalent to computing the full softmax over all logits at once), but it's important to verify empirically that the blockwise computation with $2^{m-m'}$ rescaling doesn't accumulate floating-point drift over many blocks.

Training and Configuration Details

The paper's main experiments use the OLMo2 training recipe, which provides a standardized, reproducible baseline. The key configurations are:

Model architectures:

  • 700M models: width $D = 1024$, GQA group size $G = 2$, sequence length $T = 4096$, 36 layers. Baseline OLMo2 has 669.0M parameters and 8.01T FLOPs (per sequence? the paper doesn't specify the unit precisely, but it appears to be total FLOPs per forward pass).
  • 1.5B models: same $D$, $G$, and $T$ (the paper doesn't explicitly state the architecture for 1.5B, but the parameter count implies more layers or wider FFN — likely more layers given the depth-scaling focus). Baseline OLMo2-1.5B performance is reported in Table 4 and Table 5.
  • Deeper/shallower analysis (Table 6): width $D = 384$, 6 query heads, 2 key-value heads ($H_q = 6$, $H_k = 2$, so $G = 3$). Tested at 24 layers and 48 layers.

Training data and schedule: All models are trained on a 400B-token subset of the OLMo2 dataset (which includes C4, ICE, m2d2-s2orc, Pile, Wiki-text, dolma, and other sources). Global batch size is 1024 sequences, context length is 4096 tokens. The optimizer is AdamW with the OLMo2 default hyperparameters (specific learning rates, betas, weight decay are not detailed in the paper but follow the OLMo2 paper's settings). Training uses bfloat16 precision.

Learning rate schedule: For the 700M models in Table 3, the schedule warms up to a maximum learning rate of 3e-4 over 2000 training steps, then decays to 3e-5 following a cosine schedule. For the 1.5B models, the paper doesn't specify the learning rate but states it follows the OLMo2 recipe.

MoDA variant specification (Table 3):

  • Row 3 (the minimal MoDA): reuses attention-layer KV as depth KV, no FFN KV projections. Parameter count is identical to baseline (669.0M) because no new projection weights are added; FLOPs increase by only 0.12% (from 8.01T to 8.02T).
  • Row 4 (the default MoDA used in scaling experiments): reuses attention-layer KV + adds FFN-side KV projections. Parameters increase to 705.7M (5.5% over baseline), FLOPs increase to 8.33T (4.0% over baseline). This is the configuration used for all subsequent 700M and 1.5B experiments (Tables 4, 5).
  • Row 5 (over-parameterized): Row 4 + separate attention-side depth KV projections. Parameters increase to 742.4M (11.0% over baseline), FLOPs to 8.63T (7.7% over baseline). Gains are marginal relative to the cost.

Why these specific design choices:

  • GQA group size $G=2$ for 700M models: This is relatively small (most production models use $G=8$ or higher). A smaller $G$ means fewer query heads share each key-value head, which gives more fine-grained attention but also means the group-aware optimization in the kernel provides less benefit (depth utilization is lower). The paper's efficiency analysis in Table 2 shows that $G=2$ has only 3.12% depth utilization and 27.07% extra time at T=16384 — substantially worse than $G=8$'s 12.50% utilization and 8.59% extra time. This suggests that MoDA's practical efficiency would be even better at the larger GQA group sizes typical of production models.

  • Chunk size $C=64$: This is a hardware-dependent choice that balances two factors. Smaller chunks give better depth utilization (since $\eta_{\text{depth}} = G/C$) but may under-utilize Tensor Cores (small matrix multiplications don't saturate the hardware). Larger chunks improve Tensor Core utilization but reduce depth utilization (more wasted computation on masked entries). $C=64$ is likely chosen because 64 × 64 is a common tile size for matrix multiplication on A100 GPUs, providing a good balance.

  • FFN KV projection from input rather than output: As noted earlier, this is a latency optimization. The FFN computation (typically two linear projections with an activation in between) can proceed in parallel with the lightweight KV projection, since the KV projection only needs the FFN input $X$. Projecting from the FFN output would require the FFN to complete first, adding serialization. Whether FFN-output KV would provide better depth information is an empirical question not explored in the paper.

  • Depth KV cache stores all layers unconditionally: The default MoDA keeps the full depth history. Section 6.2 discusses the memory scaling challenge this creates and sketches a bounded slot-buffer approach for industrial-scale deployment, but this is not implemented or evaluated. The practical implication is that memory for depth KV grows as $O(L T D/G)$ — for $L=64$, $T=4096$, $D=1024$, $G=8$, this is roughly 32 million elements in bf16 (~64 MB), which is manageable but would grow to concerning levels for $L=256$ or longer sequences.

The "read, operate, write" framing applied to MoDA specifically:

  • Read: MoDA reads the current hidden state $X_{l-1}$ and the historical depth KV stream $\{(K_i, V_i)\}_{i=0}^{l-1}$. The read operation is attention — both for sequence (standard) and depth (the novel component), fused into one softmax.

  • Operate: The standard Token-mixing operator $F$ (either attention or FFN) processes the attention output.

  • Write: For the attention sublayer, the current layer's KV pairs are appended to the depth stream (reusing the same KV from the sequence attention). For the FFN sublayer, a lightweight linear projection produces KV pairs that are also appended. The hidden state is updated via standard residual addition.

This decomposition makes MoDA's relationship to prior work explicit: it keeps the additive residual write from Depth Residual (for the hidden state trajectory), adds the concatenative write from Depth Dense (for the depth KV store), and replaces the identity read with attention-based read — matching the pattern that made attention successful on the sequence dimension.

4. Key Insights and Innovations

Innovation 1: Depth as a Retrieval Channel — Extending Attention's Core Principle from Sequence to Depth

The paper's most fundamental intellectual move is recognizing that the depth dimension of Transformers suffers from the same class of problem that the sequence dimension suffered from before attention — fixed-pattern, non-selective information aggregation — and that the solution should therefore be the same: data-dependent dynamic retrieval. This is not merely adding another attention operation; it is a conceptual reframing of what the depth dimension is and how information should flow through it.

Prior to this work, the dominant approach to depth information flow was the residual connection [16], which the field largely treated as a solved problem. Residuals enabled training deep networks by addressing vanishing gradients, and most subsequent architectural innovations (pre-norm vs. post-norm placement, gated residuals, hyper-connections [49]) modified how the residual write worked or where normalization was placed, without fundamentally questioning whether identity-read + additive-write was the right paradigm for depth information retrieval. Even methods that did provide cross-layer access — DenseNet-style concatenation [20, 28] — used fixed linear projections to aggregate historical states, meaning the connectivity pattern was learned during training but static at inference: layer 47's access to layer 3's output was mediated by the same trained weights regardless of what the current token actually needed.

The paper frames this limitation through an explicit analogy to the pre-attention era of sequence modeling. Before Transformers, sequence information was aggregated through fixed-pattern mechanisms — recurrence (RNNs) or convolution — where the mixing weights were either position-invariant (convolution) or determined by a learned hidden state dynamics (recurrence), but never content-dependent at the individual token-instance level. Self-attention's breakthrough was enabling each token to decide, based on its own content, which other tokens to attend to and how strongly, producing vastly more flexible and context-sensitive representations.

MoDA applies this exact principle to depth: instead of each layer passively receiving the accumulated residual sum, each token at each layer actively queries which previous layers' representations are relevant for the current computation. The query is computed from the current hidden state (what does this token need right now?), the keys from previous layers (what does each previous layer's processing of this token encode?), and the attention weights are determined dynamically for each token-layer pair. This is the depth analog of how a sequence attention query at position t might attend strongly to a noun phrase at position t-5 when resolving a pronoun — MoDA allows a query at layer 27 to attend strongly to layer 3's depth KV when layer 3's syntactic parse feature is needed for semantic composition.

What makes this a fundamental conceptual shift rather than an incremental modification is that it redefines what depth means in a Transformer. In the residual paradigm, depth is a sequential refinement process: each layer progressively transforms the hidden state, and later layers see the cumulative effect of all earlier transformations. In the MoDA paradigm, depth becomes a structured memory: earlier layers produce representations that are stored and can be selectively retrieved by any later layer, bypassing the accumulated noise of intermediate transformations. This changes depth from a pipeline (where information flows forward through every stage) to a retrieval system (where information is produced once and accessed on demand).

The evidence that this conceptual shift matters — that it's not just a theoretically elegant reframing but a practically important one — comes from the attention visualization in Figure 5. The heatmaps show that MoDA heads allocate substantial and persistent probability mass to depth-KV entries, especially in middle and late layers, confirming that the model actively uses this retrieval capability rather than ignoring it in favor of sequence attention. The complementary pattern — some heads showing sharp diagonal sequence attention while still allocating probability to depth slots — suggests that depth retrieval serves as a genuinely distinct information channel, not a redundant copy of what sequence attention already provides.

Innovation 2: The Unified Softmax — Forcing Joint Optimization of Sequence and Depth Attention

MoDA's second major innovation is fusing depth and sequence attention into a single, jointly-normalized softmax rather than computing them as separate operations. This is a design choice with deep consequences for representation learning that goes far beyond implementation convenience.

The natural alternative — which the paper's own intermediate "Depth Attention" formulation uses — is to compute depth attention and sequence attention as separate softmax operations, potentially with separate query projections. In that design, a head would first attend over depth KV to produce a depth-aware representation, and then attend over sequence KV (or vice versa) to produce the final output. Each attention operation has its own softmax normalization, meaning the probability distributions are independent: a head could allocate 100% of its depth attention to layer 5 and simultaneously 100% of its sequence attention to the current token, because these are separate probability spaces.

The unified softmax in MoDA forces these into a single probability space. For each attention head at each token at each layer, the total probability mass of 1.0 must be allocated across both sequence positions (the standard T entries) and depth positions (the l entries from previous layers). Allocating 30% of attention mass to depth layer 5 means only 70% remains for all sequence positions combined. This creates a direct competition between the two information sources, forcing the model to learn when depth information is more valuable than sequence context and vice versa.

Why does this matter? First, it imposes a representational budget constraint that encourages sparsity and selectivity. If sequence and depth attention were separate, the model could learn to always attend heavily to both, producing dense, non-selective attention patterns that don't actually prioritize information. The joint softmax prevents this: a head cannot be "good at both" without making explicit tradeoffs, which means the attention weights that do emerge are more interpretable and more likely to reflect genuine informational value.

Second, the joint softmax enables cross-modal inhibition: if a depth KV entry is highly relevant, it can actively suppress attention to less relevant sequence positions (because they compete for the same softmax denominator). This is a more powerful attention mechanism than independent softmaxes, where depth and sequence attention cannot influence each other's distributions. The model learns to use depth information not just as an additional feature source, but as a filter that reshapes how it allocates sequence attention.

Third, the shared softmax denominator means that the relative scaling of depth and sequence logits is learned jointly through the query and key projections. The model can learn to produce larger-magnitude logits for depth keys in layers where historical information is critical (effectively "turning up" the depth channel) and smaller logits where sequence context dominates. This dynamic weighting emerges from the training objective without any explicit gating mechanism.

The paper's decision to make this a unified rather than factorized attention operation reflects a deeper design principle: mixture models are more expressive than product models when the components compete for a shared resource. This is the same principle that motivates mixture-of-experts over dense ensembles — forcing specialization through competition produces more efficient and interpretable representations. The paper calls this "mixture-of-depths" attention precisely because each head produces a mixture (a convex combination) of sequence and depth information sources, with the mixing weights determined dynamically.

Evidence that this unified formulation matters comes indirectly from the paper's ablation structure. While the paper never directly compares unified vs. separate softmaxes (which would require implementing the separate Depth Attention variant at scale), the strong performance of the minimal MoDA configuration (Table 3, row 3 — reuse attention KV as depth KV with only 0.12% FLOPs overhead, improving C4 PPL by 0.11 and downstream average by 1.17) suggests that even with no additional information beyond what's already in the standard KV cache, simply restructuring how that information is accessed — through a joint softmax rather than through the residual stream — yields substantial gains. This is consistent with the unified softmax providing a better optimization landscape for learning to use depth information.

Innovation 3: FFN-Side Depth Information as a Distinct and Complementary Signal

The paper's third insight is that FFN sublayers contain depth-relevant information that is complementary to what attention sublayers produce, and that explicitly projecting FFN states into the depth memory yields meaningful gains beyond what attention-side KV alone can provide. This is a subtle but practically important finding that challenges an implicit assumption in many Transformer analyses: that attention outputs are the primary carriers of token-level representational content worth preserving across depth, while FFNs perform "local" transformations whose intermediate states don't need to be explicitly retrievable.

The evidence for this complementarity is in Table 3. Row 3 (attention KV only as depth memory) improves over the baseline by 0.11 C4 validation PPL and 1.17 downstream average — a substantial gain from simply restructuring access to information that was already being computed. Row 4 (adding FFN KV projections) provides a further 0.27 C4 PPL improvement (more than doubling the gain from depth KV alone) and 0.77 downstream average improvement. This additional gain cannot be explained by the parameter increase alone, because Row 2 (a 38-layer baseline with comparable parameters to Row 4) performs substantially worse than Row 4: train PPL 14.27 vs. 13.90, C4 PPL 18.31 vs. 18.21, downstream 57.11 vs. 58.87. The FFN depth KV is providing genuinely useful information that neither attention-side depth KV nor additional layers can substitute for.

Why would FFN depth information matter? The paper doesn't extensively theorize about this, but the result aligns with a growing understanding that FFN layers in Transformers serve as key-value memories that store factual knowledge and perform pattern completion. An FFN layer's transformation of a token representation may encode information that is qualitatively different from what attention captures — for example, retrieving a factual attribute (FFN) versus resolving a coreference (attention). By storing FFN outputs in the depth memory, MoDA allows later layers to directly access these FFN-computed features without needing to recompute them or recover them from the residual stream.

The paper also makes a negative discovery that sharpens this insight: adding separate attention-side depth KV projections (Row 5) provides only marginal gains (+0.07 train PPL, +0.10 downstream) while incurring substantial parameter overhead (742.4M vs. 705.7M). The attention KV are already informative enough for depth retrieval; the bottleneck is the FFN side, not the attention side. This negative result is valuable because it tells practitioners where to invest their parameter budget — FFN KV projections are high-return, while extra attention KV projections are saturated.

This finding reframes MoDA from "attention about attention" (retrieving what earlier attention mechanisms computed) to "attention about computation" (retrieving what earlier computational stages — both attention and FFN — produced). It suggests that the depth memory should capture all significant transformations in the forward pass, not just the attention outputs, and that different sublayer types contribute complementary information that later layers can selectively retrieve depending on their computational needs.

Innovation 4: Diagnosing Information Dilution as a Retrieval Problem Rather Than an Optimization Problem

The paper's fourth contribution is a diagnostic reframing of why deep Transformers underperform relative to their theoretical capacity. The standard narrative in the literature attributes depth-scaling difficulties primarily to optimization challenges: vanishing/exploding gradients, training instability, and the difficulty of propagating learning signals through many layers. This narrative motivated solutions like better initialization schemes (DeepNet [40]), improved normalization placement (pre-norm vs. post-norm), and gated residual variants [22, 42, 49].

MoDA's success suggests a different diagnosis: the problem is not (just) that deep networks are hard to optimize, but that the residual pathway is a lossy compression mechanism for depth history. Information that was clearly present in early layers is still technically "in" the residual stream at later layers (because addition preserves it), but it becomes increasingly difficult to extract from the accumulated superposition of dozens of subsequent updates. The signal isn't gone — it's buried under noise. This is fundamentally a retrieval problem, not an optimization problem: the information exists in the network's activations, but the architecture provides no efficient mechanism for accessing it.

This reframing has substantial implications for architecture design. If depth scaling fails primarily due to optimization difficulty, the solution is better training recipes — improved initialization, learning rate schedules, normalization schemes. But if it fails due to retrieval difficulty, the solution is architectural — providing explicit mechanisms for later layers to query earlier layers' outputs, as MoDA does. The paper's evidence supports the retrieval diagnosis: MoDA adds essentially no optimization tricks (it uses the same OLMo2 training recipe as the baseline) and achieves consistent gains purely through architectural modification of the information flow.

The post-norm vs. pre-norm finding in Table 6 provides additional evidence for this diagnosis. In the 48-layer experiments, MoDA with post-norm (Row 4) achieves substantially better validation loss than MoDA with pre-norm (Row 3): 3.3653 vs. 3.3759, a difference of 0.0106. The baseline models show the opposite pattern: pre-norm OLMo2 (Row 1, 3.3800) outperforms post-norm OLMo2 (Row 2, 3.4062) by 0.0262. Standard architectural wisdom favors pre-norm for training stability in deep networks, and the baseline confirms this. But MoDA reverses the preference: with explicit depth retrieval, post-norm becomes the better choice. This interaction suggests that post-norm's training difficulty in deep networks may stem partly from the information dilution problem that MoDA addresses — making information retrieval harder when normalization is placed after the residual addition. When MoDA provides an alternative retrieval pathway, the representational benefits of post-norm (which has been argued to produce more expressive representations [6]) can be realized without the optimization penalty.

This diagnostic contribution is significant beyond MoDA itself. It suggests that the field's focus on optimization solutions to depth scaling may be partially misdirected, and that architectural solutions to the retrieval problem — of which MoDA is one instance — may unlock depth-scaling benefits that optimization improvements alone cannot. The finding also explains a puzzle in the literature: why methods like DenseFormer [28] show promise at small scale but don't scale to LLMs (because they solve the retrieval problem but incur prohibitive computational cost), while methods like DeepNet [40] improve training stability but don't fundamentally change depth-scaling efficiency (because they solve the optimization problem but not the retrieval problem). MoDA's contribution is demonstrating that the retrieval problem can be solved at acceptable computational cost, opening a new dimension for depth scaling beyond what optimization improvements alone can achieve.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses the OLMo2 dataset [27], a 400B-token subset drawn from C4 [30], ICE [27], m2d2-s2orc [24], Pile [14], Wiki-text [27], and dolma [34] (which includes Books, Common Crawl, peS2o, Reddit, and Stack subdomains). All main experiments train on this 400B-token corpus with a global batch size of 1024 sequences and context length of 4096 tokens. For the layer-number analysis in Section 4.3.1, a separate FineWeb-Edu data pipeline is used with a held-out validation split.

  • Base model(s). The paper uses the OLMo2 architecture [27] as the baseline across two model scales: 700M parameters (width D=1024, GQA group size G=2, 36 layers) and 1.5B parameters (architecture details not explicitly stated beyond using GQA). OLMo2 is chosen because it provides a strong, reproducible open-source baseline with a standardized training recipe, making comparisons fair and replicable. The authors also train smaller models (width 384, 6 query heads, 2 key-value heads) at 24 and 48 layers for the layer-number and norm-placement analysis in Section 4.3.1.

  • Metrics. The paper evaluates on three axes: (1) Training perplexity (Train PPL) measured on the training corpus; (2) Validation perplexity (Val PPL) measured on C4 [30] and reported per-domain on C4, ICE, m2d2-s2orc, Pile, Wiki-text, Books, Common Crawl (CC), peS2o, Reddit, Stack, and dolma — with an average across all domains reported; (3) Downstream task accuracy on 10 benchmarks: PIQA [5], HellaSwag [48], WinoGrande [32], OpenBookQA [26], BoolQA [9], SciQA [3], COPA [31], MMLU [17], ARC-Easy (ARC-E) [10], and ARC-Challenge (ARC-C) [10]. The downstream average is the mean accuracy across these 10 tasks. For the FineWeb-Edu layer-number experiments, only validation loss is reported.

  • Baselines. The primary baseline is OLMo2 [27] with standard causal self-attention — this is the vanilla Transformer decoder with residual connections and no depth-aware retrieval. For parameter-matched comparisons in Table 3, the paper includes a deeper OLMo2 baseline (38 layers, 700.5M parameters) to control for the parameter increase introduced by MoDA's FFN KV projections. All baselines use the same OLMo2 training recipe, data, and hyperparameters.

  • Generation budget / compute accounting. The paper measures computational cost in two ways: (1) Parameter count (in millions, M) to quantify model size overhead; (2) FLOPs (in trillions, T) to quantify computational overhead per forward pass. The FLOPs comparison in Table 3 shows that the minimal MoDA (attention KV reuse only, Row 3) adds only 0.12% FLOPs (8.01T → 8.02T), while the default MoDA with FFN KV projections (Row 4) adds 4.0% FLOPs (8.01T → 8.33T). There is no "generation budget" sweep (as in best-of-N sampling papers) because MoDA is an architectural modification evaluated at fixed training compute. Efficiency is also measured via end-to-end kernel runtime in milliseconds (ms) for forward+backward passes in Tables 2 and 7.

  • Cross-validation / statistical protocol. The paper does not report cross-validation or statistical significance testing. Results are reported as single-run metrics after training on the full 400B-token corpus. For the layer-number analysis (Table 6), a held-out FineWeb-Edu validation split is used, but this is a single fixed split rather than k-fold cross-validation. The absence of error bars, standard deviations, or multi-seed results is a notable limitation — the reported improvements (e.g., +2.11% downstream average at 1.5B) could include run-to-run variance that is not quantified.

Main Quantitative Results

MoDA Variant Comparison at 700M Scale

Table 3 presents the core variant comparison, training 700M models on 400B tokens. The headline finding is that the simplest form of depth-aware retrieval — reusing attention KV as depth memory with no additional parameters — already provides substantial gains, and that adding FFN-side depth KV projections yields the best accuracy-efficiency tradeoff.

Minimal MoDA (Row 3 vs. Row 1). With zero additional parameters and only 0.12% extra FLOPs, reusing attention-layer KV as depth KV improves:

  • Train PPL: 14.49 → 14.08 (−0.41)
  • C4 validation PPL: 18.59 → 18.48 (−0.11)
  • Downstream average: 56.93 → 58.10 (+1.17 percentage points)

This is remarkable because the model has access to no new information — it already computed these attention KVs during the forward pass — yet simply restructuring how that information is accessed (through a joint sequence+depth softmax rather than through the residual stream alone) yields gains comparable to or exceeding what an additional 2 layers provide (Row 2, which uses 31.5M more parameters and 0.40T more FLOPs but only improves downstream average to 57.11).

Default MoDA with FFN KV (Row 4 vs. Rows 1-3). Adding lightweight FFN KV projections (parameters 705.7M, +5.5% over baseline; FLOPs 8.33T, +4.0% over baseline) yields further improvements:

  • Train PPL: 14.08 → 13.90 (−0.18 vs. Row 3; −0.59 vs. baseline Row 1)
  • C4 validation PPL: 18.48 → 18.21 (−0.27 vs. Row 3; −0.38 vs. baseline)
  • Downstream average: 58.10 → 58.87 (+0.77 vs. Row 3; +1.94 vs. baseline)

Critically, Row 4 outperforms the deeper OLMo2 baseline (Row 2) despite having similar parameters (705.7M vs. 700.5M) and lower FLOPs (8.33T vs. 8.41T): Row 4 achieves 13.90 train PPL vs. Row 2's 14.27, 18.21 C4 PPL vs. 18.31, and 58.87 downstream vs. 57.11. This means that FFN-side depth retrieval is more parameter-efficient and compute-efficient than simply adding more Transformer layers.

Over-parameterized MoDA (Row 5 vs. Row 4). Adding separate attention-side depth KV projections on top of FFN KV (parameters 742.4M, +11.0% vs. baseline; FLOPs 8.63T, +7.7%) yields marginal additional gains:

  • Train PPL: 13.90 → 13.83 (−0.07)
  • C4 validation PPL: 18.21 → 18.17 (−0.04)
  • Downstream average: 58.87 → 58.97 (+0.10)

The paper characterizes this as "overly saturated" — the attention KV already encode sufficient information for depth retrieval, and doubling the depth-specific projections provides diminishing returns. This negative result is practically important because it tells practitioners to invest their parameter budget in FFN KV projections rather than attention KV projections.

Scaling MoDA to 1.5B Parameters

Table 4 (downstream benchmarks) and Table 5 (per-domain validation perplexity) show that MoDA's gains persist and in some cases grow when scaling from 700M to 1.5B parameters under the same 400B-token training budget.

Downstream performance (Table 4). At 1.5B:

  • OLMo2 baseline: 62.28 average across 10 tasks
  • MoDA (default configuration with FFN KV): 64.39 average (+2.11 percentage points)

The gains are broad rather than concentrated on a few tasks:

  • Commonsense/causal reasoning: HellaSwag +0.38 (65.86 → 66.24), WinoGrande +2.37 (63.22 → 65.59), COPA +4.00 (81.00 → 85.00)
  • Science-oriented reasoning: OpenBookQA +2.80 (38.80 → 41.60), ARC-C +4.35 (42.47 → 46.82), SciQ +1.50 (90.60 → 92.10)
  • Broad knowledge: BoolQ +3.73 (63.61 → 67.34), MMLU +1.86 (27.73 → 29.59)
  • Neutral/small change: PIQA +0.27 (76.55 → 76.82), ARC-E −0.17 (72.98 → 72.81)

Comparing the 700M and 1.5B gains: the downstream average improvement grows from +1.76 at 700M (Table 3, Row 4 vs. Row 2) to +2.11 at 1.5B (Table 4, Row 4 vs. Row 3), suggesting that MoDA's benefits scale with model capacity. This is consistent with the intuition that deeper models have more depth history to retrieve from, and that larger models have more representational capacity to learn sophisticated depth-retrieval strategies.

Per-domain validation perplexity (Table 5). At 1.5B, MoDA improves average PPL from 13.67 to 13.47 (−0.20) and improves all ten domains without exception:

  • Largest absolute improvements: m2d2-s2orc 21.10 → 20.92 (−0.18), C4 16.16 → 15.97 (−0.19), Reddit 21.21 → 20.85 (−0.36), ICE 15.37 → 15.08 (−0.29), Wiki-text 10.41 → 10.16 (−0.25)
  • Smaller but consistent improvements: Books 8.45 → 8.33 (−0.12), CC 18.13 → 17.88 (−0.25), peS2o 8.19 → 8.09 (−0.10), Stack 3.57 → 3.52 (−0.05)

At 700M, the average PPL improvement is 15.61 → 15.46 (−0.15), which is proportionally larger relative to the baseline (roughly 1.0% reduction) compared to 1.5B (roughly 1.5% reduction). The consistency across all ten domains at both scales is strong evidence that MoDA's benefits are not domain-specific — depth retrieval helps language modeling across diverse text types including web text (C4, CC), academic papers (m2d2-s2orc, peS2o), books, code (Stack), and conversational content (Reddit).

Figure 2 provides a complementary view in terms of training dynamics for the 1.5B models. Across the full 400B-token training trajectory, MoDA maintains lower C4 validation loss than OLMo2 at every point (training curves are consistently separated by a visible gap). For downstream tasks, HellaSwag and WinoGrande show MoDA outperforming OLMo2 from early in training (~100B tokens) and maintaining or widening the gap through 400B tokens, while ARC-Challenge shows MoDA initially tracking OLMo2 closely and pulling ahead more clearly in the later stages of training. This suggests that depth retrieval helps both early in training (where the model is still developing basic language capabilities) and later (where it's refining more complex reasoning patterns).

Layer-Number and Normalization Analysis

Table 6 investigates MoDA under different depth regimes (24 and 48 layers) and normalization placements (pre-norm vs. post-norm), using smaller models (width 384, 6 query heads, 2 key-value heads) trained on FineWeb-Edu. Because these are smaller-scale experiments with different data, the absolute loss values are not directly comparable to the 700M/1.5B results, but the relative patterns are informative.

Consistency across depth (Row 8 vs. Row 7, Row 4 vs. Row 3). At both 24 and 48 layers, adding Depth KV (attention KV reuse only, no FFN KV) consistently reduces validation loss:

  • 48 layers, pre-norm: 3.3800 → 3.3759 (−0.0041, Row 1 vs. Row 3)
  • 48 layers, post-norm: 3.4062 → 3.3653 (−0.0409, Row 2 vs. Row 4)
  • 24 layers, post-norm: 3.4740 → 3.4537 (−0.0203, Row 7 vs. Row 8)

The improvement at 48 layers post-norm (−0.0409) is substantially larger than at 48 layers pre-norm (−0.0041), a 10× difference. This is the paper's key finding about normalization: MoDA interacts synergistically with post-norm in deep models. The baseline OLMo2 shows the standard pattern where pre-norm outperforms post-norm (3.3800 vs. 3.4062, a 0.0262 advantage for pre-norm), but MoDA reverses this: with Depth KV enabled, post-norm (3.3653) outperforms pre-norm (3.3759) by 0.0106. The paper interprets this as evidence that post-norm's training difficulty in deep networks stems partly from information dilution (which MoDA addresses), and that once depth retrieval is available, post-norm's representational advantages (discussed in Chen & Wei, 2026 [6]) can be realized.

FFN KV projections provide consistent additional gains (Rows 5-6 vs. Rows 3-4, Row 9 vs. Row 8). At 48 layers:

  • Pre-norm with FFN KV: 3.3759 → 3.3656 (−0.0103, Row 3 vs. Row 5)
  • Post-norm with FFN KV: 3.3653 → 3.3484 (−0.0169, Row 4 vs. Row 6)

At 24 layers post-norm: 3.4537 → 3.4338 (−0.0199, Row 8 vs. Row 9).

The best overall configuration is 48 layers, post-norm, with both Depth KV and FFN KV (Row 6, loss 3.3484), which improves over the pre-norm OLMo2 baseline (Row 1, loss 3.3800) by 0.0316 — a substantial gain at this scale. The parameter overhead from FFN KV projections is modest (128.11M vs. 123.38M, a 3.8% increase) and the FLOPs increase from 136.61G to 144.00G (5.4%), similar to the overheads observed at 700M scale.

Interaction between depth and norm placement. The results suggest that MoDA's effectiveness depends on both depth and normalization. At 24 layers (Rows 7-9), the pattern is straightforward: Depth KV helps, FFN KV helps more, both additive. At 48 layers (Rows 1-6), a more complex interaction emerges: Depth KV alone provides only marginal benefit under pre-norm but substantial benefit under post-norm; FFN KV provides further gains under both, but the combination of MoDA + post-norm (Row 6) achieves the best absolute performance. This implies that for practitioners deploying MoDA in deep networks (48+ layers), post-norm placement should be preferred even though it would be suboptimal for standard Transformers.

Ablation Studies and Robustness Checks

Kernel implementation ablations (Table 7): The paper incrementally enables optimization components to quantify each one's contribution to runtime, using a fixed configuration (B=1, T=1024, G=8, Hq=64, Hk=8, d=64, L=64, C=64). The baseline naive PyTorch implementation takes 2128.900 ms. Flash-compatible depth-KV layout alone (Row 2) reduces this to 13.102 ms — approximately 162.5× faster, demonstrating that the fundamental bottleneck in naive MoDA is non-contiguous memory access rather than compute. Adding chunk-aware layout (Row 3) further reduces runtime to 6.286 ms (a 52.0% reduction from Row 2), because the chunked organization reduces HBM traffic from masked, out-of-range depth entries. Adding group-aware indexing (Row 4) yields a further 4.31× speedup to 1.460 ms, for a total 1458× end-to-end speedup over the naive baseline. The ablation confirms that each optimization addresses a distinct bottleneck (memory coalescing, then wasted computation from low depth utilization, then GQA-aware redundancy elimination).

Kernel efficiency scaling (Table 2): The paper sweeps sequence length T, GQA group size G, and model depth L to characterize MoDA's scaling behavior relative to FlashAttention-2 Triton.

  • Scaling T (Rows 1-5): At T=4096, MoDA adds 25.86% extra time vs. FlashAttention-2 (7.970 ms → 10.750 ms). As T grows to 65536, the extra time percentage drops monotonically to 2.73% (1831.668 ms → 1883.026 ms). This is because the sequence attention cost grows as O(T²) while the depth attention cost is O(TL) — at longer sequences, the depth overhead becomes a negligible fraction of total runtime. This is the paper's strongest efficiency result: for long-context training, MoDA is essentially free.

  • Scaling G (Rows 6-10): At fixed T=16384, as G grows from 2 to 32, the extra time percentage drops from 27.07% to 2.84%. Depth utilization η_depth = G/C rises from 3.12% to 50.00%, directly reducing wasted computation. This means MoDA is most efficient at the larger GQA group sizes (G=8, 16, 32) that are standard in production LLMs — the 700M experiments with G=2 represent a relatively unfavorable efficiency regime.

  • Scaling L (Rows 11-13): At fixed T=16384 and G=8, as L grows from 64 to 256, FlashAttention-2 time remains constant (116.700 ms) while MoDA time grows from 127.661 ms to 167.958 ms, with extra time percentage rising from 8.59% to 30.52%. This is MoDA's fundamental scaling limitation: deeper models spend proportionally more time in the depth attention loop. The paper acknowledges this in Section 6.2, proposing bounded depth-KV slot caching as a mitigation for industrial-scale depths, but this mitigation is not implemented or evaluated.

Extra Attn KV Proj. saturation (Table 3, Row 5 vs. Row 4): Adding separate attention-side depth KV projections to the default MoDA (which already has FFN KV projections) yields only +0.07 train PPL, +0.04 C4 PPL, and +0.10 downstream average improvements, while increasing parameters by 5.2% (705.7M → 742.4M) and FLOPs by 3.6% (8.33T → 8.63T). This ablation establishes that attention-layer KV already contain sufficient information for depth retrieval, and that the parameter budget is better spent on FFN-side projections or not spent at all.

Parameter-matched comparison against deeper baseline (Table 3, Row 4 vs. Row 2): The default MoDA (705.7M parameters, 8.33T FLOPs) is compared against a 38-layer OLMo2 (700.5M, 8.41T) with comparable parameters and higher FLOPs. MoDA achieves better train PPL (13.90 vs. 14.27), C4 PPL (18.21 vs. 18.31), and downstream average (58.87 vs. 57.11), demonstrating that MoDA's gains are not attributable to increased parameter count — the architecture is genuinely more parameter-efficient than simply adding layers.

Post-norm + MoDA synergy (Table 6, Rows 1-6): This is not a controlled ablation in the traditional sense (since post-norm is a different training regime, not an isolated variable), but the comparison between Rows 1-2 (OLMo2 pre-norm vs. post-norm) and Rows 3-6 (MoDA pre-norm vs. post-norm) reveals that MoDA reverses the standard pre-norm advantage. The baseline: post-norm is 0.0262 worse than pre-norm. With MoDA (Depth KV only): post-norm is 0.0106 better than pre-norm. With MoDA (Depth + FFN KV): post-norm is 0.0172 better than pre-norm. This interaction is robust across both MoDA variants and suggests that depth retrieval specifically alleviates whatever mechanism makes post-norm underperform in deep standard Transformers.

Attention visualization qualitative analysis (Figure 5): While not a quantitative ablation, the attention heatmaps provide evidence that MoDA heads actively use the depth retrieval channel rather than ignoring it. Across sampled layers {0, 11, 23, 35} and randomly selected heads, substantial probability mass is allocated to the depth-KV block (right of the red dashed line). The pattern is complementary: some heads show sharp diagonal sequence attention with a smaller depth-KV allocation, while others allocate more mass to depth slots. The paper notes qualitatively that these heads appear to distribute probability more broadly across informative positions rather than concentrating on fixed attention sinks — a claim supported by the visual evidence but not quantified with attention entropy or sink-position metrics.

Critical Assessment

The experiments demonstrate that MoDA improves language modeling perplexity and downstream task performance across two model scales (700M, 1.5B) under a standardized training recipe, and that these improvements are robust across all ten validation domains and the majority of downstream benchmarks. The evidence for the paper's central architectural claim — that explicit depth-aware retrieval mitigates information dilution in deep Transformers — is strong in the sense that (a) the gains are consistent, (b) they persist and in some cases grow with scale, (c) they cannot be explained by increased parameter count alone (the parameter-matched deeper baseline performs worse), and (d) they interact with depth and normalization in ways that are consistent with the proposed mechanism (MoDA helps more in deeper models and specifically rescues post-norm in deep regimes).

However, several aspects of the experimental design limit the strength and generality of the conclusions:

The claim that MoDA "addresses the information dilution problem" is supported indirectly rather than directly. The paper never provides a direct diagnostic of information dilution — for example, measuring how well features from layer 3 can be decoded from the residual stream at layer 27 with and without MoDA, or showing that attention patterns in later layers specifically retrieve early-layer features that would otherwise be inaccessible. The evidence is all performance-based (lower perplexity, higher accuracy), which is consistent with the information dilution hypothesis but equally consistent with other mechanisms (e.g., MoDA providing a useful inductive bias for routing, acting as a form of auxiliary training signal, or simply increasing effective model capacity through the depth-KV pathway). The attention visualization in Figure 5 shows that depth retrieval is used, but does not show that it recovers diluted information specifically. A more direct test would be a probing experiment: train linear classifiers to recover specific linguistic features (part-of-speech, named entities, syntactic dependencies) from the residual stream at various depths, and show that MoDA models maintain higher probing accuracy at later layers compared to baselines.

The experimental scale is limited relative to modern LLM practice. The largest model trained is 1.5B parameters on 400B tokens — this is roughly 1-2 orders of magnitude smaller than the models where depth scaling becomes most critical (e.g., 70B+ parameter models with 80+ layers). The paper's complexity analysis shows that MoDA's FLOPs overhead is O(TL²D), which grows quadratically with depth. At L=64, the overhead is 2.73-8.59% (Table 2), which is manageable. But at L=128 or L=256 — depths that are increasingly common in frontier models — the overhead could grow to 15-30% or more based on the scaling trends in Table 2 (L=256 shows 30.52% extra time even at T=16384). The paper does not validate whether the performance gains also scale with depth at these larger regimes, or whether they saturate or diminish. This is a critical gap because the paper's stated motivation is specifically about enabling deeper Transformers.

The evaluation is limited to language modeling metrics on relatively standard benchmarks. The downstream tasks (PIQA, HellaSwag, WinoGrande, etc.) are all relatively short-context, single-turn evaluation tasks. The paper does not evaluate on tasks that specifically stress depth-dependent reasoning — for example, multi-step mathematical reasoning (where intermediate computations at different depths might be differentially important), long-document question answering (where depth retrieval across a long context might interact with sequence attention in complex ways), or code generation (where syntactic and semantic features at different depths might have clear interpretations). The absence of evaluations on these depth-stressing tasks makes it harder to attribute the observed gains specifically to improved depth information flow rather than to a general representational improvement that could come from any well-designed architectural modification.

The training budget for the 1.5B experiments (400B tokens) may be below the compute-optimal point for models of this size. The Chinchilla scaling laws [19] suggest that a 1.5B parameter model would be compute-optimally trained on roughly 30B tokens, so 400B tokens represents substantial over-training relative to that standard. However, the OLMo2 recipe is specifically designed for over-trained models (following the LLaMA paradigm), so this is not necessarily a flaw — but it means the results are specific to the over-trained regime. It's possible that MoDA's benefits are larger in over-trained settings (where depth retrieval helps extract more value from limited parameters) and would be smaller in compute-optimal training regimes, or vice versa.

The paper lacks multi-seed or statistical significance reporting. All results appear to be single training runs. At 700M and 1.5B scales with 400B tokens, run-to-run variance on downstream tasks can be non-trivial (1-2 percentage points on individual tasks, potentially 0.5-1.0 points on averages). The reported +2.11% downstream average improvement at 1.5B is large enough to likely be significant even accounting for variance, but the +1.76% at 700M is closer to the boundary where multiple seeds would be informative. The per-task results in Table 4 show some anomalies consistent with variance: ARC-E actually decreases slightly (−0.17 at 1.5B) despite all other tasks improving, and PIQA shows very small gains (+0.27 at 1.5B, and actually −0.33 at 700M in Table 3 Row 4 vs. Row 2). Without error bars, it's impossible to distinguish genuine task-specific patterns from noise.

The comparison against "adding more layers" is limited. The paper includes one deeper baseline (38-layer OLMo2 vs. 36-layer MoDA in Table 3), which is a reasonable parameter-matched comparison. But a more complete analysis would sweep over multiple depth configurations to establish the tradeoff curve: how many additional layers of standard Transformer does MoDA effectively provide? For example, does 36-layer MoDA match 38-layer, 40-layer, or 44-layer OLMo2 in performance? This "effective depth multiplier" would be a more interpretable metric for practitioners deciding whether to adopt MoDA versus simply training a deeper standard model.

Missing ablations. Several experiments would have strengthened the paper's claims:

  • MoDA with only FFN KV (no attention KV depth): The paper shows that FFN KV adds value on top of attention KV, but does not test whether FFN KV alone (without reusing attention KV) is sufficient. If FFN KV alone provided most of the gains, the mechanism would be simpler and the attention-KV reuse could be eliminated.
  • Sliding window depth (only recent layers): Section 6.2 proposes bounded depth-KV caching as future work, but a simple ablation — restricting depth attention to only the most recent K layers (e.g., last 8, last 16) rather than all preceding layers — would test whether full depth history is necessary or whether local depth context suffices.
  • MoDA with different GQA group sizes at matching model scale: The efficiency analysis (Table 2) shows MoDA is more efficient at larger G, but all training experiments use G=2 (700M) or unspecified G (1.5B). Testing MoDA at G=8 with 1.5B parameters would validate that the performance gains persist in the efficiency regime where MoDA is most practical.
  • Comparison against other depth-aware architectures: The paper only compares against OLMo2 (standard residual Transformer), but prior work like DenseFormer [28] or hyper-connections [49] represents alternative approaches to the same problem. A head-to-head comparison at small scale (e.g., the 24-layer FineWeb-Edu setting) would contextualize MoDA's efficiency-accuracy tradeoff relative to these alternatives.

The hardware efficiency claims are strong but context-dependent. The 97.3% of FlashAttention-2 efficiency is achieved at T=65536, G=8, L=64 — a specific, favorable configuration. At more common training lengths (T=4096-8192), the overhead is 8.59-25.86% (Table 2), which is still good but not "nearly free." The paper's claim that MoDA is "practical for long-context LLM training" is well-supported for the long-context regime specifically, but practitioners training at moderate sequence lengths should expect 10-25% overhead, not 3%. This is an important nuance that the abstract's "97.3% of FlashAttention-2's efficiency" headline number obscures.

The core mechanism — joint softmax competition between sequence and depth — is never ablated. The paper argues that the unified softmax is a key innovation (Section 2.2, the distinction between Depth Attention's separate softmaxes and MoDA's joint softmax), but never empirically compares unified vs. separate softmax implementations. This is a significant missing experiment because it would directly test whether the joint normalization provides representational benefits beyond simply having access to depth information. If separate softmaxes performed equally well, the unified softmax would be an implementation detail rather than a conceptual contribution.

Overall, the experiments provide convincing evidence that MoDA improves language model performance at 700M-1.5B scale under the OLMo2 training recipe, and reasonable evidence that these improvements stem from better depth information utilization (supported by the layer-number and norm-placement interactions, the attention visualization, and the parameter-matched baseline comparisons). The paper's central claims are supported for the tested regime but extrapolation to larger scales, deeper models, and production deployment remains untested. The efficiency results are robust and well-characterized across multiple scaling dimensions, though the "97.3% efficiency" figure requires careful contextualization. The most significant gaps are the absence of direct mechanistic evidence for information dilution mitigation, the lack of multi-seed statistical validation, and the limited scale of the largest training experiments relative to the depths where the problem MoDA addresses is most severe.

6. Limitations and Trade-offs

The Quadratic-Depth FLOPs Growth is Untested at Production Scale

The constraint. MoDA's FLOPs for the depth attention component scale as O(L²D) — quadratic in the number of layers. The paper's efficiency analysis in Table 2 shows this concretely: at fixed sequence length T=16384 and GQA group size G=8, increasing depth from L=64 to L=256 raises MoDA's extra time overhead over FlashAttention-2 from 8.59% to 30.52%. The paper acknowledges this is the fundamental scaling limitation, stating in Section 6.2 that "when scaling to very deep networks, caching all depth-KV states from all historical layers introduces substantial memory and bandwidth overhead. The cost grows linearly with depth, and can become the dominant bottleneck in long-context training and serving."

The consequence. The largest model trained in the paper is 1.5B parameters at an unspecified but likely modest depth (the 700M model uses 36 layers; the 1.5B model presumably uses more, but its depth is not explicitly stated). At the layer counts typical of frontier LLMs — L=80 (Llama 3 70B), L=96 (DeepSeek-V2), L=128+ (various dense models) — the depth attention overhead could grow from ~10% to 40-60% or more, potentially erasing the practical efficiency advantages that make MoDA deployable at smaller scale. This matters because the paper's stated motivation is specifically about enabling deeper Transformers: "how can a model scale depth while maintaining optimization stability and preventing information dilution?" (Section 1). If the method's cost scales quadratically with the very dimension it aims to improve, it may not be the solution it claims to be at the depths where the problem is most severe.

Evidence in the paper. Table 2, Rows 11-13: at L=64, MoDA adds 8.59% time; at L=128, 15.57%; at L=256, 30.52%. The paper does not train any models beyond L=48 (the 48-layer experiments in Table 6 use a tiny width of 384, which is not representative of production model aspect ratios). There is no experiment showing that the performance gains from MoDA continue to grow (or even persist) at L=128 or L=256 relative to the increasing overhead.

Mitigation status. The paper proposes bounded depth-KV slot caching in Section 6.2 as a future direction: "instead of storing all depth-KV entries, each query only attends to a bounded set of slots. The slot budget is fixed to S, where S << L." Two policies are sketched — dynamic utility-based selection and sliding-window recency — but neither is implemented or evaluated. The paper frames this as "future work" and acknowledges that "the key challenge is the quality of slot assignment." Until such bounded caching is validated, MoDA's deployability at industrial depth scales (L > 64) remains an open question, and the quadratic-depth FLOPs term represents a hard ceiling on naive scaling.

Difficulty Estimation or Dynamic Allocation Is Completely Absent

The constraint. MoDA treats all tokens and all layers uniformly: every token at every layer attends to depth KV from all preceding layers, with the same architecture applied identically regardless of whether depth retrieval is actually useful for that particular token or layer. The paper provides no mechanism for the model to learn when depth retrieval is beneficial versus when it adds noise, nor any difficulty-estimation or gating mechanism that could adaptively allocate the depth attention budget.

The consequence. The attention visualization in Figure 5 shows that MoDA heads do allocate probability mass to depth KV entries — but it does not show that this allocation is selective or efficient. Some heads and some layers may benefit substantially from depth retrieval, while others may perform worse because depth KV entries act as distracting noise in the joint softmax (competing for probability mass that would otherwise go to informative sequence positions). The uniform application of depth attention means that the model cannot "turn off" depth retrieval when it is harmful, only learn to assign low attention weights — which still consumes the depth-KV memory bandwidth and compute for loading and scoring depth entries, even when the resulting attention weights are near zero. More fundamentally, there is no mechanism to allocate more depth retrieval budget to layers or tokens that would benefit most (e.g., later layers recovering early-layer features for complex compositional reasoning), and less to those where depth adds no value (e.g., early layers with little depth history, or tokens where local sequence context is fully sufficient). This is the same class of limitation identified in prior work on test-time compute scaling — uniform allocation leaves efficiency on the table — but applied to the depth dimension rather than the generation budget.

Evidence in the paper. No ablation studies test whether restricting depth attention to a subset of layers (e.g., only the last N layers, or only every K-th layer having depth access) maintains or improves performance. No experiment varies the depth-KV budget per layer or per token. The attention visualization in Figure 5 shows heads with very different depth-utilization patterns (some allocate substantial mass to depth KV, others very little), confirming that depth retrieval utility varies across heads and layers, but the architecture provides no mechanism to exploit this heterogeneity.

Mitigation status. Not addressed. The paper's discussion of bounded depth-KV caching in Section 6.2 touches on slot selection (which depth entries to keep) but not on allocation (which queries should do depth attention at all, or how much of the softmax budget should go to depth vs. sequence). A dynamic gating mechanism — for example, a learned scalar gate per head that interpolates between pure sequence attention and joint sequence+depth attention, with the gate conditioned on the query representation — would be a natural extension but is not proposed or evaluated.

The Unified Softmax Innovation Is Not Empirically Validated Against Separated Alternatives

The constraint. The paper presents the unified softmax (joint normalization over sequence and depth keys) as a key architectural innovation that distinguishes MoDA from the intermediate Depth Attention formulation. The "read, operate, write" progression in Section 2.2 and Figure 3 explicitly positions the unified softmax as the upgrade: Depth Attention uses "attention to read historical depth KV pairs" but processes them separately from sequence attention, while MoDA "combines depth attention with standard sequence attention" under one softmax. The conceptual argument — that joint normalization forces competition between sequence and depth information sources, creating representational pressure for selective retrieval — is central to the paper's narrative about why MoDA works.

The consequence. Without an empirical comparison between unified and separated softmaxes, it is impossible to determine whether the joint softmax is genuinely responsible for MoDA's performance gains, or whether the gains would be achieved equally well (or better) by simply providing depth KV as an additional attention input with separate normalization. The Depth Attention formulation in Equation 5 — which computes depth attention with a separate softmax, then feeds the result into standard sequence attention — represents a natural ablation that would directly test the unified softmax claim. If Depth Attention performed comparably to MoDA, the joint softmax would be an implementation detail rather than a conceptual contribution, and the paper's framing of the "mixture" as a fundamental innovation would be overstated. If Depth Attention performed worse, the comparison would quantify how much the joint normalization matters and provide evidence for the competition-based mechanism the paper hypothesizes.

Evidence in the paper. The paper never implements or evaluates the separated Depth Attention formulation at scale. It is introduced conceptually in Section 2.2 as a "stepping stone" but immediately superseded by MoDA. The complexity analysis in Table 1 treats them as distinct rows (Depth Attention vs. MoDA) but only in terms of asymptotic complexity, not empirical performance. No experiment compares (a) separate depth softmax + separate sequence softmax, (b) depth softmax followed by sequence softmax (sequential), or (c) joint softmax (MoDA) — any of which would isolate the effect of the unified normalization.

Mitigation status. Not addressed. The paper provides no ablation on this design choice and does not acknowledge it as a missing experiment. Given that the joint vs. separated softmax distinction is one of the paper's claimed contributions (the "mixture" in Mixture-of-Depths Attention), this is a significant gap. A practitioner reading the paper cannot determine whether implementing the conceptually simpler Depth Attention with separate softmaxes would yield equivalent results with potentially simpler implementation and better hardware characteristics (separate softmaxes might allow more favorable tiling strategies).

The Single-Model-Family, Single-Benchmark-Category Evaluation Limits Generality Claims

The constraint. All training experiments use the OLMo2 architecture and training recipe, with evaluation on standard language modeling benchmarks (perplexity on C4 and related corpora) and commonsense/reasoning downstream tasks (PIQA, HellaSwag, WinoGrande, etc.). The paper states in Section 5 that "MoDA is architecture-agnostic and can be readily integrated into multimodal intelligence, visual understanding, and world models, where Transformers are increasingly adopted," but provides no empirical evidence beyond decoder-only language modeling.

The consequence. Several aspects of MoDA's design may interact with architectural choices in ways that are not captured by the OLMo2 experiments. The efficiency characteristics depend on the GQA group size G (Table 2 shows overhead drops from 27.07% at G=2 to 2.84% at G=32), meaning MoDA's practicality varies substantially with the attention configuration. Encoder-decoder architectures (where the encoder has bidirectional attention and the decoder has cross-attention to encoder outputs) would require depth-KV management for both the encoder and decoder stacks, potentially with different optimal configurations. Vision Transformers typically operate on much longer effective sequence lengths (image patches) with different attention patterns, and the depth utilization optimization (which depends on the chunk size C and GQA group size G) would need re-tuning. Multimodal models that interleave modality-specific encoders with a shared LLM backbone would require decisions about whether depth KV are shared across modalities or modality-specific. The paper's strong generality claims ("architecture-agnostic," "can be readily integrated") are supported by zero experiments outside the specific OLMo2 language modeling setting.

Evidence in the paper. All training results (Tables 3-6, Figures 2, 5) are on decoder-only language models with the OLMo2 architecture. The runtime benchmarks (Tables 2, 7) use a fixed attention configuration (Hq=64, Hk=8, d=64) that is typical for language models but not necessarily for vision or multimodal architectures. The paper does not report results on code generation, mathematical reasoning, multi-lingual tasks, or any domain that would stress test the generality of depth retrieval.

Mitigation status. The paper claims architecture-agnosticism in Section 5 without qualification and suggests future application to multimodal and visual domains, but provides no evidence and does not acknowledge this as a limitation. The absence of even a small-scale experiment on a non-language or encoder-decoder architecture makes the generality claims aspirational rather than empirically grounded.

The Claimed Mechanism (Mitigating Information Dilution) Lacks Direct Causal Evidence

The constraint. The paper's motivating problem is information dilution: "informative features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers" (Section 1). The proposed solution is that MoDA allows later layers to "adaptively read useful states from earlier layers" (Section 1), bypassing the diluted residual stream. However, the paper provides only indirect, performance-based evidence that this mechanism is actually responsible for the observed gains.

The consequence. MoDA's performance improvements are consistent with the information dilution hypothesis, but equally consistent with alternative mechanisms that have nothing to do with dilution. For example: (a) MoDA provides additional key-value parameters (the FFN KV projections) that increase model capacity, and the depth attention simply learns to use these as auxiliary memory slots; (b) the joint softmax acts as a form of regularization that improves optimization by providing gradient pathways that skip intermediate layers (analogous to how residual connections improve gradient flow, but through attention rather than addition); (c) the depth KV act as a learned positional bias that helps the model represent absolute depth position, which may matter for tasks sensitive to where in the network a computation occurs. Without direct evidence — such as probing experiments showing that specific early-layer features are more accessible at later layers in MoDA than in the baseline, or causal intervention experiments showing that disrupting early-layer features has less impact on later-layer representations in MoDA — the paper's central mechanistic claim remains a plausible hypothesis rather than an established finding.

Evidence in the paper. The attention visualization in Figure 5 shows that depth KV entries receive attention weight, but does not show what information is being retrieved from those entries or whether that information would have been inaccessible through the residual stream. The layer-number analysis (Table 6) shows that MoDA helps more in deeper models (the improvement from Depth KV is 0.0409 at L=48 vs. 0.0203 at L=24 in post-norm), which is consistent with information dilution being worse in deeper networks, but could also be explained by deeper models having more depth history to retrieve from (more layers = more depth KV entries = more capacity), regardless of dilution. The post-norm interaction — MoDA reverses the pre-norm advantage in deep models — is interpreted as evidence that MoDA addresses the specific mechanism that makes post-norm underperform in deep standard Transformers, but the paper does not establish what that mechanism is or directly show that it relates to information dilution rather than, say, gradient norm dynamics.

Mitigation status. The paper does not acknowledge this as a limitation. The abstract states that MoDA "addresses the information dilution problem" as a factual claim, and the introduction presents the information dilution diagnosis as settled. A probing experiment (e.g., training linear classifiers to decode part-of-speech tags or named entities from the residual stream at various depths, comparing MoDA vs. baseline), an intervention experiment (e.g., ablating specific early-layer attention heads and measuring the impact on later-layer representations), or even a simple representational similarity analysis (e.g., measuring CKA similarity between early-layer and late-layer representations) would substantially strengthen the mechanistic claim. Without such evidence, the paper demonstrates that MoDA improves performance but does not demonstrate how or why in a causal sense.

Training-Test Contamination from the Difficulty of Validating on Standard Benchmarks

The constraint. The paper evaluates on standard downstream benchmarks (HellaSwag, WinoGrande, ARC, MMLU, etc.) that are widely used in the LLM literature and may have been included — directly or indirectly — in the OLMo2 training corpus. The OLMo2 dataset is a 400B-token subset drawn from C4, ICE, m2d2-s2orc, Pile, Wiki-text, and dolma, which are large-scale web crawls that could contain text from or related to these benchmark datasets.

The consequence. If benchmark data is present in the training corpus, the downstream accuracy improvements may partially reflect better memorization of benchmark-specific patterns rather than genuine improvements in the underlying capabilities (commonsense reasoning, factual knowledge, etc.) that the benchmarks are designed to measure. This is a concern for any LLM evaluation, but it is particularly relevant for MoDA because depth retrieval could specifically aid in memorizing and retrieving training-time patterns: if a benchmark example appears verbatim or near-verbatim in the training data, early layers might encode the relevant information, and MoDA's depth retrieval could make it easier for later layers to access that memorized content. The per-domain perplexity improvements in Table 5 (consistent gains across all 10 domains) are less susceptible to this concern because they measure language modeling quality on held-out data rather than task-specific accuracy, but the downstream benchmark improvements could be inflated.

Evidence in the paper. The paper does not report any decontamination analysis of the training data relative to the evaluation benchmarks. The OLMo2 paper [27] likely includes decontamination results, but this paper does not reference them or conduct its own analysis. The training data description ("400B-token-subsets of OLMo2 dataset") is too vague to assess contamination risk without consulting the OLMo2 paper directly.

Mitigation status. Not addressed. The paper does not mention benchmark contamination as a potential confound. Since this is a general concern for all LLM evaluations rather than specific to MoDA, and since the primary performance metric is validation perplexity (which is measured on held-out data by construction), this limitation is less severe than the others listed here. However, the headline downstream accuracy improvements (+2.11% at 1.5B) should be interpreted with the standard caveat that benchmark scores in LLM papers may partially reflect training-data memorization rather than genuine capability improvements, especially when training on large web-scale corpora.

7. Implications and Future Directions

How This Work Changes the Landscape

MoDA represents a diagnostic reframing rather than a paradigm shift, but it is a reframing with substantial practical consequences. The paper's central contribution is not a new architecture that obsoletes the Transformer — MoDA is explicitly a drop-in replacement for standard self-attention, leaving the rest of the Transformer stack intact — but rather a reconceptualization of what the depth dimension is for and how information should flow through it. This reframing has three interconnected effects on the field.

First, it redirects depth-scaling research from optimization toward retrieval. Prior to MoDA, the dominant narrative for why deep Transformers underperform their theoretical capacity centered on optimization difficulties: vanishing/exploding gradients, training instability, and the challenge of propagating learning signals through many layers. This narrative motivated solutions like improved initialization (DeepNet), gated residual variants (hyper-connections, manifold-constrained hyper-connections), and careful normalization placement (the pre-norm vs. post-norm debate). These are all, fundamentally, solutions to optimization problems — they make training more stable and gradients better-behaved.

MoDA's success — achieving consistent gains using the identical OLMo2 training recipe, with no changes to learning rate schedules, initialization, or normalization (in the pre-norm experiments) — suggests that optimization was never the primary bottleneck. The bottleneck was retrieval: early-layer features were technically present in the residual stream at later layers, but the additive superposition accumulated so much noise that extracting them became increasingly difficult. MoDA solves this by providing an explicit, content-addressable retrieval pathway that bypasses the diluted residual stream entirely. The post-norm synergy finding in Table 6 — MoDA with post-norm outperforms both MoDA with pre-norm and OLMo2 with pre-norm, reversing the standard pre-norm advantage in deep networks — provides further evidence: post-norm's training difficulty in deep networks may itself be a symptom of information dilution, not an independent optimization pathology.

This reframing has a specific consequence for research prioritization: architectural innovations that improve depth retrieval (like MoDA) are likely to unlock more depth-scaling gains than further optimization improvements (better initialization schemes, more sophisticated normalization, etc.). The optimization community has been iterating on residual connection variants for nearly a decade since ResNet; MoDA's results suggest that the larger remaining gains lie on the retrieval side. This doesn't mean optimization research is obsolete — training stability still matters — but it does suggest that retrieval-aware architectures may be a higher-leverage investment for depth scaling specifically.

Second, it establishes depth as a first-class information dimension alongside sequence length. The standard Transformer formulation treats the sequence dimension as the primary axis of information mixing (through self-attention) and the depth dimension as a secondary axis of progressive transformation (through stacked layers with residual connections). This asymmetry — dynamic, content-based retrieval for sequence; fixed, additive accumulation for depth — is not theoretically motivated; it is a historical accident of how Transformers were designed. MoDA demonstrates that removing this asymmetry — applying the same attention-based retrieval principle to both dimensions — yields consistent improvements across model scales, domains, and evaluation metrics.

This has implications for how we think about Transformer capacity. If depth can be treated as a retrieval dimension, then the effective "memory" of a Transformer is not just the sequence length T (the number of token positions that can be attended to) but also the depth L (the number of computational stages whose outputs can be retrieved). A 48-layer MoDA model at sequence length 4096 has, in principle, T + L = 4144 retrievable memory slots per attention head — 4096 sequence positions plus 48 depth positions (the current layer plus previous layers). This is a small increase in absolute terms (about 1.2% more slots), but if depth retrieval provides qualitatively different information (computational history rather than token context), the effective representational capacity could increase more than the raw slot count suggests.

The implication for architecture design is that depth and sequence length may be partially substitutable. If depth retrieval can compensate for limited sequence context (by storing and retrieving intermediate computational results rather than keeping all relevant tokens in the active context window), or if sequence context can compensate for limited depth (by spreading computation across tokens rather than layers), then the optimal architecture for a given parameter budget might involve trading off depth against sequence length in ways that current design practices don't consider. The paper doesn't explore this tradeoff directly, but the unified softmax formulation — where depth and sequence compete for the same attention budget — provides a natural mechanism for learning the optimal allocation dynamically.

Third, it reconciles a tension in the literature between the theoretical appeal of depth scaling and its practical underperformance. Scaling laws research (Kaplan et al., 2020; Hoffmann et al., 2022) has consistently found that model width and training data scale more predictably and efficiently than model depth, leading to a practical preference for wider, shallower architectures at scale. Yet there are strong theoretical arguments for depth: deeper networks can represent more compositional functions, with each layer performing a specific transformation that builds on earlier ones, and there is evidence from computer vision that depth provides representational benefits that width cannot easily substitute for.

MoDA suggests a resolution to this tension: depth scaling underperforms in practice not because depth is inherently less valuable, but because the standard residual architecture fails to preserve and retrieve the compositional representations that depth creates. If each layer's output is immediately diluted by subsequent additive updates, the compositional structure that deeper networks are supposed to exploit — where layer 5 builds on layer 3's output in a specific, recoverable way — is lost. MoDA preserves this structure by storing layer outputs explicitly and allowing content-based retrieval, potentially enabling the theoretical benefits of depth that have been elusive in practice.

The paper's finding that MoDA's gains persist and in some cases grow when scaling from 700M to 1.5B parameters (the downstream average improvement increases from +1.76 to +2.11 points) provides tentative evidence that depth retrieval becomes more valuable at larger scales — consistent with the idea that larger models have more sophisticated compositional computations to preserve and retrieve. If this trend continues at larger scales (7B, 70B, etc.), MoDA could shift the optimal architecture frontier toward deeper, narrower configurations than current practice favors.

What this does not change: MoDA does not challenge the fundamental Transformer paradigm — it operates within the attention framework and preserves the decoder-only autoregressive structure. It does not address the quadratic sequence-length complexity of attention (the sequence attention component is unchanged). It does not provide a new theory of depth scaling or a scaling law for depth analogous to Chinchilla for training compute. And it does not solve the hardest version of the depth problem: when the base model fundamentally lacks the capability to perform a task (the analog of "difficulty bin 5" in the test-time compute scaling literature), no amount of depth retrieval will create that capability — depth retrieval amplifies existing computations but does not enable qualitatively new ones.

Follow-Up Research This Work Enables

Scaling MoDA to production depths (L=80, 128, 256) with bounded depth-KV caching. The most urgent open question is whether MoDA's performance gains persist at the layer counts where information dilution is most severe — and whether the quadratic-depth FLOPs growth can be controlled through the bounded slot-caching mechanism sketched in Section 6.2. The paper's largest depth experiment is L=48 (Table 6), which shows clear gains. A direct follow-up would train 1.5B-parameter models at L=64, L=96, and L=128, comparing MoDA with full depth caching against MoDA with sliding-window depth (only the most recent S layers' KV, where S is varied from 8 to 64) and against a standard OLMo2 baseline of equal depth. The key measurements would be: (a) Does the performance gap between MoDA and the baseline grow, shrink, or saturate as L increases? (b) What is the minimum S that recovers most of full-caching MoDA's gains — i.e., is the full depth history necessary, or is local depth context sufficient? (c) At what depth does the bounded cache become strictly better than full caching (because full caching's overhead outweighs the marginal benefit of very old depth KV)? The paper's Table 2 efficiency data (L=256 costs 30.52% extra time) makes this experiment practically urgent: if bounded caching at S=32 can recover 90% of full-caching gains at L=128 with only 10-15% overhead, MoDA becomes immediately deployable at frontier scales; if not, the quadratic-depth FLOPs term is a hard practical ceiling.

Direct causal evidence for the information dilution mitigation mechanism. The paper's central mechanistic claim — that MoDA works by allowing later layers to retrieve early-layer features that would be diluted in the residual stream — is supported only by performance improvements and attention visualizations, not by direct causal experiments. A follow-up study should design probing and intervention experiments to test this mechanism directly. One approach: train linear probes to decode specific linguistic features (part-of-speech tags, named entity types, syntactic dependency labels, coreference chains) from the residual stream and from the depth KV cache at each layer, in both MoDA and OLMo2 models at matched depths. If MoDA mitigates information dilution, the probing accuracy from the residual stream should be similar between MoDA and OLMo2 (since both use residual connections), but MoDA should show higher probing accuracy from its depth KV entries at later layers for features computed in early layers — evidence that depth retrieval successfully "rescues" diluted information. A complementary intervention experiment: ablate specific attention heads in early layers (e.g., zero out their outputs) and measure the impact on later-layer representations in MoDA vs. OLMo2. If MoDA enables later layers to retrieve ablated information from depth KV that would otherwise be lost, the impact of early-layer ablation should be smaller in MoDA (because later layers can compensate via depth retrieval) for features that have depth-KV entries, but not for features that don't.

Unified vs. separated softmax ablation to validate the "mixture" claim. The paper argues that the joint softmax over sequence and depth keys is a key innovation, but never empirically compares it against a separated-softmax baseline (Depth Attention, as defined in Equation 5). A clean ablation would train three small-scale models (e.g., the 24-layer FineWeb-Edu configuration from Table 6): (a) MoDA with joint softmax (the standard formulation), (b) Depth Attention with separate softmaxes (depth attention first, producing a depth-aware representation, then sequence attention on that representation), and (c) a parallel-attention variant (depth and sequence attention computed independently and summed, with a learned mixing weight). All three would have access to identical depth KV information; the only difference is how that information is integrated with sequence context. This experiment would answer: does the joint softmax provide representational benefits beyond simply having access to depth information? If the joint softmax outperforms, it validates the competition-based mechanism (depth and sequence competing for attention mass forces selective retrieval). If separated softmaxes perform equally well, the "mixture" aspect is an implementation detail, and practitioners could choose the variant with better hardware characteristics (separated softmaxes might allow more favorable tiling and memory access patterns, potentially reducing the overhead reported in Table 2).

MoDA for encoder-decoder and vision Transformers. The paper claims MoDA is "architecture-agnostic" (Section 5) but only evaluates it on decoder-only language models. Two natural extensions would test this claim. First, an encoder-decoder language model (e.g., T5-style) where both the encoder and decoder stacks use MoDA, with the decoder's cross-attention to encoder outputs potentially augmented with depth KV from the encoder stack. The key question is whether depth retrieval in the encoder (which processes the input bidirectionally) provides different benefits than in the decoder (which processes autoregressively), and whether cross-attention can benefit from depth-aware encoder representations. Second, a Vision Transformer (ViT) on ImageNet classification, where the "sequence length" is the number of image patches and depth retrieval might help later layers access low-level features (edges, textures) from early layers without relying on the residual stream. Vision Transformers typically have different aspect ratios than language models (shorter sequence length, often wider relative to depth), which would test whether MoDA's efficiency characteristics — particularly the depth utilization optimization that depends on chunk size C and GQA group size G — generalize beyond the language modeling configuration used in the paper.

Adaptive depth retrieval with per-head or per-layer gating. The current MoDA formulation treats depth retrieval uniformly: every head at every layer attends to depth KV from all preceding layers. However, the attention visualization in Figure 5 shows substantial heterogeneity — some heads allocate significant probability mass to depth KV, while others allocate very little. This suggests that not all heads benefit equally from depth retrieval, and the uniform application may waste compute (loading and scoring depth KV entries that receive near-zero attention weight) and potentially hurt performance (depth entries acting as distracting noise in the joint softmax for heads that don't need them). A natural extension would add a learned gating mechanism: for each attention head, a small MLP conditioned on the query representation predicts a scalar gate g ∈ [0, 1] that interpolates between pure sequence attention (g=0) and MoDA with full depth access (g=1). This gate could be trained with a sparsity penalty to encourage heads to "opt out" of depth retrieval when it's not useful, saving compute and potentially improving performance by reducing depth-key interference. An even finer-grained variant would gate per depth layer (which previous layers are relevant) rather than binary on/off, producing a sparse depth-retrieval pattern similar to the bounded slot-caching proposal in Section 6.2 but learned rather than heuristic.

MoDA as a diagnostic tool for understanding layer-wise computation in Transformers. Beyond its performance benefits, MoDA provides a unique window into which layers' outputs are useful for which downstream computations. The attention weights from later layers to earlier-layer depth KV entries are directly interpretable: if layer 27 consistently attends strongly to depth KV from layer 3 when processing a specific type of linguistic construction (e.g., relative clauses), that suggests layer 3 is computing features relevant to that construction, and layer 27 is retrieving and using them. A research program using MoDA as an analytical tool could: (a) identify functional specialization across layers by clustering depth-attention patterns across many inputs, (b) measure the "effective depth connectivity" — which layers are information sources for which other layers, and how this connectivity changes with model scale and training data, and (c) test hypotheses about Transformers as compositional computers (e.g., do later layers retrieve from early layers that perform syntactic analysis, supporting the view of Transformers as implicitly implementing a natural language processing pipeline?). This use of MoDA as an analytical instrument is independent of its performance benefits and could yield insights about Transformer internals even if MoDA itself is eventually superseded by other architectural innovations.

Practical Applications and Downstream Use Cases

Long-context LLM training and inference where depth overhead is amortized. The paper's strongest efficiency result is that at sequence length 64K, MoDA adds only 2.73% overhead over FlashAttention-2 (Table 2, row 5). This makes MoDA essentially free for long-context applications — document processing, code repository understanding, long-form dialogue — where the sequence attention cost dominates and the fixed depth overhead becomes negligible. For teams training or deploying models with 32K+ context windows, MoDA can be adopted with minimal throughput impact, and the consistent perplexity improvements across all ten validation domains (Table 5, -0.20 average PPL at 1.5B) translate directly to better language modeling quality. The key deployment consideration is that the efficiency advantage depends on the GQA group size: at G=8 (standard for many production models), the overhead is consistently under 10% for sequences above 8K tokens, while at G=2 (the configuration used in the paper's 700M training experiments), the overhead is 15-27% at typical training lengths (Table 2, rows 6-7). Practitioners should therefore evaluate MoDA at their actual GQA configuration and sequence length before committing to it for training.

Fine-tuning or continued pretraining of existing deep models to recover under-utilized depth capacity. If the paper's diagnostic reframing is correct — that deep Transformers have useful information in early layers that later layers cannot access due to dilution — then applying MoDA to an already-trained deep model via continued pretraining could "unlock" latent capabilities without training from scratch. This is speculative (the paper only trains from scratch, and fine-tuning an existing model to use depth attention would require adding the depth-KV projections and cache, which are new parameters), but it is practically compelling: many organizations have invested substantial compute in training deep models (70B+ parameters, 80+ layers) that may be under-utilizing their depth. A cost-effective strategy would be to take a pretrained checkpoint, add MoDA's depth-KV infrastructure (the FFN KV projections and the modified attention kernel), and continue training for a fraction of the original budget (e.g., 10-20% of the original tokens) to learn the depth retrieval patterns. The paper's training curves (Figure 2) show MoDA outperforming OLMo2 from early in training (~100B tokens for the 1.5B model), suggesting that depth retrieval is learned relatively quickly once the architecture enables it. If this transfers to fine-tuning, the cost of adopting MoDA for an existing model could be a small fraction of the original training cost, with immediate perplexity and downstream accuracy improvements.

Deploying deeper, narrower models for memory-constrained edge or on-device settings. The paper's finding that MoDA is more parameter-efficient than simply adding layers (Table 3: Row 4 with 705.7M parameters outperforms Row 2 with 700.5M parameters and higher FLOPs, 58.87 vs. 57.11 downstream average) has a specific deployment implication: for a fixed parameter budget (constrained by device memory), a deeper MoDA model may outperform a wider standard model. If depth provides representational benefits that width does not (compositional processing, hierarchical feature learning), and MoDA enables these benefits to be realized without the information dilution penalty, then the optimal architecture for on-device deployment might shift toward deeper configurations with MoDA rather than the wider, shallower configurations that current scaling practices favor. This is particularly relevant for applications where the model must perform multi-step reasoning (math, code, planning) — tasks where compositional depth is theoretically valuable — under tight memory constraints (smartphones, embedded systems, browser-based inference). The paper doesn't directly test this tradeoff (it compares MoDA vs. standard Transformers at fixed depth, not depth-to-width ratios), but the parameter-efficiency result provides a starting point for such exploration.

Self-improvement and knowledge distillation pipelines where deeper teacher models transfer to shallower students. In distillation settings, a deep teacher model generates training data (e.g., high-quality reasoning traces, corrected outputs) that a shallower student model learns from. If MoDA enables the teacher to utilize its depth more effectively — producing better training data through improved compositional reasoning — the student benefits indirectly even if the student doesn't use MoDA itself. Moreover, if depth retrieval helps the teacher produce more interpretable intermediate outputs (because depth-KV entries at specific layers correspond to identifiable computational stages), the distillation process could target specific depth-derived features rather than only final outputs. This application is speculative and requires validation, but it connects MoDA to the broader literature on knowledge distillation and model compression, where improving teacher quality has multiplicative benefits across all downstream student models.