xHC: остаточный поток трансформеров расширили до 16 линий

Метод Hyper-Connections (HC), способ увеличить «память» трансформера не за счёт ширины или глубины сети, а за счёт расширения остаточного потока (residual stream, канала, по которому информация проходит между слоями) в N параллельных линий. Его усовершенствованная версия mHC (Manifold-Constrained HC) стабилизирует обучение при таком расширении. Переход от N=1 к N=4 давал заметный прирост качества, и это навело исследователей на мысль, что расширение остаточного потока, перспективная ось масштабирования моделей. Но все существующие методы семейства HC на практике останавливались на N=4: дальнейшее увеличение числа линий у mHC давало всё меньший прирост качества при быстро растущей стоимости обучения.

Авторы работы объясняют затор двумя узкими местами. Во-первых, при росте числа линий не хватает информации для «обратной записи», того, что каждая линия возвращает в общий поток. Во-вторых, стоимость генерации смешивания остаточных потоков растёт кубически от N, то есть с каждой добавленной линией вычисления дорожают несоразмерно быстро.

Предложенный метод xHC (Expanded Hyper-Connections), первый в семействе HC, который добивается осмысленного расширения за пределы N=4. Он решает обе проблемы: добавляет временну́ю аугментацию признаков (temporal feature augmentation), которая обогащает информацию для обратной записи, и переходит на разреженную архитектуру остаточного потока, вместо обновления всех N=16 линий одновременно на каждом шаге активно обновляются только k=4 из них, при этом модель сохраняет плотный доступ ко всему остаточному состоянию.

На моделях типа MoE (mixture-of-experts, смесь экспертов) размером 18 и 28 миллиардов параметров xHC даёт устойчивый прирост качества. На 18-миллиардной модели средний балл по прикладным задачам вырос на 4,0 пункта по сравнению с mHC, при этом дополнительные вычисления при обучении по сравнению с обычной (vanilla) моделью оказались умеренными. В экспериментах по законам масштабирования обычной модели потребовалось в 1,50 раза больше вычислений, а mHC, в 1,19 раза больше, чтобы достичь того же значения функции потерь, что и xHC.

Отдельная проблема практического обучения при большом N, нагрузка на память из-за расширенного остаточного состояния. Для неё авторы предложили модификацию xHC-Flash: она снижает трафик обращений к памяти на подслой трансформера с 73,5 до 40 условных единиц C, это сопоставимо с 34 единицами, которые требуются mHC при N=4, и при этом сохраняет весь прирост качества полной версии xHC. Вместе xHC и xHC-Flash делают расширение остаточного потока до больших N практичным и эффективным способом предобучения больших языковых моделей.

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

Метод Hyper-Connections (HC) расширяет остаточный поток трансформера в N параллельных линий, но существующие версии упирались в потолок N=4 из-за роста стоимости обучения
Новый метод xHC впервые доводит расширение до N=16, обновляя разреженно только 4 из 16 линий за шаг, но сохраняя доступ ко всему состоянию
На 18-миллиардной MoE-модели xHC повышает средний балл по задачам на 4,0 пункта по сравнению с mHC при умеренном росте вычислений
Чтобы достичь того же качества, что и xHC, обычной модели нужно в 1,50 раза больше вычислений, а mHC, в 1,19 раза больше
Дополнение xHC-Flash снижает нагрузку на память почти вдвое (с 73,5 до 40 условных единиц на подслой) без потери прироста качества

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

Простое наращивание ширины и глубины трансформеров даёт всё меньшую отдачу, поэтому индустрия ищет новые оси масштабирования моделей. Расширение остаточного потока, один из кандидатов, но семейство методов Hyper-Connections упиралось в потолок N=4: дальше расти не получалось из-за нехватки информации для обратной записи и кубического роста стоимости смешивания потоков. xHC впервые снимает это ограничение и доводит расширение до N=16, добавляя 4,0 пункта к качеству и сокращая необходимые вычисления по сравнению и с обычной моделью (в 1,50 раза), и с предыдущей версией mHC (в 1,19 раза), то есть предлагает конкретный, измеримый путь удешевить предобучение больших моделей.

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

В первую очередь исследователям архитектур трансформеров и инженерным командам, которые предобучают крупные языковые и MoE-модели (в статье, 18 и 28 миллиардов параметров) и ищут способы снизить вычислительные затраты при сохранении или росте качества. Также полезно тем, кто занимается эффективностью обучения на уровне инфраструктуры: xHC-Flash напрямую касается управления памятью на GPU-кластерах.

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

xHC, это не готовый продукт, а изменение архитектуры трансформера на этапе предобучения: нужно встроить механизм временной аугментации признаков и разреженное обновление N=16 остаточных линий (из которых на каждом шаге активны k=4) в код тренировочного фреймворка. Для практического применения при больших N авторы рекомендуют использовать вместе с xHC модификацию xHC-Flash, которая контролирует трафик обращений к памяти и делает обучение реализуемым на существующем железе.

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

Это препринт с описанием эмпирических экспериментов на моделях 18 и 28 миллиардов параметров с конкретными числовыми результатами (прирост в 4,0 пункта, соотношения вычислений 1,50x и 1,19x, снижение памяти с 73,5 до 40 единиц), то есть заявления подкреплены измеримыми данными, а не только теоретическими рассуждениями. Над работой, 13 авторов из Shanghai Jiao Tong University, Xiaohongshu Inc., USTC, Peking University и CUHK, но она пока не прошла независимую рецензию, поэтому воспроизводимость результатов сторонними командами ещё предстоит проверить.

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

Результаты проверены только на моделях 18 и 28 миллиардов параметров, неясно, сохранится ли эффект на существенно бо́льших или меньших моделях и при дальнейшем росте N сверх 16. Кубический рост стоимости смешивания потоков, который ограничивал предыдущие методы, в принципе может вновь проявиться при ещё большем масштабировании. Как и у любой архитектурной модификации, интеграция xHC и xHC-Flash в существующие тренировочные фреймворки требует дополнительной инженерной работы и не гарантирует одинаковый выигрыш на всех типах задач и архитектур.