Atenção de Cabeças Batendo (Knocking-Heads Attention)

Resumo

A atenção multi-cabeça (MHA) tornou-se a pedra angular dos modelos de linguagem grandes modernos, aumentando a capacidade de representação por meio de cabeças de atenção paralelas. No entanto, o aumento do número de cabeças inerentemente enfraquece a capacidade individual de cada cabeça, e os mecanismos de atenção existentes - seja a MHA padrão ou suas variantes como a atenção por consulta agrupada (GQA) e a atenção agrupada e vinculada (GTA) - simplesmente concatenam as saídas de cabeças isoladas sem uma interação forte. Para superar esta limitação, propomos a atenção de cabeças intercomunicantes (KHA), que permite que as cabeças de atenção "batam" umas nas outras - facilitando interações ao nível de *features* entre cabeças antes da atenção de produto escalonado. Isto é alcançado aplicando uma matriz de projeção compartilhada e inicializada diagonalmente em todas as cabeças. A inicialização diagonal preserva a especialização específica de cada cabeça no início do treinamento, permitindo ao mesmo tempo que o modelo aprenda progressivamente representações integradas entre cabeças. A KHA adiciona apenas parâmetros e FLOPs mínimos e pode ser integrada perfeitamente na MHA, GQA, GTA e outras variantes de atenção. Validamos a KHA treinando um modelo MoE de 6,1B de parâmetros (1,01B ativados) em 1T de *tokens* de alta qualidade. Em comparação com os mecanismos de atenção de base, a KHA proporciona uma dinâmica de treinamento superior e mais estável, alcançando um melhor desempenho em tarefas subsequentes.

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.