Implementarlo entero

Attention is all you need. Episodio diez, el último: implementarlo entero. En nueve episodios armamos cada pieza por separado, y para entrenar usamos nn punto Transformer, la versión de PyTorch. Hoy, nada de eso: un transformer encoder-decoder escrito desde cero, con nuestras piezas, entrenado con la receta del paper, generando token por token. Y al final, lo que no aprendió.

Las piezas, y dónde las vimos

El plan, pieza por pieza. La multi-head attention de los episodios dos y tres. La capa de encoder, con la FFN, feed-forward network, las sumas y LayerNorm del episodio seis. La de decoder, con su máscara y la cross-attention del cinco. Los embeddings, multiplicados por raíz de d. Los senos del cuatro. La salida, que reusa la matriz de embeddings. La receta del siete. Y la decodificación greedy, que elige el token más probable en cada paso.

multi-head (ep. 3) → encoder (ep. 6) → decoder (ep. 5) → embedding · √d → senos (ep. 4) → logits atados → receta (ep. 7) → greedy

Parte 1 · Atención y capas

Primero lo que se repite: la atención, la capa de encoder y la capa de decoder. Treinta y cinco líneas en total.

Multi-head attention, en una clase

Cuatro matrices: W Q, W K, W V y W O. La query se proyecta y se parte en h heads con view y transpose, como en el episodio tres. Las queries salen de xq, y las keys y values de xkv: si son el mismo tensor, es self-attention; si xkv es la memoria del encoder, es cross-attention. La misma clase sirve para las dos. Los scores, divididos por raíz de d sub k. La máscara: menos infinito donde está prohibido. Y las heads se juntan y pasan por W O. Probamos la máscara causal cambiando el último token: las salidas de los tokens cero a tres no cambian. True.

class MultiHead(nn.Module):
    def __init__(s, d, h):
        super().__init__()
        s.h, s.dk = h, d // h
        s.wq, s.wk, s.wv, s.wo = (nn.Linear(d, d) for _ in range(4))

    def forward(s, xq, xkv, prohibido=None):       # prohibido: True = no mirar
        B, Tq, Tk = xq.shape[0], xq.shape[1], xkv.shape[1]
        q = s.wq(xq).view(B, Tq, s.h, s.dk).transpose(1, 2)    # [B, h, Tq, dk]
        k = s.wk(xkv).view(B, Tk, s.h, s.dk).transpose(1, 2)
        v = s.wv(xkv).view(B, Tk, s.h, s.dk).transpose(1, 2)
        S = q @ k.transpose(-2, -1) / s.dk ** 0.5               # [B, h, Tq, Tk]
        if prohibido is not None:
            S = S.masked_fill(prohibido, float("-inf"))
        out = S.softmax(-1) @ v                                 # [B, h, Tq, dk]
        return s.wo(out.transpose(1, 2).reshape(B, Tq, -1))

att, x = MultiHead(d=32, h=4), torch.randn(2, 5, 32)        # batch 2, 5 tokens
causal = torch.triu(torch.ones(5, 5, dtype=torch.bool), diagonal=1)
x2 = x.clone(); x2[:, 4] = torch.randn(32)                  # cambio el último
print("salida", tuple(att(x, x, causal).shape), "· tokens 0–3 iguales:",
      torch.allclose(att(x, x, causal)[:, :4], att(x2, x2, causal)[:, :4]))

salida (2, 5, 32) · tokens 0–3 iguales: True

Lo que hace la clase, en tres pasos

una tabla por head: $h = 4$, $d_k = 32/4 = 8$

la máscara causal del decoder

cada fila: un promedio de values

La clase, en las fórmulas de la serie. Los scores: Q, de T sub q por d sub k, por K transpuesta, de d sub k por T sub k. Con d sub model treinta y dos y cuatro heads, cada head compara vectores de ocho números. En el decoder, la máscara: el token i no puede mirar a ningún j mayor que i. En el dibujo, el decoder escribiendo cero, cero, dos, dos: cada fila reparte su peso solo entre lo ya escrito. El softmax de cada fila pesa los values; después se concatenan las heads y W O las mezcla.

La capa de encoder

