Как бороться с градиентным взрывом?

Градиентный взрыв: проблема и решения

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

1. Что такое градиентный взрыв

При обратном распространении ошибки (backpropagation) через много слоёв, градиенты умножаются на частные производные каждого слоя:

Входной слой → [W1 * σ'1] → [W2 * σ'2] → [W3 * σ'3] → ... → Выходной слой

Полный градиент = δL/δW = δL/δOut * dOut/dW3 * dW3/dW2 * dW2/dW1 * dW1/dInput

Если каждое умножение > 1, градиент экспоненциально растёт:

1 → 1.1 → 1.21 → 1.33 → 1.46 → ... → бесконечность (NaN)

2. Визуализация проблемы

import numpy as np
import matplotlib.pyplot as plt

# Симуляция градиентов через 100 слоёв
def simulate_gradient_flow(weight_init, num_layers=100):
    gradient = 1.0
    gradients = [gradient]
    
    for _ in range(num_layers):
        derivative = np.random.uniform(0.5, 2.0)  # Случайная производная
        gradient *= weight_init * derivative
        gradients.append(gradient)
        
        if gradient > 1e10:  # Взрыв
            break
    
    return gradients

# С большим начальным весом
grad_explode = simulate_gradient_flow(weight_init=1.5, num_layers=50)
print(f"Взрыв! Градиент вырос от {grad_explode[0]:.2e} до {grad_explode[-1]:.2e}")
# Взрыв! Градиент вырос от 1.00e+00 до 3.45e+09

3. Признаки градиентного взрыва

• Loss = NaN — веса стали бесконечными
• Loss резко скачет вверх — нестабильное обучение
• Веса модели = inf или -inf
• Очень большие значения градиентов

import tensorflow as tf
import numpy as np

# Обнаружить взрыв
model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(128, activation='relu'),
    # ...
])

# Проверить градиенты
with tf.GradientTape() as tape:
    y_pred = model(X_train)
    loss = tf.keras.losses.mse(y_train, y_pred)

gradients = tape.gradient(loss, model.trainable_weights)

for i, grad in enumerate(gradients):
    if tf.reduce_max(tf.abs(grad)) > 1e6:
        print(f"ВНИМАНИЕ! Слой {i}: max gradient = {tf.reduce_max(tf.abs(grad)):.2e}")

4. Решения

1. Gradient Clipping (Обрезание градиентов)

Ограничить максимальное значение градиента:

import tensorflow as tf

optimizer = tf.keras.optimizers.Adam(
    learning_rate=0.001,
    clipvalue=1.0  # Обрезать градиенты > 1.0
)

model.compile(optimizer=optimizer, loss='mse')
model.fit(X_train, y_train, epochs=50)

Обрезание по норме (более эффективно):

optimizer = tf.keras.optimizers.Adam(
    learning_rate=0.001,
    clipnorm=1.0  # Норма всех градиентов <= 1.0
)

# Вручную
with tf.GradientTape() as tape:
    loss = compute_loss(model, X, y)

gradients = tape.gradient(loss, model.trainable_weights)

# Обрезать по норме
clipped_gradients, _ = tf.clip_by_global_norm(gradients, clip_norm=1.0)
optimizer.apply_gradients(zip(clipped_gradients, model.trainable_weights))

2. Уменьшить Learning Rate

Меньший шаг оптимизации → градиенты не растут так быстро:

# Вместо 0.1
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)  # 100x меньше

3. Batch Normalization

Нормализует активации каждого слоя, стабилизирует градиенты:

model = tf.keras.Sequential([
    tf.keras.layers.Dense(128),
    tf.keras.layers.BatchNormalization(),  # Нормализация
    tf.keras.layers.ReLU(),
    
    tf.keras.layers.Dense(128),
    tf.keras.layers.BatchNormalization(),  # Каждый слой
    tf.keras.layers.ReLU(),
    
    tf.keras.layers.Dense(10, activation='softmax')
])

Почему помогает:

• Нормализует выходы слоёв к mean=0, std=1
• Градиенты остаются в нормальном диапазоне
• Поддерживает более высокий learning rate

4. Layer Normalization

Альтернатива BatchNorm, особенно для RNN:

