Что нужно сделать с learning rate при увеличении размера batch?

Learning Rate vs Batch Size: обратная зависимость

Это отличный вопрос об оптимизации нейросетей. Короткий ответ: learning rate нужно увеличивать при увеличении batch size. Давайте разберёмся почему и как.

Почему это связано?

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

Gradient (градиент) при большом batch размере становится более stable и reliable. Поэтому мы можем делать bigger steps (больший learning rate).

Математически

Update rule: w_new = w - lr * gradient

Что происходит:
- Большой batch → gradient более accurate (усреднение по большему числу примеров)
- Accurate gradient → можем делать bigger steps без риска diverge
- Маленький batch → noisy gradient (может указать в неправильном направлении)

Эмпирическое правило: Linear Scaling

Если увеличить batch size в k раз, увеличь learning rate в k раз:

# Baseline
batch_size_1 = 32
lr_1 = 0.001

# Если увеличиваем batch в 4 раза
batch_size_2 = 128  # 4x bigger
lr_2 = 0.004        # 4x bigger (0.001 * 4)

# Если увеличиваем batch в 8 раз
batch_size_3 = 256  # 8x bigger
lr_3 = 0.008        # 8x bigger

Это работает хорошо на практике, но есть пределы.

Почему это работает: Noise Scale Hypothesis

Шум в градиенте ∝ 1 / sqrt(batch_size)

Примеры:
- batch_size=32:  noise ∝ 1/sqrt(32) ≈ 0.177
- batch_size=128: noise ∝ 1/sqrt(128) ≈ 0.088 (в 2x меньше)
- batch_size=512: noise ∝ 1/sqrt(512) ≈ 0.044 (в 4x меньше)

Основная идея:
Если batch size растёт в k раз, шум падает в sqrt(k) раз.
Мы можем позволить себе bigger learning rate пропорционально.
Для linear scaling: if batch_size ↑ 4x, lr ↑ 4x

Практический пример

import torch
import torch.nn as nn
from torch.optim import SGD

# Модель
model = nn.Linear(10, 1)
loss_fn = nn.MSELoss()

# Scenario 1: Маленький batch
batch_size_1 = 32
lr_1 = 0.001
optimizer_1 = SGD(model.parameters(), lr=lr_1)

# Training loop
for epoch in range(10):
    for batch_idx, (X_batch, y_batch) in enumerate(train_loader):
        optimizer_1.zero_grad()
        output = model(X_batch)
        loss = loss_fn(output, y_batch)
        loss.backward()
        optimizer_1.step()
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# Scenario 2: Больший batch (4x)
batch_size_2 = 128
lr_2 = 0.004        # Увеличили в 4x
optimizer_2 = SGD(model.parameters(), lr=lr_2)

# Результат: Similar convergence speed! Оба варианта сходятся примерно одинаково
# Но Scenario 2 быстрее по wall-clock time (больше примеров за iteration)

Детальная таблица

Плюсы: быстрее per wall-clock, лучше GPU utilization

Важные ограничения

1. Linear scaling работает только до точки

# Хорошо:
batch_size 32 → lr 0.001
batch_size 512 → lr 0.016  (16x)

# Плохо (слишком большой batch):
batch_size 2048 → lr 0.064
# Модель может diverge, потому что gradient слишком stale

2. Warmup strategy

Не сразу переходи на большой LR:

# Вместо сразу lr=0.016
# Используй warmup:
epoch 0: lr = 0.001 (маленький)
epoch 1: lr = 0.004
epoch 2: lr = 0.008
epoch 3: lr = 0.016  (целевой)
epoch 4+: lr = 0.016 (hold or decay)

# Это помогает избежать divergence в начале

3. Gradient Accumulation

Если большой batch не влезает в GPU, используй накопление:

effective_batch_size = 512
accumulation_steps = 4
actual_batch_size = 128  # То, что влезает в GPU

for epoch in range(num_epochs):
    accumulated_loss = 0
    for i, (X, y) in enumerate(train_loader):
        output = model(X)
        loss = loss_fn(output, y) / accumulation_steps
        loss.backward()
        
        if (i + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            # Эффективно обработали 512 примеров

Когда НЕ увеличивать LR пропорционально batch size?

# 1. Адаптивные оптимайзеры (Adam, AdamW)
# Adam уже адаптирует LR на основе истории градиентов
# Linear scaling менее критично, но всё ещё помогает

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# batch_size 32 → lr 0.001
# batch_size 128 → lr 0.002 (может быть меньше scaling factor)

# 2. Очень большие модели (BERT, GPT)
# Там используют специальные schedules (LAMB, LARS)
lr_scaled = base_lr * sqrt(batch_size / 256)
# Масштабирование не линейное, а по sqrt

# 3. Batch Normalization
# BN заботится о variance, поэтому LR может быть меньше

Практический совет для production

# Step 1: Найти хороший LR для базового batch size
batch_size = 32
best_lr = find_lr(0.0001, 0.1, batch_size)
print(f"Best LR for batch_size={batch_size}: {best_lr}")
# Результат: best_lr = 0.001

# Step 2: Масштабировать для больших batches
for new_batch_size in [64, 128, 256, 512]:
    scaling_factor = new_batch_size / batch_size  # 2, 4, 8, 16
    new_lr = best_lr * scaling_factor              # 0.002, 0.004, 0.008, 0.016
    
    # Опционально: добавить warmup
    # Опционально: decay LR по расписанию (cosine annealing, exponential)

# Step 3: A/B test
# Вариант A: batch_size 32, lr 0.001
# Вариант B: batch_size 256, lr 0.008
# → Вариант B быстрее и точнее (больше примеров за epoch)

Исторический контекст

Это было формально показано в статье "Accurate, Large Minibatch SGD" (You, Li, Reddi, Kumar, 2017):

Key finding:
lr_new = lr_old * (batch_size_new / batch_size_old)

Это основано на theoretical анализе шума в SGD.

Вывод

TL;DR:

1. При увеличении batch size в k раз, увеличь learning rate в k раз (linear scaling)
2. Это работает потому, что gradient становится более stable
3. Используй warmup для больших LR
4. Для Adam/AdaGrad может быть меньше scaling factor
5. Profit: быстрее per wall-clock time, лучше GPU utilization

Это один из практических tricks, который может дать вам 2-3x speedup в training без loss in accuracy!