Source-linked AI summary
Mixture-of-Depths Attention
Lianghui Zhu, Yuxin Fang, Bencheng Liao, Shijie Wang, Tianheng Cheng, Zilong Huang, Chen Chen, Lai Wei, Yutao Zeng, Ya Wang, Yi Lin, Yu Li, Xinggang Wang
TL;DR
As Transformers become deeper, residual updates can dilute historical information, limiting depth scaling. MoDA lets attention retrieve sequence and preceding-layer depth information, and experiments show consistent gains with modest computational overhead. These results suggest explicit historical-depth retrieval is a practical primitive for scaling Transformer depth.
Problem
Deeper Transformers can suffer information dilution, while dense cross-layer connections preserve history at substantial parameter cost.
Method
MoDA jointly attends to current-layer sequence KV and preceding-layer depth KV, using a hardware-aware implementation for efficient execution.
Results
MoDA improves average perplexity by 0.2 across 10 validation benchmarks and average downstream performance by 2.11% on 10 tasks.
Takeaways & Limitations
Explicit retrieval of historical depth information is a practical and effective primitive for scaling Transformer depth under modest overhead.
Takeaways & Limitations
Fixed-size depth-KV slot caching remains limited by slot-assignment quality, requiring future study of selection policies and quality–latency–efficiency trade-offs.
Abstract
from arXiv · showhide
Scaling depth is a key driver for large language models (LLMs). Yet, as LLMs become deeper, they often suffer from signal degradation: informative features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers. We introduce mixture-of-depths attention (MoDA), a mechanism that allows each attention head to attend to sequence KV pairs at the current layer and depth KV pairs from preceding layers. We further describe a hardware-efficient algorithm for MoDA that resolves non-contiguous memory-access patterns, achieving 97.3% of FlashAttention-2's efficiency at a sequence length of 64K. Experiments on 1.5B-parameter models demonstrate that MoDA consistently outperforms strong baselines. Notably, it improves average perplexity by 0.2 across 10 validation benchmarks and increases average performance by 2.11% on 10 downstream tasks, with a negligible 3.7% FLOPs computational overhead. We also find that combining MoDA with post-norm yields better performance than using it with pre-norm. These results suggest that MoDA is a promising primitive for depth scaling. Code is released at https://github.com/hustvl/MoDA .
1 Introduction
MoDA addresses depth-related information dilution by dynamically retrieving historical depth information alongside current-layer sequence information. The paper combines this mechanism with a hardware-efficient implementation and reports consistent gains over strong baselines.
- Motivation: Deeper Transformers can dilute informative shallow-layer features, while dense cross-layer connections preserve history at substantial parameter cost.The introduction frames depth as under-exploited despite its representational appeal and motivates dynamic alternatives to fixed residual aggregation.
- Method: MoDA lets each attention head jointly attend to current-layer sequence KV and preceding-layer depth KV.This data-dependent retrieval is presented as a unified alternative to standard residual and dense depth connections.
- Implementation: 97.3% of FlashAttention-2 efficiency at 64K sequence length is achieved by fusing sequence and depth attention with layouts and indexing designed for efficient memory access.The implementation uses shared online-softmax states, chunk-aware depth-KV layout, and group-aware indexing.
- Results: 0.2 average perplexity improvement across 10 validation benchmarks and 2.11% average downstream improvement across 10 tasks are reported for the main 1.5B setting.The experiments use decoder-only models trained with the 400B-token OLMo2 recipe; post-norm combined with MoDA performs better than pre-norm.
- Implications: The paper presents MoDA as a practical primitive for depth scaling that improves depth-wise information aggregation without dense cross-layer overhead.The contribution is supported by the proposed formulation, fused implementation, and empirical evaluation across model scales.
2 Mixture-of-Depths Attention
MoDA frames Transformer depth stacking as a read–operate–write problem and combines adaptive depth retrieval with standard sequence attention. It targets residual signal dilution while retaining lower complexity than dense cross-layer connectivity.
- Motivation: Residual stacking repeatedly compresses the depth stream into a fixed-size tensor, diluting salient features and causing signal degradation.
- Motivation: Depth-dense connections preserve historical representations but incur dominant computation O(TL2D2), making them prohibitive for large models.
- Depth Attention: Depth Attention adaptively reads historical depth information, reducing computation to O(TL2D) compared with depth-dense connectivity.
- Mixture-of-Depths Attention: MoDA jointly attends to current-layer sequence KV and preceding-layer depth KV, normalizing both attention sources with one softmax operator.
- Mixture-of-Depths Attention: MoDA appends current-layer key-value pairs to the depth stream for subsequent layers while passing the query representation forward.
- Complexity: MoDA preserves favorable FLOPs and cache orders while reducing parameter complexity from O(LD2) to O(LD2/G) through query-projection reuse.
3 Hardware-aware efficient MoDA
MoDA’s hardware-aware implementation addresses irregular depth-KV access by reorganizing depth caches, grouping queries, and fusing sequence and depth attention. The resulting kernel preserves efficient execution across long sequences, while overhead varies predictably with sequence length, GQA grouping, and model depth.
- Motivation: Naïve MoDA suffers from non-contiguous historical-state reads, so the implementation reorganizes depth tensors and fuses computation.The design targets improved GPU utilization by replacing irregular memory access with contiguous, hardware-friendly execution.
- Depth-KV layouts: Chunk-aware layout reduces each chunk’s effective depth span from T × L to (C × L)/G by grouping queries and reusing depth-KV blocks across GQA groups.The layout pairs a query chunk with its local depth-KV region and exploits shared base-time indices among grouped query rows.
- Depth-KV layouts: Flash-compatible layout flattens each sequence’s depth cache to length T × L, enabling contiguous depth-KV block reads compatible with FlashAttention-style kernels.Each position accesses its corresponding range [tL, (t + 1)L), although dense computation still wastes work outside the valid block-diagonal region.
- Fused execution: The fused kernel reuses online-softmax states across sequence and depth loops, combining both sources into one blockwise accumulation without intermediate HBM materialization.Grouped causal and depth masks maintain consistent indexing while sequence and depth logits update shared accumulator states.
- Scaling behavior: Extra time decreases from 25.86% to 2.73% as sequence computation dominates and drops from 27.07% to 2.84% as G increases from 2 to 32.At fixed T=16384, depth utilization rises from 3.12% to 50.00% with larger G; increasing model depth instead raises extra time from 8.59% to 30.52%.
4 Experiment
The experiments assess MoDA’s expressivity and efficiency through variant comparisons on language-model training, validation, and downstream benchmarks. Table 3 focuses on 700M models trained on 400B tokens and reports both performance and computational quantities.
- Evaluation scope: The experiment section evaluates MoDA’s expressivity and efficiency on large language models.The broader study includes language-model experiments designed to assess the proposed mechanism’s empirical behavior.
- Variant comparison: Table 3 compares MoDA variants on training results, C4 validation, downstream benchmarks, parameter counts, and FLOPs.The comparison uses 700M models trained on 400B tokens under fixed width, GQA group size, and sequence length settings.
4.1 Experimental Setups
The main experiments train 700M and 1.5B decoder-only language models with GQA on 400B-token OLMo2 subsets. Evaluation covers standard downstream benchmarks and multiple perplexity-based validation measures.
- Model and training settings: Experiments use 700M and 1.5B language models with group query attention trained on 400B-token subsets of OLMo2.Models are trained in bfloat16 with global batch size 1024 and context length 4096.
- Evaluation: Evaluation includes PiQA, HellaSwag, WinoGrande, OpenBookQA, BoolQA, SciQA, COPA, MMLU, ARC-E, and ARC-C.The study also reports training perplexity, C4 validation perplexity, and per-domain validation perplexity across several corpora.
4.2 Main Results
MoDA variants consistently improve downstream and validation performance across 700M and 1.5B models, with gains depending on where depth projections are introduced. Reusing attention-side depth KV is highly efficient, while FFN-side depth KV offers the strongest accuracy-efficiency trade-off.
- MoDA Variants: With 0.12% extra FLOPs, reusing preceding-layer sequence KV as depth KV improves train PPL by 0.41, C4 validation PPL by 0.11, and downstream averaged metrics by 1.17.
- MoDA Variants: Adding FFN-side depth KV further improves train PPL by 0.18, C4 validation PPL by 0.27, and downstream averaged metrics by 0.77 over the attention-only depth-KV variant.
- MoDA Variants: Extra attention KV projection adds only 0.07 train PPL, 0.04 C4 validation PPL, and 0.10 downstream averaged metrics while increasing parameters from 705.7M to 742.4M and FLOPs from 8.33T to 8.63T.
- MoDA Variants: The experiments select the FFN-enhanced variant as default because attention-side reuse is nearly cost-free, FFN-side depth KV improves the accuracy-efficiency trade-off, and extra attention projection is marginal.
- Scaling MoDA with Model Size: MoDA improves average downstream performance by 1.76 points at 700M and 2.11 points at 1.5B under the same 400B-token training budget.
- Scaling MoDA with Model Size: MoDA lowers average validation perplexity from 15.61 to 15.46 at 700M and from 13.67 to 13.47 at 1.5B, improving all ten domains.
4.3 Analysis
Analysis shows that MoDA remains effective across layer counts, retrieves depth information persistently, and can be implemented efficiently through progressively optimized memory layouts and indexing. Its altered attention allocation is promising but not yet fully explained.
- Analyzing MoDA with Layer Number: Depth KV reduces validation loss for both 48-layer and 24-layer models, while Extra FFN KV Projection provides additional gains.
- Analyzing MoDA with Layer Number: At 48 layers, Depth KV reduces loss by 0.0409 with post-norm versus 0.0041 with pre-norm, indicating a stronger optimization impact for deeper post-norm stacks.
- Analyzing MoDA with Attention Visualization: Attention heatmaps show persistent depth-KV attention, especially in middle and late layers, alongside sequence-focused heads that still allocate probability to depth slots.
- Analyzing MoDA with Attention Visualization: MoDA distributes attention more broadly across sequence and depth slots rather than concentrating most probability on a few fixed sink positions.
- Analyzing MoDA with Attention Visualization: The precise functional role of the altered attention-sink pattern remains unclear and requires further investigation.
5 Conclusion
MoDA is presented as a depth-aware attention mechanism that improves depth-wise information aggregation while mitigating optimization and information-dilution challenges. Its hardware-aware fused implementation supports efficient execution, and experiments show consistent gains with modest overhead.
- MoDA retrieves historical depth information through a unified depth-aware attention mechanism for improving depth-wise information aggregation.
- A hardware-aware fused kernel uses unified online-softmax states, chunk-aware depth-KV layout, and group-aware indexing for efficient long-context execution.
- Experiments on 700M and 1.5B models trained with the OLMo2 recipe show consistent perplexity and downstream-performance gains under modest overhead.
- The paper positions explicit retrieval of historical depth information as a practical primitive for scaling Transformer depth.
6 Discussion
The discussion identifies scalability limits for MoDA: full historical depth-KV caching can create depth-dependent bottlenecks, while industrial deployment still requires further CUDA optimization. A bounded slot buffer offers a practical alternative but makes slot assignment quality the central challenge.
- Industrial deployment: Industrial-scale training remains beyond the current kernel’s endpoint, requiring improved memory scheduling, computation pipelining, and communication overlap.The stated examples include trillion-parameter models, where further CUDA engineering could reduce stalls and launch overhead.
- Memory bottlenecks: Full depth-KV caching incurs memory and bandwidth overhead that grows linearly with network depth.This cost can become the dominant bottleneck in long-context training and serving.
- Bounded caching: A fixed-size Depth KV slot buffer bounds each query’s attended depth memories to S slots, with S ≪ L.Dynamic selection, sliding-window retention, or a hybrid policy can determine which entries remain.
- Bounded caching: Bounded caching shifts memory and bandwidth costs from depth-dependent scaling to slot-dependent scaling and provides a stable tensor shape for fused kernels.The design changes effective depth memory from an unbounded cache to a bounded cache.
- Open challenge: The key open challenge for bounded caching is learning slot assignments while balancing quality, latency, and hardware efficiency under a fixed budget.The discussion proposes studying joint training of the selection policy with MoDA.