Mischung-der-Tiefen-Aufmerksamkeit
Mixture-of-Depths Attention
March 16, 2026
Autoren: 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
Zusammenfassung
Die Skalierung der Tiefe ist ein entscheidender Treiber für große Sprachmodelle (LLMs). Allerdings leiden LLMs mit zunehmender Tiefe häufig unter Signalverschlechterung: informative Merkmale, die in flachen Schichten gebildet werden, werden durch wiederholte Residual-Updates allmählich verdünnt, was ihre Wiederherstellung in tieferen Schichten erschwert. Wir führen Mixture-of-Depths-Attention (MoDA) ein, einen Mechanismus, der jedem Attention-Head erlaubt, auf Sequenz-KV-Paare der aktuellen Schicht und Tiefen-KV-Paare aus vorhergehenden Schichten zuzugreifen. Wir beschreiben weiterhin einen hardwareeffizienten Algorithmus für MoDA, der nicht-kontinuierliche Speicherzugriffsmuster auflöst und 97,3 % der Effizienz von FlashAttention-2 bei einer Sequenzlänge von 64K erreicht. Experimente mit 1,5B-Parameter-Modellen zeigen, dass MoDA durchgängig starke Baseline-Methoden übertrifft. Bemerkenswerterweise verbessert es die durchschnittliche Perplexität um 0,2 über 10 Validierungs-Benchmarks hinweg und steigert die durchschnittliche Leistung um 2,11 % bei 10 Downstream-Aufgaben, bei einem vernachlässigbaren Rechenaufwand von 3,7 % zusätzlichen FLOPs. Wir stellen zudem fest, dass die Kombination von MoDA mit Post-Norm eine bessere Leistung erzielt als die Verwendung mit Pre-Norm. Diese Ergebnisse deuten darauf hin, dass MoDA eine vielversprechende Grundkomponente für die Tiefenskalierung ist. Der Code ist unter https://github.com/hustvl/MoDA veröffentlicht.
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 .