Skip to content

English · Español

02 — Estabilidad del softmax y log-sum-exp

Pruébalo — temperatura y softmax

🇪🇸 La página más importante de la fase 2. Tres derivaciones, tres líneas de código, y todas las fases posteriores dependen de ellas: el truco -max, log-sum-exp, y la cross-entropy estable a partir de logits crudos. Ejemplo concreto: clasificar el token siguiente como uno de los 5 tiempos verbales del modelo (§A13).

Esta es la página de teoría de la Fase 2. Re-deriva todo lo de abajo desde una página en blanco hasta que sea mecánico. Cada fase posterior (autograd, attention, bucle de entrenamiento, sampler) se apoya en estas tres reescrituras.

El planteamiento — softmax sobre §A13

La función softmax mapea un vector de logits de longitud K x = [x_1, ..., x_K] a una distribución de probabilidad:

\[ \mathrm{softmax}(x)_i = \frac{e^{x_i}}{\sum_{j=1}^{K} e^{x_j}} \]

Instancia concreta de §A13: clasificar un token como uno de los cinco tiempos. K = 5. El modelo produce un vector de logits como

x = [ 1.2, 4.7, 3.1, 0.5, 2.9 ]    # [infinitive, present, past, past-participle, future]

El softmax convierte a:

p ≈ [ 0.018, 0.612, 0.124, 0.009, 0.237 ]

El modelo piensa que "present" es lo más probable. La cross-entropy comparará esto con la etiqueta verdadera (digamos, "past simple" — índice 2) y producirá una pérdida escalar.

Eso es las matemáticas. Ahora el modo de fallo.

El fallo de overflow

Supón que el vector de logits al inicio del entrenamiento tiene un valor enorme:

x = [ 1.2, 92.0, 3.1, 0.5, 2.9 ]

Un logit de 92 no llama la atención para salidas de modelo sin normalizar. Pero:

\[ e^{92} \approx 9.2 \times 10^{39} \]

El mayor fp32 representable es ~3.4 × 10^{38}. Así que exp(92) en fp32 devuelve +∞. Entonces exp(92) / sum(exp(x)) se convierte en inf / inf, que IEEE-754 define como NaN. El vector entero se vuelve NaN. La pérdida es NaN. El gradiente es NaN. El entrenamiento está muerto.

Este no es un ejemplo artificial — pasa siempre que los logits de un modelo no se normalizan agresivamente, que es el estado predeterminado de la mayoría de arquitecturas. El softmax ingenuo no debe existir en tu base de código. El arreglo es una línea.

El truco -max — derivación

Observación: el softmax es invariante bajo un desplazamiento constante de todos los logits.

\[ \mathrm{softmax}(x + c)_i = \frac{e^{x_i + c}}{\sum_j e^{x_j + c}} = \frac{e^c e^{x_i}}{e^c \sum_j e^{x_j}} = \mathrm{softmax}(x)_i \]

Los factores e^c se cancelan. Así que podemos desplazar x por cualquier constante sin cambiar el resultado.

Elige c = -max(x). Entonces x' = x - max(x) tiene su elemento máximo igual a 0, y todos los demás elementos son ≤ 0. Así exp(x') queda acotado por exp(0) = 1 y acotado inferiormente por cero. Sin overflow — nunca.

El softmax estable:

def stable_softmax(x):
    x_shifted = x - x.max()        # max element is now 0
    e = np.exp(x_shifted)           # all entries in (0, 1]
    return e / e.sum()

Aplicado a nuestro ejemplo adversario:

x        = [ 1.2, 92.0, 3.1, 0.5, 2.9 ]
x_shifted= [-90.8, 0.0, -88.9, -91.5, -89.1 ]
e        ≈ [ 0,    1,    0,    0,    0      ]   # all underflow to ~0 except the max
softmax  ≈ [ 0,    1,    0,    0,    0      ]

El resultado es un vector one-hot en el índice 1 (present), que es el comportamiento límite correcto del softmax cuando un logit empequeñece a los demás. Sin NaN.

El underflow es aceptable. Los valores que caen por underflow a cero habrían sido e^{-91} ≈ 4 × 10^{-40} de todas formas, que está por debajo del rango denormal y habría sido cero redondeado. El truco -max garantiza que lo que se redondeó a cero estaba correctamente cerca de cero, nunca el término dominante.

Log-sum-exp — la misma idea, aplicada al log del denominador

Muchos cómputos quieren log(sum(exp(x))) directamente (p. ej., para log-likelihood o para el término del denominador en log-softmax). Ingenuo:

