Normalización

Deep learning a secas. Episodio nueve: normalización. Una red de seis capas, entrenada cinco epochs, se queda trabada en ocho por ciento de accuracy. La misma red con una línea más por capa llega a noventa y seis. Hoy, qué hace esa línea.

Normalizar es restar y dividir

• Restar la media: centrar en 0
• Dividir por el desvío: escala 1
• BatchNorm promedia cada feature sobre el batch
• LayerNorm promedia cada ejemplo sobre sus features

Normalizar siempre es la misma cuenta, en dos pasos. Restar la media, así los números quedan centrados en cero. Dividir por el desvío estándar, así quedan con escala uno. Hacerlo antes de cada activación evita que los números crezcan o se achiquen capa tras capa. La pregunta es sobre qué números se calcula esa media. BatchNorm la calcula para cada feature, sobre los ejemplos del batch. LayerNorm, para cada ejemplo, sobre sus propias features. Todo el episodio es esa diferencia.

Parte 1 · Un batch, como matriz

Cuatro ejemplos, cuatro features. Las mismas dieciséis cuentas, agrupadas de dos maneras.

El batch: filas = ejemplos, columnas = features

El punto de partida. Una matriz de cuatro por cuatro. Cada fila es un ejemplo del batch, y cada columna una feature: una neurona de la capa, o una entrada. Las features están en escalas muy distintas: la primera va de uno a seis, la tercera de diez a cuarenta, la cuarta ronda el cero. Así, el learning rate que le sirve a una feature no le sirve a otra.

Las dos, a mano

Epsilon, un número chiquito que se suma para no dividir nunca por cero. BatchNorm: media y varianza con dim igual a cero, que recorre las filas. Una media por columna. LayerNorm: lo mismo con dim igual a uno, que recorre las columnas. Una media por fila. keepdim deja el resultado como columna, para poder restarlo a cada fila. Y medimos las medias después de normalizar. Las medias por columna son tres, cinco, veinticinco y cero. Por fila, tres coma siete cinco, nueve coma siete cinco, seis y trece coma cinco. Después, cada columna de BatchNorm y cada fila de LayerNorm tiene media cero.

X = torch.tensor([[1.0, 4.0, 10.0,  0.0],     # 4 ejemplos (filas)
                  [2.0, 6.0, 30.0,  1.0],     # × 4 features (columnas)
                  [3.0, 2.0, 20.0, -1.0],
                  [6.0, 8.0, 40.0,  0.0]])
eps = 1e-5

# BatchNorm: por COLUMNA, sobre el batch (dim=0)
mu_b, var_b = X.mean(0), X.var(0, unbiased=False)
bn = (X - mu_b) / torch.sqrt(var_b + eps)

# LayerNorm: por FILA, sobre features (dim=1)
mu_l = X.mean(1, keepdim=True)
var_l = X.var(1, unbiased=False, keepdim=True)
ln = (X - mu_l) / torch.sqrt(var_l + eps)
print("μ por columna:", mu_b.numpy())
print("μ por fila   :", mu_l.squeeze().numpy())
print("medias BatchNorm:", bn.mean(0).abs().numpy().round(2))
print("medias LayerNorm:", ln.mean(1).abs().numpy().round(2))

μ por columna: [ 3.  5. 25.  0.]
μ por fila   : [ 3.75  9.75  6.   13.5 ]
medias BatchNorm: [0. 0. 0. 0.]
medias LayerNorm: [0. 0. 0. 0.]

BatchNorm, columna 1

la media de la columna j, sobre los B ejemplos

la varianza: el promedio de las distancias al cuadrado

ε = 0,00001: para no dividir por cero

La cuenta de BatchNorm para la primera columna. Mu sub j, la letra griega mu, es la media de la columna j: sumar los B valores del batch y dividir por B. Uno más dos más tres más seis, doce, sobre cuatro: tres. En el dibujo, las cuatro barras y su promedio. Sigma sub j al cuadrado es la varianza: cada valor menos la media, al cuadrado, y promediado. Cuatro, uno, cero y nueve: tres coma cinco. Su raíz, sigma, es el desvío estándar. Y x sombrero: restar la media y dividir por la raíz de la varianza más epsilon. El seis queda en uno coma sesenta. En el dibujo, la normalización es una neurona con peso uno sobre sigma y sesgo menos mu sobre sigma.

Después de BatchNorm

