Механизм внимания с перекрёстными связями (Knocking-Heads Attention)

Аннотация

Многоголовое внимание (MHA) стало краеугольным камнем современных больших языковых моделей, повышая репрезентативную способность за счет параллельных голов внимания. Однако увеличение числа голов неизбежно ослабляет индивидуальную емкость каждой головы, а существующие механизмы внимания — будь то стандартное MHA или его варианты, такие как групповое запросное внимание (GQA) и групповое связанное внимание (GTA) — просто конкатенируют выходы изолированных голов без сильного взаимодействия. Для преодоления этого ограничения мы предлагаем механизм "стучащихся голов" (KHA), который позволяет головам внимания "стучаться" друг о друга, обеспечивая кросс-головые взаимодействия на уровне признаков до вычисления масштабированного скалярного произведения. Это достигается за счет применения общей, диагонально инициализированной проекционной матрицы ко всем головам. Диагональная инициализация сохраняет специализацию отдельных голов в начале обучения, позволяя модели постепенно изучать интегрированные межголовые представления. KHA добавляет лишь минимальное количество параметров и операций с плавающей точкой (FLOPs) и может быть бесшовно интегрировано в MHA, GQA, GTA и другие варианты внимания. Мы подтверждаем эффективность KHA, обучив MoE-модель с 6.1 млрд параметров (1.01 млрд активированных) на 1 трлн высококачественных токенов. По сравнению с базовыми механизмами внимания, KHA демонстрирует превосходную и более стабильную динамику обучения, достигая лучших результатов на последующих задачах.

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.