Atenção de Mistura de Profundidades
Mixture-of-Depths Attention
March 16, 2026
Autores: 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
Resumo
A escalagem da profundidade é um fator-chave para os grandes modelos de linguagem (LLMs). No entanto, à medida que os LLMs se tornam mais profundos, frequentemente sofrem de degradação do sinal: características informativas formadas nas camadas mais superficiais são gradualmente diluídas por atualizações residuais repetidas, tornando-as mais difíceis de recuperar nas camadas mais profundas. Apresentamos a atenção de mistura de profundidades (MoDA), um mecanismo que permite a cada cabeça de atenção atender a pares KV da sequência na camada atual e a pares KV de profundidade de camadas precedentes. Descrevemos ainda um algoritmo eficiente em hardware para MoDA que resolve padrões de acesso à memória não contíguos, atingindo 97,3% da eficiência do FlashAttention-2 em um comprimento de sequência de 64K. Experimentos em modelos de 1,5B de parâmetros demonstram que o MoDA supera consistentemente linhas de base robustas. Notavelmente, ele melhora a perplexidade média em 0,2 em 10 benchmarks de validação e aumenta o desempenho médio em 2,11% em 10 tarefas subsequentes, com uma sobrecarga computacional de FLOPs insignificante de 3,7%. Também descobrimos que combinar o MoDA com pós-normalização produz um desempenho melhor do que usá-lo com pré-normalização. Esses resultados sugerem que o MoDA é uma primitiva promissora para a escalagem de profundidade. O código foi disponibilizado em 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 .