DPO y la constitución

Alineación a bajo nivel. Episodio cinco: DPO, direct preference optimization, y la constitución. En el episodio anterior armamos RLHF, reinforcement learning from human feedback, con PPO, proximal policy optimization, con todas sus piezas. Hoy vamos a ver que buena parte de esa maquinaria se puede saltear con una cuenta de tres líneas, y después, cómo se reemplazan las etiquetas humanas por las de otro modelo.

Lo que cuesta PPO

• Cuatro redes en memoria: política, referencia, recompensa y valor
• Muestrear en cada paso: lento y ruidoso
• Muchos hiperparámetros que se rompen fácil
• DPO: el mismo objetivo, con una pérdida de clasificación

Repasemos lo que cuesta PPO. Necesitás cuatro redes en memoria: la política, la referencia congelada, el modelo de recompensa y el value head, que estima el valor. En cada paso tenés que generar texto, que es lento, y estimar ventajas, que es ruidoso. Y hay muchos hiperparámetros que, mal elegidos, rompen el entrenamiento. DPO resuelve el mismo problema de optimización, pero termina en una pérdida de clasificación sobre pares de preferencias, como el modelo de recompensa del episodio tres. No es un atajo aproximado: es una reescritura exacta del mismo objetivo. Veamos de dónde sale.

Parte 1 · La política óptima tiene fórmula

El objetivo de RLHF con penalización KL, la divergencia de Kullback-Leibler, tiene una solución exacta. No la podemos calcular, pero la podemos usar.

Derivar la política óptima

el objetivo del episodio cuatro

reacomodado: un KL contra otra distribución

el KL vale cero cuando las dos son iguales

suma sobre todas las respuestas: intratable

Partimos del objetivo del episodio cuatro. Maximizar la recompensa esperada menos beta por el KL contra la referencia. En el dibujo, para la referencia, la recompensa esperada es una suma: probabilidad por recompensa, respuesta por respuesta. Si metés la recompensa adentro del logaritmo y reacomodás, el objetivo queda como menos beta por un KL contra otra distribución, más un término que no depende de pi. Un KL es mínimo, cero, cuando las dos distribuciones son iguales. Entonces la política óptima es la referencia, multiplicada por e a la r sobre beta, y normalizada: un softmax. En el dibujo, la respuesta con recompensa tres pasa de cinco a veintinueve por ciento. El problema es Z: sumar sobre todas las respuestas posibles. Acá son cinco barras; con todas las frases que el modelo puede escribir, es imposible.

Comprobarlo con cinco respuestas

Con cinco respuestas posibles, Z sí se calcula. Multiplicamos la referencia por la exponencial de la recompensa. Normalizamos. Definimos el objetivo, recompensa esperada menos beta KL. Y probamos diez mil distribuciones al azar. La respuesta con recompensa tres tenía cinco por ciento en la referencia y pasa a veintinueve. La que no daba recompensa cae de cuarenta a once. El objetivo en pi estrella da uno coma dos cinco, exactamente beta por log Z, y la mejor al azar no llega. Fijate que no es un argmax: la óptima sigue dándole masa a todo, porque el KL no la deja alejarse.

Cada línea, su término:

• pi_opt = pi_ref * torch.exp → $\pi_{\text{ref}}(y)\,e^{r(y)/\beta}$ — repesar la referencia
• pi_opt = pi_opt / Z → $\pi^* = \pi_{\text{ref}}\,e^{r/\beta} / Z$ — con cinco respuestas, Z se calcula
• def objetivo → $J(\pi) = \mathbb{E}_\pi[r] - \beta\,\mathrm{KL}(\pi\,\|\,\pi_{\text{ref}})$ — lo que queremos maximizar
• mejor de 10000 → $J(\pi^*) = \beta\log Z \;\ge\; J(\pi)$ — ninguna otra lo supera

pi_ref = torch.tensor([0.40, 0.30, 0.15, 0.10, 0.05])  # 5 respuestas posibles
r = torch.tensor([0.0, 1.0, 2.0, 0.5, 3.0])             # recompensa de cada una
beta = 1.0

def objetivo(pi):                        # E[r] − β·KL(π ‖ π_ref)
    return (pi * r).sum() - beta * (pi * (pi / pi_ref).log()).sum()

