Entrenar bien: lotes, épocas y optimizadores

El gradiente de un lote

Con $N$ ejemplos y un error por ejemplo $\ell_i(\mathbf{w})$, el error total es $L(\mathbf{w})=\tfrac1N\sum_i \ell_i(\mathbf{w})$. Un lote $\mathcal{B}$ de $B$ ejemplos elegidos al azar da el gradiente aproximado

Es un estimador sin sesgo. Su varianza es $\frac{N-B}{N-1}\,\frac{\sigma^2}{B}$, con $\sigma^2$ la varianza del gradiente de un ejemplo: el error cae como $1/\sqrt{B}$, mientras que el costo del paso crece como $B$. Por eso, para reducir el ruido a la mitad hace falta un lote unas cuatro veces mayor. En la lección, el desvío medio de la flecha (medido con 24 lotes de cada tamaño) fue de 20 %, 10 %, 4 % y 1.7 % con lotes de 2, 8, 32 y 128: cae algo más rápido que $1/\sqrt{B}$ porque el factor $\frac{N-B}{N-1}$ pesa cuando el lote es una parte grande de los 256 ejemplos. Una época tiene $\lceil N/B\rceil$ pasos y tras $t$ pasos se han visto $tB$ ejemplos. La idea de actualizar con ruido viene de Robbins y Monro (1951).

Cerca de un mínimo con curvatura $h$, el ruido del gradiente (varianza $s^2=\sigma^2/B$) deja un suelo de error: $\mathbb{E}[L]\approx \eta\,s^2/4$ cuando $\eta h\ll 1$. Es proporcional a $\eta/B$: bajar la tasa o subir el lote lo baja. Por eso el decaimiento de la tasa ayuda a asentarse.

Momento

Con $\beta=0.9$ y un gradiente constante, $\mathbf{v}\to \mathbf{g}/(1-\beta)$: el paso efectivo es diez veces mayor; a igual paso efectivo, la tasa del momento sería unas diez veces menor que la de SGD (en la práctica la razón es menor: en el cañón de la lección, 0.02 frente a 0.07). En un cuadrático con curvaturas $\lambda_{\min}\le h\le\lambda_{\max}$ y número de condición $\kappa=\lambda_{\max}/\lambda_{\min}$, la mejor tasa fija de SGD reduce la distancia al fondo por paso en un factor $\frac{\kappa-1}{\kappa+1}$, y el momento bien ajustado en $\frac{\sqrt\kappa-1}{\sqrt\kappa+1}$ (Polyak, 1964). En el cañón de la lección, cerca del fondo, $\kappa\approx 50$: unos $0.96$ por paso frente a $0.75$ (en teoría, sin ruido).

Adam

Todo es por peso (el cuadrado y la raíz actúan componente a componente). Los valores por defecto de Kingma y Ba (2015) son $\eta=0.001$, $\beta_1=0.9$, $\beta_2=0.999$ y $\varepsilon=10^{-8}$. Las correcciones de sesgo compensan que $\mathbf{m}$ y $\mathbf{s}$ empiezan en cero. En el primer paso $\hat{\mathbf{m}}/\sqrt{\hat{\mathbf{s}}}=\operatorname{signo}(\mathbf{g})$, así que cada peso se mueve $\eta$; en régimen normal el paso de cada peso es del orden de $\eta$ como máximo. Dos consecuencias: el método es invariante al reescalar cada gradiente por separado (lo dicen sus autores), y cerca de un óptimo la señal frente al ruido baja y los pasos se encogen solos, un «recocido automático» según Kingma y Ba. Con gradientes de lotes ruidosos eso no basta: en la simulación de la lección, con $\eta$ constante el camino sigue sacudiéndose alrededor del fondo, con una amplitud que crece con $\eta$, y bajar $\eta$ al final lo calma.

Calendarios de la tasa

