Exponencial, logaritmo y softmax

Álgebra lineal para LLMs. Episodio siete: exponencial, logaritmo y softmax. Son las tres funciones que convierten los números crudos de un modelo en probabilidades. Hoy vemos por qué se usan, cómo se calculan sin que la computadora explote, y qué hace la temperatura.

Las fórmulas que nos traen

Alineación, ep. 2: la pérdida de cada token

Alineación, ep. 3: Bradley-Terry

Dos fórmulas de Alineación que hoy desarmamos. La del episodio dos: el logaritmo de la probabilidad del token correcto. ele sub k es el logit del token k, un número cualquiera que sale del modelo, positivo o negativo. Se le resta el logaritmo de la suma de e a la ele sub j, sobre todos los tokens j del vocabulario. En el dibujo, tres logits que se convierten en probabilidades. Y la del episodio tres: la probabilidad de preferir una respuesta sobre otra es la sigmoide de la diferencia de recompensas. En el dibujo, diferencia cero da un medio. Al final del episodio vas a ver que son la misma cuenta.

Parte 1 · Exponencial y logaritmo

Una convierte cualquier número en uno positivo. La otra convierte productos en sumas.

e^x y log, con números

La e es el número de Euler, dos coma siete uno ocho. e a la x para cinco valores de x. El logaritmo natural, log, es la inversa: deshace la exponencial. El logaritmo de un producto. Contra la suma de los logaritmos. Y la exponencial en float treinta y dos, el formato de treinta y dos bits de PyTorch. e a la menos dos es cero coma ciento treinta y cinco: positivo, aunque x sea negativo. e a la cero es uno; e a la diez, más de veintidós mil. Log de e al cuadrado da dos. Log de cero coma cinco por cero coma dos y la suma de los logs dan lo mismo: menos dos coma tres cero tres. Y e a la ochenta y ocho todavía entra; e a la ochenta y nueve da infinito.

for x in [-2, 0, 1, 2, 10]:
    print(f"e^{x:<3} = {math.exp(x):10.3f}")   # siempre positivo

print("log(e^2) =", math.log(math.exp(2)))     # log deshace exp
a, b = 0.5, 0.2
print(f"log(a*b) = {math.log(a * b):.3f}")
print(f"log a + log b = {math.log(a) + math.log(b):.3f}")

print("e^88 en float32:", torch.exp(torch.tensor(88.0)).item())
print("e^89 en float32:", torch.exp(torch.tensor(89.0)).item())

e^-2  =      0.135
e^0   =      1.000
e^1   =      2.718
e^2   =      7.389
e^10  =  22026.466
log(e^2) = 2.0
log(a*b) = -2.303
log a + log b = -2.303
e^88 en float32: 1.6516362661361307e+38
e^89 en float32: inf

Las tres propiedades que usamos siempre

siempre positiva; su pendiente es ella misma

log solo existe para x > 0

productos de probabilidades → sumas de logs

La exponencial convierte sumas en productos. Y su derivada, su pendiente, es ella misma. En el dibujo, la recta que toca la curva: en x igual a cero, pendiente uno; en x igual a uno, dos coma setenta y dos. Crece cada vez más rápido. El logaritmo la deshace, y solo existe para números positivos. Su pendiente es uno sobre x: en el dibujo, empinada cerca de cero y casi plana a la derecha. Y la propiedad estrella: el log de un producto es la suma de los logs. En el dibujo, las barras de log de cero coma cinco y log de cero coma dos, las dos negativas porque son menores que uno, suman menos dos coma tres cero tres.

Por qué los LLMs trabajan en log

Esa propiedad no es un capricho. Una respuesta de trescientos tokens, cada uno con una probabilidad entre cero coma dos y cero coma ocho. La probabilidad de la respuesta entera es el producto de las trescientas. Su logaritmo es la suma de los trescientos logaritmos. Y comparamos con float sesenta y cuatro, que tiene más rango. En float treinta y dos, el producto da cero exacto: underflow, el número es más chico que el mínimo representable. La suma de logs da menos doscientos treinta y seis coma nueve, sin problema. En float sesenta y cuatro el producto es uno coma dos por diez a la menos ciento tres, y e a la suma de logs da lo mismo. Por eso todo se calcula en log-probabilidades.

