Skip to content

English · Español

Lab 01 — Romper el softmax ingenuo, luego implementar la versión estable

Objetivo: ver el overflow de softmax fp32 sobre un vector de logits de clasificación de tiempos, luego implementar la versión estable y demostrar que sobrevive a entradas adversarias.

Tiempo estimado: 60–90 minutos.

Prerrequisito: teoría 02-softmax-stability.md leída.


Lo que produces

Un directorio experiments/02-softmax-stability/ que contiene:

  • naive.py — implementación ingenua exp/sum.
  • stable.py — tu implementación estable de softmax, log_sum_exp, cross_entropy.
  • compare.py — script driver que alimenta una batería de entradas adversarias a ambos y produce una tabla comparativa.
  • results.json — la tabla.
  • softmax_break.png — visualización de dónde el softmax ingenuo explota (una fila de NaN entre salidas válidas).
  • manifest.json.
  • README.md — interpretación.

Sin módulo src/ todavía. La Fase 2 se queda en experiments/. Estas funciones se re-implementarán en src/minigrad/numerics.py en la Fase 7, cuando un consumidor de autograd exista para ellas.

El encuadre de §A13

Cada vector de test representa los logits del modelo para clasificar el tiempo del próximo token entre los cinco tiempos definidos en §A13:

[infinitive, present simple, past simple, past participle, simple future]

Índices 0..4. La "etiqueta verdadera" y es el índice entero del tiempo correcto.

TODOs

Bloque A — implementación ingenua

Escribe naive.py con tres funciones:

def naive_softmax(x):
    e = np.exp(x)
    return e / e.sum()

def naive_log_sum_exp(x):
    return np.log(np.exp(x).sum())

def naive_cross_entropy(x, y):
    return -np.log(naive_softmax(x)[y])

Esta es la implementación a romper.

Bloque B — implementación estable

Escribe stable.py:

def stable_softmax(x):
    # TODO: apply the -max trick from theory/02
    ...

def log_sum_exp(x):
    # TODO: stable log-sum-exp
    ...

def stable_cross_entropy(x, y):
    # TODO: compute directly from logits via log_sum_exp(x) - x[y]
    ...

Cada función debe:

  • Manejar entrada 1D de cualquier longitud ≥ 1.
  • Manejar entrada 2D por batch (x.shape = (batch, K)) — softmax sobre el último eje.
  • Manejar entradas -inf (tratarlas como probabilidad efectivamente cero tras el desplazamiento).
  • No depender de scipy; NumPy puro.

Bloque C — entradas adversarias

En compare.py, define y ejecuta todo lo siguiente:

test_cases = [
    ("small magnitudes",     np.array([0.1, 0.2, 0.3, 0.4, 0.5])),
    ("mixed magnitudes",     np.array([-3.0, 0.0, 1.0, 2.0, 5.0])),
    ("large positive",       np.array([1.0, 92.0, 3.0, 0.0, 2.0])),  # adversarial
    ("large negative",       np.array([-100.0, -200.0, -300.0, -400.0, -500.0])),
    ("all identical",        np.array([5.0, 5.0, 5.0, 5.0, 5.0])),
    ("masked entry",         np.array([1.0, -np.inf, 3.0, 0.0, 2.0])),
    ("single element",       np.array([42.0])),
    ("verb vocabulary",      np.zeros(600)),  # uniform over §A13 vocabulary
]

Para cada caso, ejecuta tanto naive_softmax como stable_softmax. Registra:

  • Si la salida contiene algún NaN.
  • La suma de la salida (debería ser 1.0 para una distribución válida).
  • El elemento máximo de la salida.
  • Diferencia relativa elemento a elemento entre los dos (donde el ingenuo es válido).

Para la batería de cross_entropy, fija y = 1 (present simple) y ejecuta ambas versiones en cada caso de test.

Saca la tabla como results.json y una tabla markdown en README.md.

Bloque D — predice antes de ejecutar

En README.md, antes de pegar tu results.json, escribe tus predicciones para cada caso de test:

