Внимание со смешанной глубиной

Mixture-of-Depths Attention

March 16, 2026

Авторы: 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

Аннотация

Масштабирование глубины является ключевым фактором для больших языковых моделей (LLM). Однако по мере увеличения глубины LLM часто сталкиваются с проблемой деградации сигнала: информативные признаки, сформированные в поверхностных слоях, постепенно размываются из-за повторяющихся остаточных обновлений, что затрудняет их восстановление в более глубоких слоях. Мы представляем механизм внимания смешения глубин (MoDA), который позволяет каждой голове внимания обращаться к KV-парам последовательности на текущем слое и к KV-парам глубины из предыдущих слоев. Далее мы описываем аппаратно-эффективный алгоритм для MoDA, который решает проблему несмежных шаблонов доступа к памяти, достигая 97.3% эффективности FlashAttention-2 при длине последовательности 64K. Эксперименты на моделях с 1.5 млрд параметров показывают, что MoDA стабильно превосходит сильные базовые модели. В частности, он улучшает среднюю перплексию на 0.2 по 10 валидационным тестам и повышает среднюю производительность на 2.11% по 10 последующим задачам при незначительном вычислительном overhead в 3.7% FLOPs. Мы также обнаружили, что сочетание MoDA с пост-нормализацией дает лучшие результаты, чем использование с пред-нормализацией. Эти результаты свидетельствуют о том, что MoDA является перспективным примитивом для масштабирования глубины. Код доступен по адресу 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 .