Skip to content

English · Español

01 — Batching, padding, máscaras y reducción de pérdida (loss)

Aquí se decide qué cuenta como "un ejemplo". Tres convenciones triviales (mean-vs-sum, mask de padding, shift causal) determinan si el gradiente apunta hacia donde quieres ir o hacia la mediana de tus errores de implementación.

Este fichero es el más largo de la Fase 18 porque cubre en el mismo sitio los cuatro bugs más comunes de los bucles de entrenamiento. Léelo dos veces. Las matemáticas son triviales; las convenciones no.

El batch como tensor 2-D

Un batch es un tensor de forma (B, L) donde: - B = tamaño de batch (número de secuencias en este paso). - L = longitud de secuencia (número de tokens por secuencia, tras paddear al máximo del batch).

Para el corpus de gramática verbal, un único ejemplo es un string tokenizado como:

"yo trabajo / I work"          → tokens: [<bos>, "yo", "trabajo", "/", "I", "work", <eos>]
"ella va a trabajar / she is going to work"  → [<bos>, "ella", "va", "a", "trabajar", "/", "she", "is", "going", "to", "work", <eos>]

Las longitudes varían de 6 a 14 tokens. Empaquetamos el batch paddeando todas las secuencias hasta la más larga con un token especial <pad>:

input_ids: shape (B=4, L=14)
[<bos> yo  trabajo /  I  work  <eos> <pad> <pad> <pad> <pad> <pad> <pad> <pad>]
[<bos> tú  trabajas / you work  <eos> <pad> <pad> <pad> <pad> <pad> <pad> <pad>]
[<bos> él  trabaja  / he work s <eos> <pad> <pad> <pad> <pad> <pad> <pad>     ]
[<bos> ella va  a   trabajar / she is going to work <eos>                    ]

<pad> es un token real en el vocabulario, pero nunca debe contribuir ni a la pérdida ni a la atención. Para eso están la attention mask + la loss mask. Confundir las dos, o saltarse una, es el bug #1 de esta fase.

Dos máscaras, no una

Hay dos máscaras, y tienen formas distintas:

  1. La attention mask es (B, L, L) (o (B, 1, L, L) broadcasted entre cabezas). Le dice a la capa de atención "la posición de query \(i\) no puede atender a la posición de key \(j\)" por dos razones:
  2. \(j > i\) — el futuro. Máscara causal.

  3. la posición \(j\) es <pad> en esta fila del batch. Máscara de padding. Las dos se combinan con AND en una sola máscara. La atención entonces pone esas posiciones a \(-\infty\) en los logits antes del softmax.

  4. La loss mask es (B, L). Le dice a la reducción de pérdida "la posición \((b, l)\) es un token real predicho; inclúyela en la media". Es 1 en todas partes excepto donde el token target es <pad>. (Nota: el target, no el input. Veremos por qué en un momento.)

Si accidentalmente usas solo una máscara: - Saltarte la pad mask de atención → el modelo atiende a keys <pad> y aprende a copiar predicciones <pad>, contaminando otras posiciones. - Saltarte la pad mask de pérdida → la pérdida se promedia incluyendo posiciones <pad>; el optimizador recibe gradientes que empujan al modelo a predecir <pad> con más confianza, perjudicando todo lo demás.

Ambas deben existir. Ambas deben derivarse de la misma fuente de verdad: el booleano por token is_real = (input_ids != PAD_ID). Dos vistas derivadas, un hecho.

El shift causal

Para modelado de lenguaje, el modelo predice el token \(l+1\) a partir de los tokens \(1..l\). Así que el input al modelo es input_ids[:, :-1] y el target es input_ids[:, 1:]. Este es el shift causal.

input  to model:  [<bos> yo  trabajo /  I  work  <eos> <pad> <pad> ...]
                          ↓   ↓       ↓  ↓   ↓     ↓
target (next):    [yo   trabajo /     I  work <eos> <pad>  <pad> ...]

La pérdida en la posición l mide "¿predijo el modelo target[l] a partir de input[:l+1]?". Crítico: el primer token target es yo, no <bos>; la última posición de input predice <eos> a partir del contexto previo, que es un término de pérdida significativo — la colocación de <eos> es una señal aprendida. La posición donde el target es <pad> es el primer sitio que enmascaramos: ahí, se le pide al modelo "predice <pad> a partir de tokens reales", que es absurdo y no queremos que fluyan gradientes a través de eso.

La implementación estándar:

