SGD con minibatches

Optimizadores, episodio dos: SGD con minibatches. SGD es stochastic gradient descent, descenso por gradiente estocástico. En el episodio uno calculamos el gradiente con todos los datos. Nadie entrena así: con millones de ejemplos, un solo paso costaría una pasada entera. Hoy, el gradiente de un grupito al azar, el ruido que trae y por qué ese ruido no impide llegar.

Parte 1 · La pérdida es un promedio

La pérdida de un dataset es el promedio de la pérdida de cada ejemplo. Y su gradiente, también.

Del dataset al minibatch

$\ell_i$: la pérdida del ejemplo $i$

el gradiente del promedio es el promedio de los gradientes

$\mathcal{B}$: $B$ ejemplos elegidos al azar

Tres líneas, de lo que ya sabemos a lo nuevo. N es la cantidad de ejemplos; ele sub i, la pérdida de uno solo, que depende de los pesos, w. La sigma mayúscula es una suma: la pérdida total es el promedio de las N. En el dibujo, cinco ejemplos y su promedio. Derivar es lineal: la derivada de una suma es la suma de las derivadas. Entonces el gradiente de la pérdida total es el promedio de los gradientes de cada ejemplo. En el dibujo, cinco columnas, cada una pesada un quinto, que se suman en una. La idea de SGD: en vez de promediar los N, promediar solo B, elegidos al azar. Ese grupo es el minibatch, la be caligráfica. Le decimos g sub B. En el dibujo, solo cuentan dos de los cinco.

Mil ejemplos, un gradiente por grupo

Mil ejemplos; cada uno es un punto del plano, c sub i, sorteado al azar. Los corremos para que el promedio sea exactamente cero. La pérdida del ejemplo i es la distancia al cuadrado entre los pesos, p, y ese punto. Su gradiente: dos por p menos c sub i, promediado. Sumadas las mil, la pérdida es el cuenco del episodio uno. En el punto menos dos coma cinco, uno coma cinco calculamos el gradiente con todos, con minibatches de diez y con un solo ejemplo. Y el promedio de cinco mil minibatches. Con todos: menos cinco, tres. Los de diez se le parecen. Los de un ejemplo se alejan más. Pero el promedio de cinco mil da menos cinco coma cero cero tres, tres coma cero uno dos: el gradiente verdadero.

rng = np.random.default_rng(0)
C = rng.normal(0, 1, size=(1000, 2))   # 1000 ejemplos: cada uno, un punto
C = C - C.mean(0)                      # centrados: el promedio es (0, 0)

def grad(p, idx):                      # pérdida de i: |p - C[i]|²
    return 2 * (p - C[idx]).mean(0)    # gradiente promedio de esos i

p = np.array([-2.5, 1.5])
print("todos (N=1000):   ", grad(p, np.arange(1000)))
for _ in range(2):
    print("minibatch de 10:  ", grad(p, rng.choice(1000, 10, replace=False)))
for _ in range(2):
    print("un solo ejemplo:  ", grad(p, rng.choice(1000, 1)))
muchos = [grad(p, rng.choice(1000, 10, replace=False)) for _ in range(5000)]
print("promedio de 5000: ", np.mean(muchos, 0))

todos (N=1000):    [-5.  3.]
minibatch de 10:   [-4.753  3.419]
minibatch de 10:   [-4.986  2.701]
un solo ejemplo:   [-6.977  3.671]
un solo ejemplo:   [-5.164  1.58 ]
promedio de 5000:  [-5.003  3.012]

El error de cada gradiente

A cada gradiente del snippet le restamos el verdadero, el de los mil: lo que queda es su error. Los dibujamos desde el centro, al doble de largo. Un minibatch de diez: un error cortito. El otro de diez, también corto, para otro lado. Un solo ejemplo: un error mucho más largo. Y el otro ejemplo suelto, largo hacia abajo. Los errores apuntan para cualquier lado, y por eso no empujan, en promedio, hacia ninguno. Y son más largos cuanto menos ejemplos promediás. Las dos ideas de la parte que sigue.

$g_B - \nabla L$ de la salida real de ep02_ejemplos, al doble de largo

Parte 2 · Ruido sin sesgo

Cuánto se equivoca un minibatch, y cómo baja ese error cuando B crece.

El error de un minibatch, medido

