Mengsel-van-Dieptes-Aandacht
Mixture-of-Depths Attention
March 16, 2026
Auteurs: 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
cs.AI
Samenvatting
Het schalen van de diepte is een cruciale drijvende kracht voor grote taalmodellen (LLM's). Toch kampen LLM's naarmate ze dieper worden vaak met signaaldegradatie: informatieve kenmerken die in ondiepe lagen worden gevormd, worden geleidelijk verdund door herhaalde residuele updates, waardoor ze moeilijker te herstellen zijn in diepere lagen. Wij introduceren mixture-of-depths attention (MoDA), een mechanisme dat elke aandachtskop in staat stelt om aandacht te schenken aan KV-paren van de sequentie in de huidige laag en aan KV-paren van de diepte uit voorgaande lagen. Verder beschrijven we een hardware-efficiënt algoritme voor MoDA dat niet-aaneengesloten geheugentoegangspatronen oplost, waarmee 97,3% van de efficiëntie van FlashAttention-2 wordt bereikt bij een sequentielengte van 64K. Experimenten met modellen van 1,5B parameters tonen aan dat MoDA consequent sterke baseline-modellen overtreft. Opmerkelijk is dat het de gemiddelde perplexiteit met 0,2 verbetert over 10 validatiebenchmarks en de gemiddelde prestaties met 2,11% verhoogt op 10 downstreamtaken, met een verwaarloosbare rekenkostenoverhead van 3,7% FLOPs. We stellen ook vast dat de combinatie van MoDA met post-norm betere prestaties oplevert dan het gebruik ervan met pre-norm. Deze resultaten suggereren dat MoDA een veelbelovend primitief is voor het schalen van diepte. De code is vrijgegeven op https://github.com/hustvl/MoDA.
English
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 .