pi_opt = pi_ref * torch.exp(r / beta)    # π_ref · e^(r/β)
Z = pi_opt.sum()
pi_opt = pi_opt / Z                      # normalizar
print("π* =", [round(x, 3) for x in pi_opt.tolist()])
print(f"objetivo(π*) = {objetivo(pi_opt):.4f}")
print(f"β·log Z      = {beta * Z.log():.4f}")
azar = (torch.softmax(3 * torch.randn(5), 0) for _ in range(10000))
print(f"mejor de 10000 al azar = {max(objetivo(p) for p in azar):.4f}")

π* = [0.115, 0.233, 0.317, 0.047, 0.288]
objetivo(π*) = 1.2508
β·log Z      = 1.2508
mejor de 10000 al azar = 1.2368

Despejar la recompensa: Z se cancela

despejar r de la fórmula anterior

Bradley-Terry, del episodio tres

β log Z(x) aparece en los dos y se va

Acá está el truco de DPO. Despejamos la recompensa: es beta por el logaritmo del cociente entre la política y la referencia, más beta log Z. Ahora recordá Bradley-Terry: la probabilidad de preferir una respuesta es la sigmoide de la diferencia de recompensas. En el dibujo: margen cero, cincuenta por ciento; margen uno coma siete, ochenta y cinco. Y como las dos respuestas comparten el mismo prompt, el término con Z entra una vez sumando y otra restando, y se cancela: en el dibujo quedan solo las dos recompensas implícitas. Lo intratable desaparece. La preferencia queda escrita solo con probabilidades que el modelo ya sabe calcular.

Parte 2 · La pérdida DPO

La política es su propio modelo de recompensa. Para entrenarla necesitamos una sola pieza nueva: la probabilidad de una respuesta entera.

Log-prob de una respuesta

Un modelo de lenguaje da probabilidades token por token. La de una respuesta entera es el producto, o sea, la suma de logaritmos. Pasamos prompt y respuesta juntos, y descartamos el último logit, porque cada posición predice la siguiente. Con gather sacamos el log-prob del token que efectivamente vino. Una máscara se queda solo con los tokens de la respuesta, como el menos cien del episodio dos. Y sumamos. Con un GPT de juguete sin entrenar, cada token da cerca de menos tres coma nueve, que es menos log de cincuenta: el modelo es uniforme sobre su vocabulario de cincuenta. Tres tokens suman menos once coma siete cinco.

Cada línea, su término:

• logits = modelo(ids).logits[:, :-1] → $p_\theta(\,\cdot \mid x, y_{<t})$ — la posición $t$ predice el token $t+1$
• lp_tok = logp.gather → $\log \pi(y_t \mid x, y_{<t})$ — el log-prob del token que de verdad vino
• mascara[:, P - 1:] = 1 → $t \in \text{respuesta}$ — el prompt no cuenta
• return (lp_tok * mascara).sum(-1) → $\log\pi(y|x) = \textstyle\sum_t \log\pi(y_t \mid x, y_{<t})$ — producto de probabilidades = suma de logs

def logp_respuesta(modelo, prompt, respuesta):
    ids = torch.cat([prompt, respuesta], dim=1)      # [1, P + R]
    logits = modelo(ids).logits[:, :-1]              # posición t predice t+1
    logp = logits.log_softmax(-1)
    lp_tok = logp.gather(-1, ids[:, 1:, None]).squeeze(-1)
    P = prompt.shape[1]
    mascara = torch.zeros_like(lp_tok)
    mascara[:, P - 1:] = 1                           # solo la respuesta
    return (lp_tok * mascara).sum(-1), lp_tok[0, P - 1:]

with torch.no_grad():
    total, por_token = logp_respuesta(modelo, prompt, respuesta)
print("log p por token:", [round(x, 2) for x in por_token.tolist()])
print(f"log p(respuesta) = {total.item():.2f}")
print(f"p(respuesta)     = {total.exp().item():.2e}")

log p por token: [-3.97, -3.85, -3.94]
log p(respuesta) = -11.75
p(respuesta)     = 7.87e-06

La pérdida, en cuatro líneas

Con eso, la pérdida entra en cuatro líneas. La recompensa implícita de la respuesta elegida es beta por la diferencia de log-probs entre la política y la referencia. Lo mismo para la rechazada. Y la pérdida es menos log sigmoide del margen: la misma pérdida logística del modelo de recompensa, pero ahora los parámetros son los de la política. En el primer par la política subió la elegida y bajó la rechazada: la pérdida baja a cero coma cinco nueve ocho. En el segundo hizo al revés y sube. En el tercero no se movió: log de dos, cero coma seis nueve tres. Ese es siempre el punto de partida, porque al principio pi es igual a pi sub ref.

