Глава 2.4 — Нормализация и стабилизация

Содержание

1. Стек определён; можем ли мы на самом деле обучить его глубоким?
2. Зачем глубоким сетям вообще нужна нормализация
3. Механика LayerNorm и почему не BatchNorm
4. Pre-LN против Post-LN: где именно вы размещаете нормализацию, меняет всё
5. RMSNorm: нормализация с удалённым центрированием по среднему
6. Взгляд с точки зрения собеседования
7. Вопросы для самопроверки
8. Источники

1. Стек определён; можем ли мы на самом деле обучить его глубоким?

Глава 2.3 завершилась выбором в пользу стеков "только декодер" из маскированного self-attention и полносвязных подслоёв как доминирующей формы для крупномасштабных генеративных LLM, и указанием на то, что этот стек ещё не является стабильным для обучения на реальной глубине без посторонней помощи. Это утверждение стоит воспринимать буквально, а не как риторический оборот. Transformer, как и практически любая архитектура глубокой сети до него, по умолчанию вовсе не становится лучше просто от того, что вы наращиваете число слоёв друг на друга — начиная с определённой глубины наивное наращивание слоёв делает обучение сложнее, а не проще, и может привести к его полному провалу. Каждый блок Transformer, описанный в главах 2.1–2.3, был незаметно нарисован с компонентом, который эта книга ещё не рассматривала: слоем нормализации, обёрнутым вокруг каждого подслоя, без которого остаточное наращивание слоёв, делающее возможными глубокие Transformer, на практике попросту не работает.

Эта глава посвящена именно этому компоненту. Легко относиться к нормализации как к незначительной детали реализации — строчке кода, которую вы добавляете, потому что она есть в эталонной реализации, — но конкретные выборы, о которых идёт речь (по чему именно нормализовать и где именно в остаточном пути разместить нормализацию), оказываются ответственными за то, обучается ли вообще 96-слойный Transformer стабильно, и это ровно та деталь, которая отличает инженера, по-настоящему понявшего, почему архитектура работает, от того, кто её лишь использует. Как только это оказывается на столе, неразобранным остаётся лишь один компонент блока — полносвязный подслой, которым и займётся глава 2.5, завершая часть II.

2. Зачем глубоким сетям вообще нужна нормализация

Ключевая проблема, которую решает нормализация, возникла задолго до Transformer и является общим свойством глубоких сетей, обучаемых градиентным спуском: по мере того как данные проходят вперёд через множество наложенных друг на друга слоёв, а параметры каждого слоя обновляются на основе градиентов, прошедших назад через все слои выше него, распределение активаций на любом заданном слое, как правило, дрейфует и смещается на протяжении обучения, отчасти потому, что распределение входа каждого слоя зависит от параметров всех предыдущих слоёв, которые все одновременно меняются. Ioffe и Szegedy, предложившие batch-нормализацию для свёрточных сетей, описали это явление как "внутренний ковариационный сдвиг" (internal covariate shift) — каждый слой фактически пытается выучить устойчивое отображение, в то время как статистики его собственного входа постоянно смещаются под ним, что замедляет обучение и вынуждает использовать меньшие, более консервативные скорости обучения, чем хотелось бы в противном случае.

Конкретно в остаточной архитектуре — а Transformer является остаточной архитектурой насквозь, поскольку выход каждого подслоя добавляется обратно к своему входу, а не заменяет его, — стоит явно назвать ещё одну, связанную неполадку. Если наложить много остаточных блоков без какой-либо ренормализации, масштаб бегущей суммы может расти примерно пропорционально числу блоков, поскольку каждый блок добавляет свой собственный вклад поверх всё растущего накопленного сигнала — $x_L = x_0 + \sum_{l=0}^{L-1}\text{Sublayer}_l(x_l)$, сумма из $L$ слагаемых, чья норма не имеет встроенного потолка; при достаточной глубине активации, доходящие до последних слоёв, и градиенты, идущие через них обратно, могут стать численно экстремальными — слишком большими или слишком малыми, чтобы поддерживать стабильные, хорошо обусловленные обновления градиента. Слои нормализации решают обе проблемы одновременно: явно перемасштабируя активации так, чтобы они имели контролируемые статистики в различных точках сети, они удерживают сигнал, получаемый каждым слоем, в согласованном, хорошо управляемом диапазоне независимо от того, что происходит в окружающих слоях, и независимо от глубины стека, — а это ровно то свойство, которое делает наращивание очень глубоких стеков практически осуществимым, а не только теоретически возможным.

