Módulo 6 · Representaciones y atención

Lección 2 de 4

Atención y transformers

En esta páginaEl problema

El problema

Tienes una tabla de embeddings con tres palabras. Cada palabra tiene un vector de 4 números:

palabra	dim 0	dim 1	dim 2	dim 3
`perro`	1	0	1	0
`muerde`	0	1	0	1
`gato`	1	0	0	1

Con estos vectores representas dos oraciones: «perro muerde gato» y «gato muerde perro». En Embeddings y autoencoders resumiste una secuencia con la media de sus embeddings. Aquí las dos medias salen iguales: [0.667, 0.333, 0.333, 0.667]. Además, el vector de perro es el mismo en las dos oraciones, aunque en una muerde y en la otra lo muerden.

Necesitas una operación que cambie el vector de cada palabra según las otras palabras de la oración y según el lugar que ocupa.

Atención de producto punto

Cada elemento de la secuencia es un token; aquí, cada palabra. En atención, cada token produce tres vectores: una consulta q, una clave k y un valor v. En un transformer salen de tres capas nn.Linear aplicadas al embedding del token. El puntaje entre el token i y el token j es el producto punto q_i · k_j dividido entre √d_k, donde d_k es la longitud de las claves. softmax convierte cada fila de puntajes en pesos de atención. La salida del token i es la suma de los valores, cada uno multiplicado por su peso.

Con matrices, las consultas de todos los tokens forman Q, las claves forman K y los valores forman V:

Atención(Q, K, V) = softmax(Q Kᵀ / √d_k) V

Pregunta. Con las matrices de la celda, ¿qué claves reciben más peso en la fila del token 0? ¿Y en la del token 2? La salida del token 0, ¿queda más cerca de [10, 0] o de [0, 10]?

import math

import torch
import torch.nn.functional as F
from torch import nn

Q = torch.tensor([[2.0, 0.0], [0.0, 2.0], [1.0, 1.0]])
K = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
V = torch.tensor([[10.0, 0.0], [0.0, 10.0], [5.0, 5.0]])

d_k = K.shape[1]
puntajes = Q @ K.T / math.sqrt(d_k)
pesos = torch.softmax(puntajes, dim=-1)
salida = pesos @ V

print(puntajes.round(decimals=3))
print(pesos.round(decimals=3))
print(pesos.sum(dim=-1))
print(salida.round(decimals=2))
print(torch.allclose(salida, F.scaled_dot_product_attention(Q, K, V)))

Salida esperada

tensor([[1.4140, 0.0000, 1.4140],
        [0.0000, 1.4140, 1.4140],
        [0.7070, 0.7070, 1.4140]])
tensor([[0.4460, 0.1080, 0.4460],
        [0.1080, 0.4460, 0.4460],
        [0.2480, 0.2480, 0.5030]])
tensor([1., 1., 1.])
tensor([[6.6900, 3.3100],
        [3.3100, 6.6900],
        [5.0000, 5.0000]])
True

Lectura. En la fila del token 0, las claves 0 y 2 empatan y la clave 1 recibe menos peso. En la fila del token 2, la clave 2 gana porque su producto punto con [1, 1] es el mayor. Cada fila de pesos suma 1, así que cada salida es un promedio ponderado de las filas de V. La salida del token 0 queda más cerca de [10, 0]. F.scaled_dot_product_attention calcula lo mismo en una sola llamada.

La división entre √d_k controla el tamaño de los puntajes. Si q y k tienen componentes aleatorias con media 0 y varianza 1, su producto punto tiene varianza d_k.

Pregunta. Con d_k = 64 y 8 claves aleatorias, ¿qué peso máximo da softmax sin la división y cuál con ella?

Ver resultado

Cada fila compara una consulta aleatoria de 64 números con 8 claves aleatorias, primero sin la división y después con ella. Las columnas dan la desviación estándar de los 8 puntajes, el peso que softmax da a cada clave (de la clave 0 a la 7) y el mayor de esos pesos. La segunda fila dice que, con la división, la clave 1 recibe 0.223 del peso.

puntajes	desviación estándar	pesos de las 8 claves	peso máximo
`q @ k.T`	6.48	0.000, 0.162, 0.025, 0.000, 0.000, 0.000, 0.812, 0.000	0.812
`q @ k.T / √64`	0.81	0.041, 0.223, 0.177, 0.027, 0.096, 0.076, 0.273, 0.087	0.273

