Softmax y Temperatura

De Puntuaciones a Probabilidades

Las redes neuronales producen puntuaciones brutas llamadas logits — números reales sin acotar que pueden ser positivos, negativos o cero. Una red de clasificación con cuatro clases de salida podría producir logits como $[2.0, 1.0, 0.1, -1.0]$. Estos números nos dicen algo sobre las preferencias relativas, pero no son probabilidades: no suman 1 y algunos son negativos. Para interpretarlos como una distribución de probabilidad (valores positivos que sumen 1), necesitamos una función que los normalice. Softmax es esa función.

Desglosemos cada parte de esta fórmula.

$z_i$ es el logit bruto para la clase $i$. Puede ser cualquier número real: positivo, negativo o cero. Es la entrada al softmax — la puntuación que la red asigna a la clase $i$ antes de cualquier normalización.

$e^{z_i}$ es la exponencial del logit. La función exponencial hace que todos los valores sean estrictamente positivos, ya que $e^x > 0$ para todo número real $x$. Además, amplifica las diferencias: logits más grandes se vuelven exponencialmente más grandes. Un logit de 2 se mapea a $e^2 \approx 7.4$, mientras que un logit de 4 se mapea a $e^4 \approx 54.6$ — una diferencia de 2 en el espacio de logits se convierte en un factor de aproximadamente 7.4 en el espacio exponencial.

$\sum_{j=1}^{K} e^{z_j}$ es la constante de normalización. Sumamos las exponenciales de los $K$ logits y luego dividimos cada exponencial individual por esta suma. Esto garantiza que todas las salidas sumen 1, dándonos una distribución de probabilidad válida.

Rango de salida : cada $\text{softmax}(z_i) \in (0, 1)$ (estrictamente entre 0 y 1, nunca exactamente 0 ni 1) y $\sum_{i=1}^{K} \text{softmax}(z_i) = 1$.

Lo que ocurre en los extremos es instructivo. Si un logit es mucho mayor que el resto (por ejemplo, $z_1 = 10$, todos los demás $\approx 0$), entonces $\text{softmax}(z_1) \approx 1$ y todas las demás salidas son $\approx 0$. El softmax se aproxima a un hard argmax — casi toda la masa de probabilidad recae en el ganador. Si todos los logits son iguales ($z_i = c$ para todo $i$), entonces todas las exponenciales son iguales, así que $\text{softmax}(z_i) = 1/K$ para todo $i$ — una distribución perfectamente uniforme. Y si sumamos una constante $c$ a todos los logits, $\text{softmax}(z_i + c) = \text{softmax}(z_i)$, porque el factor $e^c$ se cancela entre el numerador y el denominador. Solo importan las diferencias relativas entre logits.

Veamos esto paso a paso en código.

import numpy as np

def softmax(z):
    # Subtract max for numerical stability (doesn't change result)
    z_stable = z - np.max(z)
    exp_z = np.exp(z_stable)
    return exp_z / np.sum(exp_z)

# Example: 4-class classification logits
logits = np.array([2.0, 1.0, 0.1, -1.0])

print("Step-by-step softmax:")
print(f"  Logits:        {logits}")
print(f"  Subtract max:  {logits - np.max(logits)}")
print(f"  Exponentials:  {np.exp(logits - np.max(logits)).round(4)}")
print(f"  Sum of exp:    {np.exp(logits - np.max(logits)).sum():.4f}")
print(f"  Softmax:       {softmax(logits).round(4)}")
print(f"  Sum:           {softmax(logits).sum():.6f}")
print()

# All equal -> uniform
equal = np.array([1.0, 1.0, 1.0, 1.0])
print(f"  Equal logits {equal} -> softmax {softmax(equal).round(4)} (uniform)")

# One dominant -> near argmax
dominant = np.array([10.0, 0.0, 0.0, 0.0])
print(f"  Dominant {dominant} -> softmax {softmax(dominant).round(6)} (near argmax)")

¿Por Qué Exponenciales?

Surge una pregunta natural: ¿por qué usar $e^{z_i}$ específicamente? ¿Por qué no elevar al cuadrado los logits ($z_i^2$), tomar valores absolutos ($|z_i|$), o usar alguna otra función para hacer las cosas positivas antes de normalizar? Hay tres razones convincentes por las que la exponencial es la elección correcta.

