Summary of "Mixture-of-Depths Attention"

Summary (Overview)

• Core Contribution: Introduces Mixture-of-Depths Attention (MoDA), a novel attention mechanism for Transformers that allows each query to attend to both standard sequence-level Key-Value (KV) pairs and depth KV pairs from preceding layers at the same token position. This addresses the information dilution problem in deep LLMs where informative features from shallow layers get diluted in deeper layers.
• Key Innovation: Unifies sequence and depth retrieval within a single softmax operation, enabling dynamic, data-dependent mixing of information across both dimensions. This provides a more expressive and efficient alternative to fixed cross-layer connections like residual or dense connections.
• Hardware-Efficient Implementation: Proposes a fused, hardware-aware kernel with a chunk-aware depth-KV layout and group-aware indexing. This implementation achieves 97.3% of FlashAttention-2's efficiency at a 64K sequence length, making MoDA practical for long-context training.
• Empirical Results: Demonstrates consistent performance improvements. On 1.5B-parameter models trained on 400B tokens, MoDA improves average validation perplexity by 0.2 across 10 benchmarks and boosts average downstream task performance by 2.11% with a negligible 3.7% FLOPs overhead.
• Architectural Insight: Finds that combining MoDA with post-norm yields better performance than with pre-norm, and that the primary benefit comes from incorporating depth KV from Feed-Forward Network (FFN) layers, not just attention layers.

Introduction and Theoretical Foundation

Scaling model depth is a crucial but challenging dimension for improving Large Language Models (LLMs). While deeper networks offer richer hierarchical computation, standard Transformer architectures suffer from signal degradation or information dilution: as information passes through many residual layers, salient features formed in early layers can be gradually diluted, making them harder for later layers to recover.

Existing solutions have trade-offs:

• Depth Residual (standard ResNet-style): Simple and stable for optimization but compresses all depth history into a single hidden state, leading to information loss.
• Depth Dense (DenseNet-style): Preserves all intermediate states via concatenation, mitigating dilution but incurring prohibitive quadratic parameter and computational growth ($O(L^2D^2)$) at LLM scale.

The paper posits that dynamic, data-dependent mixing (the principle behind attention) is more effective for preserving historical information than fixed aggregation patterns. This motivates extending attention from the sequence dimension to the depth dimension.

MoDA is proposed as a unified mechanism that allows each layer to adaptively "read" useful states from all preceding layers via attention, combining them with the current sequence context in a single operation. It aims to retain the expressivity of dense connections while maintaining hardware efficiency.

Methodology

Mathematical Formulation

The standard Grouped Query Attention (GQA) is defined as: Given input $X \in \mathbb{R}^{T \times D}$, it is projected into Queries ($Q$), Keys ($K$), and Values ($V$):

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

where $W_Q \in \mathbb{R}^{D \times (H_q d)}$, $W_K, W_V \in \mathbb{R}^{D \times (H_k d)}$, $H_q = G H_k$, and $D = H_q d$. The attention output for head $h$ is:

$$\text{Attention}(Q, K, V) = \text{Concat}_{h=1}^{H_q} \left( \text{softmax}\left( \frac{Q_h K_{\phi(h)}^T}{\sqrt{d}} + M \right) V_{\phi(h)} \right)$$

where $\phi(h) = \lceil h/G \rceil$ maps query heads to shared key-value heads, and $M$ is the causal mask.

MoDA Operation: At layer $l$, the query $Q_{l-1}$ attends to a concatenated KV set containing:

1. Sequence KV: The standard keys and values from the current sequence: $K_l^{seq}, V_l^{seq} \in \mathbb{R}^{T \times (H_k d)}$.
2. Depth KV: Historical key-value pairs from all preceding layers $\{K_i, V_i\}_{i=0}^{l-1}$ for the same token position, where each $K_i, V_i \in \mathbb{R}^{T \times (H_k d)}$.

The attention scores are computed over the concatenated dimension and normalized via a unified softmax. The "write" step appends the current layer's projected KV pair $(K_l, V_l)$ to the depth stream for use by subsequent layers. For FFN layers, a lightweight KV projection is applied to the FFN's input to generate its depth KV contribution.

Complexity Analysis

A key advantage of MoDA is its favorable computational profile compared to dense connections.

MoDA maintains the linear-in-width FLOPs scaling of Depth Attention ($O(T L^2 D)$) but is more parameter-efficient ($O(L D^2 / G)$) because it reuses the sequence attention's query projection, requiring only grouped depth key/value projections.

Hardware-Efficient Kernel Implementation

A naive PyTorch implementation of MoDA suffers from non-contiguous memory access. The paper introduces a fused kernel with three key optimizations (Algorithm 1):