La FFN del episodio seis: subir, ReLU, rectified linear unit, y bajar. Dos LayerNorm. Self-attention, dropout, suma y LayerNorm: post-norm, como el paper. Con dos capas no necesitamos pre-norm. Lo mismo con la FFN. Entra un batch de dos secuencias de diez tokens de treinta y dos números, y sale igual. Doce mil setecientos cuatro parámetros: la atención tiene cuatro mil doscientos veinticuatro, la FFN ocho mil trescientos cincuenta y dos, y los LayerNorm, ciento veintiocho.

def ffn(d, d_ff):
    return nn.Sequential(nn.Linear(d, d_ff), nn.ReLU(), nn.Linear(d_ff, d))

class CapaEncoder(nn.Module):
    def __init__(s, d, h, d_ff, p):
        super().__init__()
        s.att, s.ffn = MultiHead(d, h), ffn(d, d_ff)
        s.ln = nn.ModuleList(nn.LayerNorm(d) for _ in range(2))
        s.drop = nn.Dropout(p)

    def forward(s, x):                                # post-norm, como el paper
        x = s.ln[0](x + s.drop(s.att(x, x)))
        return s.ln[1](x + s.drop(s.ffn(x)))

capa = CapaEncoder(d=32, h=4, d_ff=128, p=0.1).eval()
x = torch.randn(2, 10, 32)
print("entra", tuple(x.shape), "sale", tuple(capa(x).shape),
      "· parámetros", sum(p.numel() for p in capa.parameters()))

entra (2, 10, 32) sale (2, 10, 32) · parámetros 12704

La capa de decoder

La capa de decoder: tres subcapas. Self-attention sobre lo que ya escribió, con la máscara M. La cross-attention del episodio cinco: la misma clase, pero las keys y values vienen de la memoria, m, la salida del encoder. Y la FFN. Cada una con su suma y su LayerNorm. Probamos con diez tokens de entrada y siete ya escritos: la tabla de la cross-attention es de siete por diez. Sale de dos por siete por treinta y dos, y tiene dieciséis mil novecientos noventa y dos parámetros: una atención más que el encoder.

Cada línea, su término:

• y = s.ln[0](y + s.drop(s.att(y, y, causal))) → $y' = \mathrm{LayerNorm}\big(y + \mathrm{Attn}(y, y; M)\big)$ — self-attention con máscara $M$
• y = s.ln[1](y + s.drop(s.cruzada(y, memoria))) → $y'' = \mathrm{LayerNorm}\big(y' + \mathrm{Attn}(y', m)\big)$ — queries de $y'$, keys y values de la memoria $m$
• return s.ln[2](y + s.drop(s.ffn(y))) → $\mathrm{LayerNorm}\big(y'' + \mathrm{FFN}(y'')\big)$ — tres subcapas, tres LayerNorm
• memoria, y = torch.randn(2, 10, 32), torch.randn(2, 7, 32) → $A_{\text{cruzada}}: T_{\text{dec}} \times T_{\text{enc}} = 7 \times 10$ — cada token escrito mira toda la entrada

class CapaDecoder(nn.Module):
    def __init__(s, d, h, d_ff, p):
        super().__init__()
        s.att, s.cruzada, s.ffn = MultiHead(d, h), MultiHead(d, h), ffn(d, d_ff)
        s.ln = nn.ModuleList(nn.LayerNorm(d) for _ in range(3))
        s.drop = nn.Dropout(p)

    def forward(s, y, memoria, causal):
        y = s.ln[0](y + s.drop(s.att(y, y, causal)))          # self, con máscara
        y = s.ln[1](y + s.drop(s.cruzada(y, memoria)))        # cross-attention
        return s.ln[2](y + s.drop(s.ffn(y)))

capa = CapaDecoder(d=32, h=4, d_ff=128, p=0.1).eval()
memoria, y = torch.randn(2, 10, 32), torch.randn(2, 7, 32)  # 10 de entrada, 7 escritos
causal = torch.triu(torch.ones(7, 7, dtype=torch.bool), diagonal=1)
print("sale", tuple(capa(y, memoria, causal).shape),
      "· parámetros", sum(p.numel() for p in capa.parameters()))

