Mixture-of-Depths Attention

• Standard residual paths compress depth history, while dense cross-layer connections preserve it but incur substantial parameter growth at LLM scale.The introduction motivates adaptive, data-dependent retrieval as a balance among expressivity, efficiency, and hardware friendliness.Source: §1 · p. 2
• Each MoDA head jointly attends to current-layer sequence KV and depth KV from all preceding layers, enabling data-dependent retrieval without dense cross-layer overhead.The paper analyzes depth residual, depth dense, and depth attention through a common “read, operate, write” design lens.Source: §1 · pp. 2–3
• A fused sequence-and-depth attention kernel reaches 97.3% of FlashAttention-2 efficiency at 64K sequence length.It uses shared online-softmax states, chunk-aware depth-KV layout, and group-aware indexing to improve memory access efficiency.Source: §1 · pp. 2–3
• At 1.5B, MoDA improves average perplexity by 0.2 across 10 validation benchmarks and average downstream performance by 2.11% on 10 tasks.The models use the 400B-token OLMo2 training recipe.Source: §1 · p. 2
• MoDA outperforms the strong open-source OLMo2 baseline across large-scale corpora and model scales; post-norm performs better than pre-norm, and analyses find reduced attention-sink behavior.The analyses attribute reduced attention-sink behavior to better probability allocation to informative sequence and depth KV.Source: §1 · pp. 2–3

2.1 Preliminary

• Standard self-attention projects tokens into queries, keys, and values; under GQA, groups of query heads share key-value heads.The paper uses this setup to define its attention notation.Source: §2.1 · p. 3

2.2 Stacking Transformers Along the Depth Stream

• Repeated residual updates compress the depth stream into fixed-size X_l, diluting salient features and causing signal degradation.The residual design reads the current representation and writes back by addition.Source: §2.2 · pp. 3–4
• Depth-dense methods [20] [28] connect all layers along depth and concatenate each output with historical representations, preserving them losslessly but enforcing fixed connectivity and dominant computation O(TL^2D^2), prohibitive for large models.Source: §2.2 · p. 4
• Depth Attention adaptively attends to historical key-value pairs at the same token position, with O(TL^2D) computation, a factor D below Depth Dense.It uses attention to read depth history in a data-dependent way.Source: §2.2 · pp. 4–5
• MoDA jointly normalizes attention over sequence-level and historical depth key-value pairs, then stores current-layer key-value pairs for later layers.For FFN layers, the corresponding key-value pair is obtained through a lightweight projection.Source: §2.2 · p. 5
• MoDA keeps Depth Attention’s FLOPs and cache orders while reducing parameter complexity from O(LD^2) to O(LD^2/G) by reusing sequence query projections.In GQA settings, it needs only grouped depth key-value projections.Source: §2.2 · pp. 5–6

3.1 Preliminary

• The design reorganizes depth KV layout and fuses computation to reduce non-contiguous memory access and improve effective compute utilization.Naïve PyTorch MoDA requires non-contiguous reads of historical depth states, degrading GPU utilization.Source: §3 · p. 6, §3.1 · p. 7

3.2 Hardware-aware Considerations for MoDA

• Flattening depth KV to length T × L makes each query’s depth slice a contiguous block read compatible with FlashAttention-style kernels, though most dense depth-score entries are masked.Only the depth columns corresponding to a query’s own L depth states are valid in the dense T × (T L) score matrix.Source: §3.2 · pp. 7–8
• Chunk-aware layout limits each query chunk to its local C × L depth-KV region, while group-aware reuse reduces the required span to (C/G) × L.Figure 4 depicts the reduction from a global T × L depth cache to a chunk-local span; group-aware computation reuses depth KV blocks shared by G adjacent query rows.Source: Figure 4 · p. 7, §3.2 · p. 8

3.3 Hardware-Efficient MoDA Implementation

• The kernel accumulates sequence and depth attention into shared online-softmax states, avoiding intermediate HBM materialization and normalizing the output once on chip.The sequence phase reuses the same accumulator states across blocks, which are also reused during the depth phase.Source: §3.3 · pp. 8–9

3.3.1 Efficiency Comparison

• As sequence computation becomes dominant, MoDA’s relative extra time decreases from 25.86% to 2.73%; at fixed T=16384, increasing G from 2 to 32 lowers it from 27.07% to 2.84%.At fixed T=16384 and G=8, FlashAttention-2 runtime remains 116.700 ms as MoDA runtime increases from 127.661 to 167.958 ms, raising extra time from 8.59% to 30.52%.Source: §3.3.1 · p. 9

4 Experiment

• Table 3 reports MoDA-variant performance on the training set, C4 validation set, and downstream benchmarks for 700M models trained on 400B tokens.The reported model width, GQA group size, and sequence length are 1024, 2, and 4096, respectively.Source: Table 3 · p. 10
• The evaluated variants distinguish sequence-only attention, depth KV attention, additional FFN KV projections, and separate depth KV projections.The table also reports model parameters and FLOPs.Source: Table 3 · p. 10

Model Architecture and Training Settings

• Models of 700M and 1.5B parameters use GQA [2] and train on 400B-token subsets of OLMo2 [27] in bf16, with batch size 1024 and context length 4096.Other training settings follow the OLMo2 [27] implementation.Source: §4.1 · p. 10

Evaluation Details

• Evaluation covers 10 downstream benchmarks, training and C4 validation perplexity, and per-domain validation perplexity across C4, ICE, m2d2-s2orc, Pile, Wiki-text, and dolma.The per-domain sets include Books, Common Crawl, peS2o, Reddit, and Stack.Source: §4.1 · p. 10

