English · Español

01 — `E[i] = one_hot(i) @ E`: la identidad lookup-es-matmul¶

🇪🇸 Esta es la identidad más útil de la Fase 13. Una búsqueda en una matriz de embeddings es matemáticamente idéntica a un producto matriz-vector con un one-hot. Lo importante: nunca materializamos el one-hot. La implementación es indexación; la matemática es matmul.

La identidad¶

Sea \(E \in \mathbb{R}^{V \times d}\) una matriz de embeddings e \(i \in \{0, 1, \ldots, V-1\}\) un id de token. Sea \(e_i \in \{0, 1\}^V\) el vector one-hot con un 1 en la posición \(i\).

Entonces:

\[E[i, :] = e_i^\top E\]

El lado izquierdo es indexación (tiempo constante, sin matemática). El lado derecho es un matmul (un vector fila por la matriz). Son literalmente el mismo vector. La demostración es la definición de la multiplicación de matrices:

\[\big(e_i^\top E\big)_k = \sum_{j=0}^{V-1} (e_i)_j \cdot E_{j, k} = E_{i, k}\]

ya que \((e_i)_j = 0\) para \(j \neq i\) y \((e_i)_i = 1\).

Por qué esto importa¶

Práctico: elección de implementación¶

Para el forward pass, nunca materialices el one-hot. La indexación es \(O(1)\) + copia de memoria de \(d\) floats. El matmul con un one-hot es \(O(Vd)\) y reserva un vector cero \(V\)-dimensional. Mismo resultado, coste drásticamente distinto.

# Slow, materialises one-hot:
def embed_slow(i: int, E: np.ndarray) -> np.ndarray:
    one_hot = np.zeros(E.shape[0])
    one_hot[i] = 1
    return one_hot @ E

# Fast, indexed:
def embed_fast(i: int, E: np.ndarray) -> np.ndarray:
    return E[i]

Para un batch de \(B\) ids de token: E[ids] produce forma (B, d) en tiempo \(O(Bd)\). Esto es lo que implementa el lab de la Fase 13.

Conceptual: flujo del gradiente¶

Como el lookup es matmul, el flujo del gradiente a través de un lookup queda bien definido por la regla de la cadena estándar del autograd. Si una pérdida aguas abajo \(\mathcal{L}\) depende de \(E[i]\) con gradiente \(g = \partial \mathcal{L} / \partial E[i] \in \mathbb{R}^d\), entonces \(\partial \mathcal{L} / \partial E\) es una matriz dispersa con la fila \(i\) igual a \(g\) y todas las demás filas cero.

# In autograd terms:
# Forward:  v = E[i]
# Backward: dE = scatter(g, into=row[i], shape=(V, d))
#           (all other rows zero)

Para un batch de ids \([i_0, i_1, \ldots, i_{B-1}]\) y gradientes por fila \(g_b\):

Si todos los \(i_b\) son distintos, cada gradiente cae en su propia fila de \(\partial \mathcal{L} / \partial E\).
Si dos posiciones del batch comparten el mismo id (p. ej., la palabra the aparece dos veces en una frase), sus gradientes se suman en la misma fila. Esta es la regla de la cadena estándar para "el mismo parámetro usado varias veces".

Este es el gradiente correcto (y el único). También es la razón por la que las tablas de lookup reciben optimizadores amigables con la dispersión (las tasas adaptativas por fila de Adagrad, actualizaciones SGD dispersas) en producción — la mayoría de filas tienen gradiente cero la mayoría de pasos.

Esbozo de implementación compatible con autograd¶

En src/minimodel/embedding.py:

class Embedding:
    def __init__(self, num_embeddings: int, embedding_dim: int):
        # Initialise small Gaussian, scale 0.02 (GPT-style)
        self.E = Parameter(np.random.randn(num_embeddings, embedding_dim) * 0.02)

    def __call__(self, ids: NDArray[np.int64]) -> Tensor:
        """ids: (B,) or (B, T) int → (B, d) or (B, T, d) tensor with grad."""
        return self.E[ids]   # autograd handles the gather and the scatter-back

La maquinaria del autograd (Fase 8) debe soportar indexación entera sobre un Tensor. Si tu autograd de la Fase 8 aún no la soporta, esta es la operación a añadir:

Forward: Tensor.gather(ids) devuelve las filas.
Backward: dado el gradiente aguas arriba \(g\) con forma (B, d), produce el gradiente para self.E con forma (V, d) mediante scatter: np.add.at(grad_E, ids, g). (El np.add.at es sin buffer — gestiona ids duplicados correctamente.)

np.add.at es la primitiva correcta. grad_E[ids] += g (el código obvio) silenciosamente descarta las contribuciones de gradiente de ids duplicados. Usa np.add.at(grad_E, ids, g). El lab 00 testeará esto.

Tied embeddings entrada/salida (referencia adelantada)¶

La Fase 17 atará la matriz de embeddings de entrada con la proyección de salida de la cabeza LM. Cuando están tied:

El mismo parámetro \(E\) se usa tanto para el lookup de entrada (\(E[ids]\)) como para el desembedding de salida (\(h \cdot E^\top\)).
Los gradientes de ambas contribuciones se suman en \(\partial \mathcal{L} / \partial E\).

Esto es correcto bajo autograd, pero requiere que tu librería de tensores trate "mismo parámetro reusado" del modo estándar (el gradiente es la suma de todas las contribuciones). El test de gradiente del lab 00 ejercita el caso de uso único; la Fase 17 se apoyará en la misma propiedad para los tied embeddings.

Una confusión común: embedding vs codificación one-hot para etiquetas¶

La identidad lookup-es-matmul aplica a las embeddings de entrada. Para los targets (las etiquetas en la pérdida de entropía cruzada), a veces usamos one-hot:

\[\mathcal{L} = -\sum_i y_i \log q_i \quad \text{(}y \text{ is one-hot)}\]

O, equivalentemente, simplemente \(-\log q_{y^*}\) donde \(y^*\) es el índice entero. Misma matemática, la forma entera es más barata. El cross_entropy_from_logits(z, y_star) de la Fase 05 usa la forma entera.

Así: tokens de entrada → índices enteros en \(E\) (lookup). Tokens objetivo → índices enteros en la pérdida (nunca se materializa un one-hot). Simétrico y eficiente.

Qué NO cubre este archivo¶

Embeddings posicionales. Fase 16. Aquí embeddings de token; info posicional añadida más tarde.
Embeddings multi-token (subword). Las Fases 11 / 12 trataron la tokenización; aquí cada token es un único entero.
El objetivo de entrenamiento. Siguiente archivo (CBOW).

Siguiente: 02-cbow-skipgram.md

01 — E[i] = one_hot(i) @ E: la identidad lookup-es-matmul¶