DPO: tu modelo de lenguaje es secretamente un reward model

La pregunta

En la lección anterior vimos que RLHF hace dos cosas: ajusta un reward model con Bradley-Terry y después optimiza una política con PPO bajo una restricción KL. Rafailov et al. (2023) se preguntaron: si el objetivo de RL con KL tiene solución cerrada, ¿podemos usar esa solución para escribir la pérdida de preferencias directamente en términos de la política, sin reward model explícito y sin RL?

La respuesta es sí, y el resultado se llama DPO (Direct Preference Optimization).

Paso 1: el objetivo de RL restringido por KL

Partimos del mismo objetivo que RLHF, para un prompt $x$ fijo:

elegir el modelo $\pi$ que, para el prompt $x$, dé respuestas $y$ con la mayor recompensa $r$ en promedio, sin alejarse demasiado del modelo de referencia $\pi_{\text{ref}}$. La KL mide ese alejamiento y $\beta$ dice cuánto pesa

Paso 2: la solución cerrada

Expandimos la KL y reorganizamos. Maximizar lo anterior equivale a minimizar:

lo mismo, reescrito: el cociente mide cuánto más probable hace $\pi$ una respuesta que la referencia, y le restamos la recompensa. La segunda línea junta los dos términos en un solo logaritmo

Definimos la función de partición $Z(x) = \sum_y \pi_{\text{ref}}(y \mid x) \exp\!\big(\tfrac{1}{\beta} r(x, y)\big)$ y la distribución normalizada:

la mejor política: tomar la referencia y multiplicar cada respuesta por $\exp(r/\beta)$, más peso a las que tienen más recompensa. Dividir por $Z(x)$ solo hace que todo sume 1

Sustituyendo, la expresión a minimizar queda:

minimizar esto es acercar $\pi$ a $\pi^*$ (la KL mide cuán distintas son); $\log Z(x)$ no depende de $\pi$, así que no cambia cuál es la mejor

Como $Z(x)$ no depende de $\pi$, y la KL es mínima (cero) cuando las distribuciones coinciden, la política óptima es exactamente $\pi^*$. Es una reponderación de la referencia por $\exp(r/\beta)$: las respuestas con más recompensa ganan probabilidad, con $\beta$ como temperatura.

¿Por qué no usamos esto directamente? Porque $Z(x)$ es una suma sobre todas las secuencias posibles: intratable.

Paso 3: invertir la relación

Tomamos logaritmo de la solución cerrada y despejamos la recompensa:

dado vuelta: la recompensa de una respuesta se lee en cuánto más probable la hace la política óptima que la referencia, más un término que depende solo del prompt

Esto dice: cualquier recompensa puede escribirse en términos de su política óptima, más un término que solo depende de $x$.

Paso 4: la cancelación mágica

Ahora usamos el modelo de Bradley-Terry, que solo depende de diferencias de recompensa para el mismo $x$:

la probabilidad de que una persona prefiera la respuesta ganadora $y_w$ a la perdedora $y_l$: la sigmoide $\sigma$ de la diferencia de sus recompensas

Sustituimos la expresión del paso 3. El término $\beta \log Z(x)$ aparece en ambas recompensas y se cancela:

lo mismo, con cada recompensa escrita a partir de la política: el término $\beta \log Z(x)$ estaba en las dos y se canceló

Ya no hay nada intratable. La probabilidad de preferencia queda expresada solo en términos de la política y la referencia.

Paso 5: la pérdida de DPO

Reemplazamos $\pi^*$ por una política parametrizada $\pi_\theta$ y maximizamos la verosimilitud de las preferencias observadas:

la pérdida de DPO: para cada par del dataset, premia que el modelo suba la respuesta ganadora y baje la perdedora, siempre medido contra la referencia. El signo menos y el logaritmo la vuelven algo a minimizar

Compárala con la pérdida del reward model de la lección anterior: es la misma, con la recompensa reemplazada por la recompensa implícita

la recompensa implícita: cuánto más probable hace el modelo una respuesta que la referencia, en logaritmo y escalado por $\beta$

De ahí el subtítulo del paper: el modelo de lenguaje es el reward model.

Implementación

En la práctica, $\log \pi(y \mid x)$ es la suma de log-probabilidades de los tokens de la respuesta. Necesitamos forward passes de la política (con gradiente) y de la referencia (congelada, sin gradiente).

import torch
import torch.nn.functional as F

def sequence_logprob(model, input_ids, attention_mask, response_mask):
    """Suma de log p(token) sobre los tokens de la respuesta (no del prompt)."""
    logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
    # el logit en la posición t predice el token t+1
    logits = logits[:, :-1, :]
    targets = input_ids[:, 1:]
    mask = response_mask[:, 1:].float()
    logps = torch.gather(F.log_softmax(logits, dim=-1), 2, targets.unsqueeze(-1)).squeeze(-1)
    return (logps * mask).sum(dim=-1)

def dpo_loss(policy, ref, batch, beta: float = 0.1):
    pi_w = sequence_logprob(policy, batch["chosen_ids"], batch["chosen_attn"], batch["chosen_resp_mask"])
    pi_l = sequence_logprob(policy, batch["rejected_ids"], batch["rejected_attn"], batch["rejected_resp_mask"])
    with torch.no_grad():
        ref_w = sequence_logprob(ref, batch["chosen_ids"], batch["chosen_attn"], batch["chosen_resp_mask"])
        ref_l = sequence_logprob(ref, batch["rejected_ids"], batch["rejected_attn"], batch["rejected_resp_mask"])

    # recompensas implícitas
    r_w = beta * (pi_w - ref_w)
    r_l = beta * (pi_l - ref_l)

    loss = -F.logsigmoid(r_w - r_l).mean()
    # métrica útil: % de pares donde la recompensa implícita ordena bien
    accuracy = (r_w > r_l).float().mean()
    return loss, accuracy

Observa lo que no hay: ni muestreo durante el entrenamiento, ni reward model separado, ni value network, ni clipping. Son dos modelos (uno congelado) y una pérdida de clasificación.

¿Qué controla β?

$\beta$ es el mismo coeficiente de la restricción KL. Un $\beta$ grande mantiene la política cerca de la referencia (más conservador); uno chico permite desviarse más para ajustar las preferencias. Valores en el rango de 0.1 a 0.5 son comunes en la práctica, pero depende del dataset.

Limitaciones y críticas

• Off-policy. DPO aprende de un dataset fijo de pares, generalmente muestreados de otro modelo. RLHF con PPO, en cambio, genera y evalúa respuestas de la política actual (on-policy). Varios trabajos posteriores reportan que la exploración on-policy sigue aportando, y surgieron variantes de "DPO iterativo/online" que regeneran pares.
• Probabilidad de la respuesta preferida. La pérdida solo mira la diferencia de recompensas implícitas: puede bajar las probabilidades de $y_w$ y $y_l$ al mismo tiempo, siempre que $y_l$ baje más. Esto se observó empíricamente y motivó variantes con términos adicionales.
• Sobreajuste a preferencias deterministas. Si en el dataset $y_w$ siempre gana, la pérdida empuja $\hat r_w - \hat r_l \to \infty$, lo que puede debilitar el efecto de la KL. IPO (Azar et al., 2023) analiza este problema y propone una pérdida alternativa.
• No elimina Goodhart. El reward model sigue estando, solo que implícito. Si las preferencias del dataset tienen sesgos (por ejemplo, favorecen respuestas largas), DPO los aprende igual.