log(sum(exp(x)))   # overflows the same way softmax does

Aplica el mismo desplazamiento:

\[ \log \sum_{i} e^{x_i} = \log \left( e^{m} \sum_i e^{x_i - m} \right) = m + \log \sum_i e^{x_i - m} \]

Eligiendo m = max(x), cada x_i - m ≤ 0, cada exp está en (0, 1], la suma es como mucho K, y log de algo ≤ K es finito.

def log_sum_exp(x):
    m = x.max()
    return m + np.log(np.exp(x - m).sum())

Esta es la operación canónica logsumexp. scipy.special.logsumexp es la referencia; tu implementación debería coincidir con ella hasta ε fp32 en cualquier entrada.

Cross-entropy estable a partir de logits crudos

La cross-entropy entre la distribución predicha p (sobre K clases) y la etiqueta verdadera y (un entero en [0, K)) es:

\[ \mathrm{CE}(x, y) = -\log p_y = -\log \mathrm{softmax}(x)_y = -\left( x_y - \log \sum_j e^{x_j} \right) \]

La implementación ingenua es −log(softmax(x)[y]). Esto calcula softmax (que está bien si usaste el truco -max), luego toma un log, que también está bien. El resultado también está bien — pero has hecho trabajo extra, y has convertido probabilidades que redondean a cero en log(0) = -∞.

La implementación estable va directamente de logits a pérdida:

def stable_cross_entropy(x, y):
    # x: shape (K,) logits; y: integer label
    return log_sum_exp(x) - x[y]

Una pasada, sin materializar probabilidades, sin riesgo de log(0). La probabilidad de la clase correcta está implícita en la diferencia.

Verificación: expande log_sum_exp(x) - x[y] y obtienes −log(softmax(x)[y]). Misma respuesta, distinto camino, nunca produce inf a partir de logits finitos.

Ejemplo paso a paso

Para nuestro vector de logits adversario con etiqueta verdadera y = 2 (past simple):

x = [ 1.2, 92.0, 3.1, 0.5, 2.9 ]
m = 92.0
exp(x - m).sum() = exp(-90.8) + exp(0) + exp(-88.9) + exp(-91.5) + exp(-89.1) ≈ 1.0
log_sum_exp(x) = 92.0 + log(1.0) = 92.0
x[y=2] = 3.1
CE = 92.0 - 3.1 = 88.9

El modelo está muy seguro de la respuesta incorrecta (logit 92 para "present"); la cross-entropy es enorme (88.9 nats); el gradiente será enorme; el optimizador moverá los pesos fuertemente para arreglarlo. Comportamiento correcto y bien definido — gracias al truco.

Si hubieras usado softmax ingenuo + −log p_y, el camino sería:

exp(x) → [3.3, inf, 22.2, 1.6, 18.2]
exp(x).sum() → inf
softmax → [0, NaN, 0, 0, 0]
−log p_y → −log(0) → +inf

NaN envenena cada gradiente posterior. El entrenamiento muere. Una línea de código — x = x - x.max() — separa entrenamiento muerto de entrenamiento vivo.

Equivalencia numérica — qué hay que testar

Las versiones estables deben producir la misma respuesta que las versiones ingenuas cuando las ingenuas no desbordan. Entradas de test:

  1. Magnitudes pequeñas. x = [0.1, 0.2, 0.3, 0.4, 0.5] — ambos deberían concordar hasta ε fp32.
  2. Magnitudes mixtas. x = [-3, 0, 1, 2, 5] — ambos concuerdan.
  3. Grandes positivos (adversario). x = [1, 92, 3, 0, 2] — ingenuo devuelve NaN, estable devuelve una distribución válida.
  4. Grandes negativos. x = [-100, -200, -300, -400, -500] — ingenuo devuelve NaN (exp(-300) cae por underflow a 0, 0 / 0 = NaN), estable devuelve un one-hot en el máximo (x = -100 aquí).
  5. Idénticos. x = [5, 5, 5, 5, 5] — ambos devuelven [0.2, 0.2, 0.2, 0.2, 0.2].
  6. Entradas -inf (posiciones enmascaradas). x = [1, -inf, 3, 0, 2] — ambos deberían dar probabilidad cero para la entrada -inf; el truco -max lo maneja correctamente porque -inf - max ≤ 0 y exp(-inf) = 0.

El laboratorio 01-softmax-stability.md construye cada una de estas entradas de test desde un contexto de clasificación de tiempos y pide a Borja predecir las salidas ingenuas vs estables antes de ejecutar.

