Weight normalization: перепараметризация вместо batch norm

Batch normalization ускоряет обучение сетей, но у неё есть цена: результат каждого примера зависит от того, какие ещё примеры попали в тот же батч. При батче размером 1, в рекуррентных сетях и в reinforcement learning это превращается в проблему — статистики шумные или считать их просто не по чему. Weight normalization предлагает другой подход: не нормировать активации, а перепараметризовать сами веса. Метод описали Тим Салиманс и Дидерик Кингма (тот самый автор оптимизатора Adam) в работе 2016 года, и он до сих пор живёт в стандартных слоях PyTorch и TensorFlow.

Что именно перепараметризуется

Обычный нейрон вычисляет y = φ(w · x + b), где w — вектор весов, x — вход, φ — нелинейность. Weight normalization разбивает вектор весов на две независимые части: его длину и его направление.

Формально вес представляется так:

w = g * (v / ||v||)

Здесь v — обучаемый вектор той же размерности, что и w, ||v|| — его евклидова норма, а g — обучаемый скаляр, отвечающий за длину. Направление вектора w теперь задаётся отношением v / ||v||, а его масштаб — отдельным параметром g. Градиенты по g и v считаются напрямую через это выражение, отдельного слоя с усреднением по батчу не появляется.

Ключевая идея: отвязать длину вектора весов от его направления. Оптимизатор перестаёт тратить шаги на то, чтобы одновременно крутить и растягивать один и тот же вектор.

Почему это ускоряет сходимость

В обычной параметризации норма весов и их направление жёстко сцеплены в одном векторе. Из-за этого кривизна функции потерь по разным направлениям сильно различается, и градиентному спуску тяжело подобрать общий шаг. Разделив длину и направление, weight normalization делает поверхность оптимизации ближе к сферически симметричной по масштабу — эффект, который в статье называют улучшением обусловленности задачи. На практике это означает, что можно брать больший learning rate без расхождения.

Инициализация, зависящая от данных

Одна перепараметризация масштаб активаций сама по себе не фиксирует. Поэтому авторы добавляют трюк с инициализацией: перед первым проходом обучения через сеть прогоняют один мини-батч и по нему подбирают начальные значения g и смещений b так, чтобы выход каждого слоя имел нулевое среднее и единичную дисперсию.

Дальше в обучении никакой статистики по батчу не считается — она нужна только один раз, на старте. Именно поэтому метод и работает там, где batch norm буксует.

Сравнение с другими подходами

Batch norm вносит в обучение полезный шум — активации каждого примера чуть колеблются в зависимости от соседей по батчу, и это работает как регуляризация. Weight normalization такого эффекта почти не даёт. Авторы предлагают компенсацию — mean-only batch normalization: вычитать по батчу только среднее, но не делить на дисперсию. Это возвращает часть регуляризующего шума, оставаясь дешевле полной batch norm.

Где метод особенно уместен

• Рекуррентные сети (LSTM, GRU) — batch norm по временным шагам применять неудобно, weight norm вписывается естественно.
• Reinforcement learning — батчи маленькие и нестационарные, батч-статистики шумят.
• Генеративные модели и autoencoder-ы — в исходной работе метод тестировали в том числе на них.
• Обучение с батчем 1 — batch norm здесь вырождается, weight norm работает без изменений.

Как включить на практике

В PyTorch weight normalization навешивается на существующий слой как обёртка параметра, менять архитектуру не нужно:

import torch.nn as nn
from torch.nn.utils.parametrizations import weight_norm

linear = weight_norm(nn.Linear(256, 256), name="weight", dim=0)
# теперь у слоя два параметра: weight_g (длина) и weight_v (направление)

Старый вызов torch.nn.utils.weight_norm помечен как устаревший в пользу модуля parametrizations — при обновлении версий фреймворка проверяйте, какой API актуален в вашей сборке. Параметр dim задаёт, по какой оси считать норму: для полносвязного слоя обычно dim=0 (по каждому выходному нейрону отдельно).

Что учесть при выборе

1. Если у вас крупные батчи и обычная свёрточная классификация — batch norm или её современные варианты часто дадут точность не хуже, а иногда лучше за счёт регуляризации.
2. Если батч мал, нестабилен или его нет как понятия (онлайн-обучение) — weight norm снимает основную боль.
3. Не забудьте про инициализацию по данным: без неё часть выигрыша в скорости сходимости теряется.
4. Сравнивайте на своей задаче. Универсального победителя среди схем нормализации нет — результат зависит от архитектуры и данных.

Weight normalization — не замена всему подряд, а инструмент для конкретных ситуаций, где привязка к батчу мешает. Он дешёвый по вычислениям, детерминированный между train и inference и встроен в основные фреймворки. Если ваша модель плохо обучается из-за маленьких или нестационарных батчей, это первое, что стоит попробовать перед тем, как переходить к тяжёлым схемам.