Track manipulation, locomotion and embodied systems.
Search papers, labs, and topics across Lattice.
The paper introduces Mixture-of-Depths Attention (MoDA), a novel attention mechanism that mitigates signal degradation in deep LLMs by allowing each attention head to attend to both current-layer and preceding-layer key-value pairs. A hardware-efficient implementation of MoDA is presented, achieving near-FlashAttention-2 performance. Experiments on 1.5B parameter models show that MoDA improves perplexity and downstream task performance with minimal overhead, suggesting its potential for effective depth scaling.
LLMs can now scale depth more effectively: a new attention mechanism recovers diluted features in deeper layers, boosting performance with negligible overhead.
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 .
Tsinghua AI