Atención y transformers

import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import torch
from torch import nn

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Una fila por palabra: «Ana», «sacó», «plata»
Q = torch.tensor([[0.0, 1.0],
                  [1.0, 1.0],
                  [2.0, 0.0]])
K = torch.tensor([[0.0, 2.0],
                  [1.0, 0.0],
                  [2.0, 1.0]])
V = torch.tensor([[1.0, 0.0],
                  [0.0, 1.0],
                  [1.0, 1.0]])

puntajes = Q @ K.T / math.sqrt(2)       # (3, 3): cada palabra con cada otra
pesos = torch.softmax(puntajes, dim=-1)  # cada fila suma 1
salida = pesos @ V                       # promedios ponderados de los valores

print("Puntajes:\n", puntajes.round(decimals=2))
print("Pesos:\n", pesos.round(decimals=2))
print("Suma de cada fila:", pesos.sum(dim=-1))
print("Salida:\n", salida.round(decimals=2))

for dim in [2, 64, 512]:
    q = torch.randn(1000, dim)           # 1000 consultas al azar
    k = torch.randn(8, dim)              # 8 claves al azar
    punto = q @ k.T                      # (1000, 8)
    escalado = punto / math.sqrt(dim)
    # El peso más grande de cada fila, en promedio
    p_sin = torch.softmax(punto, dim=-1).max(dim=-1).values.mean()
    p_con = torch.softmax(escalado, dim=-1).max(dim=-1).values.mean()
    print(f"d = {dim:3d}: desviación {punto.std():5.1f} -> "
          f"{escalado.std():.2f}   peso máximo {p_sin:.2f} -> "
          f"{p_con:.2f}")

palabras = ["plata", "sacó", "Ana", "del"]
s = torch.tensor([2.0, 1.0, 0.5, -1.0])   # puntajes de una fila

fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8), sharey=True)
for eje, T in zip(ejes.flat, [0.25, 0.5, 1, 4]):
    p = torch.softmax(s / T, dim=0)
    print(f"T = {T}:", np.round(p.numpy(), 2))
    eje.bar(palabras, p, color="C0")
    eje.set_title(f"T = {T}")
    eje.tick_params(axis="x", labelsize=11)
for eje in ejes[:, 0]:
    eje.set_ylabel("Peso")
plt.show()

def atencion(Q, K, V, mascara=None):
    """Atención de producto punto escalada.

    Q: (..., Lq, d)   K: (..., Lk, d)   V: (..., Lk, dv)
    mascara: True donde NO se puede mirar (o None).
    Devuelve la salida (..., Lq, dv) y los pesos (..., Lq, Lk).
    """
    d = Q.shape[-1]
    puntajes = Q @ K.transpose(-2, -1) / math.sqrt(d)
    if mascara is not None:
        # -inf: después del softmax, peso exactamente 0
        puntajes = puntajes.masked_fill(mascara, float("-inf"))
    pesos = torch.softmax(puntajes, dim=-1)
    return pesos @ V, pesos

salida_f, pesos_f = atencion(Q, K, V)
print("Igual que a mano:", torch.allclose(salida_f, salida))

# Un lote de 4 frases, 2 «cabezas», 10 palabras, vectores de 8
q = torch.randn(4, 2, 10, 8)
salida_lote, pesos_lote = atencion(q, q, q)
print("Salida:", salida_lote.shape, "  Pesos:", pesos_lote.shape)