# probabilidad de cada token de una respuesta de 300 tokens
p = torch.rand(300) * 0.6 + 0.2          # entre 0,2 y 0,8

producto = torch.prod(p)                 # P(respuesta) directo
suma_logs = torch.log(p).sum()           # log P(respuesta)

print("producto de 300 probs:", producto.item())
print("suma de 300 logs:", round(suma_logs.item(), 1))
print("en float64:", torch.prod(p.double()).item())
print("e^(suma de logs):", torch.exp(suma_logs.double()).item())

producto de 300 probs: 0.0
suma de 300 logs: -236.9
en float64: 1.252591670360075e-103
e^(suma de logs): 1.252565717604933e-103

e^x y log x: una es el espejo de la otra

Las dos curvas juntas. La exponencial: pegada a cero a la izquierda, sin tocarlo nunca, y disparada a la derecha. El logaritmo: el mismo dibujo reflejado. Muy negativo cerca de cero, y crece cada vez más despacio. El espejo es la diagonal, y igual a x. Lo que una estira, la otra comprime: por eso el logaritmo convierte probabilidades minúsculas en números manejables.

Parte 2 · Softmax

El modelo devuelve un logit por token, uno por cada palabra del vocabulario. Hay que convertirlos en probabilidades.

Softmax a mano

Una probabilidad tiene que ser positiva y todas tienen que sumar uno. Tres logits: dos, uno y cero coma uno. Primero la exponencial, que los vuelve positivos. Después dividimos por la suma. Le sumamos cien a todos los logits. Y un caso de dos opciones, contra la sigmoide. Los exponenciales son siete coma treinta y nueve, dos coma setenta y dos y uno coma once. El softmax: sesenta y seis, veinticuatro y diez por ciento, y suma uno. Igual que el de PyTorch. Sumarle cien a todos no cambia nada. Y con dos opciones, z y cero, el softmax da cero coma ocho uno siete seis: exactamente la sigmoide de z.

logits = torch.tensor([2.0, 1.0, 0.1])   # 3 tokens posibles

e = torch.exp(logits)                    # 1) exponencial: todo > 0
p = e / e.sum()                          # 2) dividir por la suma
print("e^logits:", r(e, 2))
print("softmax: ", r(p), " suma:", round(p.sum().item(), 4))
print("torch:   ", r(torch.softmax(logits, dim=0)))

p100 = torch.softmax(logits + 100, dim=0)  # sumarle 100 a todos
print("logits + 100:", r(p100))

z = 1.5                                    # dos opciones: (z, 0)
print("softmax([z, 0])[0] =", round(torch.softmax(torch.tensor([z, 0.0]), 0)[0].item(), 4))
print("sigmoid(z)        =", round(torch.sigmoid(torch.tensor(z)).item(), 4))

e^logits: [7.39, 2.72, 1.11]
softmax:  [0.659, 0.242, 0.099]  suma: 1.0
torch:    [0.659, 0.242, 0.099]
logits + 100: [0.659, 0.242, 0.099]
softmax([z, 0])[0] = 0.8176
sigmoid(z)        = 0.8176

Softmax, en fórmulas

$\ell$: vector de logits, uno por token

$e^{c}$ aparece arriba y abajo, y se cancela

dividimos arriba y abajo por $e^z$

La definición. Ele es un vector con un logit por token. El softmax del token k es e a la ele sub k, dividido por la suma de e a la ele sub j sobre todos los tokens. En el dibujo, arriba los logits, abajo las probabilidades. Si sumás la misma constante c a todos los logits, cada exponencial queda multiplicada por e a la c, arriba y abajo, y se cancela. En el dibujo, los logits suben cien y las barras de abajo no se mueven. Solo importan las diferencias entre logits. Con dos opciones, z y cero, dividimos arriba y abajo por e a la z, y queda uno sobre uno más e a la menos z: la sigmoide del episodio uno de Deep learning. La sigmoide es un softmax de dos.

