Residui di Attenzione

Abstract

Le connessioni residue con PreNorm sono standard nei moderni LLM, ma accumulano gli output di tutti i livelli con pesi unitari fissi. Questa aggregazione uniforme causa una crescita incontrollata dello stato nascosto con la profondità, diluendo progressivamente il contributo di ciascun livello. Proponiamo Attention Residuals (AttnRes), che sostituisce questo accumulo fisso con un'attenzione softmax sugli output dei livelli precedenti, consentendo a ogni livello di aggregare selettivamente le rappresentazioni precedenti con pesi appresi e dipendenti dall'input. Per affrontare l'overhead di memoria e comunicazione derivante dall'applicare l'attenzione su tutti gli output dei livelli precedenti nell'addestramento di modelli su larga scala, introduciamo Block AttnRes, che suddivide i livelli in blocchi e applica l'attenzione sulle rappresentazioni a livello di blocco, riducendo l'ingombro di memoria preservando gran parte dei vantaggi di AttnRes completo. Combinato con una comunicazione pipeline basata su cache e una strategia di calcolo a due fasi, Block AttnRes diventa una pratica sostituzione diretta per le connessioni residue standard con overhead minimo. Gli esperimenti sulle leggi di scaling confermano che il miglioramento è consistente across le dimensioni del modello, e le ablation study convalidano il beneficio della selezione depth-wise dipendente dal contenuto. Integriamo inoltre AttnRes nell'architettura Kimi Linear (48B parametri totali / 3B parametri attivati) e effettuiamo il pre-addestramento su 1.4T di token, dove AttnRes mitiga la diluizione di PreNorm, producendo magnitudini di output e distribuzioni del gradiente più uniformi lungo la profondità, e migliora le prestazioni downstream in tutti i compiti valutati.

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.