Cuatro tamaños de minibatch: uno, cuatro, dieciséis y sesenta y cuatro. Para cada uno, cuatro mil minibatches al azar en el mismo punto. El error es el gradiente del minibatch menos el verdadero. Medimos su promedio, que es el sesgo, y su desvío estándar: cuánto se aleja, típicamente, de cero. Y el desvío multiplicado por raíz de B. El sesgo da casi cero en todos los casos. El desvío: uno coma nueve nueve, cero coma nueve nueve siete, cero coma cuatro nueve cinco, cero coma dos cuatro tres. Cada vez que B se multiplica por cuatro, el desvío se divide por dos. Y desvío por raíz de B queda clavado cerca de dos.

for B in [1, 4, 16, 64]:
    gs = np.array([grad(p, rng.choice(1000, B, replace=False))
                   for _ in range(4000)])          # 4000 minibatches
    error = gs - grad(p, np.arange(1000))          # ruido: g_B - ∇L
    print(f"B = {B:2}: sesgo {error.mean(0).round(3)}"
          f"  desvío {error.std():.3f}"
          f"  desvío·√B {error.std() * np.sqrt(B):.2f}")

B =  1: sesgo [-0.044 -0.064]  desvío 1.986  desvío·√B 1.99
B =  4: sesgo [0.005 0.01 ]  desvío 0.996  desvío·√B 1.99
B = 16: sesgo [-0.011  0.003]  desvío 0.495  desvío·√B 1.98
B = 64: sesgo [0.002 0.001]  desvío 0.243  desvío·√B 1.94

Por qué raíz de B

$\xi$: el ruido del minibatch

$\sigma$: el desvío de un solo ejemplo

la cuenta coincide con lo medido

Escribimos el gradiente del minibatch como el verdadero más un ruido, xi, la letra griega. La E de doble trazo es la esperanza, de Álgebra lineal para LLMs, episodio ocho: el promedio sobre todos los minibatches posibles. Que la esperanza de xi sea cero quiere decir que el ruido no empuja para ningún lado: eso es no tener sesgo. En el dibujo, dos columnas que se suman. Sigma es cuánto varía el gradiente de un solo ejemplo. Al promediar B ejemplos independientes, el desvío se divide por raíz de B. En el dibujo, las barras del snippet. Acá sigma vale dos, así que con cuatro da uno, con dieciséis cero coma cinco y con sesenta y cuatro cero coma veinticinco. En el dibujo, desvío por raíz de B, casi igual para los cuatro. La consecuencia práctica: para reducir el ruido a la mitad hay que cuadruplicar el minibatch.

Parte 3 · SGD sobre el cuenco

El mismo cuenco del episodio uno, ahora con gradientes de minibatch.

SGD a mano

La función es descenso por gradiente con una sola diferencia. En cada paso sorteamos B ejemplos. Y el paso usa el gradiente de esos B. El resto es igual: learning rate cero coma uno, cuarenta pasos desde menos dos coma cinco, uno coma cinco. Tres corridas: con los mil ejemplos, que es el descenso del episodio uno, con diez y con uno. Medimos la distancia al mínimo en los últimos diez pasos. Con todos, cero coma cero cero uno: llegó. Con diez, cero coma uno cero uno. Con uno, cero coma dos cuatro. Llegan cerca, pero no se quedan quietos.

def sgd(B, lr=0.1, pasos=40, semilla=1):
    rng = np.random.default_rng(semilla)
    p = np.array([-2.5, 1.5]); camino = [p.copy()]
    for _ in range(pasos):
        idx = rng.choice(1000, B, replace=False)   # un minibatch al azar
        p = p - lr * grad(p, idx)                  # el paso de siempre
        camino.append(p.copy())
    return np.array(camino)

caminos = {B: sgd(B) for B in [1000, 10, 1]}
for B, c in caminos.items():
    lejos = np.linalg.norm(c[-10:], axis=1).mean()
    print(f"B = {B:4}: últimos 10 pasos, a {lejos:.3f} del mínimo")

B = 1000: últimos 10 pasos, a 0.001 del mínimo
B =   10: últimos 10 pasos, a 0.101 del mínimo
B =    1: últimos 10 pasos, a 0.240 del mínimo

Episodio uno contra SGD

Los tres caminos, cuarenta pasos cada uno. Con los mil ejemplos, la recta del episodio uno: derecho al centro, cada vez más despacio. Con diez, casi la misma recta al principio, con un temblor que se nota cuando se acerca. Con uno, el camino serpentea desde el arranque, y al llegar da vueltas alrededor de la estrella sin quedarse. Lejos del mínimo, el gradiente verdadero es grande y el ruido es chico en comparación. Cerca, el verdadero se achica y queda casi solo el ruido.

