Attention à Mélange de Profondeurs
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
Résumé
L'augmentation de la profondeur est un facteur clé pour les grands modèles de langage (LLM). Pourtant, lorsque les LLM deviennent plus profonds, ils souffrent souvent d'une dégradation du signal : les caractéristiques informatives formées dans les couches superficielles sont progressivement diluées par les mises à jour résiduelles répétées, ce qui les rend plus difficiles à récupérer dans les couches plus profondes. Nous introduisons l'attention à mélange de profondeurs (MoDA), un mécanisme qui permet à chaque tête d'attention de prendre en compte les paires KV de la séquence à la couche actuelle et les paires KV de profondeur des couches précédentes. Nous décrivons en outre un algorithme efficace matériellement pour MoDA qui résout les problèmes d'accès mémoire non contigus, atteignant 97,3 % de l'efficacité de FlashAttention-2 pour une longueur de séquence de 64K. Les expériences sur des modèles de 1,5 milliard de paramètres démontrent que MoDA surpasse constamment les bases de référence solides. Notamment, il améliore la perplexité moyenne de 0,2 sur 10 benchmarks de validation et augmente les performances moyennes de 2,11 % sur 10 tâches en aval, avec une surcharge computationnelle négligeable de 3,7 % en FLOPs. Nous constatons également que combiner MoDA avec une normalisation post-couche (post-norm) donne de meilleures performances que son utilisation avec une normalisation pré-couche (pre-norm). Ces résultats suggèrent que MoDA est une primitive prometteuse pour l'augmentation d'échelle en profondeur. Le code est disponible à l'adresse 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 .