Parte 3 · Log-sum-exp, sin que explote

e a la ochenta y nueve ya es infinito en float treinta y dos. Los logits reales pueden ser grandes.

El truco del máximo

Probemos con logits grandes. Mil, novecientos noventa y nueve y novecientos noventa y ocho. El softmax ingenuo: exponencial y dividir. La versión estable: antes de la exponencial, le restamos a todos el máximo. Por la propiedad de recién, el resultado no cambia. Lo mismo para el logaritmo de la suma de exponenciales, que se llama log-sum-exp: el máximo, más el log de la suma de lo que queda. Y la función de PyTorch. El ingenuo da nan, not a number: infinito dividido infinito. Restando el máximo, los exponenciales son uno, cero coma treinta y siete y cero coma ciento treinta y cinco, y las probabilidades salen bien. El log-sum-exp da mil coma cuatro cero ocho, igual que torch. Y log softmax da menos cero coma cuatro cero ocho para el primero.

l = torch.tensor([1000.0, 999.0, 998.0])

ingenuo = torch.exp(l) / torch.exp(l).sum()
print("ingenuo:", ingenuo.tolist())            # inf / inf

m = l.max()                                    # restar el máximo
estable = torch.exp(l - m) / torch.exp(l - m).sum()
print("e^(l - m):", r(torch.exp(l - m)))
print("estable:", r(estable))

lse = m + torch.log(torch.exp(l - m).sum())    # log-sum-exp
print("LSE a mano:", round(lse.item(), 3))
print("torch.logsumexp:", round(torch.logsumexp(l, 0).item(), 3))
print("log_softmax:", r(torch.log_softmax(l, 0)))

ingenuo: [nan, nan, nan]
e^(l - m): [1.0, 0.368, 0.135]
estable: [0.665, 0.245, 0.09]
LSE a mano: 1000.408
torch.logsumexp: 1000.408
log_softmax: [-0.408, -1.408, -2.408]

Log-sum-exp

$m = \max_j \ell_j$: el exponente más grande es $e^0 = 1$

un máximo suave: con $n = 3$, entre 1000 y 1001,1

log-probabilidad sin dividir nunca

Log-sum-exp, abreviado LSE. Llamamos m al logit más grande. Sacamos e a la m como factor común adentro del log, y por la propiedad del producto sale afuera como m. Lo que queda tiene exponentes negativos o cero: el más grande es e a la cero, uno. Nunca hay overflow. En el dibujo: uno, cero coma treinta y siete y cero coma catorce suman uno coma cinco cero tres; su log es cero coma cuatro cero ocho. LSE siempre queda entre el máximo y el máximo más log de n, con n logits: en el dibujo, mil coma cuatro, entre mil y mil coma uno. Es un máximo suave; de ahí viene el nombre softmax. Y el log del softmax es el logit menos el LSE: una resta, sin dividir nunca.

LSE es un máximo suavizado

Con dos números, x y cero. La línea punteada es el máximo exacto: cero hasta que x pasa a ser positivo, y después x. Tiene un quiebre. La curva violeta es log-sum-exp de x y cero, que es log de uno más e a la x. Pegada al máximo lejos del cero, y redondeada en la esquina, donde se separa como mucho log de dos. Esa curva suave se puede derivar en todos lados, y por eso se entrena bien.

Parte 4 · Temperatura

Al generar texto, el modelo sortea el próximo token. La temperatura decide qué tan arriesgado es el sorteo.

Dividir los logits por T

Cinco tokens candidatos. Sus logits: tres, dos, uno coma cinco, cero coma cinco y menos uno. Antes del softmax, dividimos los logits por la temperatura, T, un número positivo. Sorteamos mil tokens con esas probabilidades. Y contamos cuántas veces salió cada uno. Con T cero coma cinco, el primero tiene ochenta y cuatro por ciento, y en los sorteos sale ochocientas cuarenta y seis veces. Con T uno, cincuenta y nueve por ciento. Con T dos, la distribución se aplana: cuarenta por ciento el primero, y el último, que con T cero coma cinco no salió nunca, sale sesenta veces.