El temblor del final

el paso de SGD, separado en señal y ruido

cerca del mínimo solo queda el ruido

se achica con η más chico o B más grande

El paso de SGD, separado en dos: eta por el gradiente verdadero, y eta por el ruido. En el dibujo, la coordenada x del camino con un ejemplo: baja por la parábola, pero no en línea recta. Cerca del mínimo, el gradiente verdadero es casi cero, y el paso queda en eta por xi: un empujón al azar. Por eso el punto no se detiene. El símbolo que parece un pez es proporcional a. El tamaño del temblor crece con eta y con sigma sobre raíz de B. Hay dos perillas: bajar el learning rate, o agrandar el minibatch. En el dibujo, lo que mide el próximo snippet.

Dos perillas contra el temblor

Tres configuraciones: un ejemplo con learning rate cero coma uno, un ejemplo con cero coma cero dos, y dieciséis ejemplos con cero coma uno. Dos mil pasos cada una, así ya llegaron. Y medimos la distancia al mínimo en la segunda mitad, cuando el punto solo tiembla. Con un ejemplo y cero coma uno, cero coma cuatro cero nueve. Con un learning rate cinco veces menor, cero coma uno seis seis. Con dieciséis ejemplos, cero coma uno cero cinco. Las dos perillas funcionan. Bajar el learning rate hacia el final del entrenamiento es justo lo que hacen los schedules del episodio ocho.

for B, lr in [(1, 0.1), (1, 0.02), (16, 0.1)]:
    c = sgd(B, lr, pasos=2000)
    cola = np.linalg.norm(c[1000:], axis=1)    # ya llegó: medimos el temblor
    print(f"B = {B:2}, lr = {lr:4}: distancia media al mínimo {cola.mean():.3f}")

B =  1, lr =  0.1: distancia media al mínimo 0.409
B =  1, lr = 0.02: distancia media al mínimo 0.166
B = 16, lr =  0.1: distancia media al mínimo 0.105

Parte 4 · Minibatches de verdad

Si un gradiente con ruido es peor, ¿por qué todos usan minibatches? Porque a igual cantidad de datos, dan muchos más pasos.

Diez epochs, tres tamaños de minibatch

Ahora con datos reales: los dígitos escritos a mano de scikit-learn, como en Deep learning a secas, episodio siete. El modelo más simple: una capa lineal de sesenta y cuatro píxeles a diez clases, con entropía cruzada. Es una regresión logística. Cada epoch recorre las mil doscientas imágenes de train en minibatches de B. Contamos los pasos. Tres tamaños con el mismo learning rate, cero coma uno, y diez epochs: todo el train de una, treinta y dos, y uno. Con todo el train, diez pasos, pérdida dos coma uno uno cuatro y treinta y siete coma siete por ciento en validación. Con treinta y dos, trescientos ochenta pasos y noventa y tres por ciento. Con uno, doce mil pasos y noventa y cinco coma ocho.

def entrenar(B, lr=0.1, epochs=10):
    torch.manual_seed(0)
    modelo = nn.Linear(64, 10)                 # regresión logística
    opt = torch.optim.SGD(modelo.parameters(), lr=lr)
    hist, pasos = [], 0
    for epoch in range(epochs):
        orden = train[torch.randperm(len(train))]
        for i in range(0, len(orden), B):      # recorrer en minibatches
            b = orden[i:i + B]
            opt.zero_grad()
            F.cross_entropy(modelo(X[b]), y[b]).backward()
            opt.step(); pasos += 1
        with torch.no_grad():
            hist.append(F.cross_entropy(modelo(X[train]), y[train]).item())
    return modelo, hist, pasos

for B in [1200, 32, 1]:
    modelo, hist, pasos = entrenar(B)
    acc = (modelo(X[val]).argmax(1) == y[val]).float().mean().item()
    print(f"B = {B:4}: {pasos:5} pasos  pérdida {hist[-1]:.3f}  acc val {acc:.3f}")

B = 1200:    10 pasos  pérdida 2.114  acc val 0.377
B =   32:   380 pasos  pérdida 0.475  acc val 0.930
B =    1: 12000 pasos  pérdida 0.077  acc val 0.958

La pérdida por epoch

