Probabilidad: esperanza, entropía y KL

Álgebra lineal para LLMs. Episodio ocho: probabilidad. Esperanza, entropía, entropía cruzada y la divergencia de Kullback-Leibler, que se abrevia KL. Todo con distribuciones de cuatro opciones, para poder dibujar cada número.

Las fórmulas que nos traen

Alineación, ep. 4: el freno de RLHF

Alineación, ep. 2: ¿qué es esa constante?

Dos fórmulas de Alineación que hoy se vuelven legibles. La del episodio cuatro, el freno de RLHF, reinforcement learning from human feedback. KL de pi contra pi sub ref es una esperanza, la E de doble trazo: el promedio, sobre respuestas y que genera pi, del logaritmo del cociente entre las dos probabilidades. En el dibujo, un adelanto: los términos que vamos a calcular hoy, algunos negativos. Y la del episodio dos: la pérdida de SFT, supervised fine-tuning, es casi una KL más una constante. Hoy vemos qué es esa constante. En el dibujo, una distribución: cuatro opciones con probabilidades que suman uno.

Parte 1 · Esperanza: un promedio pesado

Una distribución asigna una probabilidad a cada resultado. La esperanza es lo que da en promedio.

Recompensa esperada, exacta y sorteada

Un modelo que puede dar cuatro respuestas. p es un vector de cuatro probabilidades: cero coma cinco, cero coma veinticinco, cero coma quince y cero coma uno. Suman uno. r es la recompensa de cada respuesta: uno, tres, cero y menos dos. La recompensa esperada: cada recompensa por su probabilidad, y sumar. Ahora sorteamos diez mil respuestas con esas probabilidades. Y promediamos sus recompensas. Los productos son cero coma cinco, cero coma setenta y cinco, cero y menos cero coma dos. Suman uno coma cero cinco. Las primeras ocho muestras tienen recompensas sueltas, menos dos, tres, uno; pero el promedio de las diez mil da uno coma cero cinco.

p = torch.tensor([0.50, 0.25, 0.15, 0.10])   # 4 respuestas posibles
r = torch.tensor([1.0, 3.0, 0.0, -2.0])      # recompensa de cada una
print("suma de p:", p.sum().item())

esperanza = (p * r).sum()                    # sum_y p(y) r(y)
print("p * r:", [round(x, 2) for x in (p * r).tolist()])
print("E[r] exacta:", round(esperanza.item(), 3))

y = torch.multinomial(p, 10000, replacement=True)  # sortear
print("primeras 8 muestras:", y[:8].tolist())
print("sus r:", r[y[:8]].tolist())
print("promedio de 10000:", round(r[y].mean().item(), 3))

suma de p: 1.0
p * r: [0.5, 0.75, 0.0, -0.2]
E[r] exacta: 1.05
primeras 8 muestras: [3, 1, 0, 3, 1, 2, 0, 0]
sus r: [-2.0, 3.0, 1.0, -2.0, 3.0, 0.0, 1.0, 1.0]
promedio de 10000: 1.05

La esperanza, dibujada

$y$: un resultado posible · $p$: su probabilidad

$y \sim p$: «y sorteada según p»

sin conocer p: sortear y promediar

Una distribución le da a cada resultado y un número p de y, nunca negativo, y todos suman uno. En el dibujo, nuestro ejemplo: arriba el log de p, abajo p. La esperanza de r, con y sorteada según p, se escribe con la E de doble trazo: la suma de cada probabilidad por su recompensa. En el dibujo, la barra roja es la de recompensa negativa. Suman uno coma cero cinco. Y la idea que usa todo RLHF: si no podés sumar sobre todos los resultados, sorteás N y promediás. En el dibujo, ocho muestras promedian cero coma seis dos cinco; con diez mil, uno coma cero cinco. Las respuestas posibles de un LLM son más que los átomos del universo: solo se puede sortear.

Parte 2 · Entropía: cuánta sorpresa

Si algo tenía probabilidad uno, verlo no sorprende. Si tenía probabilidad chiquita, sorprende mucho.

Sorpresa y entropía

