Aandachtsresiduals

Samenvatting

Residuale verbindingen met PreNorm zijn standaard in moderne LLM's, maar ze accumuleren alle laaguitvoeren met vaste eenheidsgewichten. Deze uniforme aggregatie veroorzaakt een onbeheerste groei van de verborgen toestand met de diepte, waardoor de bijdrage van elke laag progressief wordt verdund. Wij stellen Attention Residuals (AttnRes) voor, waarbij deze vaste accumulatie wordt vervangen door softmax-attentie over de uitvoeren van voorgaande lagen, waardoor elke laag eerere representaties selectief kan aggregeren met aangeleerde, invoerafhankelijke gewichten. Om het geheugen- en communicatieoverhead aan te pakken van het toepassen van aandacht op alle voorgaande laaguitvoeren voor grootschalige modeltraining, introduceren we Block AttnRes, waarbij lagen worden gepartitioneerd in blokken en aandacht wordt toegepast op blokniveau-representaties, waardoor de geheugenvoetafdruk wordt verkleind terwijl het grootste deel van de winst van volledige AttnRes behouden blijft. Gecombineerd met cache-gebaseerde pijplijncommunicatie en een tweefasen-berekeningsstrategie wordt Block AttnRes een praktische drop-in vervanging voor standaard residuale verbindingen met minimale overhead. Schaalwet-experimenten bevestigen dat de verbetering consistent is over verschillende modelgroottes, en ablatiestudies valideren het voordeel van inhoudsafhankelijke, dieptewaartse selectie. We integreren AttnRes verder in de Kimi Linear-architectuur (48B totaal / 3B geactiveerde parameters) en pre-trainen op 1.4T tokens, waarbij AttnRes PreNorm-verdunning vermindert, wat resulteert in uniformere uitvoergrootten en gradiëntverdeling over de diepte, en de downstream-prestaties verbetert voor alle geëvalueerde taken.

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.