English · Español

Lab 00 — Implementa el módulo `Embedding`¶

Lee theory/00-motivation.md y theory/01-embedding-as-lookup.md. No consultes solutions/.

Objetivo¶

Implementa un módulo Embedding respaldado por el tensor de autograd de la Fase 8. El lookup debe ser barato (sin materializar one-hot), el backward pass debe gestionar correctamente batches con ids duplicados vía np.add.at, y el módulo debe soportar save / load de embeddings entrenadas.

Setup¶

Un archivo nuevo src/minimodel/embedding.py. Tests en tests/test_phase13_embedding.py. Usa tus tipos Parameter y Tensor de las Fases ⅞.

Tareas¶

Tarea 1 — clase `Embedding`¶

class Embedding:
    """Look up dense vectors by integer id. Backed by a (V, d) parameter matrix."""

    def __init__(self, num_embeddings: int, embedding_dim: int, init_scale: float = 0.02):
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.E = Parameter(np.random.randn(num_embeddings, embedding_dim) * init_scale)

    def __call__(self, ids: NDArray[np.int64]) -> Tensor:
        """ids: (B,) or (B, T) int → (B, d) or (B, T, d) tensor with grad."""
        ...

    def save(self, path: pathlib.Path) -> None:
        """Save the embedding matrix as a .npy file plus a small JSON manifest."""

    @classmethod
    def load(cls, path: pathlib.Path) -> "Embedding":
        """Load from save()."""

Restricciones:

NumPy puro + tu autograd. Sin PyTorch.
init_scale = 0.02 coincide con la init de embedding de GPT-2.
Valida la forma de ids: array entero de cualquier dimensión con valores en [0, num_embeddings). Los ids fuera de rango deben lanzar IndexError.
save / load debe hacer round-trip exacto (bit-igual) y preservar el envoltorio Parameter del autograd.

Tarea 2 — cableado del gradiente¶

La op gather del autograd debe producir correctamente un gradiente con forma (num_embeddings, embedding_dim) a partir de un gradiente aguas arriba con forma ids.shape + (embedding_dim,). Crítico: usa np.add.at, no +=, para el scatter de vuelta.

def gather_backward(upstream_grad: NDArray, ids: NDArray, shape: tuple) -> NDArray:
    grad_E = np.zeros(shape, dtype=upstream_grad.dtype)
    np.add.at(grad_E, ids, upstream_grad)   # CRITICAL: not grad_E[ids] += upstream_grad
    return grad_E

¿Por qué np.add.at? Considera ids = [3, 3] (el mismo id dos veces) con upstream [[1, 2], [4, 5]]. Necesitamos que la fila 3 de grad_E sea [1+4, 2+5] = [5, 7]. El grad_E[ids] += upstream_grad ingenuo escribe [1, 2] y luego sobrescribe con [4, 5] por la semántica buffered de NumPy. np.add.at hace adición unbuffered y da el correcto [5, 7].

Tarea 3 — tests de propiedades¶

En tests/test_phase13_embedding.py:

Comprobación de forma. Embedding(64, 32)(np.array([1, 2, 3])) devuelve un Tensor de forma (3, 32).
Corrección del lookup. Embedding(V, d)(np.array([i])) es igual a E.value[i:i+1].
Fuera de rango lanza. Embedding(64, 32)(np.array([100])) lanza IndexError.
El gradiente fluye. Construye un grafo minúsculo: y = Embedding(64, 32)(ids).sum(); backprop; asegúrate de que E.grad es no cero exactamente en las filas de ids.
Gradiente de id duplicado (el test crítico). Con ids = [3, 3] y upstream grad_y = ones((2, d)), asegúrate de que E.grad[3] es [2, 2, ..., 2] (suma), no [1, 1, ..., 1] (last-write-wins).
Round-trip de save / load. Entrena embeddings brevemente, guarda, carga en un Embedding nuevo, comprueba bit-igual con el original.

Tarea 4 — benchmark¶

Cronometra el lookup contra el matmul ingenuo de one-hot @ matriz:

ids = np.random.randint(0, 64, size=8192)
E = Embedding(64, 32)

t0 = time.perf_counter()
for _ in range(1000):
    out = E(ids)
t1 = time.perf_counter()

# Compare to materialised one-hot:
def slow_embed(ids, E_value):
    one_hot = np.zeros((len(ids), E_value.shape[0]), dtype=np.float64)
    one_hot[np.arange(len(ids)), ids] = 1
    return one_hot @ E_value

t2 = time.perf_counter()
for _ in range(1000):
    out_slow = slow_embed(ids, E.E.value)
t3 = time.perf_counter()

Esperado: (t1 - t0) es ~100-1000× menor que (t3 - t2). Guarda en experiments/<date>-phase-13-embedding/lookup_timing.csv.

Mediciones a capturar¶

Los 6 tests de propiedades pasando.
Tiempo del lookup vs tiempo del one-hot.
Resultado del test de corrección de np.add.at.

Aceptación¶

src/minimodel/embedding.py existe; limpio con mypy --strict.
Los 6 tests de propiedades pasan.
El test de gradiente con id duplicado pasa (el bug más sutil).
El lookup es al menos 100× más rápido que la baseline one-hot.
El round-trip de save / load es bit-exacto.

Escollos esperables¶

grad_E[ids] += g silenciosamente mal para ids duplicados. Este es el bug canónico. El test de la Tarea 3.5 lo cazará; no te saltes ese test.
Desajuste de dtype float. Si tu autograd usa float32 pero la init de embedding produce float64, la acumulación de gradiente puede degradar y perder precisión. Elige un dtype (probablemente float64 para el currículo, float32 en sistemas reales) y mantén la coherencia.
load() re-randomiza. Error fácil: Embedding(num, dim) re-inicializa en __init__, y luego load se supone que debe sobrescribir self.E.value. Si olvidas la sobrescritura, obtienes una matriz aleatoria fresca en lugar de la guardada.
Manifiesto JSON vs archivo binario. Guarda la matriz float como .npy (binario, exacto) y los metadatos (forma, dtype, versión) como .json (texto). No intentes serializar la matriz a JSON — perderás precisión por el round-trip float-a-string.

Siguiente: 01-train-cbow.md

Lab 00 — Implementa el módulo Embedding¶