El batch después de BatchNorm. Mirá las columnas: ahora todas tienen media cero y varianza uno. La tercera, que iba de diez a cuarenta, quedó entre menos uno coma tres cuatro y uno coma tres cuatro, igual que la segunda. En rojo, los negativos: en cada columna hay de los dos colores. Cada feature quedó en la misma escala. Ojo con lo que esto implica: el número que recibe un ejemplo depende de los otros tres del batch.

LayerNorm, fila 1

la media de la fila i, sobre sus d features

la misma varianza, ahora a lo largo de la fila

el 10 de la fila 1

La misma cuenta, girada noventa grados. Ahora mu sub i es la media de la fila i, sobre sus d features. d es la cantidad de features, cuatro. Uno más cuatro más diez más cero: quince, sobre cuatro, tres coma setenta y cinco. En el dibujo, las barras son las features de un solo ejemplo. La varianza, a lo largo de la fila: quince coma diecinueve. Y el diez de esa fila queda en uno coma sesenta. Fijate que nada de esta cuenta mira a los otros ejemplos: cada fila se normaliza sola.

Después de LayerNorm

El batch después de LayerNorm, fila por fila. El primer ejemplo. El segundo. El tercero. Y el cuarto. Cada fila tiene media cero y varianza uno. Pero mirá las columnas: la tercera feature quedó en uno coma seis o uno coma siete en las cuatro filas, porque en cada ejemplo es la más grande. LayerNorm compara las features de un ejemplo entre sí, no los ejemplos entre ellos. Por eso tiene sentido adentro de la red, donde las features son neuronas de la misma capa.

Parte 2 · En PyTorch, con γ y β

Las capas de PyTorch hacen la misma cuenta, y agregan dos parámetros que se aprenden.

nn.BatchNorm1d y nn.LayerNorm

nn punto BatchNorm1d de cuatro: cuatro features. nn punto LayerNorm de cuatro: filas de cuatro. Las dos tienen un peso, gamma, y un sesgo, beta, por feature. Comparamos contra nuestra cuenta a mano. Y cambiamos gamma a dos y beta a cero coma cinco. gamma arranca en uno y beta en cero: al principio no hacen nada. Las dos capas coinciden con la cuenta a mano. Y con gamma dos y beta un medio, la primera columna da menos uno coma seis cuatro, menos cero coma cinco siete, cero coma cinco y tres coma siete uno: dos por x sombrero más un medio.

bn = nn.BatchNorm1d(4)      # 4 features: un γ y un β por columna
ln = nn.LayerNorm(4)        # filas de 4: un γ y un β por posición
print("γ:", bn.weight.tolist())
print("β:", bn.bias.tolist())

print("BatchNorm:", torch.allclose(bn(X), bn_mano, atol=1e-5))
print("LayerNorm:", torch.allclose(ln(X), ln_mano, atol=1e-5))

with torch.no_grad():
    bn.weight.fill_(2.0); bn.bias.fill_(0.5)   # γ = 2, β = 0,5
print("γ=2, β=0,5:", bn(X)[:, 0].detach().numpy().round(2))
print("2·x̂ + 0,5:", (2 * bn_mano[:, 0] + 0.5).numpy().round(2))

γ: [1.0, 1.0, 1.0, 1.0]
β: [0.0, 0.0, 0.0, 0.0]
BatchNorm: True
LayerNorm: True
γ=2, β=0,5: [-1.64 -0.57  0.5   3.71]
2·x̂ + 0,5: [-1.64 -0.57  0.5   3.71]

γ y β: la red puede deshacerlo

γ, β: un par por feature, aprendidos con backprop

si le conviene, vuelve a la escala original

¿Para qué gamma y beta? Después de normalizar, cada feature se multiplica por su gamma y se le suma su beta. Son parámetros como cualquier peso: los ajusta backprop. En el dibujo, el primer número normalizado, menos uno coma cero siete, por dos, más cero coma cinco: menos uno coma sesenta y cuatro, como imprimió el snippet. Y la razón de fondo: si gamma llega a valer el desvío y beta la media, la normalización se deshace exacta. En el dibujo, con gamma uno coma ocho siete y beta tres, vuelve el uno original. Normalizar no le quita nada a la red: le da un punto de partida.

Parte 3 · Train y eval no son lo mismo

BatchNorm necesita un batch para calcular la media. ¿Y cuando hay que predecir un solo ejemplo?

Las estadísticas de BatchNorm