La pérdida de train al final de cada epoch. Las tres corridas vieron exactamente las mismas imágenes, la misma cantidad de veces. Con todo el train, la pérdida apenas baja: diez pasos precisos, pero solo diez. Con treinta y dos, baja rápido: treinta y ocho pasos por epoch. Con uno, ya en la primera epoch está más abajo que las otras al final. En la práctica se elige algo intermedio, de decenas a miles, porque una GPU, graphics processing unit, procesa un minibatch entero en paralelo casi al mismo costo que un ejemplo.

El gradiente del batch es el promedio

Comprobemos la línea del principio con PyTorch. La entropía cruzada de un minibatch de ocho ya devuelve el promedio de las ocho pérdidas. Un solo backward da el gradiente del minibatch. Al lado, ocho backward, uno por imagen, sumando los gradientes. Y comparamos el del minibatch con la suma dividida por ocho. El gradiente de los pesos es de diez por sesenta y cuatro, una fila por clase. La diferencia máxima es de una centésima de millonésima: redondeo. Por eso, para SGD, un minibatch es un forward y un backward como cualquier otro: el promedio lo hace la pérdida.

Cada línea, su término:

• perdida = F.cross_entropy(modelo(X[b]), y[b]) → $L_{\mathcal{B}} = \tfrac{1}{8}\sum_{i \in \mathcal{B}} \ell_i$ — cross_entropy promedia por defecto
• perdida.backward() → $g_B = \nabla L_{\mathcal{B}}$ — un backward para los 8
• suma += modelo.weight.grad → $\textstyle\sum_{i} \nabla \ell_i$ — ocho backward, uno por ejemplo
• dif = (g_batch - suma / 8) → $g_B = \tfrac{1}{8}\sum_i \nabla \ell_i$ — la línea dos del principio

torch.manual_seed(0)
modelo = nn.Linear(64, 10)
b = torch.arange(8)                                 # un minibatch de 8

perdida = F.cross_entropy(modelo(X[b]), y[b])       # promedio de 8
perdida.backward()
g_batch = modelo.weight.grad.clone()

suma = torch.zeros_like(g_batch)
for i in b:                                         # uno por vez
    modelo.zero_grad()
    F.cross_entropy(modelo(X[i:i + 1]), y[i:i + 1]).backward()
    suma += modelo.weight.grad
dif = (g_batch - suma / 8).abs().max().item()
print(f"forma del gradiente: {tuple(g_batch.shape)}")
print(f"diferencia máxima: {dif:.1e}")

forma del gradiente: (10, 64)
diferencia máxima: 1.5e-08

Goyal et al., 2017 — Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour · arXiv:1706.02677

8192 imágenes · 256 GPUs · 1 hora · «hyper-parameter-free linear scaling rule for adjusting learning rates as a function of minibatch size»

¿Y si querés minibatches enormes para usar muchas GPU? En dos mil diecisiete, un equipo de Facebook entrenó una ResNet cincuenta sobre ImageNet con minibatches de ocho mil ciento noventa y dos imágenes, repartidas en doscientas cincuenta y seis GPU, en una hora, sin perder accuracy. La regla que usaron es lineal: si el minibatch se multiplica por k, el learning rate también. Tiene sentido con lo que vimos: más ejemplos, menos ruido, y se puede dar un paso más largo. La acompañaron con un warmup al principio, que es tema del episodio ocho.

Lo que hay que llevarse

• El gradiente de un minibatch no tiene sesgo
• Su ruido baja como $1/\sqrt{B}$
• Cerca del mínimo, SGD tiembla: bajar $\eta$ o subir $B$
• A igual cantidad de datos, más pasos ganan

Un minibatch elegido al azar da, en promedio, el gradiente verdadero. Su error baja como uno sobre raíz de B: cuadruplicar el minibatch divide el ruido por dos. Lejos del mínimo el ruido no molesta; cerca, el punto tiembla, y se calma bajando el learning rate o agrandando el minibatch. Y lo que justifica todo: con los mismos datos, muchos pasos con ruido le ganan a pocos pasos exactos.

Resumiendo

• $\nabla L$ es el promedio de los $\nabla \ell_i$
• SGD: $g_B$ promedia solo $B$ ejemplos al azar
• $g_B = \nabla L + \xi$, con $\mathbb{E}[\xi] = 0$ y desvío $\sigma/\sqrt{B}$
• Cuenco: con B = 1 llega a 0,24 y tiembla
• Dígitos, 10 epochs: B = 1200 → 37,7 %; B = 32 → 93 %
• Próximo: momentum, promediar los pasos

Papers del episodio

• Goyal et al., 2017 — Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour