SGD, батчирование и ранняя остановка

Полный градиент точно усредняет все строки, но один шаг может стоить миллионы вычислений. SGD заменяет его шумной оценкой по мини-батчу. Шум удешевляет обновления, помогает исследовать ландшафт и одновременно мешает точному финишу — а значит, обучение приходится не только запускать, но и вовремя останавливать.

Миллион примеров ради одного шага

Эмпирический риск имеет вид

$$\widehat R(\theta)=\frac1n\sum_{i=1}^n\ell_i(\theta).$$

Полный градиент

$$g(\theta)=\frac1n\sum_{i=1}^n\nabla\ell_i(\theta)$$

требует прохода по всей выборке. Если $n=10^8$, параметры обновятся только после чтения ста миллионов объектов: сутки счёта — и одна-единственная точка траектории. Мини-батч $B_t$ размера $b$ даёт

$$\widehat g_t=\frac1b\sum_{i\in B_t}\nabla\ell_i(\theta_t), \qquad \theta_{t+1}=\theta_t-\eta_t\widehat g_t .$$

При равномерном независимом выборе индексов

$$\mathbb E[\widehat g_t\mid\theta_t]=g(\theta_t).$$

Направление шумно, но в среднем верно. Это ровно та же логика, что в оценке среднего по выборке: маленькая случайная горсть говорит о целом — с ошибкой, которую можно посчитать.

Весь урок мы будем считать на реальных данных: почасовой велопрокат, $17\,379$ строк. Час суток закодирован тремя парами гармоник, как в уроке о базисных признаках, добавлены температура, влажность, ветер и признак рабочего дня — итого $10$ признаков плюс свободный член. Цель стандартизована, потери квадратичные. Точное решение нормальных уравнений даёт $\mathrm{MSE}_\star=0{,}454$, то есть $R^2=0{,}546$. Это дно, к которому будут стремиться все наши траектории; ниже него не опустится никакой оптимизатор, и всякое сравнение мы ведём именно с этим числом.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Батч — это выборка, а не приближение

Полезно раз и навсегда разделить две вещи. Мини-батч — не «грубое приближение» полного градиента в смысле округления. Это статистическая оценка: несмещённая, со своей дисперсией и своим доверительным интервалом. Для каждой координаты $j$

$$\operatorname{Var}\!\left[\widehat g_{t,j}\right] =\frac{1}{b}\operatorname{Var}\!\left[\nabla_j\ell_i(\theta_t)\right],$$

если индексы выбраны независимо. Отсюда стандартное отклонение

$$\operatorname{sd}\!\left[\widehat g_{t,j}\right] =\frac{\sigma_j}{\sqrt b},$$

где $\sigma_j$ — разброс индивидуальных градиентов по строкам. Это тот самый корень из урока о законе больших чисел, только усредняем мы теперь не отклики, а векторы производных.

Проверим закон на данных. Для батчей от $1$ до $1024$ мы измерили $\sqrt{\mathbb E\|\widehat g-g\|^2}$ по $300$ повторениям каждый. Наклон прямой в логарифмических осях получился $-0{,}498$ при теоретическом $-1/2$. При $b=1$ типичная ошибка оценки равна $2{,}64$ — сама норма полного градиента при этом порядка единицы, то есть одиночный пример указывает направление почти наугад. При $b=1024$ ошибка падает до $0{,}085$.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Закон корня: чем мы платим за тишину

Из $\sigma/\sqrt b$ следует неприятная арифметика. Чтобы уменьшить шум вдвое, надо увеличить батч вчетверо; чтобы уменьшить его в десять раз — в сто раз. Стоимость же шага растёт с $b$ примерно линейно, пока не насытится устройство. Значит,

$$\frac{\text{цена шага}}{\text{выигрыш в тишине}}\ \propto\ \frac{b}{\sqrt b}=\sqrt b .$$

Плата за каждую следующую «единицу спокойствия» растёт. Наши измерения это подтверждают: переход от $b=4$ к $b=256$ — это в $64$ раза больше вычислений на шаг и ровно в $8{,}17$ раза меньше разброса, при теоретических $\sqrt{64}=8$.

Честная ось времени

Эпоха — обработка примерно $n$ примеров; при батче $b$ число обновлений в эпохе равно