caso de test ¿predicción NaN ingenuo? ¿predicción NaN estable? ¿predicción CE?
small magnitudes No No ~1.50 (calcula a mano)
mixed magnitudes No No ...
large positive SÍ (NaN) No ~91 (calcula vía estable: log_sum_exp - x[y])
large negative SÍ (NaN: 0/0) No ~0 (¿max está en y=1? no, max está en índice 0, x[0]=-100; CE = -100 - (-100) = 0; comprueba)
... ... ... ...

Luego ejecuta, luego compara. Donde la predicción y la realidad diverjan, escribe una frase explicando por qué. Este es el paso de aprendizaje de mayor palanca del laboratorio.

Bloque E — visualización

softmax_break.png: un heatmap o visualización tipo tabla por filas mostrando, para cada caso de test, las salidas ingenuas y estables lado a lado. Las entradas NaN en rojo. La asimetría visual en la fila "large positive" es el gráfico titular de la Fase 2.

Bloque F — previsualización de gradcheck (opcional)

Verifica, en fp64, que stable_softmax(x) concuerda con scipy.special.softmax(x) con un margen de 1e-15 en todas las entradas no adversarias, y que log_sum_exp(x) coincide con scipy.special.logsumexp(x). Si scipy difiere en los casos adversarios (no debería — scipy es estable), anótalo.

Restricciones

  • NumPy puro. Nada de scipy excepto como oráculo de referencia en el Bloque F.
  • Predice primero. No ejecutes antes de haber escrito predicciones. Todo el punto es entrenar el músculo de la predicción.
  • Usa una semilla fija para cualquier entrada de test aleatoria (np.random.default_rng(42)). Decláralo en manifest.json.

Condiciones de parada

Hecho cuando:

  1. naive.py, stable.py, compare.py existen y se ejecutan.
  2. results.json muestra NaN ingenuo en al menos dos casos y NaN estable en cero casos.
  3. README.md contiene tu tabla de predicciones antes de la tabla de resultados, con explicaciones para cualquier divergencia.
  4. softmax_break.png está commiteado.
  5. Puedes recitar, en una frase, por qué el truco -max elimina el overflow sin cambiar el resultado matemático.

Escollos

  • Manejo de -inf. np.exp(-np.inf) = 0 correctamente. Pero -np.inf - (-np.inf) = nan. Si tu truco -max resta max = -inf (porque todas las entradas son -inf), obtienes NaN en todas partes. Detecta "max es -inf" y devuelve un centinela (¿uniforme? ¿NaN? documenta la elección).
  • Max por batch. x.max() sobre un array 2D da un escalar; quieres x.max(axis=-1, keepdims=True). El laboratorio está montado para pillar esto si escribes un softmax estable consciente de batches.
  • np.log(np.exp(x).sum()) para entradas desplazadas. Tras el desplazamiento -max, np.exp(x_shifted).sum() incluye el término exp(0) = 1, así que es ≥ 1, así que np.log(...) es ≥ 0. Luego suma m de vuelta para el valor final. Si olvidas sumar m, tu log_sum_exp estará mal por exactamente m.
  • Cross-entropy en el caso enmascarado. Si la etiqueta verdadera y corresponde a un logit -inf (una posición enmascarada), x[y] = -inf, y log_sum_exp(x) - x[y] = +inf. Eso es correcto: probabilidad 0 de la verdad significa pérdida infinita. Pero es una píldora envenenada para el entrenamiento — la máscara nunca debería estar en la etiqueta verdadera. Anótalo en README.md.

Cuándo consultar solutions/

Tras commitear los seis archivos. Solución en solutions/01-softmax-stability-ref.md (escrita al abrir la fase).

Pista de último recurso

Si tu softmax estable sigue dando NaN en el caso límite de "todo -inf", la opción segura es:

if np.isneginf(m):
    return np.full_like(x, 1.0 / x.size)  # uniform fallback

Discute esta elección en README.md — es defendible (la entrada no tiene señal, así que uniforme es el fallback de principios) pero oculta ligeramente la patología de la entrada.


Siguiente laboratorio: lab/02-summation-experiments.md.