En modo train, cada forward actualiza un promedio móvil de las medias, running mean, y otro de las varianzas. Comparamos la varianza con la cuenta. En modo eval, BatchNorm deja de mirar el batch y usa esos promedios. Y probamos un batch de un solo ejemplo. Después de un forward, la running mean es cero coma uno por la media del batch: arranca en cero y se mueve un diez por ciento. La varianza, uno coma tres seis siete: cero coma nueve por uno, más cero coma uno por cuatro coma seis seis siete. En eval, la misma fila da otros números. Y con un ejemplo solo, BatchNorm tira ValueError: expected more than one value per channel. LayerNorm no tiene problema.

bn = nn.BatchNorm1d(4)                # momentum=0.1 por defecto
bn(X)                                 # un forward en modo train
print("running_mean:", bn.running_mean.numpy().round(3))
print("0,1 · μ_B   :", (0.1 * X.mean(0)).numpy().round(3))
print("running_var[0]:", round(bn.running_var[0].item(), 3))
print("var no sesgada:", round(X[:, 0].var().item(), 3))

bn.eval()                             # en eval usa las running
print("eval, fila 0:", bn(X)[0].detach().numpy().round(2))

uno = X[:1]                           # un batch de un solo ejemplo
bn.train()
try:
    bn(uno)
except ValueError:
    print("BatchNorm: ValueError")
print("LayerNorm:", nn.LayerNorm(4)(uno)[0].detach().numpy().round(2))

running_mean: [0.3 0.5 2.5 0. ]
0,1 · μ_B   : [0.3 0.5 2.5 0. ]
running_var[0]: 1.367
var no sesgada: 4.667
eval, fila 0: [0.6  2.8  1.79 0.  ]
BatchNorm: ValueError
LayerNorm: [-0.71  0.06  1.6  -0.96]

El promedio móvil

momentum = 0,1 en PyTorch: el peso del batch nuevo

ojo: acá usa la varianza no sesgada, B/(B−1)

La running mean se actualiza así: cero coma nueve de lo que tenía, más cero coma uno de la media del batch. PyTorch llama momentum a ese cero coma uno, aunque no es el del optimizador. En el dibujo, las dos columnas pesadas que se suman. La varianza tiene un detalle: para normalizar el batch usa la varianza que divide por B, pero para el promedio móvil usa la que divide por B menos uno, que es la estimación sin sesgo. Tres coma cinco por cuatro tercios, cuatro coma seis seis siete. Con la varianza inicial en uno, da uno coma tres seis siete, lo que imprimió el snippet.

Ioffe y Szegedy, 2015 — Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift · arXiv:1502.03167

14× menos pasos · «Batch Normalization achieves the same accuracy with 14 times fewer training steps»

BatchNorm es de dos mil quince, de Sergey Ioffe y Christian Szegedy, en Google. Lo motivaron con lo que llamaron internal covariate shift: que la distribución de las entradas de cada capa cambia durante el entrenamiento, porque cambian los pesos de las capas anteriores. En su modelo de clasificación de imágenes, llegaron a la misma accuracy con catorce veces menos pasos. Esa explicación del porqué se discutió mucho después, y hoy no hay consenso; que acelera el entrenamiento, en cambio, se vio en muchísimos trabajos.

Parte 4 · ¿Ayuda de verdad?

Seis hidden layers, dígitos, SGD, stochastic gradient descent, sin momentum, cinco epochs. Lo único que cambia es la normalización.

La misma red, con y sin normalización

Una red de seis hidden layers de sesenta y cuatro neuronas. Si le pasamos una normalización, va entre cada capa lineal y su ReLU, rectified linear unit. Probamos sin nada, con BatchNorm y con LayerNorm. Cada una con la misma semilla, cinco epochs de SGD con learning rate cero coma uno. Sin normalizar, la pérdida de validación queda en dos coma tres cero cinco, la de no saber nada, y la accuracy en ocho coma siete por ciento: la red no aprendió. Con BatchNorm, noventa y seis por ciento. Con LayerNorm, noventa y tres.

def red(norm=None, capas=6, ancho=64):
    torch.manual_seed(0)
    bloques, d = [], 64
    for _ in range(capas):                   # 6 hidden layers
        bloques.append(nn.Linear(d, ancho))
        if norm is not None:
            bloques.append(norm(ancho))      # normalizar antes de ReLU
        bloques.append(nn.ReLU())
        d = ancho
    return nn.Sequential(*bloques, nn.Linear(d, 10))

