Skip to content

English · Español

Break 00 — Atar embeddings de entrada/salida a tamaños de vocabulario diferentes

🇪🇸 Rompemos la atadura entre embedding de entrada (V_in = 512) y la cabeza LM de salida (V_out = 256). La Fase 17 va a "atar" estas dos matrices (W_out = E.T) — pero si tienen tamaños distintos, falla en runtime con shape mismatch. Es un bug clásico que aparece al fusionar tokenizadores.

Anclas: LYNX_CORTEX.md §4 / FASE 13; teoría §01 embedding como lookup; Fase 17 §03 tied embeddings; .claude/commands/break.md.

El break

En src/minimodel/nn/embedding.py (donde vive la embedding de §A13) y src/minimodel/nn/lm_head.py:

# embedding.py
class Embedding(Module):
    def __init__(self, vocab_size: int, d_model: int) -> None:
        super().__init__()
        self.vocab_size = vocab_size
        self.weight = Parameter(np.random.uniform(...,(vocab_size, d_model)))
    ...

# lm_head.py — the "tied" version:
class LMHead(Module):
    def __init__(self, embedding: Embedding, vocab_size_out: int) -> None:
        # BUG: vocab_size_out is wrong — should be the same as embedding.vocab_size.
        super().__init__()
        self.embedding = embedding
        self.vocab_size_out = vocab_size_out  # e.g., 256 instead of 512

    def forward(self, h: Tensor) -> Tensor:
        # Untied case would be: logits = h @ W_lm.T where W_lm is (V_out, d).
        # Tied case: logits = h @ self.embedding.weight.T  → shape (B, T, V_in).
        # But the rest of the pipeline expects (B, T, V_out).
        logits = h @ self.embedding.weight.transpose((1, 0))
        return logits  # shape (B, T, V_in) ≠ (B, T, V_out) — but who is checking?

El bug es el desajuste entre embedding.vocab_size = 512 y lm_head.vocab_size_out = 256. Si los atas ingenuamente, los logits tienen forma (B, T, 512) mientras que la entropía cruzada espera (B, T, 256).

Predice, luego ejecuta

Predicciones

  1. Crash en runtime en la llamada de entropía cruzada: 
    ValueError: shape mismatch — logits (B, T, 512) vs targets (B, T) with 256 classes
  2. Si silencias la comprobación (p. ej., lookup por índice ingenuo contra vocab_size_out), el modelo entrena pero sus primeros 256 huecos de vocabulario se convierten en una quimera — en parte las primeras 256 filas de la embedding, en parte la matriz (256, d_model) de la cabeza LM.
  3. Curva de pérdida: empieza en log(256) = 5,55 (baseline aleatoria para 256 clases) y se queda ahí o crece.

Escribe las predicciones en learners/borja/phase-13/notes/breaks.md antes de ejecutar.

Observa

just exp 13-train-cbow --tag broken-tied-mismatch

Diagnósticos:

  1. Mira el traceback en el primer forward pass. El shape mismatch debería ser ruidoso.
  2. Si el mismatch está silenciado (un bug relacionado), comprueba logits.shape == (B, T, V_out) explícitamente.
  3. La pérdida debería ser log(V_out) — verifica con np.log(V_out) ≈ 5,55 para V_out = 256.

Síntoma que verá Borja

  • ValueError o RuntimeError en la llamada de entropía cruzada.
  • El traceback apunta o bien a .weight.T de la embedding o a la comprobación de forma de la función de pérdida.
  • Si está silenciado: la pérdida empieza en log(V_out) y se queda ahí.

Causa oculta (una frase)

La cabeza LM reutiliza la matriz de embeddings (tied) pero el tamaño de vocabulario de la embedding es 512 mientras que el tamaño de salida esperado por la cabeza LM es 256 — desajuste de forma en la ruta atada.

Cascada de pistas

  1. ¿Cuál es la forma de embedding.weight, y qué forma produce h @ embedding.weight.T?
  2. ¿Qué forma espera la pérdida de entropía cruzada para los logits?
  3. Compara embedding.vocab_size con el valor pasado a LMHead(vocab_size_out=...).

Diff de arreglo

class LMHead(Module):
    def __init__(self, embedding: Embedding) -> None:
        super().__init__()
        self.embedding = embedding
        self.vocab_size_out = embedding.vocab_size   # take from the source

(En el mini-GPT de la Fase 17 lo veremos codificado como model.lm_head = TiedHead(model.token_emb).)

Por qué esto enseña el concepto

El weight tying (Press & Wolf 2017, "Using the Output Embedding to Improve Language Models", e Inan et al. 2017) es un truco estándar de transformer que ata W_lm = E.T para reducir el número de parámetros a la mitad y mejorar la perplejidad en ~5%. Requiere V_in == V_out. El bug — pasar tamaños de vocabulario inconsistentes — es la forma más común en que se rompe el tying en la práctica (un refactor cambia el vocabulario del tokenizer pero no el de la cabeza LM). El lab de mini-GPT de la Fase 17 usará tied embeddings; este break asegura que Borja tenga el reflejo de disciplina de forma antes de la Fase 17.

Siguiente: El /break de la Fase 14 sobre el sigmoide de la puerta LSTM.