Softmax

Softmax es una función que convierte un vector de puntuaciones sin procesar, llamadas logits, en una distribución de probabilidad: cada valor de salida está entre 0 y 1, y la suma de las salidas es exactamente 1. Para un vector z, calcula softmax(z)ᵢ = e^(zᵢ) / Σⱼ e^(zⱼ): eleva cada puntuación a una potencia y la divide por la suma de todas las exponenciales. Softmax es la capa de salida estándar de las redes neuronales que eligen una clase entre muchas; convierte las puntuaciones de similitud en pesos dentro del mecanismo de atención de los Transformer y transforma las puntuaciones de un modelo de lenguaje sobre su vocabulario en probabilidades del siguiente token. Su nombre refleja que es una versión suave y diferenciable de argmax, que asignaría toda la probabilidad a la puntuación más alta.

Cómo funciona Softmax#

La exponencial hace que todas las salidas sean positivas y amplifica las diferencias entre puntuaciones, por lo que el logit más alto recibe una proporción desmesurada de la probabilidad. Softmax conserva el orden de los logits, de modo que la puntuación más alta siempre obtiene la probabilidad más alta. Solo importan las diferencias entre logits: sumar la misma constante a todos los logits no cambia la salida. Las bibliotecas aprovechan esta propiedad para lograr estabilidad numérica restando el logit máximo antes de calcular la exponencial, ya que e¹⁰⁰⁰ desborda la representación en coma flotante de 32 bits.

Una temperatura T controla lo concentrada que está la distribución: softmax(z / T). Los valores de T inferiores a 1 concentran la distribución en torno a la clase principal, mientras que los superiores a 1 la aplanan y la acercan a una distribución uniforme. La misma fórmula aparece en mecánica estadística como distribución de Boltzmann, donde T es literalmente la temperatura. El fragmento de código de PyTorch siguiente muestra ambos efectos:

import torch

logits = torch.tensor([2.0, 1.0, 0.1])

# Divide logits by a temperature before softmax: low sharpens, high flattens
for temperature in [0.5, 1.0, 2.0]:
    probs = torch.softmax(logits / temperature, dim=0)
    print(f"T={temperature}: {[round(p, 3) for p in probs.tolist()]}")

# Adding a constant leaves softmax unchanged, so libraries subtract the max to avoid overflow
big = logits + 1000
naive = big.exp() / big.exp().sum()
print(f"Naive: {naive.tolist()}, torch.softmax: {[round(p, 3) for p in torch.softmax(big, dim=0).tolist()]}")

Con T = 1, los logits 2.0, 1.0 y 0.1 se convierten en [0.659, 0.242, 0.099]. Al reducir la temperatura a la mitad, la distribución se concentra y queda como [0.864, 0.117, 0.019]; al duplicarla, se aplana y queda como [0.502, 0.304, 0.194]. Si se suman 1.000 a cada logit, la fórmula ingenua desborda y muestra [nan, nan, nan], mientras que torch.softmax sigue devolviendo [0.659, 0.242, 0.099].

Softmax durante el entrenamiento#

Los clasificadores se entrenan con softmax y la función de pérdida de entropía cruzada a la vez. La pérdida es −log(pᵧ), el logaritmo negativo de la probabilidad asignada a la clase correcta y, y su gradiente respecto a los logits es simplemente p − y: las probabilidades predichas menos la etiqueta one-hot. Por eso, los frameworks combinan ambas operaciones: CrossEntropyLoss de PyTorch recibe logits sin procesar y calcula internamente log-softmax, que es más estable que calcular el logaritmo de las probabilidades por separado. Un modelo entrenado de esta forma genera logits y solo aplica softmax cuando se necesitan probabilidades durante la inferencia.

Dónde se usa Softmax#

• Clasificación de imágenes: La capa final de un clasificador genera un logit por clase y softmax los convierte en las puntuaciones de confianza que ve el usuario.
• Atención: Los Transformer calculan softmax(QKᵀ / √dₖ) sobre las puntuaciones de similitud entre una consulta y cada clave, de modo que los pesos de atención de cada token suman 1. Dividir por √dₖ evita que las puntuaciones crezcan con la dimensión y saturen softmax.
• Generación de texto: Un modelo de lenguaje grande (LLM) aplica softmax sobre un vocabulario de decenas o cientos de miles de tokens en cada paso; la temperatura de muestreo es la T indicada antes. El muestreo por núcleo extrae después tokens del conjunto más pequeño cuyas probabilidades alcanzan un umbral.
• Aprendizaje por refuerzo: Softmax aplicada a los valores de las acciones genera una política que suele elegir la mejor acción, pero sigue explorando las demás; la temperatura regula el equilibrio entre exploración y explotación.

Los modelos de clasificación de YOLO26 aplican softmax en su cabeza de clasificación durante la inferencia, y el atributo probs de cada resultado contiene la distribución completa sobre las 1.000 clases de ImageNet:

from ultralytics import YOLO

# Load a YOLO26 classification model trained on ImageNet
model = YOLO("yolo26n-cls.pt")
results = model("https://ultralytics.com/images/bus.jpg")

# probs holds the softmax output of the classification head
probs = results[0].probs
for index, conf in zip(probs.top5, probs.top5conf.tolist()):
    print(f"{results[0].names[index]}: {conf:.3f}")
print(f"Sum over all 1,000 classes: {probs.data.sum().item():.3f}")

Muestra las cinco clases más probables con sus probabilidades, y la suma de las 1.000 clases da 1.000 porque la cabeza devuelve una distribución softmax.

Softmax y funciones relacionadas#

• Sigmoide: Reduce cada puntuación de forma independiente a un valor entre 0 y 1, por lo que las salidas no tienen por qué sumar 1. La sigmoide es adecuada para problemas binarios y multietiqueta, como una imagen que contiene tanto una persona como una mochila, y los detectores de objetos suelen usarla para puntuar cada clase. Para dos clases, softmax equivale exactamente a la sigmoide de la diferencia entre los dos logits.
• Argmax: Devuelve el índice de la puntuación más alta: una decisión tajante sin gradiente. Softmax conserva la misma clasificación, pero sigue siendo diferenciable, por lo que se puede entrenar; además, aplicar argmax a las probabilidades de softmax da el mismo resultado que aplicarlo a los logits.
• Log-softmax: El logaritmo de softmax, calculado directamente para mejorar la estabilidad, que utiliza la función de pérdida de entropía cruzada.
• Escalado de temperatura: Método de calibración que ajusta una única temperatura con datos de validación después del entrenamiento, para que las probabilidades de softmax reflejen la frecuencia con la que el modelo acierta realmente. No cambia la clase predicha porque dividir todos los logits por la misma T positiva conserva su orden.