La entropía, H, es menos la suma de p por el log de p. Con logaritmo natural, la unidad se llama nat. Primero la sorpresa de cada resultado: menos el log de su probabilidad. Tres distribuciones: la uniforme, donde las cuatro valen un cuarto. Y una picuda, con noventa y siete por ciento en la primera. Y la uniforme en bits, con logaritmo en base dos. Las sorpresas de nuestra p: cero coma sesenta y nueve, uno coma treinta y nueve, uno coma nueve y dos coma tres: la más rara sorprende más. Entropías: la uniforme, uno coma tres ocho seis, justo log de cuatro. La p, uno coma dos cero ocho. La picuda, cero coma uno seis ocho. Y la uniforme son dos bits: dos preguntas de sí o no.

def H(p):                                  # entropía en nats
    return -(p * torch.log(p)).sum().item()

p = torch.tensor([0.50, 0.25, 0.15, 0.10])
print("sorpresa -log p:", [round(x, 2) for x in (-torch.log(p)).tolist()])

uniforme = torch.full((4,), 0.25)
picuda = torch.tensor([0.97, 0.01, 0.01, 0.01])
for nombre, q in [("uniforme", uniforme), ("p", p), ("picuda", picuda)]:
    print(f"H({nombre:8}) = {H(q):.3f} nats")
print("log 4 =", round(torch.log(torch.tensor(4.0)).item(), 3))
print("en bits, H(uniforme) =", H(uniforme) / torch.log(torch.tensor(2.0)).item())

sorpresa -log p: [0.69, 1.39, 1.9, 2.3]
H(uniforme) = 1.386 nats
H(p       ) = 1.208 nats
H(picuda  ) = 0.168 nats
log 4 = 1.386
en bits, H(uniforme) = 2.0

La entropía es la sorpresa esperada

probabilidad 1 → sorpresa 0

cada sorpresa pesada por su probabilidad

0: seguro · log n: uniforme, con n opciones

La sorpresa de un resultado es menos el log de su probabilidad. Como p es menor que uno, el log es negativo, y el signo menos lo da vuelta. En el dibujo, la D, la más rara, es la más alta. La entropía es la esperanza de la sorpresa: cada sorpresa pesada por su probabilidad. En el dibujo, la D sorprende mucho pero sale poco, y por eso pesa menos que la A. Suman uno coma dos cero ocho. La entropía nunca es negativa. Vale cero si un resultado es seguro, y como mucho log de n, con n opciones, cuando todas son igual de probables. En el dibujo, la uniforme: log de cuatro.

Entropía de las tres distribuciones

Lado a lado. La uniforme tiene la entropía máxima: no hay forma de adivinar. Nuestra p queda cerca: una opción es más probable, pero no domina. Y la picuda casi no tiene incertidumbre. En un LLM, la entropía de la distribución del próximo token mide qué tan seguro está el modelo, y la temperatura del episodio siete la sube o la baja.

• uniforme: 1,386
• nuestra p: 1,208
• picuda: 0,168

Parte 3 · Entropía cruzada

Los datos vienen de una distribución, p. El modelo cree otra, q. ¿Cuánto le cuesta al modelo esa diferencia?

Entropía cruzada y F.cross_entropy

La entropía cruzada de p y q: menos la suma de p por el log de q. Se sortea según p, pero la sorpresa se mide con q. Un modelo uniforme. Y uno más parecido a los datos. El caso de entrenamiento: el dato es un token concreto, y p es one-hot, todo cero salvo un uno en el token que vino. Y la función de PyTorch. H de p es uno coma dos cero ocho. Con el modelo uniforme, la cruzada sube a uno coma tres ocho seis; con el mejor, uno coma dos tres uno: más cerca, pero nunca debajo de H de p. Con el one-hot da menos log de cero coma treinta, uno coma dos cero cuatro, igual que F punto cross entropy.

def CE(p, q):                              # entropía cruzada H(p, q)
    return -(p * torch.log(q)).sum().item()

p = torch.tensor([0.50, 0.25, 0.15, 0.10])   # los datos
q1 = torch.tensor([0.25, 0.25, 0.25, 0.25])  # un modelo uniforme
q2 = torch.tensor([0.40, 0.30, 0.20, 0.10])  # un modelo mejor
print(f"H(p)     = {CE(p, p):.3f}")
print(f"H(p, q1) = {CE(p, q1):.3f}")
print(f"H(p, q2) = {CE(p, q2):.3f}")

uno = torch.tensor([0.0, 1.0, 0.0, 0.0])     # one-hot: vino el token 1
logits = torch.log(q2)
print(f"one-hot: {CE(uno, q2):.3f} = -log 0,30")
print(f"F.cross_entropy: {F.cross_entropy(logits[None], torch.tensor([1])).item():.3f}")

