Feed-forward, residuales y LayerNorm

Attention is all you need. Episodio seis: feed-forward, residuales y LayerNorm. Un bloque del transformer tiene más piezas que la atención. Hoy, las otras tres, una por una, y el bloque encoder entero en PyTorch, comparado contra el de la librería. Al final, un experimento: dónde va la normalización cambia si una red de doce capas aprende o no.

Un bloque encoder del paper

El bloque, de arriba abajo. Entra x: T tokens, cada uno un vector de quinientos doce números, d sub model. La multi-head attention del episodio tres, y después dos pasos nuevos: se suma la entrada, y se normaliza con LayerNorm. Después, la FFN, feed-forward network, con la misma suma y la misma normalización. Sale otra matriz de T por quinientos doce: misma forma que la entrada. Por eso se pueden apilar seis bloques, uno arriba del otro.

x (T × 512) → multi-head attn → x + · → LayerNorm → FFN → x + · → LayerNorm → salida (T × 512)

Parte 1 · La feed-forward network

La atención mezcla información entre tokens. La FFN procesa cada token por su cuenta.

Qué hace la FFN

• Dos capas lineales con una ReLU en el medio
• Sube de 512 a 2048 y vuelve a 512
• Cada token por separado, con los mismos pesos
• Es el MLP de Alineación, ep. 1

La FFN en cuatro ideas. Son dos capas lineales, y entre ellas una ReLU, rectified linear unit: deja pasar los positivos y pone cero en los negativos. La primera capa agranda cada vector de quinientos doce a dos mil cuarenta y ocho números; la segunda lo devuelve a quinientos doce. Se aplica a cada token por separado, siempre con los mismos pesos: ningún token mira a otro acá adentro. Es lo mismo que en Alineación, episodio uno, llamamos MLP, multilayer perceptron. Gemma usa una versión con compuerta, que vemos en el episodio nueve.

La FFN, con números chicos

la ecuación (2) del paper

2048 neuronas, una por columna de $W_1$

la ReLU apaga las neuronas con $z < 0$

vuelve al tamaño del residual stream

La ecuación dos del paper. x es un token: un vector fila de quinientos doce números. W sub uno es una matriz de quinientos doce por dos mil cuarenta y ocho, y b sub uno un vector de sesgos. Una neurona sola, con tres entradas para que se pueda seguir: uno por cero coma cinco, menos dos por cero coma veinticinco, cero coma cinco por dos, más el sesgo: z igual a uno coma cinco. Es positivo, pasa entero. x por W sub uno hace eso mismo para dos mil cuarenta y ocho neuronas a la vez: en el dibujo, x por la columna j da la neurona j. Otra neurona da menos dos, y la ReLU, máximo entre cero y z, la apaga. Y W sub dos combina las dos mil cuarenta y ocho activaciones en quinientos doce números otra vez.

La FFN en PyTorch

nn punto Linear de quinientos doce a dos mil cuarenta y ocho: W sub uno y b sub uno. La ReLU. Y otra Linear de vuelta. Para ver que trabaja token por token, le pasamos el token tres solo, sin el resto de la frase. Y contamos parámetros contra las cuatro matrices de la atención, con sus sesgos. Entra un batch de dos por cinco por quinientos doce y sale igual. El token tres solo da exactamente lo mismo que dentro de la frase: True. Y la FFN tiene dos millones cien mil parámetros; la atención, un millón cincuenta mil. El doble: la mayor parte de los pesos de cada bloque está en la FFN.

d_model, d_ff = 512, 2048                     # los del paper
ffn = nn.Sequential(nn.Linear(d_model, d_ff),  # W1, b1: sube a 2048
                    nn.ReLU(),                 # max(0, ·)
                    nn.Linear(d_ff, d_model))  # W2, b2: baja a 512

x = torch.randn(2, 5, d_model)                # batch 2, 5 tokens
y = ffn(x)
print("entrada", tuple(x.shape), "→ salida", tuple(y.shape))

# cada token por separado: el token 3, solo, da lo mismo
solo = ffn(x[:, 3])
print("token 3 solo == token 3 en la frase:",
      torch.allclose(solo, y[:, 3], atol=1e-6))

n_ffn = sum(p.numel() for p in ffn.parameters())
n_att = 4 * d_model * d_model + 4 * d_model   # W^Q, W^K, W^V, W^O
print(f"parámetros FFN {n_ffn:,} · atención {n_att:,}")

entrada (2, 5, 512) → salida (2, 5, 512)
token 3 solo == token 3 en la frase: True
parámetros FFN 2,099,712 · atención 1,050,624

Vaswani et al., 2017 — Attention Is All You Need, sección 3.3 · arXiv:1706.03762

d_ff = 2048 · «This consists of two linear transformations with a ReLU activation in between.» … «Another way of describing this is as two convolutions with kernel size 1.»

