Embeddings: enteros que se vuelven geometría

Tu tokenizador produce enteros como [76, 299, 412, ...]. Un entero aislado no tiene estructura útil: el token 300 no está "más cerca" del 301 que del 7. El primer trabajo del modelo es mapear cada id a un vector denso de dimensión $d$ (n_embd) en el que la cercanía sí significa algo. A ese mapeo se lo llama embedding (incrustación).

Token embeddings: una tabla que se aprende

Un embedding es, literalmente, una tabla de búsqueda: una matriz $W_E \in \mathbb{R}^{V \times d}$, con $V$ el tamaño del vocabulario. El vector del token $t$ es la fila $t$:

el vector de un token es, literalmente, su fila en la tabla $W_E$: una lista de $d$ números

Equivale a multiplicar un vector one-hot por $W_E$, pero indexar es muchísimo más barato. En PyTorch:

import torch
import torch.nn as nn

V, d = 512, 128
tok_emb = nn.Embedding(V, d)

ids = torch.tensor([[76, 299, 412]])  # (B=1, T=3)
x = tok_emb(ids)                      # (1, 3, 128)
print(x.shape)

Al inicio la tabla es ruido aleatorio. Durante el entrenamiento, el gradiente solo actualiza las filas de los tokens que aparecieron en el batch, y tokens que se usan en contextos parecidos terminan con vectores parecidos.

El problema: la atención no sabe de orden

En la próxima lección verás que la self-attention trata su entrada como un conjunto: si permutas los tokens de entrada, las salidas se permutan igual (es equivariante a permutaciones). Sin información extra, "el perro mordió al hombre" y "el hombre mordió al perro" serían indistinguibles para la atención salvo por qué token está dónde... cosa que la atención no puede ver.

La solución es inyectar la posición en el vector. El paper original del Transformer propuso codificaciones sinusoidales fijas:

a cada posición $pos$ le toca una huella de senos y cosenos: cada par de dimensiones ($2i$ y $2i+1$) oscila a su propio ritmo, rápido en las primeras y lento en las últimas

Cada par de dimensiones oscila con una frecuencia distinta, desde muy rápida hasta muy lenta, como las agujas de un reloj con muchas escalas. Explóralo:

GPT-2, y nuestro mini-GPT, usan en cambio embeddings posicionales aprendidos: otra tabla, $W_P \in \mathbb{R}^{T_{max} \times d}$, con una fila por posición posible. El vector que entra al primer bloque es la suma:

lo que entra al modelo en la posición $i$: el vector de su token más el vector de esa posición

block_size = 64                         # largo máximo de contexto
pos_emb = nn.Embedding(block_size, d)

B, T = ids.shape
pos = torch.arange(T)                   # (T,)
h0 = tok_emb(ids) + pos_emb(pos)        # (B, T, d) + (T, d) -> broadcasting

Fíjate en el broadcasting: pos_emb(pos) tiene forma (T, d) y se suma a cada elemento del batch.

¿Sumar o concatenar?

Sumar parece destruir información: ¿cómo separa el modelo "qué token" de "qué posición" si están mezclados? En la práctica, con $d$ grande hay espacio de sobra: el modelo puede aprender a usar subespacios aproximadamente ortogonales para cada señal. Concatenar funcionaría también, pero aumenta la dimensión sin un beneficio claro. La suma es la convención.

Weight tying: la entrada y la salida comparten matriz

Al final del modelo necesitamos lo inverso: pasar de un vector $h \in \mathbb{R}^d$ a un puntaje por cada token del vocabulario (los logits). Eso es una capa lineal con matriz $V \times d$... la misma forma que $W_E$. El weight tying (atado de pesos) reutiliza la matriz de embeddings:

para predecir, comparamos el vector final $h$ con la fila de cada token del vocabulario: sale un puntaje por token

La intuición: el logit del token $t$ es el producto punto entre el estado final $h$ y el embedding de $t$. Predecir "el siguiente token es $t$" equivale a decir "el estado final apunta en la dirección de $t$". Usar el mismo espacio para leer y para escribir es una restricción razonable que además ahorra parámetros.

¿Cuántos? En GPT-2 small, $V = 50257$ y $d = 768$, así que la matriz tiene $50257 \times 768 = 38{,}597{,}376$ parámetros: unos 38,6 millones de los ~124 millones del modelo. Sin tying, serían el doble.

En PyTorch, atar pesos es una sola línea: ambos módulos apuntan al mismo tensor Parameter.

lm_head = nn.Linear(d, V, bias=False)
lm_head.weight = tok_emb.weight          # mismo objeto: los gradientes se suman

assert lm_head.weight.data_ptr() == tok_emb.weight.data_ptr()

nn.Linear(d, V) guarda su peso con forma (V, d), idéntica a la de nn.Embedding(V, d). Por eso la asignación directa funciona sin transponer.

Juntando las piezas

Esta es la "puerta de entrada" de nuestro GPT. La vamos a completar en la lección del bloque Transformer, pero ya puedes escribirla:

from dataclasses import dataclass

@dataclass
class GPTConfig:
    vocab_size: int = 512
    block_size: int = 64
    n_layer: int = 4
    n_head: int = 4
    n_embd: int = 128
    dropout: float = 0.1


class Embeddings(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.tok = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.pos = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        return self.drop(self.tok(idx) + self.pos(pos))

Lo que te llevas

• Un embedding es una tabla $V \times d$ que se aprende por gradiente.
• La atención no ve el orden: hay que sumar información posicional (sinusoidal o aprendida).
• Weight tying reutiliza $W_E$ como capa de salida: menos parámetros y un espacio compartido para leer y escribir tokens.
• Próxima lección: el corazón del modelo, la atención causal.