LayerNorm

Идея: нормализовать каждый объект

BatchNorm считает статистики «вертикально» — по объектам батча для каждого признака. LayerNorm (Ба, Кирос, Хинтон, 2016) делает то же самое «горизонтально» — по признакам каждого объекта:

Обучаемые \gamma и \beta — по одному на признак, как у BatchNorm.

import numpy as np
import nanotorch as nt

x = nt.tensor([[1.0, 2.0, 3.0, 4.0],
               [10.0, 20.0, 30.0, 40.0],
               [-5.0, 0.0, 5.0, 100.0]])
ln = nt.LayerNorm(4)
y = ln(x)
print(y.data.round(3))
print("среднее по строкам:", y.data.mean(axis=1).round(6), " std:", y.data.std(axis=1).round(3))

Первые две строки после нормализации совпали: LayerNorm не видит масштаб объекта, только «форму» вектора признаков.

Почему это удобно

• Не зависит от батча. Каждый объект нормализуется сам по себе — работает и при батче из одного объекта.
• Одинаково при обучении и применении. Нет скользящих средних, нет режимов, нет связанных с ними ошибок.
• Подходит для последовательностей. В тексте каждый токен нормализуется отдельно, длина текста и паддинг не мешают.

Поэтому LayerNorm стоит в каждом блоке трансформера — от BERT до GPT.

Эксперимент: крошечный батч

Обучим одну и ту же сеть с BatchNorm и с LayerNorm, с обычным батчем (32) и с крошечным (2). Число шагов обучения одинаковое — 600:

import numpy as np
import nanotorch as nt
import ml_data

X, y = ml_data.load_digits()
X = X / 16.0
X_train, X_test, y_train, y_test = nt.train_test_split(X, y, test_size=0.25, seed=0)
X_train, y_train = X_train[:600], y_train[:600]

for name, norm in [("BatchNorm", nt.BatchNorm1d), ("LayerNorm", nt.LayerNorm)]:
    for bs, epochs in ((32, 32), (2, 2)):          # 600 / 32 ≈ 19 шагов на эпоху, 600 / 2 = 300
        nt.manual_seed(0)
        model = nt.Sequential(nt.Linear(64, 64), norm(64), nt.ReLU(), nt.Linear(64, 64), norm(64), nt.ReLU(), nt.Linear(64, 10))
        h = nt.fit(model, X_train, y_train, nt.cross_entropy, nt.SGD(model.parameters(), lr=0.05),
                   epochs=epochs, batch_size=bs, X_val=X_test, y_val=y_test, metric=nt.accuracy, verbose=False)
        print(f"{name}, батч {bs:2}: точность на тесте {h['val_metric'][-1]:.1%}")

С батчем из 2 объектов BatchNorm почти бесполезен: после нормализации по двум объектам каждый признак принимает значения ±1 — информация о величине теряется, а накопленные скользящие средние плохо соответствуют тому, что сеть видела при обучении. LayerNorm тоже страдает от шумных градиентов маленького батча, но сама нормализация от размера батча не зависит.

BatchNorm против LayerNorm

Есть и промежуточные варианты: GroupNorm (статистики по группам каналов одного объекта) и InstanceNorm (по каждому каналу картинки отдельно) — для свёрточных сетей с маленькими батчами и для переноса стиля.

RMSNorm

RMSNorm (2019) — упрощение LayerNorm: не вычитаем среднее и не прибавляем \beta, только делим на «средний квадрат» (root mean square):

Оказалось, что центрирование почти ничего не даёт, а без него вычисления дешевле. RMSNorm используется в LLaMA, Mistral, Gemma и многих других современных языковых моделях.

Где стоит нормализация в трансформере

В блоке трансформера есть две части (внимание и полносвязная сеть) и остаточные связи x + f(x). Нормализацию можно поставить по-разному:

• Post-Norm (исходная статья 2017 года): x \leftarrow \text{LN}(x + f(x));
• Pre-Norm (GPT-2 и почти все современные модели): x \leftarrow x + f(\text{LN}(x)).

Pre-Norm устойчивее при обучении глубоких моделей: «магистраль» x проходит через все блоки без нормализации, и градиент по ней течёт беспрепятственно. Мы реализуем именно Pre-Norm, когда будем собирать энкодер трансформера.

import inspect
import nanotorch as nt

print(inspect.getsource(nt.LayerNorm))

Итог

• LayerNorm нормализует признаки каждого объекта: \gamma\frac{x - \mu}{\sqrt{\sigma^2 + \varepsilon}} + \beta, статистики по последней оси.
• Не зависит от батча, одинаково работает при обучении и применении, естественен для последовательностей.
• Стандарт в трансформерах; BatchNorm — в свёрточных сетях.
• RMSNorm — LayerNorm без центрирования и \beta, используется в современных языковых моделях.
• Pre-Norm (x + f(\text{LN}(x))) устойчивее Post-Norm.