Внимание Residuals (AttnRes) — это механизм, предложенный исследователями из Moonshot AI, который заменяет стандартное накопление остаточных соединений (Residual connections) в трансформерах.
Остаточные соединения — одна из наименее обсуждаемых частей современного дизайна трансформеров. В архитектурах PreNorm каждый слой добавляет свой выход в текущее скрытое состояние, что обеспечивает стабильность оптимизации и позволяет обучать глубокие модели. Однако исследователи Moonshot AI утверждают, что этот стандартный механизм также вводит структурную проблему: все выходы предыдущих слоёв накапливаются с фиксированными единичными весами, что приводит к росту величины скрытого состояния с глубиной и постепенно ослабляет вклад любого отдельного слоя.
Почему стандартные остаточные соединения становятся узким местом
Исследовательская группа выявила три проблемы со стандартным накоплением остаточных соединений:
1. Нет избирательного доступа: все слои получают одно и то же агрегированное состояние, даже если слои внимания, прямого распространения или MoE могут выиграть от разных смесей более ранней информации.
2. Необратимая потеря: как только информация смешивается в единый остаточный поток, более поздние слои не могут выборочно восстановить конкретные более ранние представления.
3. Рост выходных данных: более глубокие слои стремятся производить более крупные выходные данные, чтобы оставаться влиятельными внутри постоянно растущего накопленного состояния, что может дестабилизировать обучение.
Full AttnRes: внимание ко всем предыдущим слоям
В Full AttnRes каждый слой вычисляет веса внимания по всем предыдущим источникам глубины. По умолчанию дизайн не использует входной запрос, зависящий от условий. Вместо этого каждый слой имеет изученный специфичный для слоя псевдозапрос $w_l ∈ R^d$, в то время как ключи и значения поступают из встраивания токена и выходов предыдущего слоя после RMSNorm.
Шаг RMSNorm важен, поскольку он предотвращает доминирование выходных данных большого размера над весами внимания по глубине.
Full AttnRes прост, но увеличивает затраты. На токен требуется $O(L^2 d)$ арифметических операций и $(O(Ld))$ памяти для хранения выходных данных слоя.
Block AttnRes: практический вариант для больших моделей
Чтобы сделать метод масштабируемым, исследовательская группа Moonshot AI представляет Block AttnRes. Вместо того чтобы учитывать каждый предыдущий выход слоя, модель разбивает слои на $N$ блоков. В каждом блоке выходы накапливаются в единое блочное представление, и внимание применяется только к этим блочным представлениям плюс встраивание токена. Это сокращает накладные расходы на память и коммуникации с $O(Ld)$ до $O(Nd)$.
Исследовательская группа описывает кэшированную конвейерную связь и стратегию двухфазных вычислений, которые делают Block AttnRes практичным при распределённом обучении и выводе. Это приводит к менее чем 4% накладным расходам на обучение при конвейерном параллелизме, в то время как в репозитории сообщается о менее чем 2% накладных расходах на задержку вывода при типичных рабочих нагрузках.
Результаты масштабирования
Исследовательская группа оценивает пять размеров моделей и сравнивает три варианта в каждом размере: базовый уровень PreNorm, Full AttnRes и Block AttnRes примерно с восемью блоками. Все варианты в каждой группе размеров имеют одинаковые гиперпараметры, выбранные на основе базового уровня, что делает сравнение консервативным.
Отчёты о масштабировании представлены в виде:
* Базовый уровень: $L = 1,891 × C^{-0,057}$.
* Block AttnRes: $L = 1,870 × C^{-0,058}$.
* Full AttnRes: $L = 1,865 × C^{-0,057}$.
Практический вывод заключается в том, что AttnRes достигает более низкого уровня проверки потерь в протестированном диапазоне вычислений, а Block AttnRes соответствует потерям базового уровня, обученного примерно с $1,25×$ большим количеством вычислений.
Интеграция в Kimi Linear
Moonshot AI также интегрирует AttnRes в Kimi Linear, свою архитектуру MoE с $48$ миллиардами общих параметров и $3$ миллиардами активированных параметров, и предварительно обучает её на $1,4$ триллионах токенов.
Согласно исследовательской статье, AttnRes смягчает разбавление PreNorm, ограничивая величину выходных данных по глубине и более равномерно распределяя градиенты по слоям.
Основные выводы
* Attention Residuals заменяет фиксированное накопление остаточных соединений с помощью softmax-внимания к предыдущим слоям.
* По умолчанию дизайн AttnRes использует изученный специфичный для слоя псевдозапрос, а не входной запрос, зависящий от условий.
* Block AttnRes делает метод практичным, сокращая глубину памяти и коммуникации с $O(Ld)$ до $O(Nd)$.
* Исследовательская группа Moonshot сообщает о более низких потерях при масштабировании, чем базовый уровень PreNorm, при этом Block AttnRes соответствует примерно $1,25×$ большему объёму вычислений базового уровня.
* В Kimi Linear AttnRes улучшает результаты по тестам на рассуждения, кодирование и оценку с ограниченными накладными расходами.
1. Какие проблемы исследователи из Moonshot AI обнаружили в стандартных остаточных соединениях в трансформерах?
В стандартных остаточных соединениях исследователи из Moonshot AI обнаружили три проблемы:
* Нет избирательного доступа: все слои получают одно и то же агрегированное состояние, даже если слои внимания, прямого распространения или MoE могут выиграть от разных смесей более ранней информации.
* Необратимая потеря: как только информация смешивается в единый остаточный поток, более поздние слои не могут выборочно восстановить конкретные более ранние представления.
* Рост выходных данных: более глубокие слои стремятся производить более крупные выходные данные, чтобы оставаться влиятельными внутри постоянно растущего накопленного состояния, что может дестабилизировать обучение.
2. В чём заключается основное отличие механизма Attention Residuals (AttnRes) от стандартных остаточных соединений?
Основное отличие механизма Attention Residuals (AttnRes) от стандартных остаточных соединений заключается в том, что AttnRes заменяет фиксированное накопление остаточных соединений с помощью softmax-внимания к предыдущим слоям. В то время как стандартные остаточные соединения накапливают выходы предыдущих слоёв с фиксированными единичными весами, AttnRes позволяет каждому слою вычислять веса внимания по всем предыдущим источникам глубины.
3. Какие два варианта реализации Attention Residuals (AttnRes) предложены в статье?
В статье предложены два варианта реализации Attention Residuals (AttnRes):
* Full AttnRes: каждый слой вычисляет веса внимания по всем предыдущим источникам глубины.
* Block AttnRes: модель разбивает слои на N блоков. В каждом блоке выходы накапливаются в единое блочное представление, и внимание применяется только к этим блочным представлениям плюс встраивание токена.
4. Какой вывод можно сделать о практической применимости и эффективности Block AttnRes на основе информации из статьи?
На основе информации из статьи можно сделать вывод, что Block AttnRes является практичным вариантом для больших моделей, так как сокращает накладные расходы на память и коммуникации. Исследовательская группа описывает кэшированную конвейерную связь и стратегию двухфазных вычислений, которые делают Block AttnRes практичным при распределённом обучении и выводе. Также сообщается о менее чем 4% накладных расходах на обучение при конвейерном параллелизме и менее чем 2% накладных расходах на задержку вывода при типичных рабочих нагрузках.
5. Какие результаты были получены при сравнении базового уровня PreNorm, Full AttnRes и Block AttnRes?
При сравнении базового уровня PreNorm, Full AttnRes и Block AttnRes было получено, что все варианты в каждой группе размеров имеют одинаковые гиперпараметры, выбранные на основе базового уровня. Отчёты о масштабировании представлены в виде:
* Базовый уровень: $L = 1,891 × C^{-0,057}$.
* Block AttnRes: $L = 1,870 × C^{-0,058}$.
* Full AttnRes: $L = 1,865 × C^{-0,057}$.
Практический вывод заключается в том, что AttnRes достигает более низкого уровня проверки потерь в протестированном диапазоне вычислений, а Block AttnRes соответствует потерям базового уровня, обученного примерно с $1,25×$ большим количеством вычислений.