BatchNorm

Проблема

Каждый слой сети учится на выходах предыдущего. Но предыдущий слой тоже учится — и распределение его выходов постоянно меняется: сдвигается среднее, растёт или падает разброс. Следующему слою приходится подстраиваться под «движущуюся цель». Кроме того, если активации какого-то слоя оказались слишком большими или маленькими, мы получаем знакомые проблемы — насыщение, затухание, взрыв.

Нормализация входов модели (стандартизация признаков) помогает первому слою. Идея BatchNorm (Иоффе и Сегеди, 2015): давайте нормализовать входы каждого слоя — прямо внутри сети, на каждом шаге обучения.

Алгоритм

Для батча из N объектов и каждого признака (нейрона) отдельно:

1. Считаем среднее и дисперсию признака по объектам батча.
2. Нормализуем: теперь у признака среднее 0 и дисперсия 1 (\varepsilon \approx 10^{-5} защищает от деления на ноль).
3. Умножаем на обучаемый масштаб \gamma и прибавляем обучаемый сдвиг \beta.

Третий шаг важен: жёстко зафиксировать среднее 0 и дисперсию 1 — лишнее ограничение (например, для сигмоиды это означало бы работу только в её почти линейной середине). С \gamma и \beta сеть сама выбирает удобный масштаб — но теперь он задаётся двумя явными параметрами, а не случайно складывается из всех весов предыдущих слоёв.

import numpy as np
import nanotorch as nt

rng = np.random.default_rng(0)
x = nt.tensor(rng.normal([10, -3, 0.5], [5, 0.1, 2], size=(256, 3)))
print("до:    среднее", x.data.mean(0).round(2), " std", x.data.std(0).round(2))

bn = nt.BatchNorm1d(3)
y = bn(x)
print("после: среднее", y.data.mean(0).round(2), " std", y.data.std(0).round(2))
print("обучаемые параметры:", [name for name, _ in bn.named_parameters()], " всего:", bn.num_params())

Обучение и применение

Во время обучения статистики считаются по текущему батчу. Но при применении модели:

• объект может быть один — статистику по батчу не посчитать;
• ответ модели для объекта не должен зависеть от того, с какими соседями он попал в батч.

Поэтому во время обучения BatchNorm копит скользящие средние \mu и \sigma^2:

а в режиме применения использует их. Переключение — model.train() и model.eval().

import numpy as np
import nanotorch as nt

rng = np.random.default_rng(0)
bn = nt.BatchNorm1d(1)
for step in range(100):
    bn(nt.tensor(rng.normal(7, 2, size=(32, 1))))        # обучение: копим статистики
print("накопленные: mean =", bn.running_mean.round(2), " var =", bn.running_var.round(2))

bn.eval()
print("применение к одному объекту 9.0 →", bn(nt.tensor([[9.0]])).data.round(3), "(≈ (9 − 7) / 2)")
bn.train()

Где ставить

Обычно BatchNorm ставят после линейного (или свёрточного) слоя и перед активацией:

Сдвиг b у линейного слоя перед BatchNorm бесполезен (его всё равно вычтет нормализация, а роль сдвига играет \beta), поэтому его часто отключают: Linear(..., bias=False).

Что даёт BatchNorm

• Можно брать больший шаг обучения: масштаб активаций больше не зависит от масштаба весов, и большие шаги не приводят к взрыву.
• Меньше зависимость от инициализации.
• Быстрее сходимость — часто в разы меньше эпох.
• Небольшая регуляризация: статистики батча немного шумят, и это мешает переобучению.

Почему именно BatchNorm так хорошо работает — до сих пор предмет исследований. Исходное объяснение («уменьшение внутреннего сдвига распределений») оказалось неполным; сейчас считается, что главное — он делает поверхность потерь более гладкой, и градиентный спуск ведёт себя предсказуемее.

import numpy as np
import nanotorch as nt
import ml_data
import vis

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)

def make(bn):
    layers, d = [], 64
    for _ in range(5):
        layers.append(nt.Linear(d, 64))
        if bn:
            layers.append(nt.BatchNorm1d(64))
        layers.append(nt.ReLU())
        d = 64
    return nt.Sequential(*layers, nt.Linear(64, 10))

curves = {}
for bn in (False, True):
    for lr in (0.05, 0.1):
        nt.manual_seed(0)
        model = make(bn)
        h = nt.fit(model, X_train, y_train, nt.cross_entropy, nt.SGD(model.parameters(), lr=lr),
                   epochs=12, batch_size=64, X_val=X_test, y_val=y_test, metric=nt.accuracy, verbose=False)
        name = f"{'с BatchNorm' if bn else 'без BN'}, lr = {lr}"
        curves[name] = [v if np.isfinite(v) else np.nan for v in h["val_metric"]]
        print(f"{name:22}: точность на тесте {h['val_metric'][-1]:.1%}")
vis.line(curves, title="Точность на тесте по эпохам (5 скрытых слоёв)", xlabel="эпоха")

С BatchNorm та же сеть с тем же шагом обучается быстрее и ровнее. Но чудес не бывает: при слишком большом шаге (попробуйте 0.3) разваливается и она — веса меняются так быстро, что накопленные скользящие средние не успевают за ними.

Недостатки

• Зависимость от батча. При маленьком батче (2–8 объектов) статистики очень шумные, и качество падает. Это проблема для больших моделей, где в память помещается лишь несколько объектов.
• Разное поведение при обучении и применении — источник ошибок.
• Неудобен для последовательностей (тексты разной длины) и рекуррентных сетей.

Поэтому в трансформерах используют другую нормализацию — LayerNorm, которая считает статистики по признакам одного объекта. О ней — следующий урок.

Итог

• BatchNorm нормализует каждый признак по объектам батча: \hat x = (x - \mu)/\sqrt{\sigma^2 + \varepsilon}, затем \gamma\hat x + \beta.
• При обучении — статистики батча и накопление скользящих средних; при применении (eval()) — накопленные средние.
• Ставят между линейным слоем и активацией; позволяет брать больший шаг, ускоряет и стабилизирует обучение.
• Недостатки: зависимость от размера батча, разные режимы работы, неудобство для последовательностей.