shifted_inputs  = input_ids[:, :-1]          # (B, L-1)
shifted_targets = input_ids[:, 1:]           # (B, L-1)
loss_mask       = (shifted_targets != PAD_ID).astype(np.float32)  # (B, L-1)

logits = model(shifted_inputs)               # (B, L-1, V)
per_token_loss = cross_entropy(logits, shifted_targets)  # (B, L-1)
loss = (per_token_loss * loss_mask).sum() / loss_mask.sum()  # scalar

El denominador loss_mask.sum() es el número de tokens reales predichos en el batch. Esta es la media a nivel de token — la única reducción que hace que la LR sea significativa entre batches con distribuciones de longitud de secuencia distintas. Derivamos el porqué a continuación.

Reducción de pérdida (loss): la trampa

Puedes reducir el tensor de pérdida por token (B, L-1) a un escalar de tres maneras:

Reducción Fórmula Escala del gradiente por paso
Media a nivel de token \(\frac{1}{\sum_{b,l} M_{b,l}} \sum_{b,l} M_{b,l} \cdot \ell_{b,l}\) \(O(1)\) — invariante a \(B, L\)
Media a nivel de secuencia (suma por secuencia, luego media por batch) \(\frac{1}{B} \sum_b \sum_l M_{b,l} \ell_{b,l}\) \(O(L)\) — crece con la longitud de secuencia
Suma \(\sum_{b,l} M_{b,l} \ell_{b,l}\) \(O(B \cdot L)\) — crece con ambos

Si usas "sum" con lr=1e-3, entonces doblar el tamaño de batch reduce a la mitad la LR efectiva por token, porque los gradientes se suman pero la actualización es grad × lr. Tendrías que re-tunear lr para cada tamaño de batch. Si usas "sequence mean", re-tuneas para cada cambio en la longitud media de secuencia. La media a nivel de token es la única reducción cuyo lr no depende de las elecciones de batching.

Convención de Borja para la Fase 18 en adelante: media a nivel de token. Fíjala. Tésteala. No la cambies sin escribir el cambio en el manifest.json.

Derivación: por qué la media por token es la elección correcta

El modelo produce una distribución de probabilidad por token. El log-likelihood del corpus es:

\[\log P(\text{corpus}) = \sum_{\text{real tokens}} \log P(t \mid \text{context}(t))\]

El entrenamiento por máxima verosimilitud maximiza esta suma. El gradiente de la suma es la suma de los gradientes. Pero SGD/Adam da un paso proporcional a la media sobre el batch — eso es lo que mide lr: el tamaño de la actualización por "unidad de trabajo". Si la unidad de trabajo es "la verosimilitud de un token", entonces la media por token es lo que quieres. La media por secuencia trata cada frase como una unidad sin importar la longitud, lo que cuenta doble las frases cortas frente a las largas. La suma trata cada batch como una unidad, lo que hace que lr dependa de B.

Esto no es solo contabilidad. Un modelo entrenado con pérdida por media de secuencia sobreenfatiza sistemáticamente las frases cortas. Para las conjugaciones de verbos, las frases cortas son las regulares en presente simple (I work); las frases largas son las formas perifrásticas de futuro simple (he is going to work). Una pérdida por media de secuencia sesga el modelo hacia los casos fáciles — exactamente lo contrario de lo que queremos, que es que los casos difíciles de cola larga irregular reciban más atención, no menos.

Sequence packing (opcional, diferido)

Una optimización común es concatenar varias secuencias cortas en una secuencia larga empaquetada, con una máscara de "sequence ID" en la atención que impide la atención entre secuencias. Esto sube la utilización de la GPU en modelos de contexto largo. No hacemos esto en la Fase 18. El corpus tiene 600 formas; el padding es barato. El sequence packing introduce sus propios bugs de máscara off-by-one que no merecen depurarse a esta escala. Las notas de modern-attention de la Fase 27 tocan packing para FlashAttention.

Ensamblándolo: el paso de batching canónico

def make_batch(examples: list[list[int]], pad_id: int) -> tuple[ndarray, ndarray, ndarray]:
    """examples: list of token-id sequences (variable length)."""
    B = len(examples)
    L = max(len(x) for x in examples)
    input_ids = np.full((B, L), pad_id, dtype=np.int32)
    for i, x in enumerate(examples):
        input_ids[i, :len(x)] = x
    # Causal shift produces input/target of length L-1.
    inp = input_ids[:, :-1]              # (B, L-1)
    tgt = input_ids[:, 1:]               # (B, L-1)
    # Loss mask: 1 where target is not pad.
    loss_mask = (tgt != pad_id).astype(np.float32)   # (B, L-1)
    # Attention mask: 1 where input is not pad (used by attention layer).
    attn_pad_mask = (inp != pad_id).astype(np.float32)  # (B, L-1)
    return inp, tgt, loss_mask, attn_pad_mask