for nombre, norm in [("sin norma", None), ("BatchNorm", nn.BatchNorm1d),
                     ("LayerNorm", nn.LayerNorm)]:
    lv, acc = entrenar(red(norm))            # 5 epochs, SGD lr=0,1
    print(f"{nombre:10} val {lv:.3f}  acc {acc:.3f}")

sin norma  val 2.305  acc 0.087
BatchNorm  val 0.155  acc 0.960
LayerNorm  val 0.268  acc 0.930

Accuracy en validación, 5 epochs

En barras. Sin normalizar, ocho coma siete por ciento: peor que adivinar uno de diez. BatchNorm, noventa y seis. LayerNorm, noventa y tres. ¿Por qué la red sin normalizar ni siquiera arranca? Con la inicialización por defecto, la señal se achica un poco en cada capa, y después de seis capas casi no llega nada, ni para adelante ni para atrás. Medir eso capa por capa es el episodio diez. La normalización lo tapa: cada capa vuelve a escala uno.

• sin norma: 8,7 %
• BatchNorm: 96,0 %
• LayerNorm: 93,0 %

Ba, Kiros y Hinton, 2016 — Layer Normalization · arXiv:1607.06450

un solo ejemplo · «Layer normalization performs exactly the same computation at training and test times.»

LayerNorm llegó un año después, de Jimmy Ba, Jamie Kiros y Geoffrey Hinton. Partieron justamente del problema que vimos: el efecto de BatchNorm depende del tamaño del batch. Su propuesta es transponer BatchNorm: calcular la media y la varianza sobre todas las neuronas de una capa, en un solo ejemplo de entrenamiento. Así, hace exactamente la misma cuenta en train y en eval. Es la normalización que adoptaron los transformers.

La prima de LayerNorm: RMSNorm

Un hilo hacia Alineación. En el episodio uno de esa serie, el bloque de Gemma normalizaba con layer norm antes de la atención, y en el siete usamos la RMSNorm, root mean square norm, antes del unembedding. Una LayerNorm sin gamma ni beta. Y una RMSNorm. RMSNorm divide por la raíz cuadrática media: la raíz del promedio de los cuadrados. Comparamos las medias. LayerNorm deja la fila centrada: media cero. RMSNorm no resta la media: solo divide por cinco coma cuatro tres, y la media queda en cero coma siete cuatro. Es una cuenta menos, y Gemma, como muchos modelos recientes, la usa.

x = torch.tensor([1.0, 4.0, 10.0, 1.0])     # una fila, 4 features

ln = nn.LayerNorm(4, elementwise_affine=False)
rms = nn.RMSNorm(4, elementwise_affine=False)
print("LayerNorm:", ln(x).numpy().round(2))
print("RMSNorm  :", rms(x).numpy().round(2))

raiz = x.pow(2).mean().sqrt()                # raíz cuadrática media
print("x / raíz :", (x / raiz).numpy().round(2))
print("raíz =", round(raiz.item(), 2))
print("medias:", round(ln(x).mean().item(), 2), round(rms(x).mean().item(), 2))

LayerNorm: [-0.82  0.    1.63 -0.82]
RMSNorm  : [0.18 0.74 1.84 0.18]
x / raíz : [0.18 0.74 1.84 0.18]
raíz = 5.43
medias: 0.0 0.74

BatchNorm y LayerNorm, lado a lado

Para cerrar, lado a lado. BatchNorm promedia cada feature sobre el batch, así que el resultado de un ejemplo depende de con quién vino; necesita estadísticas móviles para eval, y no funciona con un batch de uno. Es la clásica de redes convolucionales, las de imágenes, y redes densas. LayerNorm promedia cada ejemplo sobre sus features: cada uno se normaliza solo, con la misma cuenta en train y en eval. Es la que usan los transformers: entre otras cosas, no le importa el tamaño del batch.

Resumiendo

• Normalizar: $\hat x = (x - \mu)/\sqrt{\sigma^2 + \varepsilon}$
• BatchNorm: $\mu$ por columna (dim=0); LayerNorm: por fila (dim=1)
• $\gamma$, $\beta$ aprendidos: la red puede deshacerla
• BatchNorm en eval usa running mean y running var
• 6 capas, 5 epochs: 8,7 % sin norma, 96 % con BatchNorm
• Próximo: profundidad, inicialización y residuales

Papers del episodio

• Ioffe y Szegedy, 2015 — Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• Ba, Kiros y Hinton, 2016 — Layer Normalization