sale (2, 7, 32) · parámetros 16992

Parte 2 · El modelo

Las capas apiladas, dos de cada lado, en vez de las seis del paper, más lo que pasa en la entrada y en la salida.

Una matriz de embeddings, tres usos

la fila $t$ de $E$, escalada, más la posición

la salida usa la misma $E$, transpuesta

probabilidades del próximo token

La entrada y la salida comparten una sola matriz, E: once filas, una por token, de treinta y dos números. En la entrada, cada token elige su fila de E, se multiplica por raíz de d, y se le suman los senos de su posición: el positional encoding del episodio cuatro. En el dibujo, BOS, cero, cero, dos. BOS es el token de arranque del decoder. En la salida, el vector final h se multiplica por E transpuesta: el logit de cada dígito es el producto punto de h con la fila de ese dígito. Y el softmax los vuelve probabilidades. El encoder, el decoder y la salida usan la misma E.

El transformer entero

Una sola matriz de embeddings. La entrada: embedding por raíz de d, más los senos, con dropout. codificar: la entrada pasa por las capas del encoder, y queda la memoria. decodificar arma la máscara causal, pasa por las capas del decoder, y multiplica por la misma matriz transpuesta. forward es decodificar de codificar. Cincuenta y nueve mil setecientos cuarenta y cuatro parámetros. La matriz de embeddings son trescientos cincuenta y dos números, y se cuentan una sola vez aunque se usen tres veces. El base del paper tenía sesenta y cinco millones.

class Transformer(nn.Module):
    def __init__(s, d=32, h=4, N=2, d_ff=128, p=0.1):
        super().__init__()
        s.d, s.emb = d, nn.Embedding(V, d)           # una sola matriz
        s.register_buffer("pe", pe(64, d))           # senos del episodio 4
        s.enc = nn.ModuleList(CapaEncoder(d, h, d_ff, p) for _ in range(N))
        s.dec = nn.ModuleList(CapaDecoder(d, h, d_ff, p) for _ in range(N))
        s.drop = nn.Dropout(p)
    def entrada(s, t):                               # embedding · √d + posición
        return s.drop(s.emb(t) * math.sqrt(s.d) + s.pe[:t.shape[1]])
    def codificar(s, x):
        m = s.entrada(x)
        for capa in s.enc: m = capa(m)
        return m                                     # la memoria
    def decodificar(s, y_in, memoria):
        T = y_in.shape[1]
        causal = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
        y = s.entrada(y_in)
        for capa in s.dec: y = capa(y, memoria, causal)
        return y @ s.emb.weight.T                    # logits: la misma matriz
    def forward(s, x, y_in):
        return s.decodificar(y_in, s.codificar(x))

m = Transformer()
print(f"parámetros: {sum(p.numel() for p in m.parameters()):,}"
      f" · embeddings: {m.emb.weight.numel()} (una sola vez)")

parámetros: 59,744 · embeddings: 352 (una sola vez)

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

E se usa 3 veces · × √d_model · «we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation … In the embedding layers, we multiply those weights by √d_model.»

Las dos decisiones vienen de la sección tres punto cuatro. Una sola matriz para los embeddings del encoder, los del decoder y la transformación lineal antes del softmax. Y en los embeddings, los pesos se multiplican por raíz de d sub model. Funciona porque en el paper los dos idiomas comparten un vocabulario de unos treinta y siete mil tokens; en el nuestro, los dígitos son los mismos a la entrada y a la salida.

Parte 3 · Entrenar

La tarea de los episodios siete y ocho: entra una secuencia de diez dígitos, sale ordenada.

El loop, con la receta del paper

El modelo: d sub model treinta y dos, cuatro heads, dos capas de cada lado, FFN de ciento veintiocho, dropout cero coma uno. Adam, adaptive moment estimation, con los betas del paper. El schedule con warmup. Label smoothing. El decoder recibe la salida correcta corrida un lugar, con BOS adelante: todo en paralelo, como en el episodio cinco. Y guardamos los pesos. Mil doscientos pasos, unos diez segundos. La pérdida baja de cero coma nueve dos ocho a cero coma cinco siete nueve. Y no va a bajar mucho más: ahí hay un piso.