Lectura. Sin la división, la desviación estándar de los puntajes es 6.48, del mismo orden que √64 = 8. softmax da 0.812 a una clave y 0 a cinco de las ocho, redondeado a tres decimales. Una clave con peso casi nulo también recibe un gradiente casi nulo. Con la división, la desviación estándar baja a 0.81, ninguna clave pasa de 0.273 y todas reciben algo de peso.

Máscara causal

Un modelo que genera texto palabra por palabra no puede usar las palabras que todavía no escribió. Una máscara causal impide que el token i atienda a los tokens j > i. Antes del softmax, los puntajes bloqueados se reemplazan por -inf. La exponencial de -inf es 0.

torch.triu(..., diagonal=1) marca con True las posiciones que quedan por encima de la diagonal.

Pregunta. Escribe la primera fila de pesos_causales. ¿Cuántos pesos distintos de 0 tiene la fila del token 1?

mascara = torch.triu(torch.ones(3, 3, dtype=torch.bool), diagonal=1)
pesos_causales = torch.softmax(puntajes.masked_fill(mascara, float("-inf")), dim=-1)

print(mascara)
print(pesos_causales.round(decimals=3))
print(
    torch.allclose(
        pesos_causales @ V, F.scaled_dot_product_attention(Q, K, V, is_causal=True)
    )
)

Salida esperada

tensor([[False,  True,  True],
        [False, False,  True],
        [False, False, False]])
tensor([[1.0000, 0.0000, 0.0000],
        [0.1960, 0.8040, 0.0000],
        [0.2480, 0.2480, 0.5030]])
True

Lectura. El token 0 solo se atiende a sí mismo: su fila es [1, 0, 0]. El token 1 reparte el peso entre los tokens 0 y 1. Cada fila sigue sumando 1, pero solo entre las posiciones permitidas. is_causal=True aplica la misma máscara.

Pregunta. Cambias el valor del último token por [-100, -100]. ¿Qué salidas cambian con la máscara causal?

Ver resultado

Cada fila es un token. Las columnas dan su salida de atención causal con V y con V_cambiado, que es V con la fila del token 2 reemplazada por [-100, -100]. La fila del token 1 dice que su salida es [1.96, 8.04] con los dos juegos de valores.

token	salida con `V`	salida con `V_cambiado`
0	`[10.00, 0.00]`	`[10.00, 0.00]`
1	`[1.96, 8.04]`	`[1.96, 8.04]`
2	`[5.00, 5.00]`	`[-47.87, -47.87]`

Lectura. Las salidas de los tokens 0 y 1 no cambian. Ninguno de los dos puede atender al token 2. Solo cambia la salida del token 2, que sí atiende a su propio valor.

Explorador

Atención con máscara causal

Elige la palabra que consulta con un clic o con las flechas y aplica la máscara causal; observa qué pesos pasan a 0 y la columna Σ fila.

Qué probar

Mueve la temperatura y compara el mayor peso de la fila.

ParámetrosTemperatura τ = 1.00

Con τ menor que 1 la atención se concentra; con τ mayor que 1 se reparte. Pesos = softmax(QKᵀ / (τ·√4)), con Q y K de 4 números.

Consulta «la», sin máscara: ve los 7 tokens; el mayor peso va a «alfombra» (0.74). Todas las filas suman 1.00.

Atención multi-cabeza

Una sola atención produce un juego de pesos por token. La atención multi-cabeza repite el cálculo h veces en paralelo. Cada cabeza tiene sus propias proyecciones para Q, K y V, de tamaño d / h. Las h salidas se concatenan y una capa lineal final las mezcla en un vector de tamaño d. Cada cabeza puede asignar pesos distintos a los mismos tokens.

nn.MultiheadAttention(embed_dim, num_heads) implementa ese cálculo. Recibe query, key y value. En autoatención, los tres son la misma secuencia x. Con batch_first=True, x tiene forma (batch, tokens, embed_dim). La llamada devuelve la salida y los pesos.

Aquí x es la oración «perro muerde gato» con la tabla de El problema.

Pregunta. Con embed_dim=4 y num_heads=2, escribe la forma de x, de la salida y de los pesos con average_attn_weights=False. ¿Cuántos parámetros tiene el módulo?

tabla = {
    "perro": [1.0, 0.0, 1.0, 0.0],
    "muerde": [0.0, 1.0, 0.0, 1.0],
    "gato": [1.0, 0.0, 0.0, 1.0],
}


