Atención de Cabezas Golpeantes

Resumen

La atención multicabezal (MHA) se ha convertido en la piedra angular de los modelos de lenguaje modernos a gran escala, mejorando la capacidad de representación mediante cabezales de atención paralelos. Sin embargo, aumentar el número de cabezales debilita inherentemente la capacidad individual de cada uno, y los mecanismos de atención existentes -ya sea la MHA estándar o sus variantes como la atención por consultas agrupadas (GQA) y la atención agrupada y vinculada (GTA)- simplemente concatenan las salidas de cabezales aislados sin una interacción sólida. Para abordar esta limitación, proponemos la atención de cabezales interconectados (KHA), que permite a los cabezales de atención "golpearse" entre sí, facilitando interacciones a nivel de características entre cabezales antes de la atención de producto escalado. Esto se logra aplicando una matriz de proyección compartida e inicializada diagonalmente a través de todos los cabezales. La inicialización diagonal preserva la especialización específica de cada cabezal al inicio del entrenamiento, mientras permite que el modelo aprenda progresivamente representaciones integradas entre cabezales. KHA agrega solo parámetros y operaciones de punto flotante (FLOPs) mínimos, y puede integrarse perfectamente en MHA, GQA, GTA y otras variantes de atención. Validamos KHA entrenando un modelo MoE de 6.1B parámetros (1.01B activados) con 1 billón de tokens de alta calidad. En comparación con los mecanismos de atención de referencia, KHA aporta una dinámica de entrenamiento superior y más estable, logrando un mejor rendimiento en diversas tareas posteriores.

English

Multi-head attention (MHA) has become the cornerstone of modern large language models, enhancing representational capacity through parallel attention heads. However, increasing the number of heads inherently weakens individual head capacity, and existing attention mechanisms - whether standard MHA or its variants like grouped-query attention (GQA) and grouped-tied attention (GTA) - simply concatenate outputs from isolated heads without strong interaction. To address this limitation, we propose knocking-heads attention (KHA), which enables attention heads to "knock" on each other - facilitating cross-head feature-level interactions before the scaled dot-product attention. This is achieved by applying a shared, diagonally-initialized projection matrix across all heads. The diagonal initialization preserves head-specific specialization at the start of training while allowing the model to progressively learn integrated cross-head representations. KHA adds only minimal parameters and FLOPs and can be seamlessly integrated into MHA, GQA, GTA, and other attention variants. We validate KHA by training a 6.1B parameter MoE model (1.01B activated) on 1T high-quality tokens. Compared to baseline attention mechanisms, KHA brings superior and more stable training dynamics, achieving better performance across downstream tasks.

Atención de Cabezas Golpeantes

Knocking-Heads Attention

Resumen

Support