Что такое warm-up при обучении?

Warm-up при обучении нейронных сетей

Warm-up (прогрев) — это техника, при которой в начале обучения используется малый learning rate, который затем постепенно увеличивается до целевого значения. Это улучшает стабильность обучения, особенно в больших моделях и с большим batch size.

Проблема без warm-up

import torch
import torch.nn as nn

# Без warm-up: большой learning rate с самого начала
model = nn.Linear(10, 5)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)  # большой lr

for epoch in range(100):
    output = model(x)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()  # большие скачки в начале -> нестабильность
    optimizer.zero_grad()

Что происходит без warm-up:

1. Веса инициализированы случайно
2. Градиенты в начале очень шумные (высокая дисперсия)
3. Большой learning rate + шумные градиенты = большие скачки в весах
4. Модель может diverge (расходиться) или попасть в плохой local minima
5. Обучение становится нестабильным, loss может спайковаться

Как работает warm-up

import torch
import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR

# Стратегия: линейный warm-up от 0 до целевого lr за K шагов
def linear_warmup_schedule(step, warmup_steps, total_steps):
    """
    step: текущий номер шага (0, 1, 2, ...)
    warmup_steps: количество шагов для прогрева (например, 1000)
    total_steps: общее количество шагов обучения
    """
    if step < warmup_steps:
        # Linear increase from 0 to 1
        return float(step) / float(max(1, warmup_steps))
    else:
        # После warm-up: постепенное уменьшение
        return max(0.0, float(total_steps - step) / float(max(1, total_steps - warmup_steps)))

model = nn.Linear(10, 5)
base_lr = 1e-3
optimizer = optim.Adam(model.parameters(), lr=base_lr)

warmup_steps = 1000
total_steps = 10000

# Используем LambdaLR с функцией расписания
scheduler = LambdaLR(
    optimizer,
    lr_lambda=lambda step: linear_warmup_schedule(step, warmup_steps, total_steps)
)

# Обучение
for step in range(total_steps):
    output = model(x)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    scheduler.step()  # обновляем learning rate
    
    if step % 100 == 0:
        current_lr = optimizer.param_groups[0]['lr']
        print(f"Step {step}: LR = {current_lr:.6f}, Loss = {loss:.4f}")

Различные стратегии warm-up

1. Линейный warm-up

def linear_warmup(step, warmup_steps):
    if step < warmup_steps:
        return step / warmup_steps
    return 1.0

# LR = base_lr * (step / warmup_steps) для step < warmup_steps
# Плюсы: просто, популярно
# Минусы: скорость увеличения постоянная

2. Экспоненциальный warm-up

import math

def exponential_warmup(step, warmup_steps):
    if step < warmup_steps:
        return math.exp(step / warmup_steps - 1)
    return 1.0

# LR медленно растёт сначала, быстро в конце
# Плюсы: более плавное начало
# Минусы: параметр a нужно подобрать

3. Cosine annealing с warm-up

def cosine_warmup_schedule(step, warmup_steps, total_steps):
    if step < warmup_steps:
        return step / warmup_steps
    else:
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * progress))

# 1. Линейный рост на warmup_steps
# 2. Затем плавное косинусное уменьшение
# Популярно в трансформер-моделях

4. Cyclic Learning Rate с warm-up

def cyclic_warmup_schedule(step, warmup_steps, cycle_steps, num_cycles):
    if step < warmup_steps:
        return step / warmup_steps
    else:
        cycle_position = ((step - warmup_steps) % cycle_steps) / cycle_steps
        return 0.5 * (1 + math.cos(math.pi * cycle_position))

# Несколько циклов увеличения/уменьшения LR
# Помогает выходить из local minima

Практический пример: обучение трансформера

import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import get_linear_schedule_with_warmup
from transformers import AutoModel

model = AutoModel.from_pretrained('bert-base-uncased')

# Гиперпараметры
num_epochs = 3
train_dataloader_len = 1000  # количество батчей
warmup_ratio = 0.1  # 10% от всех шагов

total_steps = num_epochs * train_dataloader_len
warmup_steps = int(total_steps * warmup_ratio)

optimizer = AdamW(model.parameters(), lr=2e-5)
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps
)

# Обучение
for epoch in range(num_epochs):
    for step, batch in enumerate(train_dataloader):
        outputs = model(**batch)
        loss = outputs.loss
        
        loss.backward()
        optimizer.step()
        scheduler.step()  # Update LR
        optimizer.zero_grad()
        
        if step % 100 == 0:
            current_lr = scheduler.get_last_lr()[0]
            print(f"Epoch {epoch}, Step {step}: LR = {current_lr:.6f}")

Параметры warm-up

Количество шагов warm-up

# Рекомендации:
# - Для BERT/GPT: 10% от total_steps
# - Для ResNet: 5% от total_steps
# - Для маленьких данных: 1000-5000 шагов
# - Для больших данных: зависит от размера батча

warmup_steps = int(0.1 * total_steps)  # 10%

Начальный learning rate

# Обычно используют 0 или очень маленький lr (1e-7)
# и увеличивают до целевого (1e-3, 2e-5 для BERT)

# Если начинаем с 0:
warmup_factor = 1e-7  # tiny learning rate

# Если начинаем с целевого:
warmup_factor = 1.0  # уже целевой, потом уменьшаем

Когда использовать warm-up

# ИСПОЛЬЗУЙ warm-up:
# 1. Большой batch size (>256)
batch_size = 512  # -> используй warm-up

# 2. Глубокие сети
model = DeepResNet(depth=200)  # -> используй warm-up

# 3. Transfer learning
model = AutoModel.from_pretrained('bert-base')  # -> используй warm-up

# 4. Низкий learning rate
lr = 2e-5  # -> часто нужен warm-up для стабильности

# НЕ ИСПОЛЬЗУЙ warm-up:
# 1. Малый batch size (<32)
# 2. Простые модели
# 3. Когда уже есть другой scheduler (например, ReduceLROnPlateau)

Визуализация эффекта warm-up

import matplotlib.pyplot as plt

# График learning rate с warm-up
warmup_steps = 1000
total_steps = 10000
lrs = []

for step in range(total_steps):
    if step < warmup_steps:
        lr = (step / warmup_steps) * 1e-3
    else:
        lr = 1e-3  # constant after warmup
    lrs.append(lr)

plt.figure(figsize=(12, 4))
plt.plot(lrs)
plt.xlabel('Training Step')
plt.ylabel('Learning Rate')
plt.title('Linear Warm-up Schedule')
plt.axvline(x=warmup_steps, color='r', linestyle='--', label='Warmup End')
plt.legend()
plt.show()

# График loss с и без warm-up
plt.figure(figsize=(12, 4))
plt.plot(loss_with_warmup, label='With Warm-up')
plt.plot(loss_without_warmup, label='Without Warm-up')
plt.xlabel('Training Step')
plt.ylabel('Loss')
plt.legend()
plt.title('Effect of Warm-up on Training Loss')
plt.show()

Выводы

1. Warm-up улучшает стабильность обучения больших моделей
2. Рекомендуется использовать с большим batch size или низким learning rate
3. Линейный warm-up — самый популярный и простой подход
4. 10% от total_steps — хороший стартовый параметр
5. Комбинируй с другими schedulers (cosine annealing, step decay)

Warm-up — это простая но мощная техника, которая должна быть в арсенале каждого практика глубокого обучения.