1. Flash-Compatible Depth-KV Layout: Flattens the depth cache along a single axis of length $T \times L$, enabling contiguous block reads.
2. Chunk-Aware Depth-KV Layout: Groups queries into chunks of size $C$. Each chunk only accesses a local depth-KV region of size $C \times L$, reducing the effective depth span scanned and improving depth utilization from $1/T$ to $1/C$.
3. Group-Aware Indexing: Leverages the GQA structure where $G$ adjacent query rows share the same base-time index $t_{base}(i_q) = \lfloor i_q / G \rfloor$. This allows reusing the same depth KV blocks within a group, further improving utilization to $G/C$.

The kernel fuses sequence and depth attention in one forward pass with shared online-softmax states, minimizing HBM traffic. The mask for depth attention ensures a query only attends to its corresponding depth slots:

$$\text{mask}(i_q, j_d) = \mathbb{1}\left[ \left\lfloor \frac{i_q}{G} \right\rfloor = \left\lfloor \frac{j_d}{L} \right\rfloor \right]$$

Empirical Validation / Results

Main Results (Model Scaling)

Models were trained on 400B tokens following the OLMo2 recipe.

Table 4: Downstream Task Performance (Accuracy %)

• MoDA provides stable gains across model scales (+1.76 avg at 700M, +2.11 avg at 1.5B).
• Improvements are broad-based, seen across commonsense, reasoning, and knowledge tasks.

Table 5: Per-Domain Validation Perplexity (Lower is Better)

• MoDA consistently lowers perplexity across all ten validation domains at both model sizes.

Ablation Studies

Table 3: MoDA Variants (700M Model) The study compares different design choices for incorporating depth information:

1. Baseline (Vanilla Attention): Only sequence KV.
2. + Depth KV: Reuses preceding attention layer's KV as depth KV. Provides significant gains (+1.17 avg downstream) with only 0.12% extra FLOPs.
3. + Extra FFN KV Proj.: Adds lightweight KV projections for FFN layers. Yields the best accuracy-efficiency trade-off, improving over the baseline by +1.76 avg downstream.
4. + Extra Attn KV Proj.: Adds separate depth KV projections for attention layers. Offers only marginal gains (+0.10) for substantial parameter/FLOPs overhead, indicating saturation.

Key Finding: The primary benefit comes from incorporating FFN layers' depth information. The default MoDA variant uses Sequence KV + Depth KV + Extra FFN KV Projection.

Kernel Efficiency

Table 2: Kernel Efficiency vs. FlashAttention-2 The hardware-aware MoDA kernel achieves high efficiency, especially for long sequences and large GQA group sizes $G$.

• The relative overhead of MoDA decreases as sequence length $T$ increases, dropping to ~2.7% at $T=65536$.
• Larger group sizes $G$ improve depth utilization ($\eta_{depth}=G/C$) and reduce overhead.
• Increasing model depth $L$ linearly increases MoDA's runtime, as expected.

Table 7: Kernel Optimization Ablation

The combined optimizations yield a ~1458x speedup over the naive implementation.

Additional Analyses

• Layer-Number Analysis: MoDA improves validation loss for both shallow (24-layer) and deep (48-layer) models. Gains are more pronounced with post-norm than pre-norm in deeper models.
• Attention Visualization: Heatmaps (Figure 5) show that models allocate non-trivial and persistent attention mass to the depth-KV block, confirming the mechanism's active use. The attention pattern also appears to differ from typical "attention sink" behavior, distributing probability more broadly.

Theoretical and Practical Implications

• Theoretical: MoDA provides a principled, data-dependent mechanism for mitigating information dilution in deep networks, bridging the gap between the stability of residual connections and the expressivity of dense connections. The unified softmax over sequence and depth creates a coherent representation space for cross-layer information retrieval.
• Practical: The method offers a cost-effective way to improve LLM performance. With only a ~3.7% FLOPs overhead, it delivers consistent gains in perplexity and downstream task accuracy. The provided hardware-efficient kernel makes it feasible for integration into large-scale training pipelines.
• Architectural Guidance: The finding that post-norm works better with MoDA than pre-norm offers a new perspective on normalization choices in depth-scaled Transformers. The result that FFN depth KV is crucial highlights the importance of non-attention features in the depth stream.

Conclusion

Mixture-of-Depths Attention (MoDA) is an effective and efficient architectural innovation for scaling Transformer depth. By enabling dynamic, attentive retrieval of information from all preceding layers, it alleviates the information dilution problem in deep LLMs. The proposed hardware-aware implementation ensures the approach is practical for long-context training. Empirical results across multiple scales and tasks confirm its robustness and performance benefits. MoDA establishes depth-aware attention as a promising primitive for future LLM architecture design.

Future Directions include scaling MoDA for industrial training via advanced CUDA engineering and investigating bounded depth-KV slot caching to mitigate memory bottlenecks when scaling to extreme depths. The mechanism is architecture-agnostic and could benefit multimodal and other Transformer-based models.