Normas y normalización

Álgebra lineal para LLMs. Episodio cuatro: normas y normalización. Medir el largo de un vector parece trivial, pero hay varias maneras, y cada una se usa para algo distinto. Hoy, las tres normas que aparecen en Alineación, y las dos normalizaciones que tiene adentro cada capa de un modelo.

De dónde venimos

Alineación, ep. 9: la pérdida de un SAE

Alineación, ep. 1: normalizar antes de leer

Alineación, ep. 7: el logit lens

En Alineación, episodio nueve, la pérdida de un SAE, sparse autoencoder: reconstruir bien x, más lambda por la norma uno de f. Las dos barras dobles son una norma, el largo de un vector; el subíndice dice cuál. En el dibujo, un f con pocas coordenadas distintas de cero. En el episodio uno, cada capa empieza con LN, layer norm. En el dibujo, el primer paso: el promedio de las coordenadas de x. Y en el siete, el logit lens: antes de multiplicar por W sub U, RMSNorm, root mean square norm. En el dibujo, normalizar y después puntuar. Hoy, qué hace cada una.

Parte 1 · Tres maneras de medir

L2, L1 y L0: tres respuestas distintas a la pregunta «qué tan grande es este vector».

Las tres normas de x = (4, 3, 0)

Un vector de tres coordenadas: cuatro, tres y cero. La norma L2 es la de Pitágoras, la del episodio uno: raíz de la suma de los cuadrados. La L1 suma los valores absolutos, sin elevar a nada. Y la L0 ni siquiera mira cuánto valen: cuenta cuántas coordenadas no son cero. NumPy las calcula con norm, cambiando el segundo argumento. Y la L2 al cuadrado es x punto x: el producto punto de x consigo mismo. L2, cinco. L1, siete. L0, dos. Las mismas con NumPy. Y la L2 al cuadrado, veinticinco.

x = np.array([4.0, 3.0, 0.0])

l2 = np.sqrt(np.sum(x ** 2))          # Pitágoras
l1 = np.sum(np.abs(x))                # sumar valores absolutos
l0 = np.count_nonzero(x)              # contar los que no son cero
print("L2 =", l2, "  L1 =", l1, "  L0 =", l0)

print("np.linalg.norm:", np.linalg.norm(x, 2),
      np.linalg.norm(x, 1), np.linalg.norm(x, 0))
print("L2 al cuadrado =", x @ x)      # producto punto consigo mismo

L2 = 5.0   L1 = 7.0   L0 = 2
np.linalg.norm: 5.0 7.0 2.0
L2 al cuadrado = 25.0

En línea recta o por las calles

En el plano, con las dos primeras coordenadas. La L2 es la distancia en línea recta, la del vuelo de un pájaro: cinco. La L1 es la distancia por las calles de una ciudad en cuadrícula: primero cuatro cuadras hacia un lado, después tres hacia arriba. Siete. Por eso a la L1 también se le dice distancia Manhattan. Y la L0 no es una distancia: es un conteo de coordenadas activas.

Las tres normas, en fórmula

L2: raíz de la suma de cuadrados

L1: suma de valores absolutos

L0: cuántas coordenadas no son cero

La L2: raíz de la suma, sobre i, de x sub i al cuadrado. Dieciséis más nueve más cero, veinticinco; la raíz, cinco. En el dibujo, las tres barras son los cuadrados. La L1: la suma de los valores absolutos. Las barras verticales simples son el valor absoluto, el número sin signo. Cuatro más tres más cero: siete. La L0: el numeral es «cantidad de», y las llaves, un conjunto: cuántos i cumplen que x sub i es distinto de cero. En el dibujo, un uno por cada coordenada activa. Suma dos.

Parte 2 · Por qué L1 deja ceros

Un SAE quiere pocas features activas, o sea, L0 chica. Pero castiga L1. ¿Por qué esa y no L2?

Castigar con L2² o con L1

Un experimento chico. Cinco números a, lo que cada peso w querría valer si nadie lo frenara: tres, cero coma cuatro, menos cero coma dos, uno coma cinco y cero coma uno. Una función que busca w con gradientes, partiendo de cero. La pérdida tiene dos partes: parecerse a a, más lambda por un castigo. Primero, el castigo es la L2 al cuadrado. Después, la L1. Con L2 al cuadrado, todos los pesos se achican a la mitad, y ninguno llega a cero. Con L1, los grandes bajan una cantidad fija, de tres a dos coma cinco, y los tres chicos quedan exactamente en cero. Tres ceros contra ninguno.

a = torch.tensor([3.0, 0.4, -0.2, 1.5, 0.1])   # lo que «querría» valer w
lam = 1.0