tokens = ["el", "un", "la", "este", "mi"]
logits = torch.tensor([3.0, 2.0, 1.5, 0.5, -1.0])

for T in [0.5, 1.0, 2.0]:
    p = torch.softmax(logits / T, dim=0)       # dividir por T
    muestras = torch.multinomial(p, 1000, replacement=True)
    cuenta = torch.bincount(muestras, minlength=5)
    print(f"T={T}: p      = {r(p)}")
    print(f"       sorteos = {cuenta.tolist()}")

T=0.5: p      = [0.84, 0.11, 0.04, 0.01, 0.0]
       sorteos = [846, 96, 51, 7, 0]
T=1.0: p      = [0.59, 0.22, 0.13, 0.05, 0.01]
       sorteos = [599, 228, 111, 51, 11]
T=2.0: p      = [0.4, 0.24, 0.19, 0.11, 0.05]
       sorteos = [395, 246, 178, 121, 60]

La temperatura, dibujada

$T = 1$: el softmax de siempre

más concentrada en el mejor

$T \to 0$: siempre el máximo · $T \to \infty$: uniforme

La fórmula: p sub k es el softmax de los logits divididos por T. Con T igual a uno, el softmax de siempre: en el dibujo, cincuenta y nueve por ciento para el primero. Con T menor que uno, dividir agranda los logits y las diferencias entre ellos. En el dibujo, los logits se estiran al doble y la probabilidad se concentra: ochenta y cuatro por ciento. Con T mayor que uno, los logits se achican y se juntan, y la distribución se aplana. En el límite: T tendiendo a cero da siempre el token de logit máximo, y T muy grande da a todos la misma probabilidad.

La probabilidad de cada token según T

Barriendo la temperatura de cero coma quince a tres coma cinco, con los mismos logits. El primer token: cerca de uno con T chica, y baja. El segundo arranca en cero, sube rápido y se estaciona. Y el último, el de logit menos uno, arranca en cero y sube despacio. Si seguimos subiendo T, las cinco curvas se juntan en un quinto: cinco tokens, todos con la misma probabilidad.

Hinton, Vinyals y Dean, 2015 — Distilling the Knowledge in a Neural Network · arXiv:1503.02531

T = 1 · «T is a temperature that is normally set to 1. Using a higher value for T produces a softer probability distribution over classes.»

La temperatura como perilla del softmax se popularizó con este paper. Hinton, Vinyals y Dean la usan para destilar: entrenar un modelo chico para que imite las probabilidades de uno grande. Escriben que T normalmente vale uno, y que una T más alta da una distribución más suave. Subir la temperatura deja ver cuánto se cree el modelo grande las opciones que no ganan.

Volvemos a las fórmulas de Alineación

log-softmax: un logit menos un log-sum-exp estable

Bradley-Terry: un softmax de dos respuestas

Ahora las dos fórmulas se leen enteras. La pérdida de cada token, en Alineación episodio dos, es el logit del token correcto menos un log-sum-exp. Eso es log softmax, y así lo calcula F punto cross entropy por dentro: restando el máximo, sin dividir y sin overflow. Y Bradley-Terry, del episodio tres: la sigmoide de la diferencia de recompensas es el softmax de dos respuestas, donde las recompensas hacen de logits. Por eso sumarle una constante a todas las recompensas no cambia nada: es la invariancia del softmax.

Resumiendo

• $e^x$ siempre positiva; $\log$ la invierte y convierte productos en sumas
• Probabilidades de secuencias: se suman $\log p$, nunca se multiplican $p$
• $\mathrm{softmax}(\ell)_k = e^{\ell_k} / \sum_j e^{\ell_j}$; solo importan diferencias
• $\mathrm{LSE} = m + \log \sum_j e^{\ell_j - m}$: restar el máximo evita el overflow
• Temperatura: $\mathrm{softmax}(\ell/T)$; $T<1$ concentra, $T>1$ aplana
• Próximo: probabilidad, entropía y KL

Papers del episodio

• Hinton, Vinyals y Dean, 2015 — Distilling the Knowledge in a Neural Network