Что такое Batch Normalization и как она работает?

«Что такое Batch Normalization и как она работает?» — вопрос из категории Нейронные сети и Deep Learning, который задают на 26% собеседований Data Scientist / ML Инженер. Ниже — развёрнутый ответ с разбором ключевых моментов.

Ответ

Batch Normalization (BN) — это метод, который стабилизирует и ускоряет обучение глубоких нейронных сетей за счёт нормализации активаций каждого слоя по мини-батчу во время обучения.

Проблема, которую она решает: Internal Covariate Shift — изменение распределения входных данных для каждого последующего слоя в процессе обучения, что замедляет сходимость и требует осторожного подбора learning rate.

Алгоритм (на этапе обучения): Для каждого признака (канала/нейрона) в активациях слоя:

  1. Вычислить среднее и дисперсию по текущему мини-батчу. (muB = frac{1}{m} sum{i=1}^{m} x_i) (sigmaB^2 = frac{1}{m} sum{i=1}^{m} (x_i - mu_B)^2)
  2. Нормализовать данные: (hat{x}_i = frac{x_i - mu_B}{sqrt{sigma_B^2 + epsilon}})
  3. Масштабировать и сдвинуть с помощью обучаемых параметров (gamma) и (beta): (y_i = gamma hat{x}_i + beta). Это позволяет сети при необходимости "отключить" нормализацию.

На этапе инференса используются скользящие средние значений (mu) и (sigma), накопленные во время обучения.

Реализация в PyTorch:

import torch.nn as nn

# Для полносвязных слоев (после активации или перед? Чаще перед.)
model_fc = nn.Sequential(
    nn.Linear(784, 256),
    nn.BatchNorm1d(256),  # Нормализует по батчу для 256 признаков
    nn.ReLU(),
    nn.Linear(256, 10)
)

# Для сверточных слоев
model_conv = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=3),
    nn.BatchNorm2d(64),  # Нормализует по батчу для 64 каналов
    nn.ReLU(),
    nn.MaxPool2d(2)
)

Преимущества:

  • Позволяет использовать более высокие learning rates.
  • Снижает чувствительность к инициализации весов.
  • В некоторой степени действует как регуляризатор (за счёт шума от статистик батча).

Недостатки/нюансы:

  • Поведение ухудшается при очень маленьких размерах батча (неточная оценка статистик).
  • Увеличивает потребление памяти.
  • В таких архитектурах, как RNN или Transformer, чаще используют Layer Normalization.