def oracion(texto):
    return torch.tensor([[tabla[palabra] for palabra in texto.split()]])


x = oracion("perro muerde gato")

torch.manual_seed(0)
atencion = nn.MultiheadAttention(embed_dim=4, num_heads=2, batch_first=True)

with torch.no_grad():
    salida_mha, pesos_mha = atencion(x, x, x, average_attn_weights=False)

print(x.shape, salida_mha.shape, pesos_mha.shape)
print(pesos_mha[0].round(decimals=2))
print(sum(p.numel() for p in atencion.parameters()))

Salida esperada

torch.Size([1, 3, 4]) torch.Size([1, 3, 4]) torch.Size([1, 2, 3, 3])
tensor([[[0.3900, 0.3000, 0.3100],
         [0.2000, 0.5400, 0.2600],
         [0.2000, 0.5100, 0.2800]],

        [[0.3800, 0.2800, 0.3400],
         [0.3800, 0.2900, 0.3300],
         [0.3900, 0.2800, 0.3300]]])
80

Lectura. La salida conserva la forma de x: un vector de 4 números por token. Los pesos tienen forma (1, 2, 3, 3): una matriz de 3×3 por cabeza. Las dos matrices son distintas, así que cada cabeza mezcla los tokens de otra manera. Los 80 parámetros salen de las proyecciones Q, K y V (3 × (4 × 4 + 4) = 60) y de la capa final (4 × 4 + 4 = 20).

nn.MultiheadAttention acepta una máscara en attn_mask. En este módulo, True significa «bloquear».

Pregunta. Con attn_mask=mascara, ¿qué entradas de la matriz de pesos promediada quedan en 0?

with torch.no_grad():
    _, pesos_promedio = atencion(x, x, x, attn_mask=mascara)

print(pesos_promedio.shape)
print(pesos_promedio[0].round(decimals=2))

Salida esperada

torch.Size([1, 3, 3])
tensor([[1.0000, 0.0000, 0.0000],
        [0.4200, 0.5800, 0.0000],
        [0.2900, 0.4000, 0.3100]])

Lectura. Las entradas por encima de la diagonal valen 0, igual que en pesos_causales. Sin average_attn_weights=False, el módulo promedia las cabezas y devuelve una sola matriz de 3×3 por oración.

Bloque transformer

Un bloque transformer del encoder combina cuatro piezas:

autoatención multi-cabeza
una conexión residual, que suma la entrada de una subcapa a su salida
normalización de capa (nn.LayerNorm), que resta al vector de cada token la media de sus propios números y lo divide entre su desviación estándar; después lo multiplica por una escala y le suma un desplazamiento, ambos aprendidos
un MLP que se aplica a cada token por separado, aquí 4 → 16 → 4 con ReLU

El transformer original aplica x = LayerNorm(x + Atención(x)) y después x = LayerNorm(x + MLP(x)). nn.TransformerEncoderLayer sigue ese orden por defecto. bloque_a_mano repite el cálculo con las subcapas del mismo objeto.

Pregunta. ¿Qué forma tiene la salida del bloque para x? Con los parámetros iniciales de LayerNorm (escala 1, desplazamiento 0), ¿qué media y qué desviación estándar tiene el vector de cada token a la salida?

torch.manual_seed(0)
capa = nn.TransformerEncoderLayer(
    d_model=4, nhead=2, dim_feedforward=16, dropout=0.0, batch_first=True
)
capa.eval()


def bloque_a_mano(x):
    atendido, _ = capa.self_attn(x, x, x)
    x = capa.norm1(x + atendido)
    mlp = capa.linear2(torch.relu(capa.linear1(x)))
    return capa.norm2(x + mlp)


with torch.no_grad():
    y_capa = capa(x)
    y_mano = bloque_a_mano(x)

print(y_capa.shape)
print(torch.allclose(y_capa, y_mano, atol=1e-5))
print(y_capa.mean(dim=-1).abs().round(decimals=3))
print(y_capa.std(dim=-1, unbiased=False).round(decimals=3))

Salida esperada

torch.Size([1, 3, 4])
True
tensor([[0., 0., 0.]])
tensor([[1., 1., 1.]])

Lectura. La salida tiene la misma forma que la entrada, (1, 3, 4). Por eso los bloques se pueden apilar: la salida de uno es la entrada del siguiente. Cada vector de token sale con media 0 y desviación estándar 1 porque la última operación es LayerNorm. bloque_a_mano coincide con nn.TransformerEncoderLayer: el bloque no hace nada más que esas cuatro piezas.