El paper, en la sección tres punto tres. Dos transformaciones lineales con una ReLU en el medio, aplicada a cada posición por separado e idénticamente. La capa de adentro mide dos mil cuarenta y ocho. Y lo describen también como dos convoluciones de tamaño uno: el mismo filtro en cada lugar, sin mirar a los vecinos.

Parte 2 · Residual y LayerNorm

Cada subcapa se envuelve igual: se suma la entrada y se normaliza. Dos ideas de Deep learning a secas, ahora en su lugar.

El residual stream

Esto ya lo dibujamos en Alineación, episodio uno. Cada subcapa, atención o FFN, es un bloque al costado. Lee el vector, calcula algo, y lo suma al que venía, en vez de reemplazarlo. Así el vector de cada token es una corriente que va juntando lo que cada subcapa le agrega: el residual stream. Por esa corriente pasa todo lo que el modelo sabe de ese token. En Deep learning a secas, episodio diez, medimos por qué ayuda: con veinte capas, la red sin residuales no aprendió nada.

Sumar en vez de reemplazar

$F$: la atención o la FFN

el gradiente tiene un camino directo: la identidad $I$

sin normalizar en el medio, la salida es una suma

Las cuentas. x prima es x más F de x, donde F es la subcapa. En el dibujo, la columna x pasa entera y se le suma otra. La derivada de x prima respecto de x es la identidad, I, la matriz que deja todo igual, más la derivada de F. Hacia atrás, el gradiente siempre tiene un camino directo, aunque F tenga derivada chica. Ese uno más algo es lo que en Deep learning a secas evitó que el gradiente se apague. Y si apilás tres subcapas sin nada en el medio, la salida es x cero más lo que aportó cada una: una suma. Guardá esta línea: el paper no la cumple, porque normaliza en el medio.

LayerNorm, con cuatro números

la media de un token, sobre sus $d$ números

la varianza: el promedio de los desvíos al cuadrado

$\varepsilon = 10^{-5}$ evita dividir por cero

$\gamma$, $\beta$: vectores de $d$ números que se aprenden

LayerNorm la vimos en Deep learning a secas, episodio nueve; acá, un token con d igual a cuatro: dos, cuatro, seis y ocho. mu es la media de sus números: cinco. sigma al cuadrado, la varianza: el promedio de cada número menos la media, al cuadrado. Nueve, uno, uno y nueve: cinco. Restamos la media y dividimos por la raíz de la varianza más épsilon, un número chiquito para no dividir por cero. Queda menos uno coma treinta y cuatro, menos cero coma cuarenta y cinco, y los simétricos: media cero y desvío uno. Después, gamma y beta, dos vectores que se aprenden, reescalan y corren. El círculo con un punto es el producto elemento a elemento. En el dibujo, con gamma dos y beta uno.

LayerNorm a mano y con nn.LayerNorm

La media del token. Restar y dividir, como en la fórmula. Y la de PyTorch, que arranca con gamma igual a uno y beta igual a cero. Después, cinco tokens de quinientos doce números con media tres y desvío diez, bien desordenados. Media cinco, varianza cinco, y los mismos cuatro números del dibujo. Igual a nn punto LayerNorm: True. Y los cinco tokens salen con media cero y desvío uno, cada uno por su cuenta. Lo importante: LayerNorm normaliza cada token sobre sus propios números, sin mirar a los otros tokens ni al resto del batch.

x = torch.tensor([2., 4., 6., 8.])            # un token, d = 4
mu = x.mean()
var = ((x - mu) ** 2).mean()                  # varianza sin corregir
x_hat = (x - mu) / torch.sqrt(var + 1e-5)
print("media", mu.item(), "· varianza", var.item())
print("normalizado", [round(v, 2) for v in x_hat.tolist()])

ln = nn.LayerNorm(4)                          # gamma = 1, beta = 0 al inicio
print("igual a nn.LayerNorm:", torch.allclose(ln(x), x_hat, atol=1e-6))

X = 3 + 10 * torch.randn(5, 512)              # 5 tokens desordenados
Y = nn.LayerNorm(512)(X)
print("media por token", [round(v, 3) + 0.0       # +0.0: sin -0.0
                         for v in Y.mean(-1).tolist()])
print("desvío por token", [round(v, 3) for v in Y.std(-1, unbiased=False).tolist()])

media 5.0 · varianza 5.0
normalizado [-1.34, -0.45, 0.45, 1.34]
igual a nn.LayerNorm: True
media por token [0.0, 0.0, 0.0, 0.0, 0.0]
desvío por token [1.0, 1.0, 1.0, 1.0, 1.0]

Parte 3 · El bloque entero

Atención, FFN, dos sumas y dos LayerNorm. Y una decisión: normalizar antes o después de la subcapa.