model = tf.keras.Sequential([
    tf.keras.layers.Dense(128),
    tf.keras.layers.LayerNormalization(),  # Вместо BatchNorm
    tf.keras.layers.ReLU(),
])

5. Инициализация весов (Xavier / He initialization)

Правильная инициализация уменьшает риск взрыва:

model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, 
                         kernel_initializer='glorot_uniform',  # Xavier
                         activation='tanh'),
    
    tf.keras.layers.Dense(128,
                         kernel_initializer='he_normal',  # He initialization
                         activation='relu'),  # He лучше для ReLU
])

Почему:

• Xavier: инициализирует веса в диапазоне sqrt(6 / (n_in + n_out))
• He: инициализирует для ReLU сетей
• Случайная инициализация может привести к взрыву

6. Residual Connections (Skip Connections)

Позволяют градиентам обходить слои без умножения:

def residual_block(x):
    y = tf.keras.layers.Dense(128, activation='relu')(x)
    y = tf.keras.layers.Dense(128)(y)
    return x + y  # Skip connection!

inputs = tf.keras.Input(shape=(128,))
x = residual_block(inputs)
model = tf.keras.Model(inputs=inputs, outputs=x)

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

• Градиент может потечь напрямую через identity
• Глубокие сети (100+ слоёв) становятся обучаемы

7. Использовать стабильные активации

Избегать функции, которые усиливают градиенты:

# ПЛОХО: sigmoid может привести к взрыву
model.add(tf.keras.layers.Dense(128, activation='sigmoid'))

# ХОРОШО: ReLU стабильнее
model.add(tf.keras.layers.Dense(128, activation='relu'))

# ХОРОШО: GELU (современный вариант)
model.add(tf.keras.layers.Dense(128, activation='gelu'))

5. Практический пример: RNN (часто страдает от взрыва)

import tensorflow as tf

# RNN без защиты
model_bad = tf.keras.Sequential([
    tf.keras.layers.LSTM(128, return_sequences=True),
    tf.keras.layers.LSTM(128),
    tf.keras.layers.Dense(10, activation='softmax')
])

model_bad.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),  # Большой LR
    loss='sparse_categorical_crossentropy'
)
model_bad.fit(X_train, y_train, epochs=10)  # Может упасть!

# RNN со защитой
model_good = tf.keras.Sequential([
    tf.keras.layers.LSTM(128, return_sequences=True),
    tf.keras.layers.LSTM(128),
    tf.keras.layers.Dense(10, activation='softmax')
])

model_good.compile(
    optimizer=tf.keras.optimizers.Adam(
        learning_rate=0.001,
        clipnorm=1.0  # Обрезание градиентов
    ),
    loss='sparse_categorical_crossentropy'
)
model_good.fit(X_train, y_train, epochs=10)  # Будет работать!

6. Мониторинг градиентов

import tensorflow as tf
from tensorboard.plugins import projector

class GradientMonitor(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        # Вычислить градиенты
        with tf.GradientTape() as tape:
            loss = self.model(X_val, training=True)
        
        grads = tape.gradient(loss, self.model.trainable_weights)
        
        # Проверить на взрыв
        for i, grad in enumerate(grads):
            max_grad = tf.reduce_max(tf.abs(grad))
            if max_grad > 100:
                print(f"ВНИМАНИЕ! Слой {i}: gradient magnitude = {max_grad:.2e}")

model.fit(X_train, y_train, callbacks=[GradientMonitor()])

Выводы

Основные методы борьбы:

1. Gradient Clipping — самый простой, срабатывает во всех случаях
2. Batch Normalization — стандарт для современных сетей
3. Learning Rate scheduling — постепенное снижение LR
4. Residual connections — для очень глубоких сетей
5. Правильная инициализация — Xavier/He инициализация
6. Мониторинг — отслеживать max gradient

Лучшая практика:

optimizer = tf.keras.optimizers.Adam(learning_rate=0.001, clipnorm=1.0)
model.add(tf.keras.layers.BatchNormalization())
model.add(tf.keras.layers.Dense(128, activation='relu', kernel_initializer='he_normal'))

Это комбинация обрезания + нормализации + правильной инициализации обеспечивает стабильное обучение даже в глубоких сетях.