Ahora las dos oraciones de El problema pasan por el mismo bloque. perro está en la posición 0 de la primera y en la posición 2 de la segunda.

Pregunta. ¿El vector de perro a la salida del bloque es el mismo en las dos oraciones? ¿Las medias de las dos salidas son iguales?

x_a = oracion("perro muerde gato")
x_b = oracion("gato muerde perro")

with torch.no_grad():
    y_a = capa(x_a)
    y_b = capa(x_b)

print(x_a[0, 0], y_a[0, 0].round(decimals=3))
print(x_b[0, 2], y_b[0, 2].round(decimals=3))
print(torch.allclose(y_a[0, 0], y_b[0, 2], atol=1e-5))
print(torch.allclose(y_a.mean(dim=1), y_b.mean(dim=1), atol=1e-5))

Salida esperada

tensor([1., 0., 1., 0.]) tensor([ 0.4380, -0.9030,  1.4340, -0.9680])
tensor([1., 0., 1., 0.]) tensor([ 0.4380, -0.9030,  1.4340, -0.9680])
True
True

Lectura. El vector de perro ya no es su embedding: el bloque lo mezcló con muerde y gato. Pero es el mismo en las dos oraciones. La atención compara cada consulta con todas las claves y suma los valores ponderados. El MLP y LayerNorm actúan sobre cada token por separado. Ninguno de esos pasos usa la posición. Si se reordenan los tokens de entrada, las salidas se reordenan igual, y la media no cambia.

Codificación posicional

Una codificación posicional es un vector por posición que se suma al embedding de cada token antes del primer bloque. El transformer original usa senos y cosenos con frecuencias distintas. Para la posición pos y el par de dimensiones (2i, 2i + 1):

PE[pos, 2i] = sin(pos / 10000^(2i / d))
PE[pos, 2i + 1] = cos(pos / 10000^(2i / d))

Pregunta. Escribe la fila de la posición 0. Después de sumar pe a las dos oraciones, ¿el vector de perro a la salida del bloque sigue siendo el mismo? ¿Y las medias?

def posicional_senoidal(longitud, d):
    posiciones = torch.arange(longitud, dtype=torch.float32).unsqueeze(1)
    i = torch.arange(0, d, 2, dtype=torch.float32)
    angulos = posiciones / 10000 ** (i / d)
    pe = torch.zeros(longitud, d)
    pe[:, 0::2] = torch.sin(angulos)
    pe[:, 1::2] = torch.cos(angulos)
    return pe


pe = posicional_senoidal(3, 4)
print(pe.round(decimals=3))

with torch.no_grad():
    z_a = capa(x_a + pe)
    z_b = capa(x_b + pe)

print(z_a[0, 0].round(decimals=3))
print(z_b[0, 2].round(decimals=3))
print(torch.allclose(z_a[0, 0], z_b[0, 2], atol=1e-5))
print(torch.allclose(z_a.mean(dim=1), z_b.mean(dim=1), atol=1e-5))

Salida esperada

tensor([[ 0.0000,  1.0000,  0.0000,  1.0000],
        [ 0.8410,  0.5400,  0.0100,  1.0000],
        [ 0.9090, -0.4160,  0.0200,  1.0000]])
tensor([-1.5980,  0.6030, -0.0440,  1.0400])
tensor([ 0.4870, -1.6880,  0.2920,  0.9100])
False
False

Lectura. La posición 0 suma [0, 1, 0, 1], porque sin(0) = 0 y cos(0) = 1. Con pe, perro en la posición 0 entra al bloque con un vector distinto del de perro en la posición 2. Sus salidas ya no coinciden y las medias de las dos oraciones tampoco. El bloque es el mismo; la diferencia viene solo de la posición.

Otra opción es aprender la codificación posicional: una tabla nn.Embedding(longitud_máxima, d) que se indexa con torch.arange(longitud) y se entrena con el resto del modelo. Dos modelos conocidos usan esta opción:

BERT apila bloques del encoder como el de esta lección. BERT-base tiene 12 bloques, vectores de 768 números y 12 cabezas. No usa máscara causal: cada token atiende a los tokens de su izquierda y de su derecha.
ViT (Vision Transformer) aplica bloques del encoder a imágenes. A diferencia del bloque de esta lección, normaliza antes de cada subcapa: x = x + Atención(LayerNorm(x)). Divide una imagen de 224×224 píxeles en parches de 16×16, así que obtiene 14 × 14 = 196 parches. Cada parche se convierte en un vector con una capa lineal y funciona como un token.