$$T_{\text{эпоха}}=\left\lceil\frac{n}{b}\right\rceil,$$

то есть для нашего велопроката $17\,379/64\approx272$ шага. Сравнивать два режима по числу шагов нечестно: один мог прочитать в сотни раз больше данных. Сравнивать по эпохам честнее, но и это ответ лишь на один из трёх вопросов. Их три:

$$\text{шаги}\ \longleftrightarrow\ \text{сколько раз мы двинули веса},$$ $$\text{примеры}\ \longleftrightarrow\ \text{сколько данных прочитано},$$ $$\text{секунды}\ \longleftrightarrow\ \text{сколько это стоило.}$$

Мы дали трём режимам одинаковый бюджет данных — $400\,000$ прочитанных примеров — и один и тот же шаг $\eta=0{,}35/L$. Батч $8$ успел сделать $50\,000$ обновлений, батч $64$ — $6250$, батч $512$ — всего $781$. После первой тысячи примеров лидирует мелкий батч ($0{,}518$ против $0{,}672$ у крупного): дешёвые шаги важнее их точности, пока мы далеко от минимума. Но к концу бюджета порядок меняется на противоположный: $0{,}535$ у $b=8$ против $0{,}455$ у $b=512$ при дне $0{,}454$. Мелкий батч быстро добежал до окрестности минимума и остался в ней дрожать.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Порог устойчивости

Для квадратичной задачи полный градиентный спуск с постоянным шагом ведёт себя как умножение ошибки в каждом собственном направлении на множитель $1-\eta\lambda$. Значит, сходимость требует $|1-\eta\lambda|<1$ для всех $\lambda$, то есть

$$0<\eta<\frac{2}{L},\qquad L=\lambda_{\max}.$$

На нашей матрице признаков $L=1{,}45$ и $\mu=0{,}31$, число обусловленности $\kappa=L/\mu=4{,}6$, а порог $2/L=1{,}38$. Шаг $0{,}5/L$ спокойно приводит к минимуму, $1{,}9/L$ уже качает траекторию поперёк долины, а $2{,}05/L$ — уносит потери за $10^3$ за сотню шагов. Порог здесь не метафора: это точное число, которое считается из данных до всякого обучения.

Статья Роббинса и Монро — точка рождения всей темы. В ней нет ни нейросетей, ни изображений: речь о том, как найти корень уравнения, если каждое измерение зашумлено. Ответ — двигаться маленькими, постепенно затухающими шагами. Через семьдесят лет тот же рецепт крутит миллиардные модели.

Пол из шума

Постоянный шаг в SGD не приводит к минимуму даже в идеальной выпуклой задаче. Причина видна из простой оценки: около минимума полезный сигнал $g(\theta)\to0$, а шум $\sigma/\sqrt b$ никуда не девается. Траектория стабилизируется в «шаровом облаке» вокруг минимума, радиус которого растёт с $\eta$ и падает с $b$. Для квадратичной задачи асимптотический избыток риска ведёт себя как

$$\mathbb E\left[\widehat R(\theta_t)\right]-\widehat R_\star \ \approx\ \frac{\eta\,\sigma^2}{2b}\cdot\frac{1}{1-\eta L/2},$$

то есть примерно линейно по $\eta$ вдали от порога и взрывается у него. Измерения на велопрокате при $b=32$: шаг $0{,}05/L$ даёт финальные потери $0{,}4555$ (избыток $0{,}0017$ над дном), шаг $1{,}6/L$ — $0{,}633$ (избыток $0{,}179$). Увеличение шага в $32$ раза подняло избыток в $106$ раз: линейная часть объясняет $32$, остальное добавил множитель близости к порогу.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Расписание шага

Классические условия Роббинса—Монро требуют

$$\sum_{t}\eta_t=\infty,\qquad \sum_{t}\eta_t^2<\infty .$$

Первая сумма запрещает остановиться слишком рано: суммарного «пути» должно хватить, чтобы добраться до минимума из любой начальной точки. Конечность второй ограничивает накопленный шум и заставляет пол опускаться к нулю. Последовательность $\eta_t=1/t$ удовлетворяет обоим условиям, $\eta_t=1/\sqrt t$ — только первому, $\eta_t=1/t^2$ — только второму.

