Учёные представили Multi-Head Attention Residuals, раздельное внимание к истории слоёв трансформера

В обычном трансформере информация идёт по слоям через единый остаточный поток: каждый подслой видит только состояние с предыдущего шага. Более новый подход, «attention residuals» (остаточные связи с механизмом внимания), даёт подслою возможность через обучаемый softmax самому выбирать, из каких предыдущих слоёв читать информацию. Проблема в том, что этот выбор делается одним общим запросом (query) на всю ширину модели, то есть все группы признаков (подпространства) вынуждены пользоваться одним и тем же распределением внимания к истории слоёв, даже если им на самом деле нужны разные слои. Чем шире модель, тем сильнее подпространства расходятся во мнениях, какие слои им важны, и тем дороже обходится этот вынужденный компромисс.

Авторы предлагают Multi-Head Attention Residuals (MHAR): тот же самый маршрутизирующий запрос разбивается на H отдельных голов, по одной на подпространство признаков, и каждая голова считает свой собственный softmax по истории слоёв независимо от остальных. Технически чтение становится блочно-диагональным. Важно, что это разбиение не добавляет ни одного нового параметра и почти не увеличивает объём вычислений, а при H = 1 метод в точности превращается в исходные attention residuals, то есть MHAR является обобщением, а не заменой.

Метод обучали с нуля на очищенном от дублей корпусе на основе Nemotron, отфильтрованном по качеству и смещённом в сторону STEM и кода, на трёх масштабах, 100 миллионов, 350 миллионов и 1 миллиард параметров. На всех трёх MHAR снизил ошибку на валидации относительно обычного трансформера (на -0,061, -0,149 и -0,140 соответственно) и оказался лучшим из четырёх сравниваемых методов на каждом масштабе, причём выигрыш рос при переходе от 100M к более крупным моделям.

Число голов H, не формальность, а реальный параметр архитектуры: зависимость ошибки от H имеет форму буквы U с плоским минимумом при H = 4 или H = 8 (устойчиво на всех масштабах); авторы рекомендуют H = 8 для крупных моделей. Дальнейшее дробление до H = 16 частично сводит выигрыш на нет. Прямая проверка обученных запросов подтвердила механизм: подпространства признаков действительно расходятся в том, какие слои им нужны, и именно это расхождение, причина выигрыша MHAR.

Отдельно авторы ускорили обучение с помощью специальных Triton-ядер для маршрутизации: скорость обучения attention residuals выросла с 0,2, 0,5x от скорости базового трансформера до 0,55, 0,88x при почти той же пиковой памяти. Наконец, авторы показали способ перевести уже обученную модель на эту архитектуру без обучения с нуля: преобразование, сохраняющее тождественную функцию модели, переводит её в «дельта»-версию attention residuals; на модели с 8 миллиардами параметров такое промежуточное дообучение дало прирост +3,2 пункта на бенчмарке GSM8K и +3,1 пункта на GPQA.

Ключевые факты

MHAR заменяет один общий запрос внимания к истории слоёв на H независимых голов, по одной на группу признаков (подпространство), каждая со своим softmax по прошлым слоям
На масштабах 100M, 350M и 1B параметров MHAR снизил ошибку на валидации относительно стандартного трансформера на 0,061, 0,149 и 0,140 соответственно и стал лучшим из четырёх сравниваемых методов на каждом масштабе
Оптимальное число голов, H = 4 или H = 8, и это устойчиво на всех масштабах; дальнейшее дробление до H = 16 частично отдаёт выигрыш обратно
Разбиение на головы не добавляет параметров и почти не увеличивает вычисления; специальные Triton-ядра подняли скорость обучения attention residuals с 0,2, 0,5x до 0,55, 0,88x от скорости базового трансформера при сопоставимой пиковой памяти
Уже обученную модель на 8 миллиардов параметров перевели на новую архитектуру без обучения с нуля через промежуточное дообучение, это дало +3,2 на GSM8K и +3,1 на GPQA

Почему это важно

Стандартный трансформер передаёт информацию между слоями через один поток, и даже более гибкие «attention residuals» упираются в то, что весь выбор, из каких слоёв читать историю, идёт через единственный общий запрос на всю ширину модели. Из-за этого разные группы признаков вынуждены довольствоваться одним и тем же компромиссным распределением внимания, хотя им реально нужны разные слои, и чем шире модель, тем дороже этот компромисс. MHAR устраняет именно это архитектурное узкое место, причём почти бесплатно с точки зрения параметров и вычислений, что делает его логичным следующим шагом для целого семейства архитектур с «остаточными связями через внимание».

Кому это важно

Прежде всего, исследователям и инженерам, занимающимся архитектурами трансформеров и предобучением больших языковых моделей, особенно тем, кто уже экспериментирует с альтернативами обычным остаточным связям. Также релевантно командам, которые обучают модели на масштабах от сотен миллионов до нескольких миллиардов параметров и считают компромисс между качеством и скоростью обучения.

Как это применить

Технически изменение сводится к тому, чтобы разбить единый маршрутизирующий запрос на H голов, каждая со своим softmax по истории слоёв; авторы советуют H = 8 для крупных моделей. Чтобы получить конкурентную скорость обучения, нужны специализированные Triton-ядра для маршрутизации, без них attention residuals заметно медленнее базового трансформера. Отдельный практический путь, не обучать модель заново, а перевести уже готовую (в статье, на 8 миллиардов параметров) на архитектуру через промежуточное дообучение с сохраняющим тождество преобразованием в «дельта»-версию attention residuals.

Можно ли доверять

Методология выглядит основательно: метод обучали с нуля на реальном (не игрушечном) корпусе на основе Nemotron на трёх разных масштабах, 100M, 350M и 1B параметров, сравнивали с четырьмя альтернативами, провели чёткую абляцию по числу голов с воспроизводимым U-образным минимумом и напрямую проверили обученные запросы, подтвердив причинный механизм выигрыша (расхождение подпространств), а не просто постфактум-объяснение. Ограничение в том, что это одна препринтная работа на arXiv без указаний на рецензирование, а обучение с нуля проверено только до 1 миллиарда параметров, тест на 8 миллиардах сделан через дообучение уже готовой модели, а не обучение с нуля.

Риски и подводные камни

Главный риск, ограниченный масштаб: обучение с нуля проверено только до 1 миллиарда параметров, а результаты на моделях в десятки и сотни миллиардов параметров (где реально работают современные крупные модели) пока не показаны. Метод также требует специальных Triton-ядер для маршрутизации, без них варианты attention residuals работают заметно медленнее базового трансформера (0,2, 0,5x скорости). Наконец, MHAR добавляет архитектурную сложность: новый гиперпараметр (число голов H), который нужно подбирать под масштаб, и собственный набор кастомных ядер для эффективной реализации, то есть выигрыш в качестве достаётся не совсем бесплатно с точки зрения инженерной сложности.