Primero, positividad . Necesitamos que todos los valores sean positivos para formar una distribución de probabilidad válida. La exponencial satisface esto: $e^x > 0$ para todo $x$, incluyendo valores negativos. Elevar al cuadrado también hace las cosas positivas, pero...

Segundo, monotonía . La exponencial es monótonamente creciente — logits más grandes siempre producen exponenciales más grandes, lo que significa probabilidades más altas. Esto preserva el ranking: si la red cree que la clase A es más probable que la clase B (logit más alto), la clase A recibe una probabilidad mayor. Elevar al cuadrado rompería esto: $z^2$ mapea $z = -5$ a 25 y $z = 2$ a 4, invirtiendo completamente el ranking. Una función como $|z|$ tiene el mismo problema.

Tercero, propiedades del gradiente . Cuando el softmax se combina con la pérdida de entropía cruzada (la pérdida estándar para clasificación, que cubriremos en el siguiente artículo), el gradiente se simplifica elegantemente a $p_i - y_i$ — la probabilidad predicha menos la etiqueta verdadera. Este gradiente limpio proviene específicamente de la familia exponencial de distribuciones y hace que la optimización sea estable y eficiente. Otras funciones que hacen los valores positivos producirían gradientes mucho más complicados.

Temperatura: Controlando la Nitidez

A veces no queremos la distribución estándar del softmax. Podríamos querer una distribución más nítida (más confiada, más determinista) o una más plana (más incierta, más exploratoria). El escalado por temperatura nos da este control dividiendo los logits por un parámetro $\tau$ (tau) antes de aplicar softmax.

El parámetro $\tau$ se llama temperatura , tomando prestada la terminología de la mecánica estadística donde la temperatura controla la aleatoriedad de los estados de las partículas. Veamos qué hacen los diferentes valores de temperatura.

$\tau = 1$ : softmax estándar. Sin cambios — dividir por 1 no tiene efecto.

$\tau \to 0^+$ (temperatura baja): dividir por un número positivo muy pequeño hace que todos los logits sean enormes en magnitud, amplificando las diferencias entre ellos. El softmax se aproxima a un hard argmax — casi toda la masa de probabilidad se concentra en el logit más grande. El modelo se vuelve muy confiado y determinista.

$\tau \to \infty$ (temperatura alta): dividir por un número enorme hace que todos los logits se aproximen a 0, borrando las diferencias entre ellos. El softmax se aproxima a una distribución uniforme — el modelo se vuelve máximamente incierto y aleatorio, asignando probabilidad igual a cada clase.

En resumen: $\tau < 1$ (temperatura baja) produce distribuciones más nítidas y confiadas. $\tau > 1$ (temperatura alta) produce distribuciones más planas e inciertas. La temperatura es una única perilla que interpola suavemente entre "siempre elegir la mejor opción" y "elegir uniformemente al azar."

El siguiente gráfico lo hace concreto. Tomamos el mismo conjunto de cinco logits y aplicamos softmax a cinco temperaturas diferentes, mostrando cómo la distribución de probabilidad cambia de casi one-hot ($\tau$ baja) a casi uniforme ($\tau$ alta).

import numpy as np
import json
import js

logits = np.array([2.0, 1.0, 0.5, -0.5, -1.0])
classes = ["A", "B", "C", "D", "E"]

def softmax_temp(z, tau):
    z_t = z / tau
    z_t = z_t - np.max(z_t)
    exp_z = np.exp(z_t)
    return exp_z / np.sum(exp_z)

temperatures = [0.2, 0.5, 1.0, 2.0, 5.0]
colors = ["#ef4444", "#f59e0b", "#3b82f6", "#10b981", "#8b5cf6"]

lines = []
for tau, color in zip(temperatures, colors):
    probs = softmax_temp(logits, tau)
    lines.append({"label": f"\u03c4 = {tau}", "data": probs.tolist(), "color": color})

plot_data = [{
    "title": "Softmax with Different Temperatures",
    "x_label": "Class",
    "y_label": "Probability",
    "x_data": classes,
    "lines": lines
}]
js.window.py_plot_data = json.dumps(plot_data)

Temperatura en la Práctica

La temperatura no es solo una curiosidad teórica — se usa en todas partes del aprendizaje automático moderno, a menudo como uno de los hiperparámetros más importantes. Estas son tres aplicaciones principales.