На практике используют прогрев (warmup), ступенчатое уменьшение, косинусное угасание

$$\eta_t=\frac{\eta_0}{2}\left(1+\cos\frac{\pi t}{T}\right)$$

и one-cycle. Формально косинусное расписание конечно и потому не проверяет условия Роббинса—Монро — оно просто доводит шаг до нуля к концу заданного бюджета. Этого достаточно: при $b=32$ и бюджете $300\,000$ примеров постоянный шаг $1{,}2/L$ оставил нас на $0{,}5516$, а тот же шаг с косинусным угасанием — на $0{,}4542$ при дне $0{,}4538$. Избыток уменьшился в $235$ раз без единого дополнительного прочитанного примера.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Русская линия: квазиградиент Ермольева

Пока в статистике развивалась стохастическая аппроксимация, в Киеве складывалась своя школа. Юрий Михайлович Ермольев в Институте кибернетики АН УССР занимался задачами стохастического программирования — оптимизацией, где и целевая функция, и ограничения заданы средними по случайности. Его книга «Методы стохастического программирования» (М.: Наука, 1976) ввела в обиход понятие стохастического квазиградиента: направления, которое не обязано быть несмещённой оценкой градиента и даже не обязано существовать как градиент — достаточно, чтобы в среднем оно составляло острый угол с направлением убывания.

Формально требуется лишь

$$\mathbb E\left[\xi_t\mid\theta_t\right]=g(\theta_t)+r_t, \qquad \|r_t\|\to0,$$

и тогда при условиях Роббинса—Монро метод $\theta_{t+1}=\Pi\left(\theta_t-\eta_t\xi_t\right)$ сходится. Эта постановка охватывает случаи, где мини-батч смещён: недифференцируемые потери, конечно-разностные оценки, вложенные средние, проекция на множество ограничений из урока о задачах с ограничениями. Современные приёмы, которые формально ломают несмещённость — обрезка градиента, hard example mining, аппроксимация внутреннего ожидания одной выборкой, — живут именно в этой теории, а не в исходной схеме Роббинса—Монро.

Рядом стоит вторая ветвь той же школы: Аркадий Немировский и Давид Юдин в конце 1970-х описали, каких скоростей в принципе нельзя превзойти для выпуклых задач, и придумали метод зеркального спуска, где выбор геометрии подстраивается под форму допустимого множества. Из их нижних оценок следует, что $O(1/\sqrt T)$ для общих выпуклых задач с шумом — не слабость SGD, а свойство самой задачи. Стоит помнить об этом, читая очередное обещание «ускорить обучение в десять раз».

Лаборатория батча и расписания

Бюджет в лаборатории фиксирован: восемь тысяч примеров, как их ни дели. Начните с постоянного шага и переберите батчи: при $b=1$ траектория — клубок, при $b=256$ — гладкая дуга, которая не успевает дойти. Затем доводите шаг до порога $2/L$ и наблюдайте, как аккуратная дуга превращается в поперечные скачки, а потом в расходимость. Наконец включите косинусное угасание и посмотрите на золотой пунктир «пола»: он опускается, хотя прочитано ровно столько же данных.

Ранняя остановка: где разворачивается validation

Оптимизация и обучение — разные цели. Оптимизатор минимизирует $\widehat R$ на обучающей выборке; нас же интересует истинный риск. После некоторого момента дальнейшая минимизация обучающего среднего перестаёт улучшать будущее качество и начинает его ухудшать: модель дописывает в веса шум конкретных меток.

Мы поставили честный эксперимент на реальных данных. Двести случайных часов велопроката, из них $50$ в обучении и $150$ в validation; $53$ признака — шестнадцать гармоник часа плюс погодные и их взаимодействия; к целям добавлен дополнительный шум с фиксированным seed, чтобы имитировать грязные метки. Полный градиентный спуск, тридцать тысяч шагов.

Результат: validation достигает минимума $1{,}76$ на шаге $54$, а к концу поднимается до $3{,}33$ — в $1{,}89$ раза хуже. Обучающая ошибка за то же время падает с $0{,}632$ до $0{,}081$. Тот, кто смотрел только на train, увидел бы восьмикратное улучшение и не заметил, что модель испортилась вдвое.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Почему остановка похожа на ridge

