Что такое loss landscape LLM и как оно влияет на обучение (sharp vs flat minima)?

Краткий тезис

Loss landscape — это многомерная поверхность ошибки (loss) в пространстве параметров модели. Для LLM с миллиардами параметров этот ландшафт чрезвычайно сложен и содержит как sharp minima (узкие, крутые впадины), так и flat minima (широкие плато). Flat minima обеспечивают лучшее обобщение на новых данных, тогда как sharp minima ведут к переобучению. Понимание ландшафта помогает выбирать оптимизаторы (Adam, SGD) и техники регуляризации, которые направляют обучение в плоские области.

1. Определение loss landscape

Loss landscape (ландшафт функции потерь) — это графическое представление значения функции потерь (например, кросс-энтропии) в зависимости от всех весов модели. Поскольку весов миллиарды, ландшафт невозможно визуализировать напрямую; его изучают через проекции на 2D-плоскости (метод Li et al., 2018) или через анализ кривизны (собственные числа гессиана).

• Ось X, Y — направления в пространстве параметров (случайные или главные компоненты).
• Ось Z — значение loss.
• Минимумы — точки, где градиент равен нулю; модель обучена.

Для LLM ландшафт невыпуклый, с множеством локальных минимумов, седловых точек и плато.

2. Sharp minima vs flat minima

Почему flat minima лучше обобщают

• PAC-Bayes — плоские минимумы имеют меньшую «сложность» (описательную длину), что даёт более тесные границы обобщения.
• Шум в данных — при тестовых данных, слегка отличающихся от тренировочных, flat minima остаются в области низкого loss, а sharp — «вылетают» в высокий loss.
• Сглаживание — flat minima соответствуют решениям, устойчивым к возмущениям весов (например, при квантовании или прунинге).

3. Как loss landscape влияет на обучение LLM

3.1 Выбор оптимизатора

• SGD (с momentum) часто находит sharp minima, если learning rate мал. При большом LR SGD может «перепрыгивать» узкие впадины и оседать в плоских областях.
• Adam (адаптивный LR) склонен находить flat minima благодаря адаптивному шагу и моментуму. Эмпирически Adam даёт лучшее обобщение для LLM, чем SGD.
• Sharpness-Aware Minimization (SAM) — специальный метод, который явно ищет flat minima, минимизируя loss в окрестности точки. SAM улучшает обобщение для BERT и GPT, но требует вдвое больше вычислений.

3.2 Learning rate schedule

• Cosine annealing с warm restarts (SGDR) помогает выходить из sharp minima, периодически увеличивая LR.
• One-cycle policy — сначала разогрев, потом снижение — также способствует flat minima.

3.3 Batch size

• Маленький batch size (шумный градиент) → больше шума → выход из sharp minima → flat minima.
• Большой batch size (точный градиент) → риск застрять в sharp minima. Для LLM часто используют большие batch (например, 1024+), поэтому добавляют шум (gradient noise) или используют Adam.

3.4 Регуляризация

• Weight decay (L2) — штрафует большие веса, что сглаживает ландшафт и способствует flat minima.
• Dropout — стохастическая регуляризация, эквивалентная ансамблированию, также ведёт к плоским минимумам.
• Label smoothing — заменяет жёсткие метки на мягкие, уменьшает остроту минимумов.

4. Методы анализа loss landscape для LLM

4.1 Визуализация (Li et al., 2018)

Выбираются два случайных направления в пространстве параметров, и loss вычисляется вдоль их линейной комбинации. Результат — контурный график.

import torch
import numpy as np
import matplotlib.pyplot as plt