Un bloque encoder, contra PyTorch

El bloque entero en una clase. La versión del paper: la atención, dropout, la suma con x, y LayerNorm después de sumar. Por eso se llama post-norm. Dropout apaga al azar una fracción de los números al entrenar; lo vemos en el episodio siete. Lo mismo con la FFN. Y la otra versión, pre-norm: normalizamos la entrada de la subcapa, y la suma queda sin nada encima. Es la línea de la suma de tres subcapas de hace un rato. Copiamos los pesos de nn punto TransformerEncoderLayer, que por defecto es post-norm, y apagamos dropout. La diferencia máxima es nueve coma cinco por diez a la menos siete: es el mismo bloque.

Cada línea, su término:

• x = s.ln1(x + s.drop(s.attn(x, x, x)[0])) → $x' = \mathrm{LayerNorm}\big(x + \mathrm{Attn}(x)\big)$ — post-norm: el paper (dropout sobre Attn)
• return s.ln2(x + s.drop(s.ffn(x))) → $y = \mathrm{LayerNorm}\big(x' + \mathrm{FFN}(x')\big)$ — la misma receta con la FFN
• x = x + s.drop(s.attn(h, h, h)[0]) → $x' = x + \mathrm{Attn}\big(\mathrm{LayerNorm}(x)\big)$ — pre-norm: la suma queda limpia
• print("dif. máx con PyTorch:" → $\max |y - y_{\text{PyTorch}}| \approx 10^{-6}$ — redondeo de punto flotante

class Bloque(nn.Module):
    def __init__(s, d=512, h=8, d_ff=2048, p=0.1, pre=False):
        super().__init__()
        s.attn = nn.MultiheadAttention(d, h, dropout=p, batch_first=True)
        s.ffn = nn.Sequential(nn.Linear(d, d_ff), nn.ReLU(), nn.Linear(d_ff, d))
        s.ln1, s.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        s.drop, s.pre = nn.Dropout(p), pre

    def forward(s, x):
        if s.pre:                                   # pre-norm (hoy)
            h = s.ln1(x)
            x = x + s.drop(s.attn(h, h, h)[0])
            return x + s.drop(s.ffn(s.ln2(x)))
        x = s.ln1(x + s.drop(s.attn(x, x, x)[0]))   # post-norm (paper)
        return s.ln2(x + s.drop(s.ffn(x)))

ref = nn.TransformerEncoderLayer(512, 8, 2048, batch_first=True)
nuestro = Bloque(); copiar(nuestro, ref)            # mismos pesos
nuestro.eval(); ref.eval()                          # sin dropout
x = torch.randn(2, 7, 512)
print("dif. máx con PyTorch:", f"{(nuestro(x) - ref(x)).abs().max():.1e}")

dif. máx con PyTorch: 9.5e-07

Dónde va la normalización

Las dos variantes, lado a lado. Post-norm, la del paper: normaliza después de sumar, así que el residual stream pasa por un LayerNorm en cada subcapa, y cada vez se reescala. Y el paper entrena con un warmup del learning rate: arranca muy chico y lo sube de a poco. Pre-norm normaliza solo lo que entra a la subcapa: la corriente es una suma limpia, y hace falta un LayerNorm al final, antes de la salida. Gemma va un paso más allá: normaliza la entrada y también la salida de cada subcapa antes de sumarla. Lo vemos en el episodio nueve.

Parte 4 · Post-norm contra pre-norm

Una pila de doce bloques, las dos variantes, con la misma inicialización. Medimos los gradientes y entrenamos.

La tarea de juguete

• Diez tokens, cada uno un dígito del 0 al 9
• Para cada token: ¿cuántas veces más aparece?
• Necesita comparar con todos: atención
• No necesita posición: no hay positional encoding

El experimento necesita una tarea chica. Una secuencia de diez dígitos al azar. Para cada token, el modelo tiene que decir cuántas veces más aparece ese dígito en la secuencia: cero, uno, dos. Para eso cada token tiene que compararse con todos los otros: es un trabajo para la atención. Y como contar no depende del orden, no le damos positional encoding. Doce bloques, d sub model de treinta y dos, cuatro heads.

El gradiente de cada capa, al inicio

La misma semilla: las dos pilas arrancan con los mismos pesos. norm first es el interruptor de PyTorch: falso es post-norm, verdadero es pre-norm. La pre-norm lleva su LayerNorm final. Un solo backward, sin entrenar, y medimos la norma del gradiente de la primera matriz de la FFN en cada una de las doce capas: el tamaño del empujón que recibiría. En post-norm crece hacia la salida: cero coma dieciséis abajo, cero coma cincuenta y siete en la capa doce, tres veces y media más. En pre-norm no: arranca en cero coma treinta y cuatro y las capas de arriba quedan en cero coma trece.

def pila(pre, N=12, d=32):
    torch.manual_seed(0)                        # misma init para las dos
    capa = nn.TransformerEncoderLayer(d, 4, 4 * d, dropout=0.0,
                                      batch_first=True, norm_first=pre)
    enc = nn.TransformerEncoder(capa, N, enable_nested_tensor=False,
                                norm=nn.LayerNorm(d) if pre else None)
    return nn.Sequential(nn.Embedding(10, d), enc, nn.Linear(d, 10))

torch.manual_seed(1); x, y = datos(256)
for pre in (False, True):
    m = pila(pre)
    F.cross_entropy(m(x).flatten(0, 1), y.flatten()).backward()
    g = [c.linear1.weight.grad.norm().item() for c in m[1].layers]
    print("pre " if pre else "post", " ".join(f"{v:.2f}" for v in g))

post 0.16 0.14 0.14 0.15 0.16 0.19 0.22 0.26 0.30 0.33 0.37 0.57
pre  0.34 0.27 0.24 0.22 0.19 0.17 0.15 0.14 0.13 0.13 0.13 0.13

Norma del gradiente por capa (salida real)

Las dos curvas. Post-norm: plana abajo, y sube cerca de la salida, con un salto en la última capa. Pre-norm: baja suave y queda plana arriba. Con un learning rate fijo, las capas de arriba de la post-norm reciben los pasos más grandes, justo al principio, cuando los pesos todavía son ruido.

Xiong et al., 2020 — On Layer Normalization in the Transformer Architecture · arXiv:2002.04745

Post-LN: gradientes grandes cerca de la salida · «for the original-designed Post-LN Transformer, which places the layer normalization between the residual blocks, the expected gradients of the parameters near the output layer are large.»

Eso mismo es lo que demuestra este paper de dos mil veinte. En el transformer original, post layer norm, al inicio los gradientes esperados de los parámetros cerca de la salida son grandes, y usar un learning rate grande con esos gradientes vuelve inestable el entrenamiento. Su conclusión: el warmup es lo que lo salva, y con pre-norm se puede entrenar sin warmup.

Entrenar doce capas sin warmup

Ahora entrenamos. Doscientos pasos, learning rate de uno por diez a la menos tres, fijo, sin warmup. Adam, adaptive moment estimation, con los betas del paper; lo vemos en el próximo episodio. Medimos el acierto en dos mil secuencias nuevas. Y la referencia tonta: el acierto de responder siempre cero. Post-norm: pérdida uno coma veinticinco, acierto treinta y nueve por ciento. Pre-norm: pérdida cero coma veintidós, noventa y dos por ciento. Y responder siempre cero da treinta y nueve: la post-norm no aprendió nada, se quedó en la respuesta más común.

def entrenar(pre, pasos=200, lr=1e-3):          # 12 capas, sin warmup
    m = pila(pre)
    opt = torch.optim.Adam(m.parameters(), lr=lr, betas=(0.9, 0.98))
    torch.manual_seed(1)
    for _ in range(pasos):
        x, y = datos(64)
        loss = F.cross_entropy(m(x).flatten(0, 1), y.flatten())
        opt.zero_grad(); loss.backward(); opt.step()
    x, y = datos(2000)
    acc = (m(x).argmax(-1) == y).float().mean()
    return loss.item(), acc.item(), (y == 0).float().mean().item()

for pre in (False, True):
    loss, acc, base = entrenar(pre)
    print("pre " if pre else "post", f"pérdida {loss:.2f} · acierto {acc:.0%}")
print(f"decir siempre «0»: acierto {base:.0%}")

post pérdida 1.25 · acierto 39%
pre  pérdida 0.22 · acierto 92%
decir siempre «0»: acierto 39%

Acierto tras 200 pasos, 12 capas (salida real)

En barras. Responder siempre cero: treinta y nueve por ciento. La post-norm, lo mismo: se trabó. La pre-norm, noventa y dos. Ojo: con seis capas, la post-norm también aprende. El problema aparece con profundidad y pasos grandes; por eso el paper usó warmup.

• siempre «0»: 39 %
• post-norm: 39 %
• pre-norm: 92 %

Resumiendo

• FFN: $\max(0, xW_1 + b_1)W_2 + b_2$, de 512 a 2048 y de vuelta
• La FFN tiene el doble de parámetros que la atención
• Residual: $x + F(x)$, el gradiente tiene un camino directo
• LayerNorm: media 0 y desvío 1 por token, después $\gamma$, $\beta$
• 12 capas sin warmup: post-norm 39 %, pre-norm 92 %
• Próximo: cómo se entrenó, warmup y label smoothing

Papers del episodio

• Vaswani et al., 2017 — Attention Is All You Need
• Xiong et al., 2020 — On Layer Normalization in the Transformer Architecture