3. Механика LayerNorm и почему не BatchNorm

Batch-нормализация, более ранняя и устоявшаяся техника, нормализует каждую активацию по среднему и дисперсии, вычисленным по измерению батча — для данного признака вы вычисляете его среднее и дисперсию по всем примерам в текущем минибатче и используете эти статистики для нормализации этого признака у каждого примера. Это хорошо работает для фиксированных по размеру, плотно заполненных карт признаков, типичных для свёрточных сетей, но плохо подходит для последовательных данных переменной длины по конкретной, ощутимой причине: батч последовательностей переменной длины требует дополнения (padding) до общей длины, а вычисление постатейной статистики по измерению батча в заданной позиции последовательности означает смешивание в одной и той же статистике позиций реальных токенов в одних последовательностях и токенов-заполнителей в других (или, в случае авторегрессивного обучения, статистик, которые должны были бы различаться в зависимости от того, насколько далеко модель продвинулась в генерации) — что плохо соответствует тому, что на самом деле требуется вычислению на уровне позиции и токена, каким является внимание. Статистики батча также ведут себя по-разному во время обучения (вычисляются из текущего батча) и во время инференса (обычно используется скользящее среднее), что вносит собственное рассогласование между обучением и инференсом и становится ещё более неудобным, когда инференс происходит по одному токену за раз, как это бывает при авторегрессивной генерации, где может вовсе не быть осмысленной структуры "батча".

Нормализация слоя (layer normalization), предложенная Ba, Kiros и Hinton, обходит всё это, нормализуя по измерению признаков, а не по измерению батча, вычисляя статистику независимо для каждого отдельного токена. Для вектора активации одного токена $x \in \mathbb{R}^d$ LayerNorm вычисляет среднее $\mu = \frac{1}{d}\sum_{k=1}^d x_k$ и дисперсию $\sigma^2 = \frac{1}{d}\sum_{k=1}^d (x_k - \mu)^2$ по собственным $d$ признакам этого токена, а затем нормализует: $$\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad \text{LN}(x) = \gamma \odot \hat{x} + \beta$$ где $\gamma$ и $\beta$ — обучаемые постатейные параметры масштаба и сдвига (позволяющие сети частично или полностью отменить нормализацию, если это оказывается полезным), а $\epsilon$ — малая константа, добавляемая для численной устойчивости. Поскольку эта статистика вычисляется на уровне токена, независимо от всех остальных токенов последовательности и независимо от всех остальных примеров в батче, она не требует никакой бухгалтерии на уровне батча, ведёт себя одинаково во время обучения и инференса и обрабатывает последовательности переменной длины и однотокенную авторегрессивную генерацию вообще без каких-либо особых случаев — именно те свойства, которые действительно нужны последовательной модели с входами переменной длины и инференсом по одному токену за раз.

4. Pre-LN против Post-LN: где именно вы размещаете нормализацию, меняет всё

Знание того, что в каждый подслой Transformer нужна некоторая форма нормализации, всё ещё оставляет открытым важное проектное решение: где именно, относительно остаточного соединения, размещается нормализация? Исходный Transformer, как описано у Vaswani et al., использует то, что сейчас называется конфигурацией Post-LN: подслой (self-attention или полносвязный) применяется к входу, добавляется обратно через остаточное соединение, и только после этого нормализуется — $$x_{l+1} = \text{LN}\big(x_l + \text{Sublayer}(x_l)\big)$$ так что нормализация оказывается прямо на основном остаточном пути, после каждого сложения. У этого есть реальный, эмпирически наблюдаемый недостаток в очень глубоких стеках: поскольку нормализация применяется после остаточного сложения, сам остаточный путь не остаётся в своём сыром, нетронутом виде — выход каждого блока ренормализуется перед передачей следующему блоку, а значит, почти-тождественный путь, который остаточные соединения призваны предоставлять (позволяя градиентам протекать по сути неизменными через сколь угодно много слоёв), периодически нарушается операцией нормализации, находящейся прямо на этом пути.