m = Transformer(d=32, h=4, N=2, d_ff=128, p=0.1)
opt = torch.optim.Adam(m.parameters(), lr=2e-3, betas=(0.9, 0.98), eps=1e-9)
sch = LambdaLR(opt, lambda s: forma(s + 1, warmup=100))       # ecuación (3)
m.train()
for paso in range(1, 1201):
    x, y, y_in = datos(64)                       # ordenar 10 dígitos
    logits = m(x, y_in)                          # [64, 10, 11]
    loss = F.cross_entropy(logits.flatten(0, 1), y.flatten(),
                           label_smoothing=0.1)
    opt.zero_grad(); loss.backward(); opt.step(); sch.step()
    if paso % 300 == 0:
        print(f"paso {paso:4}: pérdida {loss.item():.3f}")
torch.save(m.state_dict(), CKPT)

paso  300: pérdida 0.928
paso  600: pérdida 0.674
paso  900: pérdida 0.607
paso 1200: pérdida 0.579

El piso de la pérdida

label smoothing con $\varepsilon = 0{,}1$ y $V = 11$

la entropía de $q$: lo mejor posible es $p = q$

lo que todavía le falta aprender

¿Por qué no llega a cero? Con label smoothing, el objetivo no es un uno: es cero coma novecientos nueve para el correcto y cero coma cero cero nueve uno para cada uno de los otros diez tokens. La pérdida es mínima cuando el modelo predice exactamente q, y ese mínimo es la entropía de q: menos la suma de q por el logaritmo de q. El correcto aporta cero coma cero ocho siete, los otros diez juntos, cero coma cuatro dos siete. Total, cero coma quinientos catorce. Estamos en cero coma quinientos setenta y nueve: sesenta y cinco milésimas arriba del piso.

La pérdida y su piso (salida real)

En la curva. La pérdida cae rápido al principio y después se aplana. La línea punteada es el piso de cero coma cinco uno cuatro: lo que queda entre las dos es lo que el modelo todavía duda. Es lo mismo que vimos en el episodio siete: label smoothing no deja que el modelo esté del todo seguro. La pérdida sola no dice si ordena bien. Para eso hay que hacerlo generar.

Parte 4 · Generar y medir

Al entrenar, el decoder recibía la respuesta correcta. Al usarlo, tiene que leer lo que él mismo escribió.

Decodificación greedy

El encoder corre una sola vez: la memoria no cambia mientras se genera. En cada paso, el decoder corre sobre todo lo escrito y nos quedamos con el último vector. Elegimos el token más probable, el argmax, y lo pegamos al final. Eso es greedy: nunca vuelve atrás. El paper usaba beam search, que mantiene cuatro candidatos. Y comparamos con teacher forcing: darle la respuesta correcta como entrada. Greedy: noventa y nueve coma cinco por ciento de secuencias perfectas en mil nuevas. Igual que con teacher forcing. Por ejemplo, dos, nueve, dos, cero, cero, dos, seis, siete, nueve, cuatro sale cero, cero, dos, dos, dos, cuatro, seis, siete, nueve, nueve.

@torch.no_grad()
def greedy(modelo, x):
    memoria = modelo.codificar(x)                    # el encoder corre una vez
    y = torch.full((x.shape[0], 1), BOS)
    for _ in range(x.shape[1]):
        logits = modelo.decodificar(y, memoria)[:, -1]   # solo el último
        logits[:, BOS] = float("-inf")               # BOS nunca es salida
        y = torch.cat([y, logits.argmax(-1, keepdim=True)], dim=1)
    return y[:, 1:]                                  # sin el BOS

m = Transformer(); m.load_state_dict(torch.load(CKPT)); m.eval()
x, y, y_in = datos(1000, torch.Generator().manual_seed(123))   # 1000 nuevas
pred = greedy(m, x)
forzado = m(x, y_in).argmax(-1)               # con la respuesta como entrada
for nombre, p in [("greedy", pred), ("teacher forcing", forzado)]:
    print(f"{nombre:15}: secuencias {(p == y).all(1).float().mean():.1%}"
          f" · dígitos {(p == y).float().mean():.1%}")