H(p)     = 1.208
H(p, q1) = 1.386
H(p, q2) = 1.231
one-hot: 1.204 = -log 0,30
F.cross_entropy: 1.204

Entropía cruzada = entropía + KL

sortear con p, medir sorpresa con q

la pérdida de cada token de un LLM

$H(p)$ no depende del modelo

La entropía cruzada: el promedio, sorteando según p, de la sorpresa que asigna q. En el dibujo, los cuatro términos con q sub dos: suman uno coma dos tres uno. Cuando p es one-hot en el token k, queda un solo término: menos log de q sub k, la pérdida de cada token. En el dibujo, el modelo le daba treinta por ciento a la B, y la pérdida es uno coma dos cero cuatro. Y la descomposición clave: la entropía cruzada es la entropía de p más la KL de p contra q. El primer término es de los datos; el segundo mide cuánto se equivoca el modelo. En el dibujo, uno coma dos cero ocho más cero coma cero dos tres.

Parte 4 · KL: cuánto se aleja una de otra

La KL mide cuánto difieren dos distribuciones. No es una distancia: depende del orden.

p y q₂, barra por barra

Las dos distribuciones juntas. En celeste, p, los datos; en ámbar, q sub dos, el modelo. En la A, el modelo le da menos de lo que debería: cuarenta contra cincuenta. En la B y la C, de más. En la D, coinciden. La KL va a sumar esas diferencias, pesadas por p.

• A · p: 0,50
• A · q₂: 0,40
• B · p: 0,25
• B · q₂: 0,30
• C · p: 0,15
• C · q₂: 0,20
• D · p: 0,10
• D · q₂: 0,10

KL exacta, sorteada y en PyTorch

La KL de p contra q: la suma de p por el log del cociente p sobre q. Miramos cada término por separado. Y las dos direcciones. La función de PyTorch: ojo, F punto kl div recibe primero el log de q y después p. Y la versión sorteada: el log del cociente para cada muestra de p, como en el episodio cuatro de Alineación. Los términos: cero coma ciento doce, y dos negativos, menos cero coma cero cuatro seis y menos cero coma cero cuatro tres. La suma es cero coma cero dos dos ocho, positiva. Al revés da cero coma cero dos tres cero: distinto. KL de p contra p, cero. PyTorch coincide. Y el promedio de cien mil muestras, cero coma cero dos dos siete.

def KL(p, q):
    return (p * torch.log(p / q)).sum().item()

p = torch.tensor([0.50, 0.25, 0.15, 0.10])
q = torch.tensor([0.40, 0.30, 0.20, 0.10])
print("términos p log(p/q):", [round(x, 3) for x in (p * torch.log(p / q)).tolist()])
print(f"KL(p||q) = {KL(p, q):.4f}   KL(q||p) = {KL(q, p):.4f}")
print(f"KL(p||p) = {KL(p, p):.4f}")
print(f"F.kl_div: {F.kl_div(torch.log(q), p, reduction='sum').item():.4f}")

y = torch.multinomial(p, 100000, replacement=True)   # muestras de p
k = torch.log(p[y] / q[y])                            # log p/q por muestra
print("k de 6 muestras:", [round(x, 3) for x in k[:6].tolist()])
print(f"promedio de k: {k.mean().item():.4f}")

términos p log(p/q): [0.112, -0.046, -0.043, 0.0]
KL(p||q) = 0.0228   KL(q||p) = 0.0230
KL(p||p) = 0.0000
F.kl_div: 0.0228
k de 6 muestras: [0.0, -0.182, 0.223, 0.0, -0.182, -0.288]
promedio de k: 0.0227

La KL, dibujada

términos negativos sí; total negativo nunca

el estimador por muestra de Alineación, ep. 4

no es simétrica: el orden importa

Cada término es la probabilidad según p por el log del cociente. Si q le da menos que p a un resultado, el término es positivo; si le da más, negativo. En el dibujo, la A suma y la B y la C restan. El total nunca es negativo: desigualdad de Gibbs. Se estima sorteando: el promedio de k, el log del cociente por muestra. En el dibujo, seis muestras dan menos cero coma cero siete: con pocas, puede dar negativo. Con cien mil, cero coma cero dos dos siete. Nunca es negativa, vale cero solo si p es igual a q, y el orden cambia el resultado: en el dibujo, cero coma cero dos dos ocho contra cero coma cero dos tres cero.