Анализ этой проблемы у Xiong et al., опубликованный в 2020 году, формулирует аргумент точно: они показывают, что в конфигурации Post-LN ожидаемая величина градиентов на слоях, ближайших ко входу, растёт или иным образом плохо контролируется по мере увеличения числа слоёв, что и является той самой нестабильностью, из-за которой очень глубоким Post-LN Transformer требуются тщательно подобранные графики разогрева скорости обучения (learning-rate warmup), чтобы вообще обучаться, и которая может привести к невозможности обучения после определённой глубины даже при наличии разогрева. Pre-LN переставляет те же два компонента так, чтобы подслой работал с нормализованной копией входа, а остаточное сложение происходило после, на сыром (ненормализованном) пути: $$x_{l+1} = x_l + \text{Sublayer}\big(\text{LN}(x_l)\big)$$ Теперь основной остаточный путь — последовательность сложений, несущая сигнал от самого первого слоя до самого последнего, — вообще не проходит через операцию нормализации; нормализуется лишь вход во внутреннее вычисление каждого подслоя. Это сохраняет свойство, которое изначально сделало остаточные соединения эффективными для очень глубоких сетей: по сути беспрепятственный аддитивный путь от входа к выходу, вдоль которого градиенты могут протекать с величиной, не зависящей от глубины сети так, как это происходит у Post-LN. Теоретический и эмпирический анализ Xiong et al. показывает, что Pre-LN Transformer можно обучать стабильно без графика разогрева скорости обучения, который обычно требуется Post-LN Transformer, и можно успешно обучать на глубинах, при которых конфигурации Post-LN становятся нестабильными.

Эта стабильность не даётся полностью бесплатно. В литературе, посвящённой этому сравнению, есть свидетельства того, что улучшенная обучаемость Pre-LN сопровождается некоторой ценой в итоговой представительной силе по сравнению с моделью Post-LN, которую действительно удаётся обучить до сходимости на той же глубине — нормализация, применяемая перед каждым подслоем на каждом слое, несколько сильнее ограничивает диапазон функций, которые может выразить каждый подслой, чем её применение лишь один раз после остаточной суммы на блок. На практике эта цена признана вполне оправданной в масштабе современных LLM: чуть менее выразительная, но надёжно обучаемая конфигурация побеждает теоретически более выразительную, которую не удаётся обучить после умеренной глубины без масштабной настройки, и Pre-LN (или её близкие варианты) сейчас является стандартным выбором практически в каждой крупной LLM типа "только декодер".

5. RMSNorm: нормализация с удалённым центрированием по среднему

Среднеквадратичная нормализация (root mean square normalization) Zhang и Sennrich задаёт другой вопрос о LayerNorm: учитывая, что LayerNorm делает две отдельные вещи — перецентрирование активаций путём вычитания среднего и перемасштабирование их путём деления на стандартное отклонение, — действительно ли шаг перецентрирования необходим для стабилизирующего эффекта, который даёт LayerNorm, или же практически всю полезную работу выполняет само перемасштабирование? Их ответ, подкреплённый как рассуждением, так и экспериментом, состоит в том, что именно перемасштабирование — контроль общей величины, или среднеквадратичного значения, вектора активации — является компонентом, действительно отвечающим за стабилизирующий эффект LayerNorm, и шаг перецентрирования по среднему можно опустить с минимальными потерями в качестве модели или вовсе без них, при этом сэкономив реальные вычисления.

Конкретно, RMSNorm заменяет нормализацию LayerNorm на $$\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{k=1}^d x_k^2}, \qquad \text{RMSNorm}(x) = \gamma \odot \frac{x}{\text{RMS}(x)}$$ что перемасштабирует вектор активации напрямую по его среднеквадратичной величине, без какого-либо вычитания среднего где бы то ни было, и обычно также без обучаемого параметра сдвига $\beta$ (сохраняется только обучаемый постатейный масштаб $\gamma$). На бумаге это небольшое изменение, но вычислительно оно значимо: удаление вычисления среднего и связанного с ним вычитания устраняет полную операцию редукции из каждого вызова нормализации, в каждом подслое, на каждом слое, для каждого токена — в масштабе моделей с десятками или сотнями миллиардов параметров, обрабатывающих триллионы обучающих токенов, это складывается в реальное, измеримое снижение стоимости обучения и инференса, причём Zhang и Sennrich сообщают об ускорениях порядка 10-40% для самого вычисления нормализации в зависимости от оборудования, практически без потерь в качестве модели на задачах, которые они оценивали. Этот компромисс "эффективность в обмен на пренебрежимо малую цену в качестве" — ровно тот тип аргумента, который, как правило, побеждает в масштабе LLM, поэтому RMSNorm, обычно сочетаемая с размещением Pre-LN, является схемой нормализации, используемой в большинстве современных крупных языковых моделей с открытыми весами.