def ajustar(castigo):
    w = torch.zeros(5, requires_grad=True)
    opt = torch.optim.SGD([w], lr=0.002)
    for _ in range(5000):
        perdida = ((w - a) ** 2).sum() + lam * castigo(w)
        opt.zero_grad(); perdida.backward(); opt.step()
    return w.detach()

w2 = ajustar(lambda w: (w ** 2).sum())     # castigo L2²
w1 = ajustar(lambda w: w.abs().sum())      # castigo L1
print("con L2²:", w2.numpy().round(2))
print("con L1 :", w1.numpy().round(2))
ceros = lambda w: int((w.abs() < 0.01).sum())
print("ceros:", ceros(w2), "con L2²,", ceros(w1), "con L1")

con L2²: [ 1.5   0.2  -0.1   0.75  0.05]
con L1 : [ 2.5  0.  -0.   1.  -0. ]
ceros: 0 con L2², 3 con L1

El castigo cerca de cero

El porqué está en la forma del castigo. w al cuadrado es una parábola: cerca de cero es casi plana. Un peso de cero coma uno paga cero coma cero uno: casi nada. Bajarlo más no gana nada, así que el gradiente deja de empujar. El valor absoluto es una ve: tiene la misma pendiente en todos lados. Un peso de cero coma uno paga cero coma uno, y el empujón hacia cero sigue igual de fuerte hasta que llega. Por eso L1 produce ceros de verdad.

La pendiente de cada castigo

se apaga cerca de cero

siempre el mismo empujón

el SAE: reconstruir, con pocas features

La derivada, la pendiente del episodio uno de Deep learning, escrita con la d redondeada: cuánto cambia el castigo si movés w. Para lambda w al cuadrado es dos lambda w: proporcional a w. En el dibujo, la recta tangente: empinada en uno, casi horizontal en cero coma uno. Para lambda por valor absoluto de w es lambda por el signo de w: más uno o menos uno. En el dibujo, la ve: pendiente menos uno a la izquierda y más uno a la derecha, igual en uno que en cero coma uno. Por eso la pérdida del SAE de Alineación castiga la L1 de las features: la L0 es la que queremos, pero no tiene pendiente, y la L1 es la que más se le parece y sí la tiene.

Parte 3 · Dividir por la norma

Antes de normalizar capas, normalizar un vector: dejarle la dirección y sacarle el largo.

Vectores unitarios

Dividir x por su norma da un vector que apunta igual y mide uno: un vector unitario, x sombrero. El coseno del episodio uno. Es exactamente el producto punto de los dos unitarios: el coseno es el producto punto después de sacarle el largo a cada uno. Y si multiplicamos x por diez, el producto punto se multiplica por diez. El coseno, no. x sombrero, cero coma ocho, cero coma seis, de largo uno. Coseno y producto de unitarios, cero coma ochocientos noventa y cuatro, los dos. Diez x punto y da cien, contra diez. El coseno sigue en cero coma ochocientos noventa y cuatro.

x = np.array([4.0, 3.0])
x_hat = x / np.linalg.norm(x)              # mismo sentido, largo 1
print("x̂ =", x_hat, "  ‖x̂‖ =", np.linalg.norm(x_hat))

y = np.array([1.0, 2.0])
y_hat = y / np.linalg.norm(y)
cos = (x @ y) / (np.linalg.norm(x) * np.linalg.norm(y))
print("coseno   =", round(cos, 3))
print("x̂ · ŷ    =", round(x_hat @ y_hat, 3))  # el mismo número

print("(10x)·y  =", (10 * x) @ y, "  vs x·y =", x @ y)
z = 10 * x
print("cos(10x, y) =", round((z / np.linalg.norm(z)) @ y_hat, 3))

x̂ = [0.8 0.6]   ‖x̂‖ = 1.0
coseno   = 0.894
x̂ · ŷ    = 0.894
(10x)·y  = 100.0   vs x·y = 10.0
cos(10x, y) = 0.894

Parte 4 · LayerNorm

Adentro de un modelo, los vectores del residual stream crecen capa a capa. Antes de leerlos, cada bloque los normaliza.

LayerNorm a mano y con PyTorch

LayerNorm normaliza cada token por separado, mirando sus d coordenadas. Primero el promedio, mu: la suma de las coordenadas, dividida por d. Después la varianza, sigma al cuadrado: el promedio de la distancia al cuadrado a mu. Restamos mu y dividimos por la raíz de la varianza; épsilon es un número chiquito para no dividir por cero. Y nn punto LayerNorm, además, multiplica por gamma y suma beta: dos vectores de d números que se aprenden, y arrancan en uno y en cero. Mu, tres; sigma al cuadrado, tres coma cinco. A mano y con PyTorch: menos uno coma cero siete, menos cero coma cincuenta y tres, cero, uno coma seis. Media cero, desvío uno.

Cada línea, su término:

• mu = x.mean() → $\mu = \tfrac{1}{d}\textstyle\sum_i x_i$ — el promedio de las $d$ coordenadas
• var = ((x - mu) ** 2).mean() → $\sigma^2 = \tfrac{1}{d}\textstyle\sum_i (x_i - \mu)^2$ — la varianza
• a_mano = (x - mu) / torch.sqrt(var + 1e-5) → $\hat x_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$ — centrar y escalar
• ln = nn.LayerNorm(4) → $y_i = \gamma_i\, \hat x_i + \beta_i$ — $\gamma$ y $\beta$ se entrenan

x = torch.tensor([1.0, 2.0, 3.0, 6.0])     # un token, d = 4

mu = x.mean()                               # promedio
var = ((x - mu) ** 2).mean()                # varianza
a_mano = (x - mu) / torch.sqrt(var + 1e-5)
print("μ =", mu.item(), "  σ² =", var.item())
print("a mano   :", a_mano.numpy().round(3))

ln = nn.LayerNorm(4)                        # γ = 1, β = 0 al inicio
print("LayerNorm:", ln(x).detach().numpy().round(3))
print("γ:", ln.weight.data.numpy(), " β:", ln.bias.data.numpy())
print("media =", round(a_mano.mean().item(), 3),
      " desvío =", round(a_mano.std(unbiased=False).item(), 3))

μ = 3.0   σ² = 3.5
a mano   : [-1.069 -0.535  0.     1.604]
LayerNorm: [-1.069 -0.535  0.     1.604]
γ: [1. 1. 1. 1.]  β: [0. 0. 0. 0.]
media = 0.0  desvío = 1.0

LayerNorm con x = (1, 2, 3, 6)

el promedio

distancias a 3, al cuadrado

centrado y escalado

Las mismas cuentas, dibujadas. El promedio de uno, dos, tres y seis: doce sobre cuatro, tres. En el dibujo, cuatro barras y su promedio. Las distancias a tres son menos dos, menos uno, cero y tres. Al cuadrado: cuatro, uno, cero, nueve. Su promedio, tres coma cinco: la varianza. Su raíz, uno coma ochenta y siete: el desvío. Restar tres y dividir por uno coma ochenta y siete da el x normalizado. En el dibujo, las barras quedan a los dos lados del cero: la media es cero. El vector perdió su posición y su escala; le queda la forma: cuál coordenada es grande y cuál chica.

Parte 5 · RMSNorm, la de Gemma

Gemma no usa LayerNorm: usa una versión más simple, que no resta el promedio.

RMSNorm a mano y la de Gemma 2

RMS es root mean square: la raíz del promedio de los cuadrados. Fijate que es la norma L2 dividida por raíz de d. RMSNorm divide x por eso, y listo: no resta el promedio. La clase de Gemma dos, de transformers. Multiplica, elemento a elemento, por uno más w; ese círculo con punto es el producto elemento a elemento. Con w en cero, como al crearla, es dividir y nada más. Y leemos el w de la norma final de Gemma, la que viene antes de W sub U. RMS, tres coma cincuenta y cuatro. x dividido por eso, a mano y con la clase de Gemma: lo mismo. En la norma final entrenada, uno más w va de cero coma veintidós a diez coma cuarenta y cinco: algunas coordenadas se amplifican diez veces.

Cada línea, su término:

• rms = torch.sqrt((x ** 2).mean() + 1e-6) → $\mathrm{RMS}(x) = \sqrt{\tfrac{1}{d}\textstyle\sum_i x_i^2 + \epsilon}$ — raíz del promedio de los cuadrados
• print("x / RMS → $\hat x = x \,/\, \mathrm{RMS}(x)$ — sin restar nada
• norma = Gemma2RMSNorm(4) → $y = \hat x \odot (1 + w)$ — Gemma multiplica por $1 + w$
• w = f.get_tensor(nombre) → $w \in \mathbb{R}^{2304}$ — la norma final de Gemma, real

x = torch.tensor([1.0, 2.0, 3.0, 6.0])

rms = torch.sqrt((x ** 2).mean() + 1e-6)      # raíz del promedio de x²
print("RMS =", round(rms.item(), 3))
print("x / RMS      :", (x / rms).numpy().round(3))

norma = Gemma2RMSNorm(4)                      # la de Gemma 2
print("Gemma2RMSNorm:", norma(x).detach().numpy().round(3))
print("weight inicial:", norma.weight.data.numpy())   # usa 1 + w

nombre = "model.norm.weight"                   # la norma final de Gemma
with safe_open(hf_hub_download(REPO, mapa[nombre]), "pt") as f:
    w = f.get_tensor(nombre)
print("norma final:", tuple(w.shape), " 1 + w va de",
      round(1 + w.min().item(), 2), "a", round(1 + w.max().item(), 2))

RMS = 3.536
x / RMS      : [0.283 0.566 0.849 1.697]
Gemma2RMSNorm: [0.283 0.566 0.849 1.697]
weight inicial: [0. 0. 0. 0.]
norma final: (2304,)  1 + w va de 0.22 a 10.45

RMSNorm con x = (1, 2, 3, 6)

promedio de los cuadrados

la forma se mantiene; el promedio no es 0

todos los tokens al mismo largo

El promedio de los cuadrados: uno, cuatro, nueve y treinta y seis, sobre cuatro: doce coma cinco. La raíz, tres coma cincuenta y cuatro. En el dibujo, las cuatro barras y su promedio. Dividimos: cero coma veintiocho, cero coma cincuenta y siete, cero coma ochenta y cinco, uno coma setenta. Las cuatro siguen positivas: no se centró nada. En el dibujo, los cuadrados nuevos: su promedio es uno. Y como la L2 es raíz de d por el RMS, después de RMSNorm todo vector mide raíz de d. En Gemma, cuarenta y ocho. El embedding real de perro medía setenta y tres coma seis; normalizado, cuarenta y ocho, como vemos en el código que sigue.

Qué ignora cada una

Cada normalización ignora algo, a propósito. LayerNorm de x más diez es igual a LayerNorm de x: no le importa correr todo. Y tampoco escalar por cinco. RMSNorm de cinco x es igual a RMSNorm de x: tampoco le importa la escala. Pero sumarle diez sí lo cambia. Y el embedding real de perro, multiplicado por cuarenta y ocho como hace Gemma. True, True, True, y False. Y el embedding pasa de setenta y tres coma seis a cuarenta y ocho, raíz de dos mil trescientos cuatro.

ln, rms = nn.LayerNorm(4), Gemma2RMSNorm(4)
x = torch.tensor([1.0, 2.0, 3.0, 6.0])
igual = lambda a, b: torch.allclose(a, b, atol=1e-4)
print("LayerNorm(x + 10) == LayerNorm(x):", igual(ln(x + 10), ln(x)))
print("LayerNorm(5 x)    == LayerNorm(x):", igual(ln(5 * x), ln(x)))
print("RMSNorm(5 x)      == RMSNorm(x)  :", igual(rms(5 * x), rms(x)))
print("RMSNorm(x + 10)   == RMSNorm(x)  :", igual(rms(x + 10), rms(x)))

(i,) = tok(" perro", add_special_tokens=False).input_ids
x0 = W_E[i] * 2304 ** 0.5                     # embedding real de Gemma
y = Gemma2RMSNorm(2304)(x0)
print("‖x0‖ =", round(x0.norm().item(), 1), " → ‖RMSNorm(x0)‖ =",
      round(y.norm().item(), 1), " (√2304 = 48)")

LayerNorm(x + 10) == LayerNorm(x): True
LayerNorm(5 x)    == LayerNorm(x): True
RMSNorm(5 x)      == RMSNorm(x)  : True
RMSNorm(x + 10)   == RMSNorm(x)  : False
‖x0‖ = 73.6  → ‖RMSNorm(x0)‖ = 48.0  (√2304 = 48)

LayerNorm vs. RMSNorm

LayerNorm resta el promedio, divide por el desvío, y tiene dos vectores entrenables. Le dan igual los corrimientos y la escala. RMSNorm hace la mitad: no centra, divide por la RMS, y tiene un solo vector de escala. Le da igual la escala, pero no los corrimientos. Gemma y Llama usan esta. La apuesta es que centrar no hacía falta, y que sacarlo ahorra cuentas.

Zhang y Sennrich, 2019 — Root Mean Square Layer Normalization · arXiv:1910.07467

7 % – 64 % · «We hypothesize that re-centering invariance in LayerNorm is dispensable» … RMSNorm «reduces the running time by 7%~64% on different models».

La apuesta tiene paper. En dos mil diecinueve, Biao Zhang y Rico Sennrich propusieron RMSNorm con la hipótesis de que la invariancia a centrar de LayerNorm es prescindible, y conservaron solo la invariancia a escalar. Reportan rendimiento comparable y entre siete y sesenta y cuatro por ciento menos de tiempo de cómputo, según el modelo.

Resumiendo

• L2 = $\sqrt{\sum x_i^2}$, L1 = $\sum |x_i|$, L0 = cuántas no son cero
• L1 empuja siempre igual: deja ceros exactos (el $\lambda\lVert f\rVert_1$ del SAE)
• Coseno = producto punto de vectores unitarios $x / \lVert x\rVert$
• LayerNorm: $(x - \mu)/\sigma$, después $\gamma$ y $\beta$
• RMSNorm: $x / \mathrm{RMS}(x) \odot (1 + w)$; en Gemma, todo mide 48
• Próximo: proyecciones y bases ortogonales

Papers del episodio

• Zhang y Sennrich, 2019 — Root Mean Square Layer Normalization