Práctica

Traza guiada

Con estas matrices, escribe antes de procesar: los puntajes Q Kᵀ / √2, la clave que más pesa para cada consulta, la suma de cada fila de pesos y si la salida del token 0 queda más cerca de [4, 0] o de [0, 4]. Repite la traza con is_causal=True.

Q_traza = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
K_traza = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
V_traza = torch.tensor([[4.0, 0.0], [0.0, 4.0]])

Comprueba con torch.softmax y con F.scaled_dot_product_attention. Con la máscara causal, la salida del token 0 debe ser exactamente [4, 0]. Si no coincide, revisa en qué dimensión aplicaste softmax.

Número de cabezas

Crea nn.MultiheadAttention(embed_dim=4, num_heads=h, batch_first=True) con h igual a 1, 2 y 4. Antes de procesar, predice el número de parámetros y la forma de los pesos con average_attn_weights=False para la oración x. Comprueba ambas cosas con assert. Después intenta num_heads=3 y relaciona el error con el tamaño d / h de cada cabeza.

Codificador diminuto

Escribe un módulo CodificadorDiminuto que:

reciba índices de palabras con forma (batch, tokens) y los convierta con nn.Embedding(3, 4)
sume posicional_senoidal(tokens, 4)
aplique un nn.TransformerEncoderLayer con d_model=4, nhead=2, dim_feedforward=16, dropout=0.0 y batch_first=True
promedie los vectores de los tokens y devuelva 2 logits con nn.Linear(4, 2)

Con torch.manual_seed(0), pasa las oraciones [0, 1, 2] y [2, 1, 0] en un solo batch. Comprueba con assert que la salida tiene forma (2, 2) y que los logits de las dos oraciones son distintos. Quita el paso 2 y comprueba con torch.allclose que ahora son iguales.

Reto opcional

Genera 2000 secuencias de 5 tokens con valores de 0 a 9 y la etiqueta 1 si el primer token es mayor que el último. Adapta CodificadorDiminuto a un embedding de 10 palabras y d_model=16, y entrénalo con y sin codificación posicional. Compara la exactitud de validación de las dos versiones y explica la diferencia con lo que viste en «Codificación posicional».

Errores que revisar

Si las filas de pesos no suman 1, revisa la dimensión de softmax: debe ser la de las claves, dim=-1, no la de las consultas.
Si los pesos de nn.MultiheadAttention tienen forma (tokens, batch, batch) en vez de (batch, tokens, tokens), falta batch_first=True. Sin él, el módulo lee la entrada como (tokens, batch, embed_dim): la salida conserva la forma, pero la atención mezcla las oraciones del batch en vez de los tokens.
Si un token atiende a tokens futuros, o solo a ellos, revisa el significado de True en la máscara booleana. En nn.MultiheadAttention, True bloquea la posición. En F.scaled_dot_product_attention, True la permite.
Si nn.MultiheadAttention lanza AssertionError: embed_dim must be divisible by num_heads, elige un número de cabezas que divida a embed_dim.
Si al reordenar los tokens de entrada las salidas solo se reordenan, falta sumar la codificación posicional antes del primer bloque.

Cierre

Explica por qué el vector de perro sale igual en «perro muerde gato» y «gato muerde perro» después del bloque sin codificación posicional, y qué cambia cuando se suma pe.

Lista de comprobación

¿Qué representa cada fila de la matriz softmax(Q Kᵀ / √d_k) y por qué suma 1?
¿Qué pasa con los pesos de atención si los puntajes no se dividen entre √d_k y d_k es grande?
¿Qué salidas cambian si se modifica el último token con máscara causal?
¿Qué forma tienen la entrada y la salida de un bloque transformer y por qué permite apilar bloques?
¿Por qué la atención sin codificación posicional no distingue el orden de los tokens?

Documentación oficial

torch.nn.functional.scaled_dot_product_attention: fórmula de la atención, attn_mask e is_causal.
torch.nn.MultiheadAttention: cabezas, batch_first, attn_mask y pesos devueltos.
torch.nn.TransformerEncoderLayer: bloque del encoder con atención, residual, LayerNorm y MLP.
BERT en Hugging Face Transformers: arquitectura, configuración y uso del modelo preentrenado.
ViT en Hugging Face Transformers: división en parches, codificación posicional y uso del modelo preentrenado.

Cargando progreso…