Cada línea, su término:

• r_w = beta * (lp_w - ref_w) → $\hat r_w = \beta \log\frac{\pi(y_w|x)}{\pi_{\text{ref}}(y_w|x)}$ — recompensa implícita de la elegida
• r_l = beta * (lp_l - ref_l) → $\hat r_l = \beta \log\frac{\pi(y_l|x)}{\pi_{\text{ref}}(y_l|x)}$ — y de la rechazada
• return -F.logsigmoid(margen) → $\mathcal{L}_{\text{DPO}} = -\log\sigma\big(\hat r_w - \hat r_l\big)$ — la misma pérdida logística del episodio tres

def perdida_dpo(lp_w, lp_l, ref_w, ref_l, beta=0.1):
    r_w = beta * (lp_w - ref_w)        # recompensa implícita de la elegida
    r_l = beta * (lp_l - ref_l)        # recompensa implícita de la rechazada
    margen = r_w - r_l
    return -F.logsigmoid(margen), r_w, r_l

loss, r_w, r_l = perdida_dpo(lp_w, lp_l, ref_w, ref_l)
for i in range(3):
    print(f"par {i}: r_w {r_w[i]:+.2f}  r_l {r_l[i]:+.2f}  "
          f"loss {loss[i]:.3f}")
print(f"log 2 = {torch.log(torch.tensor(2.)):.3f}")

par 0: r_w +0.10  r_l -0.10  loss 0.598
par 1: r_w +0.00  r_l +0.10  loss 0.744
par 2: r_w +0.00  r_l +0.00  loss 0.693
log 2 = 0.693

Qué hace el gradiente

sube una, baja la otra

peso grande cuando el par está mal ordenado

chico: se aleja más de π_ref; grande: se queda cerca

Derivando, el gradiente tiene una forma muy legible. Empuja hacia arriba el log-prob de la elegida y hacia abajo el de la rechazada. Pero pesado por la sigmoide del margen al revés. En el dibujo: un par que el modelo ya ordena bien queda a la izquierda, con peso cero coma uno cinco, y casi no aprende de él; uno mal ordenado queda a la derecha, cerca de uno, y aprende mucho. Y beta hace de correa: es el mismo beta del KL, ahora escondido adentro de la pérdida.

Entrenar con DPO

Ahora entrenamos de verdad, con preferencias sintéticas: preferimos las respuestas hechas de tokens del cero al veinticuatro. La referencia es una copia congelada de la política. Sus log-probs se calculan sin gradiente. La pérdida es la de recién. No hay modelo de recompensa, ni muestreo, ni value head. La pérdida baja de cero coma seis nueve tres a cero coma uno seis dos. Mirá las recompensas implícitas: la elegida sube apenas a más cero coma tres tres, la rechazada se hunde a menos uno coma cuatro uno. DPO gana más bajando lo rechazado que subiendo lo elegido. Y en quinientos pares nuevos, la recompensa implícita ordena bien el cien por ciento.

ref = copy.deepcopy(pi).eval()           # π_ref congelada
opt = torch.optim.AdamW(pi.parameters(), lr=1e-3)
for paso in range(61):
    p, yw, yl = lote(32)
    with torch.no_grad():
        ref_w, ref_l = logp(ref, p, yw), logp(ref, p, yl)
    r_w = beta * (logp(pi, p, yw) - ref_w)
    r_l = beta * (logp(pi, p, yl) - ref_l)
    loss = -F.logsigmoid(r_w - r_l).mean()
    opt.zero_grad(); loss.backward(); opt.step()
    if paso % 20 == 0:
        print(f"paso {paso:2d}  loss {loss:.3f}  "
              f"r_w {r_w.mean():+.2f}  r_l {r_l.mean():+.2f}")
with torch.no_grad():                    # pares nuevos, nunca vistos
    p, yw, yl = lote(500)
    rw = logp(pi, p, yw) - logp(ref, p, yw)
    rl = logp(pi, p, yl) - logp(ref, p, yl)
print(f"r_w > r_l en 500 pares nuevos: {(rw > rl).float().mean():.0%}")