Более современная, дополняющая техника нацелена на иной режим отказа, чем те, что решают RMSNorm и размещение Pre-LN: сами логиты внимания становятся достаточно большими — особенно на длинных контекстах или большом масштабе, — чтобы дестабилизировать обучение, даже когда сам остаточный поток ведёт себя хорошо. QK-norm, используемая в моделях вроде Gemma 2, нормализует векторы запроса и ключа (обычно с помощью RMSNorm) перед вычислением их скалярного произведения внутри внимания, что ограничивает величину скалярного произведения независимо от того, насколько большими могли бы иначе стать отдельные компоненты запроса и ключа, — воздействуя напрямую на вычисление оценки внимания, а не на остаточный поток, на который нацелено обычное размещение нормализации. Родственный приём, мягкое ограничение логитов (logit soft-capping), ограничивает логиты внимания (а отдельно — и финальные выходные логиты перед softmax, порождающим вероятности следующего токена) гладкой ограниченной функцией вместо того, чтобы оставлять их неограниченными, не давая редким экстремальным значениям добраться до численно хрупких участков softmax. Ни одна из этих техник не заменяет RMSNorm или Pre-LN; обе — узко нацеленные меры стабилизации, надстроенные поверх, и это само по себе паттерн, который стоит заметить: по мере того как модели подталкивают ко всё большей глубине, масштабу и длине контекста, продолжают всплывать новые, конкретные численные режимы отказа, которые исходный набор инструментов LayerNorm/Pre-LN не предвидел, — и ответ области стабильно состоит в том, чтобы добавить небольшой, точечный фикс, а не переизобретать нормализацию с нуля.

Вместе взятые, решения, рассмотренные в этой главе, — постатейные, по измерению признаков статистики LayerNorm вместо батчевых статистик BatchNorm, размещение Pre-LN для сохранения чистого остаточного пути через произвольную глубину и упрощение RMSNorm, оставляющее лишь важное перемасштабирование, — это то, что на самом деле делает возможным наращивание архитектуры "только декодер" из главы 2.3 до тех глубин, которые используют современные передовые модели, причём стабильно и эффективно. Это оставляет сам блок почти полностью собранным: внимание как ключевой механизм, позиционные схемы для восстановления информации об упорядоченности, которую внимание отбрасывает, выбор архитектурной формы, подходящей для генеративного моделирования, и нормализацию, делающую глубокие стеки этой формы действительно обучаемыми. Единственная часть, которая до сих пор описывалась лишь мимоходом, — второй подслой, сидящий внутри каждого блока: позиционный полносвязный слой, который в большинстве изложений удостаивается одного предложения, хотя держит большинство параметров плотной модели. Глава 2.5 уделяет ему то внимание, которого он заслуживает, и на этом часть II завершается.

6. Взгляд с точки зрения собеседования

"Почему batch-нормализация плохо работает с Transformer механистически?" Сильный ответ определяет, что BatchNorm вычисляет статистики по измерению батча для каждого признака, что требует чётко определённого, согласованно заполненного батча примеров в этой точке вычисления. При последовательностях переменной длины дополнение (padding) загрязняет эти статистики, а однотокенный авторегрессивный инференс во время генерации может вообще не иметь осмысленной структуры батча. Также стоит упомянуть рассогласование статистики обучения/инференса (статистики батча против скользящих средних) как вторую, независимую проблему, которую вносит BatchNorm и которой LayerNorm избегает, вычисляя статистики на уровне токена, независимо от батча.

"Объясните, с реальными уравнениями, разницу между Pre-LN и Post-LN и почему одна из них обучается более стабильно на глубине." Сильный ответ выписывает оба правила обновления — $x_{l+1} = \text{LN}(x_l + \text{Sublayer}(x_l))$ для Post-LN против $x_{l+1} = x_l + \text{Sublayer}(\text{LN}(x_l))$ для Pre-LN — и объясняет, что в Pre-LN сам остаточный путь (бегущая сумма слагаемых $x_l$) никогда не проходит через операцию нормализации, сохраняя почти-тождественный путь градиента, который призваны предоставлять остаточные соединения. В Post-LN нормализация находится прямо на этом пути после каждого сложения, что, как показывают Xiong et al., приводит к тому, что величины градиентов на ранних слоях плохо масштабируются с глубиной, требуя тщательного разогрева и ограничивая, насколько глубокой сеть может быть на практике.

