BatchNormalization

BatchNormalization (BN) — это техника нормализации входов слоев нейронной сети, которая позволяет ускорить обучение, стабилизировать градиенты и улучшить общую производительность модели. В TensorFlow.js BN реализуется через слой tf.layers.batchNormalization(), который можно интегрировать в последовательные (Sequential) или функциональные (Functional) модели.

Основная идея

BatchNormalization выполняет нормализацию входного тензора по мини-батчу. Для каждого признака (x_i) вычисляются среднее и дисперсия по батчу:

[ B = {i=1}^{m} x_i, B^2 = {i=1}^{m} (x_i - _B)^2]

где (m) — размер батча. Затем нормализуются значения:

[ _i = ]

где () — маленькое значение для предотвращения деления на ноль. После этого вводятся обучаемые параметры γ (gamma) и β (beta) для масштабирования и сдвига:

[ y_i = _i + ]

Эта операция позволяет сети сохранять представительную мощность, несмотря на нормализацию.

Синтаксис в TensorFlow.js

const batchNormLayer = tf.layers.batchNormalization({
  axis: -1,        // ось нормализации (обычно -1 для последнего измерения)
  momentum: 0.99,  // коэффициент для скользящего среднего
  epsilon: 1e-3,   // предотвращение деления на ноль
  center: true,    // добавление β (сдвиг)
  scale: true      // добавление γ (масштаб)
});
  • axis — ось, по которой выполняется нормализация. Для изображений это обычно ось каналов.
  • momentum — определяет скорость обновления скользящего среднего на протяжении обучения.
  • epsilon — малое значение для числовой стабильности.
  • center и scale — включение или отключение параметров β и γ.

Пример использования в Sequential модели

const model = tf.sequential();

model.add(tf.layers.dense({ units: 128, inputShape: [64] }));
model.add(tf.layers.batchNormalization());
model.add(tf.layers.activation({ activation: 'relu' }));

model.add(tf.layers.dense({ units: 64 }));
model.add(tf.layers.batchNormalization());
model.add(tf.layers.activation({ activation: 'relu' }));

model.add(tf.layers.dense({ units: 10, activation: 'softmax' }));

В этом примере после каждого полносвязного слоя применяется BatchNormalization, что помогает стабилизировать распределение входов для последующей функции активации.

BatchNormalization в функциональном API

Функциональный API предоставляет возможность более гибкой архитектуры:

const inputs = tf.input({ shape: [64] });
let x = tf.layers.dense({ units: 128 })(inputs);
x = tf.layers.batchNormalization()(x);
x = tf.layers.activation({ activation: 'relu' })(x);

x = tf.layers.dense({ units: 64 })(x);
x = tf.layers.batchNormalization()(x);
x = tf.layers.activation({ activation: 'relu' })(x);

const outputs = tf.layers.dense({ units: 10, activation: 'softmax' })(x);

const model = tf.model({ inputs, outputs });

Использование функционального подхода удобно при построении моделей с разветвлениями, объединениями и skip-соединениями.

Особенности обучения и инференса

BatchNormalization ведёт себя по-разному на этапе обучения и при инференсе:

  • Обучение: нормализация выполняется по текущему батчу, а скользящие среднее и дисперсия обновляются для последующего инференса.
  • Инференс: используется накопленное скользящее среднее и дисперсия, а не параметры текущего батча.

TensorFlow.js автоматически управляет этим через флаг training в методах model.fit() и model.predict().

Влияние на скорость обучения и регуляризацию

  1. Ускорение сходимости: нормализация входов снижает риск «взрыва» или «затухания» градиентов, позволяя использовать более высокие значения learning rate.
  2. Смягчение зависимости от инициализации: BN уменьшает чувствительность к начальным весам.
  3. Регуляризационный эффект: небольшое шумовое влияние при использовании мини-батчей действует как лёгкая регуляризация, снижая переобучение.

Советы по использованию

  • BatchNormalization эффективен в сочетании с функциями активации, особенно ReLU и её модификациями.
  • Размещение BN перед или после активации зависит от выбранной архитектуры, но наиболее распространён вариант: слой → BN → активация.
  • Для небольших батчей (например, <16) скользящее среднее может быть нестабильным; в этом случае иногда используют LayerNormalization.

Производительность в TensorFlow.js

BatchNormalization в TensorFlow.js оптимизирован для выполнения на GPU через WebGL/WebGPU. Однако при больших моделях и больших батчах может наблюдаться значительное потребление памяти, поэтому важно контролировать размер батча и архитектуру слоёв.