Совпадение с ridge из урока о регуляризации не случайно. Разложим линейную задачу по собственным направлениям матрицы $\frac1nX^\top X$ с собственными числами $\lambda_k$. Градиентный спуск из нуля после $t$ шагов восстанавливает $k$-ю компоненту решения с множителем

$$f^{\text{GD}}_k=1-(1-\eta\lambda_k)^{t},$$

а ridge со штрафом $\alpha$ — с множителем

$$f^{\text{ridge}}_k=\frac{\lambda_k}{\lambda_k+\alpha}.$$

Обе функции равны нулю при $\lambda_k\to0$ и стремятся к единице при больших $\lambda_k$; при малых $\eta\lambda t$ обе ведут себя как $\eta\lambda_k t$. Отсюда знаменитое соответствие

$$\alpha\approx\frac{1}{\eta t}.$$

Дольше учим — слабее штраф. При $\eta=0{,}5$ и $t=50$ направление с $\lambda=1{,}5$ выучено полностью ($1{,}000$), направление с $\lambda=0{,}02$ — на $0{,}395$ (ridge при $\alpha=1/(\eta t)=0{,}04$ даёт $0{,}333$), а направление с $\lambda=0{,}001$ — всего на $0{,}025$. Максимальное расхождение двух кривых на сетке равно $0{,}209$: соответствие качественное, а не тождество.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Заодно понятно, почему ранняя остановка часто помогает именно шумным задачам: направления с малым $\lambda$ плохо определены данными, и именно в них селится шум. Спуск добирается до них последними. Эту же спектральную картину мы видели в PCA — только там мы отбрасывали слабые направления явно, а здесь просто не успеваем их выучить.

Patience, checkpoint и цена подглядывания

Реальная кривая validation не гладкая: она шумит, потому что validation конечен. Отсюда правило patience — ждать $p$ проверок без улучшения, прежде чем остановиться. Оно сглаживает случайные провалы, но не создаёт новой независимости.

Если проверять validation каждые десять шагов и выбирать лучший из тысяч замеров, то сам минимум становится оптимистичным: мы выбрали максимум по шумной выборке. Для $m$ независимых замеров с шумом $\tau$ типичный оптимизм растёт как

$$\mathbb E\left[\max_{j\le m}\varepsilon_j\right]\ \approx\ \tau\sqrt{2\ln m}.$$

При $m=1000$ множитель $\sqrt{2\ln1000}=3{,}72$: почти четыре стандартных отклонения бесплатного «улучшения». Это та же цена поиска, которую разбирает урок про oracle inequality, и лечится она независимым test.

Перемешивание, редкие группы и стратификация

Если данные отсортированы по классу или по времени, последовательные батчи систематически смещены: градиент подолгу смотрит на один режим. Shuffle возвращает оценкам почти равномерность. Но для временных рядов бездумное перемешивание разрушает содержательный порядок оценки — модель начинает подглядывать в будущее, и это уже утечка, а не оптимизация.

Отдельная опасность — редкая группа. Пусть её доля $p$; тогда вероятность, что в батче размера $b$ не окажется ни одного её представителя, равна

$$\mathbb P(\text{группа не попала})=(1-p)^b .$$

Возьмём реальные числа: в SMS Spam Collection $747$ спам-сообщений из $5574$, то есть $p=0{,}134$. При $b=8$ доля батчей без единого спама равна $0{,}866^8=0{,}316$ — почти треть шагов не содержит информации о редком классе вовсе. При $b=32$ это уже $0{,}0100$, один шаг из ста.

Стратифицированный батч, где доля редких фиксирована, уменьшает дисперсию — если затем правильно взвесить слагаемые. Простое принудительное равенство классов меняет саму задачу: оптимизируется перевзвешенный риск, и выходные вероятности требуют коррекции prior и калибровки, как в уроке о порождающих классификаторах. Это допустимый выбор, но он должен быть записан явно, а не спрятан в сэмплере.

Микробатчи и эффективный размер

Память устройства ограничена, и большой батч часто собирают из микробатчей: несколько прямых-обратных проходов, градиенты складываются, обновление одно.

$$\widehat g=\frac1M\sum_{m=1}^{M}\widehat g^{(m)}, \qquad b_{\text{eff}}=M\,b_{\text{micro}} .$$