"Что именно RMSNorm убирает из LayerNorm и почему это безопасно убрать?" Ответ должен точно утверждать, что RMSNorm убирает шаг вычитания среднего (перецентрирование), сохраняя только перемасштабирование по среднеквадратичной величине, и что аргументы и эксперименты Zhang и Sennrich указывают на то, что именно перемасштабирование действительно отвечает за стабилизирующий эффект на обучение, тогда как перецентрирование вносит сравнительно небольшой вклад. Мотивацию по эффективности стоит сформулировать конкретно: удаление одной операции редукции из каждого вызова нормализации, по каждому токену, на каждом слое, на каждом шаге обучения — это значимая экономия в масштабе LLM, даже если эффект на один вызов выглядит небольшим.

"Если Pre-LN стабильнее, почему её не использовали с самого начала? В чём компромисс?" Сильный ответ отмечает, что исходный Transformer использовал Post-LN и что выигрыш в стабильности Pre-LN не даётся полностью бесплатно — в литературе есть некоторые свидетельства того, что сети Pre-LN, слой за слоем, несколько менее представительно выразительны, чем сеть Post-LN того же размера, которую действительно удаётся обучить до сходимости, поскольку нормализация перед вычислением каждого подслоя несколько сильнее ограничивает функции, которые может выразить каждый подслой, чем единственная нормализация после остаточной суммы. Сильный ответ верно формулирует фактический отраслевой выбор: на глубинах и масштабах, актуальных для современных LLM, обучаемость вообще надёжно побеждает теоретическое преимущество в выразительности, которое трудно реализовать на практике.

"Внутренний ковариационный сдвиг был изначально описан для свёрточных сетей. Применимо ли это понятие к Transformer так же чисто, и важно ли это для того, как вы бы объяснили роль нормализации здесь?" Хороший ответ трактует это как нюансированный вопрос, а не как да/нет: общая интуиция (статистики активаций дрейфуют от слоя к слою и в ходе обучения, что затрудняет оптимизацию) переносится концептуально, но конкретный механизм, вокруг которого была построена BatchNorm (статистики батча для карт признаков свёрточных сетей), вовсе не переносится на последовательные данные переменной длины, что и объясняет, почему перепроектирование LayerNorm вокруг постатейной статистики на уровне токена было необходимо, а не сводилось к простому повторному использованию BatchNorm. Уместно, и, пожалуй, верно, отметить, что понимание областью того, почему именно нормализация помогает, со временем эволюционировало и не полностью устоялось, и представить это как по-настоящему открытый исследовательский вопрос, а не притворяться, что существует единая, полная теоретическая картина.

7. Вопросы для самопроверки

1. Какие две различные проблемы решает нормализация в очень глубокой остаточной сети, и почему простое наращивание остаточных блоков без нормализации усугубляет обе проблемы с ростом глубины?
2. Почему постатейная, побатчевая статистика BatchNorm плохо подходит именно для последовательных данных переменной длины и однотокенного авторегрессивного инференса?
3. Выпишите формулу LayerNorm и объясните, с точки зрения того, что она вычисляет и по какому измерению, почему она избегает конкретных проблем, названных вами в вопросе 2.
4. Запишите уравнения обновления Post-LN и Pre-LN рядом друг с другом и точно объясните, какой именно вычислительный путь между ними различается.
5. Почему сохранение остаточного пути свободным от операций нормализации (как это делает Pre-LN) помогает градиентам лучше протекать через очень глубокие стеки, связав это обратно с тем, что остаточные соединения изначально призваны обеспечивать?
6. Какой конкретный компонент LayerNorm убирает RMSNorm и в чём состоит аргумент, почему его удаление существенно не вредит качеству модели?
7. Оглядываясь на главы 2.1–2.4, перечислите конкретную проблему, которую решала центральная идея каждой главы (последовательное узкое место рекуррентности, инвариантность внимания к перестановкам, несоответствие архитектурной формы и задачи и обучаемость глубокого стека), и объясните, как необходимость каждого решения причинно следовала из проектного выбора предыдущей главы.

8. Источники

• Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450. https://arxiv.org/abs/1607.06450
• Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., & Liu, T.-Y. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020. arXiv:2002.04745. https://arxiv.org/abs/2002.04745
• Zhang, B., & Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS 2019. arXiv:1910.07467. https://arxiv.org/abs/1910.07467
• Henry, A., Dachapally, P. R., Pawar, S. S., & Chen, Y. (2020). Query-Key Normalization for Transformers. Findings of EMNLP 2020. arXiv:2010.04245. https://arxiv.org/abs/2010.04245
• Gemma Team, Riviere, M., et al. (2024). Gemma 2: Improving Open Language Models at a Practical Size (QK-norm и мягкое ограничение логитов в продакшене). arXiv:2408.00118. https://arxiv.org/abs/2408.00118