paso  0  loss 0.693  r_w +0.00  r_l +0.00
paso 20  loss 0.455  r_w +0.21  r_l -0.34
paso 40  loss 0.273  r_w +0.29  r_l -0.87
paso 60  loss 0.162  r_w +0.33  r_l -1.41
r_w > r_l en 500 pares nuevos: 100%

Parte 3 · PPO contra DPO

Ya tenemos las dos recetas. Veamos qué ganás y qué perdés con cada una.

PPO contra DPO

Pongámoslos lado a lado. PPO entrena un modelo de recompensa aparte, genera texto en cada paso, y carga cuatro redes. A cambio, puede explorar respuestas que no estaban en los datos. DPO usa la política como su propio modelo de recompensa y aprende de pares fijos, con dos redes. Es más simple y más estable, pero solo ve lo que hay en el dataset. La elección práctica depende de eso: si ya tenés pares buenos, o si necesitás que el modelo genere y explore.

DPO en la práctica

• $\pi_{\text{ref}}$ = el modelo SFT; si no hay, SFT sobre las elegidas
• $\beta = 0{,}1$ por defecto en el paper ($0{,}5$ para resúmenes)
• Sin el peso $\sigma(\hat r_l - \hat r_w)$, el modelo degenera
• Los pares tienen que parecerse a lo que genera $\pi_{\text{ref}}$

Cuatro detalles del paper que importan cuando lo corrés. La referencia es el modelo después del SFT, supervised fine-tuning. Si no lo tenés, primero hacés SFT sobre las respuestas elegidas, para que la referencia se parezca a la que generó los datos. Beta vale cero coma uno por defecto, y cero coma cinco en resúmenes. Los autores probaron una versión sin el peso de la sigmoide, y el modelo degeneraba: ese peso no es un adorno. Y como DPO es offline, todo depende de que los pares cubran lo que el modelo realmente genera. Lo que no está en los datos, no se corrige.

Rafailov et al., 2023 — Direct Preference Optimization: Your Language Model is Secretly a Reward Model · arXiv:2305.18290

≥ PPO · DPO supera a RLHF con PPO en controlar el sentimiento de lo generado, e iguala o mejora la calidad en resumen y diálogo de un turno, siendo bastante más simple de implementar y entrenar.

El paper original lo resume así, comparando contra RLHF con PPO en las mismas tareas. DPO supera a RLHF con PPO en controlar el sentimiento de los textos generados, e iguala o mejora la calidad en resumen y en diálogo de un turno, siendo bastante más simple de implementar y de entrenar.

Parte 4 · Que la IA ponga las etiquetas

Todo lo anterior necesita pares de preferencias. La pregunta siguiente: ¿quién los etiqueta?

Constitutional AI, en dos fases

Constitutional AI, de Anthropic, reemplaza las etiquetas humanas de daño por una lista de principios escritos: la constitución. Arranca con un prompt problemático y una respuesta del modelo. En la fase supervisada, el mismo modelo critica su respuesta según un principio, la revisa, y se hace ajuste supervisado sobre las revisiones. En la fase de RL, el modelo genera dos muestras, otro modelo elige la mejor, con eso se entrena un modelo de preferencias, y se optimiza con RL, reinforcement learning. Eso es RLAIF: RL a partir de feedback de IA.

Prompt de red team → Respuesta inicial → Crítica → Revisión → SFT → Dos muestras → Juez IA → Modelo de pref. → RL (RLAIF)

El bucle de crítica y revisión

Este es el bucle, a escala juguete: en el paper, criticar y revisar son llamadas al mismo modelo con un prompt; acá son reglas, para que corra. La crítica recibe la respuesta y un principio, y devuelve los problemas. En cada ronda se sortea un principio, como en el paper. La revisión reescribe la respuesta. Y el par prompt, respuesta revisada, va al dataset de ajuste supervisado. La primera ronda encuentra dos certezas inventadas, la segunda ya no encuentra nada, la tercera saca la culpa. La respuesta final dice quizás llueve, a veces pasa.

def criticar(resp, principio):           # en CAI: "¿la respuesta viola…?"
    return [p for p in CONSTITUCION[principio] if p in resp]

resp = "seguro que llueve, siempre pasa. es tu culpa."
inicial, datos_sft = resp, []
for ronda in range(3):
    principio = random.choice(list(CONSTITUCION))   # se sortea por ronda
    problemas = criticar(resp, principio)
    print(f"ronda {ronda}: {principio:20} → {problemas}")
    resp = revisar(resp, problemas)                  # en CAI: otro prompt