class AtencionMulticabeza(nn.Module):
    def __init__(self, d_modelo, cabezas):
        super().__init__()
        assert d_modelo % cabezas == 0, "d_modelo debe dividirse"
        self.cabezas = cabezas
        # Una sola capa calcula Q, K y V juntos: 3 · d_modelo salidas
        self.qkv = nn.Linear(d_modelo, 3 * d_modelo)
        self.salida = nn.Linear(d_modelo, d_modelo)

    def forward(self, x, mascara=None):
        B, L, D = x.shape                 # lote, largo, d_modelo
        q, k, v = self.qkv(x).chunk(3, dim=-1)

        def separar(t):
            # (B, L, D) -> (B, cabezas, L, D / cabezas)
            t = t.view(B, L, self.cabezas, D // self.cabezas)
            return t.transpose(1, 2)

        o, self.pesos = atencion(separar(q), separar(k), separar(v),
                                 mascara)
        # Pegar de nuevo las cabezas: (B, L, D)
        o = o.transpose(1, 2).reshape(B, L, D)
        return self.salida(o)


torch.manual_seed(0)
capa = AtencionMulticabeza(d_modelo=16, cabezas=4)
x = torch.randn(2, 5, 16)                 # 2 frases de 5 palabras
print("Entrada:", x.shape, "  Salida:", capa(x).shape)
print("Pesos guardados:", capa.pesos.shape)

oficial = nn.MultiheadAttention(16, 4, batch_first=True)
with torch.no_grad():
    oficial.in_proj_weight.copy_(capa.qkv.weight)
    oficial.in_proj_bias.copy_(capa.qkv.bias)
    oficial.out_proj.weight.copy_(capa.salida.weight)
    oficial.out_proj.bias.copy_(capa.salida.bias)

nuestra = capa(x)
# Autoatención: la misma x hace de consultas, claves y valores
suya, pesos_suyos = oficial(x, x, x, average_attn_weights=False)
print("Salidas iguales:", torch.allclose(nuestra, suya, atol=1e-6))
print("Pesos iguales:  ",
      torch.allclose(capa.pesos, pesos_suyos, atol=1e-6))

frase = ["Ana", "sacó", "plata", "del", "banco"]
L = len(frase)
# True arriba de la diagonal: las palabras del futuro
causal = torch.triu(torch.ones(L, L, dtype=torch.bool), diagonal=1)
print(causal.int())

torch.manual_seed(1)
q_frase = torch.randn(L, 8)
k_frase = torch.randn(L, 8)
_, pesos_causal = atencion(q_frase, k_frase, k_frase, causal)
print(pesos_causal.round(decimals=2))

# Blanco para el 0 y tinta para el 1
blanco_tinta = LinearSegmentedColormap.from_list("bt", ["white", "C0"])
# Las casillas del futuro quedan en blanco, sin número
tabla = np.ma.masked_where(causal.numpy(), pesos_causal.numpy())

fig, eje = plt.subplots()
eje.imshow(tabla, cmap=blanco_tinta, vmin=0, vmax=1)
for i in range(L):
    for j in range(i + 1):
        p = pesos_causal[i, j].item()
        eje.text(j, i, f"{p:.2f}", ha="center", va="center",
                 fontsize=11, color="white" if p > 0.5 else "C0")
eje.set_xticks(range(L), frase, rotation=45, ha="right")
eje.set_yticks(range(L), frase)
eje.grid(False)                         # sin rejilla sobre la tabla
eje.set_xlabel("Palabra que se mira")
eje.set_ylabel("Palabra que mira")
eje.set_title("Pesos con máscara causal")
plt.show()

torch.manual_seed(0)
x = torch.randn(1, 5, 16)
al_reves = torch.tensor([4, 3, 2, 1, 0])
y = capa(x)
y_de_la_frase_al_reves = capa(x[:, al_reves])
print("Mismo resultado, en otro orden:",
      torch.allclose(y_de_la_frase_al_reves, y[:, al_reves], atol=1e-6))

def posicion_sinusoidal(largo, d_modelo):
    pos = torch.arange(largo).unsqueeze(1)          # (largo, 1)
    i = torch.arange(0, d_modelo, 2)                # 0, 2, 4, ...
    frecuencia = 1 / 10000 ** (i / d_modelo)        # (d_modelo / 2,)
    pe = torch.zeros(largo, d_modelo)
    pe[:, 0::2] = torch.sin(pos * frecuencia)       # columnas pares
    pe[:, 1::2] = torch.cos(pos * frecuencia)       # columnas impares
    return pe

pe = posicion_sinusoidal(50, 32)
print(pe.shape)
print(pe[:3, :4].round(decimals=2))

plt.imshow(pe.T, cmap=blanco_tinta, aspect="auto")
plt.grid(False)
plt.colorbar(label="Valor (de −1 a 1)")
plt.xlabel("Posición en la frase")
plt.ylabel("Columna del vector")
plt.title("Codificación posicional sinusoidal")
plt.show()

ln = nn.LayerNorm(4)
palabra = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
print(ln(palabra).detach().round(decimals=2))

from matplotlib.patches import FancyBboxPatch

fig, eje = plt.subplots(figsize=(4.8, 4.8))
cajas = [  # (texto, centro y), de abajo hacia arriba
    ("Embeddings + posición", 0.5),
    ("LayerNorm", 1.6),
    ("Atención multicabeza", 2.7),
    ("LayerNorm", 4.9),
    ("MLP", 6.0),
    ("Salida: mismo tamaño", 8.2),
]
sumas = [3.8, 7.1]                      # los dos signos +
for texto, y in cajas:
    clave = texto.startswith("Atención")
    eje.add_patch(FancyBboxPatch(
        (0.8, y - 0.3), 2.8, 0.6, boxstyle="round,pad=0.05",
        facecolor="C3", edgecolor="C1" if clave else "C0",
        linewidth=2.5 if clave else 1.2))
    eje.text(2.2, y, texto, ha="center", va="center")
for y in sumas:
    eje.scatter([2.2], [y], s=700, facecolor="white", edgecolor="C0",
                zorder=3)
    eje.text(2.2, y, "+", ha="center", va="center", fontsize=18,
             zorder=4)
# Flechas del camino principal, de centro a centro
centros = sorted([y for _, y in cajas] + sumas)
for y0, y1 in zip(centros[:-1], centros[1:]):
    eje.annotate("", (2.2, y1), (2.2, y0),
                 arrowprops=dict(arrowstyle="->", color="C0",
                                 shrinkA=15, shrinkB=15))
# Los atajos residuales: la entrada de cada parte va directo al +
for y_desde, y_hasta in [(1.05, 3.8), (4.35, 7.1)]:
    eje.plot([2.2, 0.35, 0.35], [y_desde, y_desde, y_hasta],
             color="C2", linewidth=1.2)
    eje.annotate("", (2.2, y_hasta), (0.35, y_hasta),
                 arrowprops=dict(arrowstyle="->", color="C2",
                                 shrinkA=0, shrinkB=14))
    eje.text(0.15, (y_desde + y_hasta) / 2, "residual", rotation=90,
             color="C2", fontsize=12, ha="center", va="center")
eje.set_xlim(-0.1, 3.8)
eje.set_ylim(0, 8.7)
eje.axis("off")
eje.set_title("Un bloque transformer")
plt.show()

class BloqueTransformer(nn.Module):
    def __init__(self, d_modelo, cabezas, d_mlp):
        super().__init__()
        self.norma1 = nn.LayerNorm(d_modelo)
        self.atencion = AtencionMulticabeza(d_modelo, cabezas)
        self.norma2 = nn.LayerNorm(d_modelo)
        self.mlp = nn.Sequential(nn.Linear(d_modelo, d_mlp), nn.ReLU(),
                                 nn.Linear(d_mlp, d_modelo))

    def forward(self, x, mascara=None):
        # Las palabras se miran entre sí, y se suma (residual)
        x = x + self.atencion(self.norma1(x), mascara)
        # Cada palabra se procesa por separado, y se suma
        x = x + self.mlp(self.norma2(x))
        return x


def contar(modelo):
    return sum(p.numel() for p in modelo.parameters())

bloque = BloqueTransformer(d_modelo=64, cabezas=4, d_mlp=256)
print("Forma:", bloque(torch.randn(2, 10, 64)).shape)
print("Parámetros del bloque:", contar(bloque))
print("  atención:", contar(bloque.atencion),
      "  MLP:", contar(bloque.mlp),
      "  LayerNorm:", contar(bloque.norma1) + contar(bloque.norma2))

# El bloque equivalente que trae PyTorch
suyo = nn.TransformerEncoderLayer(64, 4, dim_feedforward=256,
                                  dropout=0.0, batch_first=True,
                                  norm_first=True)
print("Parámetros del de PyTorch:", contar(suyo))

for L_texto in [100, 1_000, 10_000, 100_000]:
    casillas = L_texto ** 2
    # 4 bytes por número (float32)
    print(f"{L_texto:>7,} palabras: {casillas:>15,} pesos "
          f"= {casillas * 4 / 1e6:>9,.2f} MB por cabeza")

from torch.utils.data import DataLoader, TensorDataset

LARGO, DIGITOS = 8, 10
g = torch.Generator().manual_seed(0)
X_ent = torch.randint(0, DIGITOS, (20000, LARGO), generator=g)
X_val = torch.randint(0, DIGITOS, (2000, LARGO), generator=g)
# La respuesta: la misma secuencia, al revés
Y_ent, Y_val = X_ent.flip(dims=[1]), X_val.flip(dims=[1])
print("Entrada: ", X_ent[0].tolist())
print("Objetivo:", Y_ent[0].tolist())

class TransformerJuguete(nn.Module):
    def __init__(self, d_modelo=32, cabezas=2, bloques=1,
                 con_posicion=True):
        super().__init__()
        self.fichas = nn.Embedding(DIGITOS, d_modelo)
        self.posicion = (nn.Embedding(LARGO, d_modelo)
                         if con_posicion else None)
        self.bloques = nn.ModuleList(
            [BloqueTransformer(d_modelo, cabezas, 4 * d_modelo)
             for _ in range(bloques)])
        self.norma = nn.LayerNorm(d_modelo)
        self.salida = nn.Linear(d_modelo, DIGITOS)

    def forward(self, x):                 # x: (B, L) con dígitos
        z = self.fichas(x)                # (B, L, d_modelo)
        if self.posicion is not None:
            z = z + self.posicion(torch.arange(x.shape[1],
                                               device=x.device))
        for bloque in self.bloques:
            z = bloque(z)
        return self.salida(self.norma(z))  # (B, L, 10)

torch.manual_seed(0)
modelo = TransformerJuguete()
print("Parámetros:", contar(modelo))

def exactitud(modelo, X, Y):
    """Fracción de dígitos y de secuencias enteras acertados."""
    modelo.eval()
    with torch.no_grad():
        pred = modelo(X.to(device)).argmax(dim=-1).cpu()
    acierto = pred == Y
    por_digito = acierto.float().mean().item()
    por_secuencia = acierto.all(dim=1).float().mean().item()
    return por_digito, por_secuencia


def entrenar_secuencias(modelo, X, Y, X_v, Y_v, epocas, tasa=1e-3):
    cargador = DataLoader(TensorDataset(X, Y), batch_size=64,
                          shuffle=True,
                          generator=torch.Generator().manual_seed(0))
    optimizador = torch.optim.Adam(modelo.parameters(), lr=tasa)
    funcion_perdida = nn.CrossEntropyLoss()
    modelo.to(device)
    for epoca in range(1, epocas + 1):
        modelo.train()
        suma = 0.0
        for x, y in cargador:
            x, y = x.to(device), y.to(device)
            optimizador.zero_grad()
            logits = modelo(x)                         # (B, L, 10)
            perdida = funcion_perdida(logits.transpose(1, 2), y)
            perdida.backward()
            optimizador.step()
            suma += perdida.item() * len(y)
        digitos, secuencias = exactitud(modelo, X_v, Y_v)
        # Exactitud de validación: por dígito y por secuencia entera
        print(f"Época {epoca}: pérdida {suma / len(X):.3f}   "
              f"dígitos {digitos:.1%}   secuencias {secuencias:.1%}")


entrenar_secuencias(modelo, X_ent, Y_ent, X_val, Y_val, epocas=2)

ejemplo = X_val[:1].to(device)
modelo.eval()
with torch.no_grad():
    pred = modelo(ejemplo).argmax(dim=-1)
print("Entrada:   ", ejemplo[0].tolist())
print("Predicción:", pred[0].tolist())

# Los pesos de la cabeza 0 que guardó la última pasada
mapa = modelo.bloques[0].atencion.pesos[0, 0].cpu()
fig, eje = plt.subplots()
eje.imshow(mapa, cmap=blanco_tinta, vmin=0, vmax=1)
eje.grid(False)
entrada = ejemplo[0].tolist()
eje.set_xticks(range(LARGO),
               [f"{d}\n{i}" for i, d in enumerate(entrada)])
eje.set_yticks(range(LARGO),
               [f"{i}: {d}" for i, d in enumerate(pred[0].tolist())])
eje.set_xlabel("Posición que se mira (dígito y posición)")
eje.set_ylabel("Posición que escribe")
eje.set_title("¿A dónde mira cada posición?")
plt.show()

torch.manual_seed(0)
sin_posicion = TransformerJuguete(con_posicion=False)
entrenar_secuencias(sin_posicion, X_ent, Y_ent, X_val, Y_val,
                    epocas=2)

from sklearn.datasets import load_sample_image

foto = load_sample_image("china.jpg")[::4, ::4][:96, :160]
imagen = torch.tensor(foto).permute(2, 0, 1) / 255   # (3, 96, 160)
P = 16                                               # lado del parche
cortes = imagen.unfold(1, P, P).unfold(2, P, P)
print("Después de unfold:", cortes.shape)   # (C, filas, columnas, P, P)
# Una fila por parche, con sus 3 · 16 · 16 números aplanados
parches = cortes.permute(1, 2, 0, 3, 4).reshape(-1, 3 * P * P)
print("Secuencia de parches:", parches.shape)

# El parche 11 es fila 1, columna 1: píxeles 16 a 31 en las dos
igual = torch.equal(parches[11].view(3, P, P),
                    imagen[:, 16:32, 16:32])
print("El parche 11 es el recorte correcto:", igual)

filas, columnas = cortes.shape[1], cortes.shape[2]
fig, eje = plt.subplots()
eje.imshow(foto)
for k in range(filas * columnas):
    f, c = divmod(k, columnas)
    eje.text(c * P + P / 2, f * P + P / 2, str(k), ha="center",
             va="center", fontsize=9, color="white",
             bbox=dict(boxstyle="round,pad=0.15", facecolor="C0",
                       edgecolor="none"))
for f in range(filas + 1):
    eje.axhline(f * P - 0.5, color="white", linewidth=1)
for c in range(columnas + 1):
    eje.axvline(c * P - 0.5, color="white", linewidth=1)
eje.axis("off")
eje.set_title("60 parches de 16 × 16 píxeles")
plt.show()

torch.manual_seed(0)
d_vit = 64
proyeccion = nn.Linear(3 * P * P, d_vit)
cls = nn.Parameter(torch.zeros(1, 1, d_vit))
pos_vit = nn.Parameter(torch.zeros(1, 1 + len(parches), d_vit))

fichas_img = proyeccion(parches.unsqueeze(0))       # (1, 60, 64)
fichas_img = torch.cat([cls, fichas_img], dim=1)    # (1, 61, 64)
z = BloqueTransformer(d_vit, 4, 4 * d_vit)(fichas_img + pos_vit)
print("Tokens de la imagen:", fichas_img.shape)
print("Después del bloque: ", z.shape)
print("Resumen ([CLS]):    ", z[:, 0].shape)