Dos recetas reales: el transformer original (Vaswani et al., 2017) usó Adam con $\beta_2=0.98$ y calentó durante 4000 pasos; GPT-3 (Brown et al., 2020) usó Adam con $\beta_2=0.95$, calentamiento lineal durante los primeros 375 millones de tokens y decaimiento coseno hasta el 10 % del valor máximo a lo largo de 260 mil millones de tokens. Además, su lote creció de 32 mil tokens hasta el valor final (3.2 millones en el modelo de 175 mil millones de parámetros) durante los primeros 4 a 12 mil millones de tokens, según el tamaño del modelo, y de su texto del rastreo web vio solo 0.44 épocas (tabla 2.2). Al aumentar el lote, Goyal et al. (2017) proponen multiplicar la tasa por el mismo factor y calentar: con esa regla entrenaron ResNet-50 con lotes de 8192 imágenes en 256 GPU, en una hora, sin pérdida de exactitud.

Hiperparámetros

Los pesos $\mathbf{w}$ se aprenden con $\nabla L$; los hiperparámetros ($\eta$, $B$, $\beta$, $T_w$, el optimizador…) fijan cómo se aprende y no tienen gradiente a su alcance. Se eligen probando: en espacios de varias dimensiones, probar combinaciones al azar suele rendir más por prueba que recorrer una cuadrícula (Bergstra y Bengio, 2012).

En código

import math

# Ocho ejemplos fijos de y ≈ 2x + 1. Los lotes se toman en orden para que el
# resultado sea reproducible (en la práctica se baraja en cada época).
X = [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]
Y = [1.2, 1.8, 3.1, 4.1, 4.9, 6.2, 6.9, 8.1]

def grad(w, b, idx):                       # gradiente del error cuadrático medio del lote
    n = len(idx)
    gw = sum(2 * (w * X[i] + b - Y[i]) * X[i] for i in idx) / n
    gb = sum(2 * (w * X[i] + b - Y[i]) for i in idx) / n
    return [gw, gb]

def entrenar(opt, eta, B, epocas, beta=0.9, b1=0.9, b2=0.999, eps=1e-8):
    p = [0.0, 0.0]                         # pesos: w y b
    v = [0.0, 0.0]; m = [0.0, 0.0]; s = [0.0, 0.0]; t = 0
    for _ in range(epocas):                # una época = recorrer todos los lotes
        for ini in range(0, len(X), B):
            g = grad(p[0], p[1], range(ini, min(ini + B, len(X))))
            t += 1
            if opt == "sgd":
                paso = [eta * gi for gi in g]
            elif opt == "momento":
                v = [beta * vi + gi for vi, gi in zip(v, g)]
                paso = [eta * vi for vi in v]
            else:                          # Adam
                m = [b1 * mi + (1 - b1) * gi for mi, gi in zip(m, g)]
                s = [b2 * si + (1 - b2) * gi * gi for si, gi in zip(s, g)]
                paso = [eta * (mi / (1 - b1 ** t)) / (math.sqrt(si / (1 - b2 ** t)) + eps)
                        for mi, si in zip(m, s)]
            p = [pi - di for pi, di in zip(p, paso)]
    return [round(pi, 3) for pi in p]

# la solución exacta de mínimos cuadrados es [1.998, 1.042]
for opt, eta in [("sgd", 0.05), ("momento", 0.005), ("adam", 0.3)]:
    print(opt, entrenar(opt, eta, B=4, epocas=50))
# sgd [1.999, 1.033]
# momento [2.008, 1.043]
# adam [1.993, 1.028]

En un problema tan pequeño y bien condicionado, SGD ya basta: las diferencias aparecen en paisajes como el cañón, y cada optimizador necesitó su propia tasa.

Lo que viene

Bajar el error de entrenamiento no basta: hay que medir con datos que el modelo no vio. De eso trata la lección 10 (generalizar y evaluar), que añade la parada temprana y el decaimiento de pesos como formas de regularizar.