print("entrada:", *x[0].tolist()); print("salida: ", *pred[0].tolist())

greedy         : secuencias 99.5% · dígitos 99.9%
teacher forcing: secuencias 99.5% · dígitos 99.9%
entrada: 2 9 2 0 0 2 6 7 9 4
salida:  0 0 2 2 2 4 6 7 9 9

Lo que no aprendió

La prueba incómoda. Lo entrenamos solo con secuencias de diez dígitos. ¿Qué pasa con seis, ocho, doce o catorce? Generamos con greedy. Con diez, en otras mil secuencias, noventa y nueve coma tres por ciento. Con doce, cuarenta y cinco. Con ocho, nueve. Con catorce y con seis, prácticamente cero. Hasta las más cortas fallan. Y no es por la posición: los senos de las posiciones once a catorce existen y el modelo los puede leer. Lo que nunca vio son secuencias de ese largo.

for largo in (6, 8, 10, 12, 14):                 # entrenado solo con 10
    x, y, _ = datos(1000, torch.Generator().manual_seed(5), largo=largo)
    pred = greedy(m, x)
    print(f"largo {largo:2}: secuencias {(pred == y).all(1).float().mean():6.1%}"
          f" · dígitos {(pred == y).float().mean():.1%}")

largo  6: secuencias   0.0% · dígitos 28.6%
largo  8: secuencias   9.0% · dígitos 57.0%
largo 10: secuencias  99.3% · dígitos 99.9%
largo 12: secuencias  45.4% · dígitos 82.1%
largo 14: secuencias   0.3% · dígitos 37.2%

Secuencias perfectas según el largo (salida real)

En barras. Aprendió a ordenar secuencias de diez, no a ordenar. Una explicación posible: con diez dígitos al azar, la salida empieza con cero dos de cada tres veces y termina con nueve otras tantas, y el modelo pudo aprender esas costumbres, que con otro largo dejan de valer. Generalizar fuera de lo que se vio es uno de los problemas abiertos de estos modelos, y los de lenguaje no son la excepción.

• 6: 0 %
• 8: 9 %
• 10 (entrenado): 99,3 %
• 12: 45,4 %
• 14: 0,3 %

La serie, en diez pasos

La serie entera, en un camino. Arrancamos con las RNN, recurrent neural networks, y la atención de Bahdanau. Después la scaled dot-product attention, las heads, la posición con senos, el encoder y el decoder, el resto del bloque, la receta de entrenamiento, los resultados en BLEU, bilingual evaluation understudy, el salto a Gemma, y hoy, todo junto. Un transformer de sesenta mil parámetros que ordena, y uno de dos mil seiscientos millones que escribe, son la misma arquitectura.

RNN y Bahdanau → dot-product → multi-head → posición → encoder-decoder → FFN y LayerNorm → entrenamiento → BLEU y ablaciones → Gemma → desde cero

De predecir el token siguiente a seguir instrucciones

• Pre-entrenar: predecir el token siguiente en mucho texto
• Fine-tuning con demostraciones escritas por personas
• Un modelo que compara respuestas y las puntúa
• Refuerzo para que el modelo prefiera las mejores

Lo que falta para llegar a un asistente. Un transformer que predice el token siguiente completa texto, pero no sigue instrucciones. Primero se pre-entrena en muchísimo texto. Después, fine-tuning con ejemplos de cómo responder, escritos por personas. Se entrena otro modelo que compara dos respuestas y dice cuál es mejor. Y con aprendizaje por refuerzo, el modelo aprende a preferir las respuestas que ese juez puntúa alto. Ese camino es el paper de InstructGPT, y es la próxima serie.

Resumiendo

• MultiHead: una clase para self y cross-attention
• Encoder y decoder post-norm: 12.704 y 16.992 parámetros por capa
• Una matriz $E$: entrada × √d y salida $hE^\top$; 59.744 en total
• 1200 pasos: pérdida 0,579, piso 0,514; greedy 99,5 %
• Largo 10: 99,3 %; largo 14: 0,3 %
• Próximo: la serie «InstructGPT»

Papers del episodio

• Vaswani et al., 2017 — Attention Is All You Need