Softmax

El modelo ha puntuado tres posibles continuaciones de una frase con los números 2, 1 y 0. Quieres sortear una respuesta según sus probabilidades. Las puntuaciones brutas no son probabilidades: pueden ser negativas y no suman uno.

Softmax convierte esas puntuaciones, llamadas logits, en proporciones positivas que suman uno. Primero aplica la función exponencial —una transformación que crece con la puntuación— y después divide cada valor por la suma de todos.

Para las puntuaciones 2, 1, 0, obtenemos aproximadamente 67%, 24%, 9%. El primer candidato tiene la mayor probabilidad, pero el sorteo todavía puede elegir el segundo o el tercero. Softmax no realiza el sorteo ni añade texto por sí misma.

Los mismos cálculos pueden determinar las contribuciones de los fragmentos en attention. En ese caso las proporciones se refieren a la mezcla de información, no a elegir el siguiente token. Ni siquiera una probabilidad del 67% para un token evalúa la veracidad del texto completo.

Mecanismo y detalles

$$p_i=\frac{\exp(z_i)}{\sum_{j=1}^{n}\exp(z_j)}$$

$z_i$ es el logit del candidato $i$, $n$ el número de candidatos y $p_i$ su probabilidad. En un tensor hay que indicar el eje de normalización: la suma es uno para cada grupo considerado a lo largo de ese eje. La documentación de PyTorch: Softmax presenta la fórmula y el significado del parámetro dim.

Un pequeño ejemplo

Supongamos tres logits: $[\ln 2;0;0]$. Después de exponenciarlos obtenemos $[2;1;1]$ y, al dividir por su suma, $[0{,}5;0{,}25;0{,}25]$. Es un ejemplo de cálculo propio. El primer candidato tiene el doble de probabilidad que cada uno de los demás.

En un modelo de lenguaje, la cabeza del modelo de lenguaje calcula los logits del vocabulario. Softmax los convierte en una distribución, pero por sí solo no elige un token. En el ejemplo, elegir el máximo siempre señala al primer candidato; muestrear de la distribución puede producir cualquiera de los tres.

En la atención por producto escalar escalado, la misma operación normaliza los resultados de comparar las consultas con las claves. En ese caso, las componentes son pesos de posiciones, y no probabilidades de tokens del vocabulario. Attention Is All You Need, §3.2.1 y §3.4 describe ambas aplicaciones.

Un valor de 0,5 en la distribución de tokens no significa que haya un 50% de posibilidades de que toda la frase sea verdadera. El resultado se refiere a la elección entre candidatos concretos con una entrada determinada, no a una evaluación de la veracidad de lo dicho.