Resíduos de Atenção

Resumo

As conexões residuais com PreNorm são padrão em LLMs modernos, mas acumulam todas as saídas das camadas com pesos unitários fixos. Esta agregação uniforme causa um crescimento não controlado do estado oculto com a profundidade, diluindo progressivamente a contribuição de cada camada. Propomos os *Attention Residuals* (AttnRes), que substituem este acúmulo fixo por atenção softmax sobre as saídas das camadas precedentes, permitindo que cada camada agregue seletivamente representações anteriores com pesos aprendidos e dependentes da entrada. Para lidar com a sobrecarga de memória e comunicação ao atender sobre todas as saídas de camadas precedentes no treinamento de modelos em larga escala, introduzimos o *Block AttnRes*, que particiona as camadas em blocos e atende sobre representações a nível de bloco, reduzindo a pegada de memória enquanto preserva a maior parte dos ganhos do AttnRes completo. Combinado com comunicação de pipeline baseada em *cache* e uma estratégia de computação em duas fases, o Block AttnRes torna-se uma substituição prática e direta (*drop-in*) para as conexões residuais padrão, com sobrecarga mínima. Experimentos com leis de escala confirmam que a melhoria é consistente em diferentes tamanhos de modelo, e *ablation studies* validam o benefício da seleção dependente de conteúdo ao longo da profundidade. Integramos ainda o AttnRes na arquitetura Kimi Linear (48B no total / 3B parâmetros ativados) e realizamos pré-treinamento em 1.4T de *tokens*, onde o AttnRes mitiga a diluição do PreNorm, resultando em magnitudes de saída e distribuição de gradiente mais uniformes ao longo da profundidade, e melhora o desempenho *downstream* em todas as tarefas avaliadas.

English

Residual connections with PreNorm are standard in modern LLMs, yet they accumulate all layer outputs with fixed unit weights. This uniform aggregation causes uncontrolled hidden-state growth with depth, progressively diluting each layer's contribution. We propose Attention Residuals (AttnRes), which replaces this fixed accumulation with softmax attention over preceding layer outputs, allowing each layer to selectively aggregate earlier representations with learned, input-dependent weights. To address the memory and communication overhead of attending over all preceding layer outputs for large-scale model training, we introduce Block AttnRes, which partitions layers into blocks and attends over block-level representations, reducing the memory footprint while preserving most of the gains of full AttnRes. Combined with cache-based pipeline communication and a two-phase computation strategy, Block AttnRes becomes a practical drop-in replacement for standard residual connections with minimal overhead. Scaling law experiments confirm that the improvement is consistent across model sizes, and ablations validate the benefit of content-dependent depth-wise selection. We further integrate AttnRes into the Kimi Linear architecture (48B total / 3B activated parameters) and pre-train on 1.4T tokens, where AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient distribution across depth, and improves downstream performance across all evaluated tasks.