Muestreo en LLMs. Cuando un modelo de lenguaje grande genera texto, produce logits sobre todo el vocabulario en cada paso. La temperatura controla cómo se muestrea el siguiente token de la distribución resultante. Una temperatura alta (1.0-1.5) fomenta texto creativo y diverso al distribuir la probabilidad entre muchos tokens. Una temperatura baja (0.1-0.5) produce salidas más factuales y deterministas al concentrar la probabilidad en los tokens más probables. Temperatura = 0 es equivalente a la decodificación voraz (greedy decoding): siempre elegir el token más probable, sin ninguna aleatoriedad.

Destilación de conocimiento. Al entrenar un modelo "estudiante" más pequeño para imitar a un "profesor" más grande, Hinton et al. (2015) demostraron que usar una temperatura alta (típicamente $\tau = 4$ a $\tau = 20$) en las salidas softmax de ambos modelos es crucial. A $\tau = 1$, las predicciones del profesor suelen ser casi one-hot — la mayor parte de la probabilidad está en una clase — así que el estudiante solo aprende "la respuesta es la clase 3." A temperatura alta, la distribución se suaviza, revelando las clasificaciones relativas de todas las clases. El profesor podría mostrar que la clase 5 es más plausible que la clase 7, aunque ambas tengan probabilidad mínima a $\tau = 1$. Este "conocimiento oscuro" en la cola de la distribución contiene información rica sobre similitudes entre clases que ayuda al estudiante a generalizar mejor.

Aprendizaje contrastivo. En modelos como CLIP , se usa una temperatura baja ($\tau \approx 0.07$) en la pérdida contrastiva. La pérdida empuja a los pares imagen-texto coincidentes a tener alta similitud y a los pares no coincidentes a tener baja similitud. Una temperatura baja hace que esta pérdida sea más nítida, forzando al modelo a discriminar más agresivamente entre pares coincidentes y no coincidentes. La temperatura es un parámetro aprendido en CLIP, comenzando alrededor de 0.07 y adaptándose durante el entrenamiento.

Log-Softmax y Estabilidad Numérica

En la práctica, casi siempre necesitamos $\log(\text{softmax}(z_i))$ en lugar de $\text{softmax}(z_i)$ directamente. La razón es que la función de pérdida estándar para clasificación — la entropía cruzada, que cubriremos en el siguiente artículo — involucra el logaritmo de la probabilidad predicha. Calcular el softmax primero y luego tomar el logaritmo es numéricamente peligroso: el softmax puede producir valores extremadamente cercanos a 0 (por ejemplo, $10^{-45}$), y $\log(0) = -\infty$. Incluso valores que son simplemente muy pequeños pueden perder precisión al almacenarse en punto flotante.

La solución es calcular el log-softmax directamente, sin materializar nunca los valores intermedios del softmax.

Esto se obtiene tomando el logaritmo de la fórmula del softmax: $\log(e^{z_i} / \sum_j e^{z_j}) = z_i - \log \sum_j e^{z_j}$. El término clave es $\log \sum_j e^{z_j}$, conocido como log-sum-exp . Se puede calcular de forma estable factorizando el logit máximo $m = \max_j z_j$:

Después de restar $m$, todos los exponentes son $\leq 0$, así que ninguna exponencial se desborda. La mayor exponencial es $e^0 = 1$, que es perfectamente segura. Por eso PyTorch proporciona F.log_softmax y F.cross_entropy (que fusiona internamente log-softmax con el cálculo de la pérdida) — usan esta identidad internamente para evitar el paso intermedio peligroso.

El siguiente código demuestra por qué esto importa.

import numpy as np

logits = np.array([100.0, 101.0, 102.0])  # large logits

# Naive: overflow!
try:
    exp_z = np.exp(logits)
    naive_softmax = exp_z / np.sum(exp_z)
    naive_log_softmax = np.log(naive_softmax)
    print(f"Naive log-softmax: {naive_log_softmax}")
except:
    print("Naive approach: overflow or nan!")

# Stable: subtract max, then use log-sum-exp identity
m = np.max(logits)
log_sum_exp = m + np.log(np.sum(np.exp(logits - m)))
stable_log_softmax = logits - log_sum_exp
print(f"Stable log-softmax: {stable_log_softmax.round(4)}")
print(f"Sum of exp(log-softmax): {np.exp(stable_log_softmax).sum():.6f} (should be 1.0)")

Quiz

Pon a prueba tu comprensión de softmax y temperatura.