print("inicial:", inicial)
print("final:  ", resp)
datos_sft.append(("¿llueve hoy?", resp))   # el SFT se entrena con esto

ronda 0: no inventes certezas → ['seguro que', 'siempre']
ronda 1: no inventes certezas → []
ronda 2: no culpes al usuario → ['es tu culpa']
inicial: seguro que llueve, siempre pasa. es tu culpa.
final:   quizás llueve, a veces pasa. le pasa a muchos.

El juez: etiquetas blandas

La fase de RL necesita un juez. En el paper, el juez lee los dos candidatos como una pregunta de opción múltiple, A o B, y un principio sorteado entre dieciséis. Se toman los log-probs de las opciones y se normalizan. Esa probabilidad se usa como etiqueta blanda, no como cero o uno: el paper reporta que funcionó mucho mejor. Y cada comparación se guarda como un par. Cuando A viola el principio, la probabilidad de A cae a cero coma uno dos; si son iguales, queda en cero coma cinco. Esos pares pueden entrenar un modelo de preferencias, o ir directo a DPO.

def juez(a, b, principio):               # P(A mejor), normalizada entre A y B
    la, lb = logp_opcion(a, principio), logp_opcion(b, principio)
    return math.exp(la) / (math.exp(la) + math.exp(lb))

pares = []
for i, (prompt, a, b) in enumerate(muestras):
    principio = list(CONSTITUCION)[i % 2]
    p_a = juez(a, b, principio)                 # etiqueta blanda, no 0/1
    elegida, rechazada = (a, b) if p_a >= 0.5 else (b, a)
    pares.append(dict(prompt=prompt, y_w=elegida, y_l=rechazada, p=p_a))
    print(f"{prompt:17} P(A)={p_a:.2f}  y_w={elegida!r}")
# pares → modelo de preferencias + RL (RLAIF), o directo a DPO

¿llueve hoy?      P(A)=0.12  y_w='quizás llueve.'
¿llego a tiempo?  P(A)=0.12  y_w='depende del tránsito.'
¿qué hora es?     P(A)=0.50  y_w='son las tres.'

Bai et al., 2022 — Constitutional AI: Harmlessness from AI Feedback · arXiv:2212.08073

0 etiquetas · Entrenan un asistente inofensivo sin ninguna etiqueta humana que identifique respuestas dañinas: la única supervisión humana es una lista de reglas o principios.

Lo notable del resultado. Entrenaron un asistente inofensivo sin ninguna etiqueta humana que marcara respuestas dañinas. La única supervisión humana fue una lista de principios, escrita en lenguaje natural. Además, las dos fases pueden usar cadena de pensamiento, lo que hace más transparente por qué el juez decide lo que decide. Y el asistente resultante no es evasivo: ante pedidos dañinos, explica por qué objeta.

Qué cambia y qué no

• La especificación pasa a ser texto: se lee y se audita
• El juez es un modelo: hereda sus sesgos
• DPO y RLAIF siguen optimizando un proxy
• Del proxy al problema: el próximo episodio

Qué ganamos y qué no. La especificación pasa a ser un texto que cualquiera puede leer y discutir, en vez de miles de etiquetas sueltas. Pero el juez es un modelo, y hereda sus sesgos y sus puntos ciegos. Y tanto DPO como RLAIF siguen optimizando un proxy de lo que queremos. Un juez que prefiere respuestas largas, o que coinciden con el usuario, va a producir un modelo que hace exactamente eso. Qué pasa cuando optimizás un proxy con demasiada fuerza es justo el tema del próximo episodio.

Lo que te llevás

• Con KL, la óptima es $\pi^* \propto \pi_{\text{ref}}\,e^{r/\beta}$
• Recompensa implícita: $\beta\log(\pi/\pi_{\text{ref}})$; $Z$ se cancela
• log p(respuesta) = suma de log-probs de sus tokens
• $\mathcal{L}_{\text{DPO}} = -\log\sigma(\hat r_w - \hat r_l)$: dos redes, sin muestrear
• Constitutional AI: criticar, revisar, y un juez IA con etiquetas blandas
• Próximo: qué sale mal al optimizar un proxy (Goodhart, adulación, engaño)

Papers del episodio

• Rafailov et al., 2023 — Direct Preference Optimization: Your Language Model is Secretly a Reward Model
• Bai et al., 2022 — Constitutional AI: Harmlessness from AI Feedback