La capa de atención entonces combina internamente attn_pad_mask con una máscara triangular causal. El Lab 00 implementa esto.

El split held-out para gramática verbal

El corpus de la Fase 12 es 20 × 5 × 3 = 300 formas en inglés + 300 formas en español = 600 en total, más la estructura de pares. Necesitamos un split held-out (val).

Tres políticas de split, todas válidas:

  1. Dejar fuera 4 verbos por completo. Train sobre 16 verbos × 5 tiempos × 3 personas = 240 formas; val sobre 4 verbos × 5 tiempos × 3 personas = 60 formas. Prueba generalización a verbo nuevo: ¿puede el modelo conjugar un verbo que nunca vio?
  2. Dejar fuera 1 tiempo por completo. Train sobre 20 verbos × 4 tiempos × 3 personas = 240 formas; val sobre 20 verbos × 1 tiempo × 3 personas = 60 formas. Prueba generalización a tiempo nuevo: ¿puede el modelo manejar un patrón de tiempo que nunca vio?
  3. Dejar fuera 1 persona por completo. Train sobre 20 × 5 × 2 = 200; val sobre 20 × 5 × 1 = 100. Prueba generalización a persona nueva.

El default de la Fase 18 es la política 1: dejar fuera 4 verbos — 2 regulares (look, like) + 2 irregulares (see, eat). Esto aísla "¿puede el modelo aprender el patrón regular/irregular y aplicarlo a lexemas no vistos?". La baseline de la Fase 14 se recalcula sobre el mismo split, así que la comparación es apples-to-apples.

Un holdout uniforme aleatorio 80/20 está prohibido. Con 600 formas y un modelo de ~103k parámetros, un holdout uniforme aleatorio es trivialmente memorizable; el modelo puede interpolar celdas (verbo, tiempo, persona) promediando las formas vecinas en el espacio de embedding sin aprender la regla subyacente. Eso no es generalización — es un lookup de hoja de cálculo de alta dimensión. La política "dejar fuera dimensiones completas" fuerza un aprendizaje real de reglas.

Problemas de drill (hacer antes del lab 00)

  1. El batch tiene 4 secuencias de longitudes [6, 11, 9, 14]. Tras empaquetar y aplicar el shift causal, ¿cuál es la forma de input_ids[:, :-1]? ¿Cuántas posiciones <pad> hay en loss_mask (es decir, ceros)?
  2. Olvidaste enmascarar <pad> en la pérdida, y el 30% de los tokens del batch son <pad>. La pérdida reportada es 0.7× la pérdida real. Demuestra por qué.
  3. Pones lr = 1e-3 y entrenas con batch size 32 usando reducción token-mean. Doblas el batch a 64. ¿Necesitas re-tunear la LR? ¿Por qué o por qué no?
  4. El corpus de la Fase 12 tiene 600 formas. Dejas fuera 4 verbos. El training set tiene 240 formas. Tu val PPL es 9.0; tu train PPL es 4.5. La val PPL de la baseline de la Fase 14 es 13.0. ¿Está el modelo en sobreajuste (overfitting)? ¿Está batiendo la baseline?

Si puedes responder los cuatro, sigue. Si no, relee la sección relevante.

Recap de un párrafo

Un batch es (B, L) de token IDs, paddeados hasta la secuencia más larga. Dos máscaras derivadas de input != PAD: una attention mask (combinada con el triángulo causal) de forma (B, L, L), y una loss mask de forma (B, L-1) aplicada tras el shift causal. La pérdida se reduce como media a nivel de token para que lr sea invariante a la forma del batch. El split held-out deja fuera verbos enteros, no formas aleatorias, para forzar aprendizaje de reglas sobre memorización. Tres bugs off-by-one se esconden aquí: atención sin la pad mask, pérdida sin la pad mask, y reducción sum/sequence-mean reescalando silenciosamente el gradiente. Lee este fichero dos veces antes de lab/00.

Lo que esta sección NO cubre

  • Sequence packing para throughput. Fase 27.
  • Curriculum learning / dynamic batching. Fuera de alcance; misma config fija durante todo el run.
  • Augmentation (back-translation, token dropout). Fase 28 si acaso.

Siguiente: theory/02-optimizer-and-schedule.md.