Attenzione a Miscela di Profondità
Mixture-of-Depths Attention
March 16, 2026
Autori: 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
Abstract
La scalabilità in profondità è un fattore chiave per i grandi modelli linguistici (LLM). Tuttavia, man mano che gli LLM diventano più profondi, spesso soffrono di degradazione del segnale: le caratteristiche informative formate negli strati superficiali vengono gradualmente diluite da ripetuti aggiornamenti residui, rendendole più difficili da recuperare negli strati più profondi. Introduciamo l'attenzione a miscela di profondità (MoDA), un meccanismo che consente a ogni testa di attenzione di focalizzarsi sulle coppie chiave-valore (KV) della sequenza nello strato corrente e sulle coppie KV di profondità dagli strati precedenti. Descriviamo inoltre un algoritmo efficiente per l'hardware per MoDA che risolve i modelli di accesso alla memoria non contigui, raggiungendo il 97,3% dell'efficienza di FlashAttention-2 con una lunghezza di sequenza di 64K. Esperimenti su modelli da 1,5 miliardi di parametri dimostrano che MoDA supera costantemente i baseline di riferimento. In particolare, migliora la perplexity media di 0,2 su 10 benchmark di validazione e aumenta le prestazioni medie del 2,11% su 10 task downstream, con un trascurabile overhead computazionale del 3,7% in termini di FLOP. Troviamo anche che combinare MoDA con la normalizzazione post-layer (post-norm) produce prestazioni migliori rispetto all'uso con la normalizzazione pre-layer (pre-norm). Questi risultati suggeriscono che MoDA è un primitivo promettente per il scaling in profondità. Il codice è rilasciato all'indirizzo 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 .