Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift

Разбор статьи

Проблема и мотивация

Во время оптимизации глубоких сетей с помощью SGD изменение параметров ранних слоев меняет распределение входов для последующих слоев. Авторы называют это явление Internal Covariate Shift (внутренний ковариантный сдвиг).

Внутренний ковариантный сдвиг порождает две основные проблемы:

1. Подсети вынуждены непрерывно адаптироваться к постоянно меняющимся распределениям активаций.
2. Активации легко смещаются в области насыщения нелинейностей (например, сигмоид), что приводит к затуханию градиентов и требует низких скоростей обучения и тщательной инициализации параметров.

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

Ключевые идеи

• Дифференцируемый слой нормализации: Вместо стандартизации активаций вне оптимизационного процесса нормализация встраивается прямо в архитектуру модели, что позволяет backpropagation учитывать зависимость между активациями слоя и статистиками мини-батча.
• Покомпонентная нормализация: Вместо совместного многомерного обеления каждый скалярный признак нормализуется независимо до нулевого среднего и единичной дисперсии.
• Сохранение репрезентативной емкости: Чтобы нормализация не ограничивала выразительную способность слоя (например, не переводила сигмоиду исключительно в линейный режим работы), вводятся обучаемые параметры $\gamma$ (масштаб) и $\beta$ (сдвиг), позволяющие слою при необходимости восстановить тождественное преобразование.
• Детерминированный инференс: Во время обучения статистики рассчитываются по каждому мини-батчу $\mathcal{B}$. Во время инференса статистики мини-батча заменяются фиксированными генеральными статистиками (вычисленными через скользящее среднее по обучающим батчам), что гарантирует детерминированный вывод для одиночных примеров.
• Инвариантность к масштабу параметров: Масштабирование весов слоя удовлетворяет свойству $\text{BN}((aW)\mathrm{u})=\text{BN}(W\mathrm{u})$, что означает: большие веса приводят к меньшим эффективным градиентам, предотвращая взрыв градиентов.

Метод (с ключевыми формулами)

Для мини-батча $\mathcal{B}=\{x_{1\dots m}\}$, соответствующего одной размерности активации:

1. Вычисление среднего по мини-батчу: $${\mu }_{\mathcal{B}}\leftarrow \frac{1}{m}\sum _{i=1}^m{x}_i$$
2. Вычисление дисперсии по мини-батчу: $${\sigma }_{\mathcal{B}}^2\leftarrow \frac{1}{m}\sum _{i=1}^m(x_i-{\mu }_{\mathcal{B}}{)}^2$$
3. Нормализация активаций (с константой стабильности $ϵ$): $${\hat{x}}_i\leftarrow \frac{x_i-{\mu }_{\mathcal{B}}}{\sqrt{{\sigma }_{\mathcal{B}}^2+ϵ}}$$
4. Масштабирование и сдвиг с помощью обучаемых параметров $\gamma$ и $\beta$: $$y_i\leftarrow \gamma {\hat{x}}_i+\beta \equiv {\text{BN}}_{\gamma ,\beta }(x_i)$$

Применение к слоям

• Расположение: Применяется непосредственно перед нелинейностями: $z=g(\text{BN}(Wu))$. Исходный свободный член (bias) $b$ исключается, поскольку вычитание среднего его нивелирует (его роль берет на себя $\beta$).
• Сверточные слои: Чтобы сохранить трансляционную эквивариантность, нормализация применяется к картам признаков целиком, а не к отдельным пространственным точкам. Для размера батча $m$ и пространственного разрешения $p\times q$ эффективный размер батча равен $m^{\prime }=m\cdot p\cdot q$, а одна пара параметров $(\gamma ,\beta )$ разделяется между всеми элементами карты признаков.

Преобразование на инференсе

Во время тестирования преобразование сводится к статичной линейной операции с использованием генеральных статистик $\mathrm{E}[x]$ и $\mathrm{Var}[x]$: $$y=\frac{\gamma }{\sqrt{\mathrm{Var}[x]+ϵ}}\cdot x+\left(\beta -\frac{\gamma \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+ϵ}}\right)$$

Результаты (с важнейшими показателями)

• MNIST с сигмоидами: Сети с Batch Normalization сохраняли стабильные распределения активаций на протяжении всего обучения, ускоряя сходимость и достигая более высокой точности на тесте по сравнению с ненормализованными сетями.
• Скорость обучения на ImageNet (одна сеть):
  • Базовая сеть Inception без нормализации достигла точности валидации 72,2% за $31,0\times 10^6$ шагов.
  • Baseline Inception с Batch Normalization достиг 72,2% за $13,3\times 10^6$ шагов.
  • С увеличенной в 5 раз скоростью обучения (BN-x5) модель достигла 72,2% за $2,1\times 10^6$ шагов — в 14 раз быстрее, показав итоговый максимум 73,0%.
  • При увеличении скорости обучения в 30 раз (BN-x30) сеть достигла максимальной точности 74,8% за $6,0\times 10^6$ шагов.
• Насыщающиеся нелинейности: В то время как исходная Inception с сигмоидными активациями не смогла обучиться (оставаясь на уровне случайного угадывания ~0,1%), BN-x5-Sigmoid показала точность top-1 в 69,8%.
• SOTA-ансамбль на ImageNet: Ансамбль из 6 сетей с батч-нормализацией показал ошибку top-5 на валидации 4,9% и 4,82% на тестовом сервере ILSVRC, превзойдя уровень человека.

Ограничения и открытые вопросы

• Зависимость от размера батча: Поскольку нормализация опирается на выборочные статистики мини-батча, слишком маленькие размеры батчей могут ухудшать оценку среднего и дисперсии.
• Рекуррентные сети: В работе рассматриваются полносвязные и сверточные архитектуры, а рекуррентные нейронные сети (RNN) упомянуты как открытая область, где ковариантный сдвиг также может быть критичен.
• Теоретическое обоснование: Авторы предполагают, что Batch Normalization удерживает сингулярные числа якобианов слоев близкими к 1, однако строгое теоретическое доказательство поведения градиентов оставлено для будущих исследований.