Por qué el truco generaliza — y dónde más vive

El truco -max es una instancia de un patrón más amplio: escala las entradas a un régimen donde la operación se comporta bien, luego compensa. Ejemplos que Borja verá después:

  • Layer normalization (Fase 10) — restar la media, dividir por la desviación típica antes de aplicar transformaciones. Misma idea, distinta motivación (estabilidad de gradiente más que overflow).
  • RMSNorm (Fase 10) — dividir solo por RMS. Los LLM modernos lo prefieren (más barato, igual de estable).
  • Attention scores (Fase 15) — dividir QK^T por √d_k antes del softmax. Esto no es numérico; es un argumento de preservación de varianza, pero tiene el efecto colateral de empujar los logits hacia el régimen donde el softmax está bien condicionado.
  • Gradient clipping (Fase 18) — reescalar gradientes cuya norma supere un umbral, para que el paso del optimizador esté bien condicionado.
  • Mixed precision training (Fase 26) — escalar la pérdida por 2^k para que los gradientes aterricen en el rango representable de fp16, luego desescalar antes de aplicar el paso del optimizador.

Cada uno es "escala a un régimen bueno, opera, compensa". El truco -max es la primera instancia que conocerás. Memorízala.

Advertencias honestas

  1. El truco -max no arregla todo. Arregla el overflow. No arregla la cancelación catastrófica en log(1 - p) cuando p ≈ 1 (usa log1p(-p) en su lugar) ni la pérdida de precisión en exp(x) - 1 para x diminuto (usa expm1(x)).
  2. El truco no es único. Podrías desplazar por min(x), por 0, por la media — cualquier constante. Se elige max porque hace que el mayor exponente sea exp(0) = 1, que es la cota más ajustada posible sobre los valores post-desplazamiento.
  3. Vectorizado sobre batches. Para un batch de vectores K, max debe calcularse por fila, no sobre el tensor entero. x.max(axis=-1, keepdims=True). El lab 01 te hace caer en esta trampa a propósito.
  4. max de todos -inf es -inf. Caso límite si cada entrada es -inf (una fila de attention totalmente enmascarada). El softmax estable debería devolver nan o un centinela aquí — discutido en el lab 01.

Problemas de práctica (trabájalos antes del laboratorio)

Soluciones en solutions/02-softmax-stability-ref.md (escritas al abrir la fase, no visibles durante la pre-escritura). Inténtalos razonando, no ejecutando.

  1. Para x = [10, 20, 30], calcula softmax(x) ingenuo y softmax(x) estable a mano con 4 cifras significativas. ¿Concuerdan? (Sí — el ingenuo no desborda con estas magnitudes.)
  2. Para x = [100, 100, 100], calcula ambos. El estable da [1/3, 1/3, 1/3]. ¿Qué da el ingenuo en fp32? ¿En fp64?
  3. Cross-entropy para x = [1.2, 4.7, 3.1, 0.5, 2.9], y = 1 (present). Calcula a mano usando log_sum_exp(x) - x[y]. Verifica contra −log(stable_softmax(x)[y]).
  4. Muestra que log_sum_exp(x) ≥ max(x) siempre, con igualdad si y solo si una entrada domina estrictamente.
  5. ¿Cuál es stable_softmax([0, 0, 0, 0, 0])? ¿Cuál es log_sum_exp([0, 0, 0, 0, 0])? Ambos tienen formas cerradas; derívalas.

Si puedes responder las cinco de memoria + papel, pasa al lab 01. Si alguna te tambalea, relee.

Recapitulación en un párrafo

El softmax ingenuo desborda siempre que algún logit sea lo bastante grande para que exp exceda el rango representable (fp32: ~89). La solución es desplazar los logits por -max(x) antes de exponenciar — el softmax es invariante bajo desplazamientos constantes, pero el desplazamiento mueve los valores post-exp a (0, 1], eliminando el overflow. El mismo desplazamiento convierte log(sum(exp(x))) en el estable m + log(sum(exp(x - m))). La cross-entropy a partir de logits crudos es entonces log_sum_exp(x) - x[y] — una pasada, sin overflow, sin log(0). Estas tres transformaciones son no negociables en toda fase posterior que toque una distribución de probabilidad.

Lo que esta página NO cubre

  • Gradiente del softmax — Fase 4.
  • Uso específico del softmax en multi-head attention — Fase 15.
  • Softmax distribuido / online (FlashAttention) — Fase 27.

Siguiente: theory/03-summation-and-cancellation.md.