Внимание с остаточными связями

Аннотация

Остаточные связи с PreNorm являются стандартом в современных больших языковых моделях, однако они суммируют выходы всех слоев с фиксированными единичными весами. Такая равномерная агрегация приводит к неконтролируемому росту скрытых состояний с увеличением глубины, постепенно размывая вклад каждого слоя. Мы предлагаем Attention Residuals (AttnRes), которые заменяют это фиксированное суммирование на механизм внимания softmax над выходами предыдущих слоев, позволяя каждому слою выборочно агрегировать предыдущие представления с обучаемыми, зависящими от входных данных весами. Для решения проблемы высоких затрат памяти и коммуникации при обработке всех предыдущих слоев в обучении крупномасштабных моделей мы вводим Block AttnRes, который разбивает слои на блоки и применяет внимание к представлениям на уровне блоков, сокращая объем используемой памяти при сохранении большей части преимуществ полного AttnRes. В сочетании с конвейерной коммуникацией на основе кэширования и двухфазной стратегией вычислений Block AttnRes становится практичной заменой стандартных остаточных связей с минимальными накладными расходами. Эксперименты с законами масштабирования подтверждают, что улучшение согласованно для моделей разных размеров, а ablation-исследования подтверждают преимущество содержательно-зависимого глубинного выбора. Мы также интегрировали AttnRes в архитектуру Kimi Linear (48B общих / 3B активируемых параметров) и провели предварительное обучение на 1.4T токенах, где AttnRes смягчает размытие PreNorm, обеспечивая более равномерные амплитуды выходов и распределение градиентов по глубине, а также улучшает производительность на всех оцениваемых downstream-задачах.

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.