La KL no es simétrica

Dos opciones, A y B. p le da treinta por ciento a la A y setenta a la B; movemos q de casi cero a casi uno en la A. KL de p contra q, en celeste: cero cuando q es igual a p, y se dispara a la derecha, cuando q le deja casi nada a la B, que p usa mucho. Le exige a q cubrir todo lo que p usa. KL de q contra p, en ámbar: otra curva, más baja a la derecha, con el mismo cero. Por eso el orden en la fórmula de RLHF no es un detalle: pi contra pi sub ref castiga que pi ponga probabilidad donde la referencia casi no ponía.

Shannon, 1948 — A Mathematical Theory of Communication · fuente

H = −Σ pᵢ log pᵢ · «Quantities of the form H = −Σ pᵢ log pᵢ […] play a central role in information theory as measures of information, choice and uncertainty.»

La entropía tiene dueño. En mil novecientos cuarenta y ocho, Claude Shannon escribió que las cantidades de la forma menos suma de p por log de p tienen un papel central como medidas de información, de elección y de incertidumbre. Kullback y Leibler definieron su divergencia tres años después, en mil novecientos cincuenta y uno.

La KL de best-of-n, comprobada

Una fórmula de Alineación, episodio seis: best-of-n, sortear n respuestas y quedarse con la de mayor recompensa. Su KL contra el modelo base sería log de n menos n menos uno sobre n. Cien mil respuestas, ordenadas por recompensa. El modelo base las elige a todas con la misma probabilidad. La respuesta i gana si las n sorteadas tienen rango i o menos, sin que todas sean menores que i. La KL exacta, con la fórmula de hoy. Y la fórmula cerrada. Con n igual a uno, cero. Con cuatro, cero coma seis tres seis las dos. Con dieciséis, uno coma ocho tres cinco. Con sesenta y cuatro, tres coma uno siete cinco. Coinciden.

N = 100_000                                # respuestas, ordenadas por recompensa
i = torch.arange(1, N + 1, dtype=torch.float64)
ref = torch.full((N,), 1 / N, dtype=torch.float64)  # el modelo base: uniforme

for n in [1, 4, 16, 64]:
    # best-of-n: sortear n y quedarse con la de mayor recompensa
    bon = (i / N) ** n - ((i - 1) / N) ** n
    kl = (bon * torch.log(bon / ref)).sum().item()
    formula = math.log(n) - (n - 1) / n
    print(f"n={n:3}: KL medido {round(kl, 3) + 0.0:.3f}   log n - (n-1)/n = {formula:.3f}")

n=  1: KL medido 0.000   log n - (n-1)/n = 0.000
n=  4: KL medido 0.636   log n - (n-1)/n = 0.636
n= 16: KL medido 1.835   log n - (n-1)/n = 1.835
n= 64: KL medido 3.175   log n - (n-1)/n = 3.175

Volvemos a las fórmulas de Alineación

esperanza sobre lo que genera π: se estima sorteando

la constante es la entropía de los datos

Alineación, ep. 6: n = 4 → 0,636 · 16 → 1,835 · 64 → 3,175

Ahora las fórmulas se leen enteras. La KL del freno de RLHF es una esperanza sobre lo que genera la política: por eso se estima sorteando respuestas y promediando el log del cociente, token por token. Con pocas muestras, es ruidoso. La pérdida de SFT es la entropía cruzada entre las demostraciones y el modelo. Se parte en la entropía de las demostraciones, que el modelo no puede tocar, y la KL. Esa es la constante. Y la KL de best-of-n crece como log de n: en el dibujo, cero coma seis, uno coma ocho, tres coma dos. Multiplicar n por cuatro suma apenas uno coma tres.

Resumiendo

• Esperanza: $\mathbb{E}_{y\sim p}[f] = \sum_y p(y) f(y)$; sin $p$, sortear y promediar
• Entropía: $H(p) = -\sum_y p \log p$, la sorpresa esperada; entre 0 y $\log n$
• Entropía cruzada $H(p,q)$: con one-hot es $-\log q_k$, la pérdida de cada token
• $H(p, q) = H(p) + \mathrm{KL}(p\|q)$: la constante de SFT es $H(p)$
• KL $\ge 0$, cero solo si $p = q$, y no es simétrica
• Próximo: gradientes y la regla de la cadena

Papers del episodio

• Shannon, 1948 — A Mathematical Theory of Communication