4.2.1 MoDA Variants

• Reusing preceding-layer attention KV as depth KV improves performance at 0.12% extra FLOPs without adding parameters.Against row 1, row 3 improves train PPL by 0.41, C4 validation PPL by 0.11, and downstream averaged metrics by 1.17.Source: §4.2.1 · p. 10

4.2.2 Scaling MoDA with Model Size

• MoDA raises downstream averages from 57.11 to 58.87 (+1.76) at 700M and from 62.28 to 64.39 (+2.11) at 1.5B.Gains span HellaSwag (+0.42 at 700M; +0.38 at 1.5B), ARC-C (+1.34; +4.35), BoolQ (+3.09; +3.73), and MMLU (+0.92; +1.86).Source: Table 4 · p. 11, §4.2.2 · pp. 11–12
• MoDA lowers average validation PPL from 15.61 to 15.46 at 700M and from 13.67 to 13.47 at 1.5B, improving all ten domains at both scales.At 1.5B, PPL falls on Reddit from 21.21 to 20.85, ICE from 15.37 to 15.08, and Wiki-text from 10.41 to 10.16.Source: Table 5 · p. 11, §4.2.2 · p. 12

4.2.1 MoDA Variants

• Adding FFN-side depth KV yields the best accuracy-efficiency trade-off and is adopted as MoDA’s default variant.Against row 3, row 4 improves train PPL by 0.18, C4 validation PPL by 0.27, and downstream averaged metrics by 0.77; against the similarly sized row 2, it improves them by 0.37, 0.10, and 1.76, respectively.Source: §4.2.1 · pp. 10–11
• Extra attention KV projection adds only marginal gains while increasing parameters from 705.7M to 742.4M and FLOPs from 8.33T to 8.63T.Comparing rows 4 and 5, it improves train PPL by 0.07, C4 validation PPL by 0.04, and downstream averaged metrics by 0.10.Source: §4.2.1 · p. 11

4.3.1 Analyzing MoDA with Layer Number

• Validation loss falls with Depth KV at both 48 layers (3.3800 to 3.3759 pre-norm; 3.4062 to 3.3653 post-norm) and 24 layers (3.4740 to 3.4537).At 48 layers, the reduction is larger with post-norm (0.0409) than pre-norm (0.0041).Source: §4.3.1 · p. 12
• Adding Extra FFN KV Projection further reduces validation loss at 48 layers (3.3759 to 3.3656 pre-norm; 3.3653 to 3.3484 post-norm) and 24 layers (3.4537 to 3.4338).The results show that FFN-side depth information brings additional gains when compute budget allows.Source: §4.3.1 · pp. 12–13

4.3.3 Analyzing MoDA with Efficiency

• 1458×: combining Flash-compatible layout, chunk-aware layout, and group-aware indexing speeds end-to-end runtime over the naive PyTorch baseline.At T=1024, the successive steps reduce runtime from 2128.900 ms to 13.102 ms, then to 6.286 ms, and finally to 1.460 ms.Source: Table 7 · p. 13, §4.3.3 · pp. 13, 15

4.3.2 Analyzing MoDA with Attention Visualization

• Unlike typical attention sinks, visualized MoDA heads distribute attention more broadly across sequence and depth slots, though the functional role of this pattern remains unclear.The visualization suggests probability mass may shift from fixed sink positions toward potentially useful locations; the precise role warrants further investigation.Source: §4.3.2 · p. 13
• Attention heatmaps show persistent depth-KV attention, especially in middle and late layers; sharper diagonal heads also use depth slots, while broader heads rely on them more heavily.The combined-softmax visualization includes both attention KV and FFN KV in the depth-KV block.Source: §4.3.2 · p. 13, Figure 5 · p. 14

• Experiments on 700M and 1.5B models trained with the OLMo2 recipe show consistent gains in perplexity and downstream performance under modest overhead.Source: §5 · p. 15
• MoDA combines historical depth-information retrieval with a hardware-aware fused kernel for efficient long-context execution.The kernel uses unified online-softmax states, a chunk-aware depth-KV layout, and group-aware indexing.Source: §5 · p. 15
• The authors propose MoDA as an architecture-agnostic primitive for scaling Transformer depth, with potential integration into multimodal intelligence, visual understanding, and world models.They plan to release the full implementation and hope it will support stronger open-source language models.Source: §5 · p. 15

6.1 Scaling MoDA for Industrial Training via Advanced CUDA Engineering

• Industrial-scale MoDA training still needs CUDA optimizations for scheduling, pipelining, and communication overlap to reduce stalls and overhead and improve throughput.These optimizations preserve MoDA’s algorithmic behavior; the current kernel is not yet considered sufficient for trillion-parameter production runs.Source: §6.1 · p. 15

6.2 Mitigating Memory Bottlenecks with Bounded Depth-KV Slot Caching

• Full depth-KV caching grows linearly with network depth, creating memory and bandwidth overhead that can dominate long-context training and serving.The discussion identifies this as increasingly difficult to sustain at industrial scale.Source: §6.2 · p. 15
• A fixed-size buffer of S slots, where S ≪ L, could bound depth-KV memory using dynamic top-S selection, a recent-entry sliding window, or a hybrid policy.A hybrid can reserve slots for recent entries and others for high-score global memories.Source: §6.2 · p. 15
• Bounded slots shift memory and bandwidth scaling from depth to slot count and stabilize fused-kernel tensor shapes, but slot assignment and quality-latency-efficiency tradeoffs need study.Future work should explore training the selection policy jointly with MoDA under a fixed slot budget.Source: §6.2 · p. 15