Суть эксперимента
Автор использовал доступ к бесплатным вычислительным ресурсам (4 GPU NVIDIA H200) для обучения модели GPT-2 (11 слоёв, ширина 768) на фиксированном наборе из 800 миллионов токенов из датасета FineWeb. Ключевое отличие от стандартного обучения: вместо мини-батчей использовался полный батч (full-batch), то есть градиент вычислялся по всем 800M токенов одновременно. Это устранило шум, возникающий из-за случайной выборки подмножества данных, но сохранило шум дискретизации шага обучения.
Результаты: переобучение против обобщения
Эксперимент наглядно показал разрыв между качеством подгонки под обучающую выборку и способностью модели работать с новыми данными. Полнобатчное обучение достигло лучшей ошибки на тренировочных данных, но провалилось на валидации. В то же время стандартное обучение с мини-батчами показало более сбалансированные результаты.
| Конфигурация | Размер батча | Шаги | Training Loss | Validation Loss |
|---|---|---|---|---|
| Малый батч (стандарт) | 393K токенов | 2,372 | 3.186 | 3.248 |
| Полный батч (эксперимент) | 800M токенов | 350 | 3.094 | 3.521 (лучший: 3.498) |
| Свежие данные (без повторов) | 50M токенов | 350 | 3.228 | 3.230 |
Как видно из таблицы, при сопоставимом уровне ошибки на тренировке (3.186 против 3.094), полнобатчная модель имеет валидационную ошибку на 0.27 nats выше, чем стандартная. Это означает, что шум от выбора мини-батчей не является «помехой», а выполняет функцию регуляризации, предотвращающую переобучение.
Динамика оптимизации: край стабильности
Полнобатчное обучение велось в режиме, который авторы называют «краем стабильности» (edge-of-stability). Градиенты на соседних шагах были практически антинаправленны (медианное косинусное сходство -0.945). На 98% шагов каждый микро-батч был анти-выровнен с полным градиентом предыдущего шага. Модель постоянно «перелетала» минимум функции потерь, увеличивая ошибку примерно на трети шагов, при этом собственное значение гессиана (Hessian eigenvalue) росло с 48 до 113 к 300-му шагу.
Технические детали и выводы
Для обеспечения чистоты эксперимента были удалены хэшированные n-граммные таблицы эмбеддингов (которые вызывали взрыв валидационной ошибки при повторении токенов) и sampled softmax. Использовался оптимизатор семейства Muon (ANVIL) для матриц и Adam для остальных параметров. Эксперимент занял около 100 GPU-часов и пережил 106 прерываний вычислений (preemptions), что потребовало частого чекпоинтинга.
Результат подтверждает гипотезу: шум оптимизации, создаваемый мини-батчами, необходим для нахождения широких минимумов в ландшафте потерь, которые обеспечивают хорошее обобщение. Чистый градиентный спуск на полном наборе данных приводит к переобучению, даже если модель формально лучше «запоминает» обучающую выборку.
Источник: LessWrong ↗