При неизменных параметрах внутри цикла это в точности воспроизводит большой батч. Но эквивалентность легко потерять: BatchNorm считает статистики внутри микробатча, случайная аугментация даёт разные преобразования, обрезка градиента, применённая к каждому микробатчу отдельно, — это уже другая операция, чем обрезка суммы.

Слагаемые этой суммы не обязаны лежать на одной машине. Если каждый микробатч живёт у своего владельца и не покидает его устройства, усреднение градиентов превращается в федеративное обучение: формула та же, а вот равенство долей и одинаковость распределений — уже нет.

Ходовое эвристическое правило переноса рецепта — линейное масштабирование шага вместе с батчем:

$$\eta(b)=\eta_0\cdot\frac{b}{b_0}, \qquad \frac{\eta(b)}{b}=\frac{\eta_0}{b_0}=\text{const},$$

то есть неизменной сохраняется именно комбинация $\eta/b$, задающая пол из шума. Правило работает лишь пока $\eta(b)$ остаётся заметно ниже порога $2/L$; дальше его страхуют прогревом.

Отсюда практическое правило отчётности: указывать микробатч, эффективный батч и число устройств. При переносе рецепта на вдвое большее число ускорителей эффективный батч удваивается, шум падает в $\sqrt2$ раз — и шаг обычно поднимают, добавляя прогрев, чтобы не задеть порог устойчивости.

Шум как неявная регуляризация — осторожно

Популярное объяснение: мелкий батч чаще покидает узкие минимумы, а плоские области лучше обобщают. Первая часть верна — амплитуда дрожания растёт с $\eta/b$, и из узкой ямы шум действительно выбивает. Вторая часть требует осторожности: «плоскость» зависит от параметризации. В сети с ReLU можно умножить веса одного слоя на $c$, а следующего — на $1/c$; функция не изменится, а кривизна изменится сильно. Значит, само по себе значение гессиана не является характеристикой качества.

Честная формулировка проще: шум SGD — часть определения алгоритма, и решение, к которому он приходит, зависит от $\eta$, $b$ и расписания. Меняя их, мы меняем не скорость, а конечную модель. Именно поэтому связка «батч + шаг + расписание» проверяется по эмпирическому риску на нескольких seed, а не по одной удачной кривой.

Журнал обучения

Что писать в журнал, чтобы запуск можно было понять постфактум:

$$\text{train loss},\quad \text{validation loss},\quad \eta_t,\quad \|\widehat g_t\|,\quad \text{время},\quad \text{seed}.$$

Норма градиента — самый дешёвый детектор беды. Резкий всплеск означает нестабильность или испорченный батч; норма, ушедшая в ноль при высокой потере, — насыщение активаций либо ошибка кода. Перед долгим запуском градиент проверяют конечными разностями на малой гладкой модели:

$$\frac{f(\theta+he_j)-f(\theta-he_j)}{2h} \approx\frac{\partial f}{\partial\theta_j},$$

помня, что слишком большое $h$ даёт ошибку аппроксимации, а слишком малое — ошибку округления. Такой тест ловит неверный знак и пропущенное слагаемое, но ничего не говорит о скорости SGD — это проверка кода, а не оптимизатора.

Сборка

Мини-батч — это выборка, а не приближение: несмещённая оценка со стандартным отклонением $\sigma/\sqrt b$. Отсюда всё остальное. Шум делает шаги дешёвыми и потому быстрыми в начале, но ставит пол, который при постоянном шаге не опускается: $0{,}5516$ вместо $0{,}4542$ на нашем велопрокате. Порог устойчивости $2/L$ считается заранее и не обсуждается. Расписание опускает пол бесплатно. Ранняя остановка — не хитрость, а регуляризатор, почти совпадающий с ridge: $1{,}760$ против $1{,}753$ на одном и том же разбиении. А сравнение двух рецептов честно только при равном бюджете прочитанных данных, показанном расписании и нескольких seed.

Дальше — урок 62: momentum накапливает направление, AdaGrad и Adam дают каждой координате свой масштаб. Ландшафт при этом не меняется — меняется траектория и то, в какую из многих одинаково хороших точек мы придём.

Задачи