def plot_loss_landscape(model, loss_fn, data_loader, directions, n_points=20):
    # directions: два вектора той же размерности, что и параметры
    # нормализуем направления
    dir1 = directions[0] / np.linalg.norm(directions[0])
    dir2 = directions[1] / np.linalg.norm(directions[1])
    # сетка
    alphas = np.linspace(-2, 2, n_points)
    betas = np.linspace(-2, 2, n_points)
    losses = np.zeros((n_points, n_points))
    # сохраняем исходные веса
    orig_params = [p.clone().detach() for p in model.parameters()]
    for i, a in enumerate(alphas):
        for j, b in enumerate(betas):
            # устанавливаем веса = orig + a*dir1 + b*dir2
            idx = 0
            for p, orig in zip(model.parameters(), orig_params):
                p.data = orig + a * dir1[idx] + b * dir2[idx]
                idx += p.numel()
            # считаем loss
            total_loss = 0
            for x, y in data_loader:
                out = model(x)
                loss = loss_fn(out, y)
                total_loss += loss.item()
            losses[i, j] = total_loss / len(data_loader)
    # восстанавливаем веса
    for p, orig in zip(model.parameters(), orig_params):
        p.data = orig
    # рисуем
    plt.contourf(alphas, betas, losses, levels=50)
    plt.colorbar()
    plt.xlabel('Direction 1')
    plt.ylabel('Direction 2')
    plt.title('Loss landscape')
    plt.show()

4.2 Анализ гессиана

Собственные числа гессиана (матрицы вторых производных) показывают кривизну. Большие положительные собственные числа → sharp minima. Для LLM гессиан огромен, поэтому используют приближения (Hutchinson trace estimator, Lanczos).

4.3 Метрики sharpness

• Sharpness measure (Keskar et al., 2017): максимальное изменение loss при возмущении весов в окрестности радиуса epsilon.
• PAC-Bayes bound — оценка обобщения через flatness.

5. Эмпирические наблюдения для LLM

• LLM имеют очень сложный ландшафт с множеством седловых точек и плато. Градиент часто мал, что замедляет обучение.
• Adam эффективнее SGD для LLM, так как адаптивный LR помогает преодолевать плато и избегать sharp minima.
• Pre-training часто заканчивается в относительно плоском минимуме, что позволяет успешно fine-tune на новых задачах (transfer learning).
• Fine-tuning может «свалиться» в sharp минимум, если learning rate слишком велик или мало данных. Поэтому для fine-tuning LLM используют маленький LR (1e-5 – 5e-5) и early stopping.
• Large batch training (например, 4096 для GPT-3) требует дополнительных техник (gradient accumulation, loss scaling), чтобы не уйти в sharp minima.

6. Связь с другими аспектами обучения

7. Практические рекомендации для обучения LLM

1. Используйте AdamW (Adam с weight decay) — он сочетает адаптивность и регуляризацию, способствуя flat minima.
2. Применяйте cosine annealing с warm restarts — помогает выходить из sharp minima.
3. Добавляйте gradient noise (например, Gaussian noise к градиенту) при большом batch size.
4. Контролируйте sharpness через validation loss — если train loss падает, а validation растёт, возможно, модель попала в sharp минимум; увеличьте LR или добавьте dropout.
5. Для fine-tuning используйте маленький LR и early stopping, чтобы не покинуть flat область pre-trained весов.

Пет-проект для закрепления

Задача: Визуализировать loss landscape для небольшой модели (2-слойный перцептрон на MNIST) и сравнить sharpness после обучения с SGD и Adam.

Инструменты: PyTorch, matplotlib, numpy.

Шаги:

1. Обучите модель двумя способами: SGD (LR=0.01, momentum=0.9) и Adam (LR=0.001).
2. Сохраните финальные веса для каждого оптимизатора.
3. Реализуйте функцию plot_loss_landscape, как в разделе 4.1.
4. Сгенерируйте два случайных направления (нормализованных) и постройте контурные графики для обеих обученных моделей.
5. Вычислите метрику sharpness: среднее изменение loss при возмущении весов на epsilon=0.1.
6. Сравните: у SGD loss landscape должен быть более «острый» (sharp), у Adam — более плоский.

Ожидаемый результат: Контурные графики, показывающие разницу в форме минимумов, и численное значение sharpness (у Adam меньше). Вывод: Adam находит более плоские минимумы, что объясняет его лучшее обобщение.

Связь с другими вопросами

Навигация

• Предыдущий: 669
• Следующий: 671
• Индекс: 00. Индекс разборов