Atención de Mezcla de Profundidades

Resumen

La escalabilidad en profundidad es un factor clave para los grandes modelos de lenguaje (LLM). Sin embargo, a medida que los LLM se vuelven más profundos, a menudo sufren de degradación de la señal: las características informativas formadas en las capas superficiales se diluyen gradualmente por las actualizaciones residuales repetidas, dificultando su recuperación en las capas más profundas. Introducimos la atención de mezcla de profundidades (MoDA), un mecanismo que permite a cada cabeza de atención atender a los pares clave-valor (KV) de la secuencia en la capa actual y a los pares KV de profundidad de las capas precedentes. Además, describimos un algoritmo eficiente en hardware para MoDA que resuelve los patrones de acceso a memoria no contiguos, alcanzando un 97.3% de la eficiencia de FlashAttention-2 con una longitud de secuencia de 64K. Los experimentos en modelos de 1.5B de parámetros demuestran que MoDA supera consistentemente a los baselines sólidos. En particular, mejora la perplejidad promedio en 0.2 puntos a lo largo de 10 benchmarks de validación y aumenta el rendimiento promedio en un 2.11% en 10 tareas posteriores, con un overhead computacional negligible de solo un 3.7% en FLOPS. También encontramos que combinar MoDA con post-norm produce un mejor rendimiento que usarlo con pre-norm. Estos resultados sugieren que MoDA es una primitiva prometedora para el escalado en profundidad. El código se ha publicado en 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 .

Atención de Mezcla de Profundidades

Mixture-of-Depths Attention

Resumen

Support