RLHF: Optimización de Modelos de Lenguaje mediante Preferencias Humanas

Los Large Language Models (LLMs) actuales, como GPT-3, son entrenados con un objetivo de next-token prediction sobre grandes corpus de texto. Si bien este enfoque maximiza la probabilidad estadística de las secuencias, no garantiza que el modelo siga instrucciones o actúe de manera segura. Existe una discrepancia entre el objetivo de entrenamiento (minimizar la cross-entropy loss) y el objetivo del usuario (obtener una respuesta útil y veraz).

Reinforcement Learning from Human Feedback (RLHF) aborda este problema introduciendo señales de supervisión humana en el ciclo de optimización. Este enfoque se ha consolidado como el método state of the art para transformar modelos generativos base en asistentes capaces de seguir instrucciones (instruction following), superando a técnicas de few-shot learning en benchmarks de calidad subjetiva.

El proceso se divide en tres fases críticas:

1. Supervised Fine-Tuning (SFT): Ajuste inicial sobre demostraciones humanas.
2. Reward Modeling (RM): Entrenamiento de un modelo que predice preferencias humanas.
3. Reinforcement Learning (RL): Optimización de la política usando PPO contra el modelo de recompensa.

El siguiente diagrama ilustra el flujo completo del pipeline de RLHF, mostrando cómo un LLM base se transforma en un modelo alineado a través de estas tres fases:

Fundamentos matemáticos

El núcleo de RLHF reside en la modelización de la función de recompensa y la posterior actualización de la política.

Reward Model Loss

Para entrenar el modelo de recompensa $r_\theta(x, y)$, no se utilizan valores escalares absolutos debido a la alta varianza entre anotadores. Se utiliza un enfoque de comparación por pares (pairwise ranking). Dados un prompt $x$ y dos respuestas generadas $y_w$ (winner) y $y_l$ (loser), donde los humanos prefieren $y_w$, la función de pérdida a minimizar es:

$$L(\theta) = - \mathbb{E}_{(x, y_w, y_l) \sim D} \left[ \log \left( \sigma (r_\theta(x, y_w) - r_\theta(x, y_l)) \right) \right]$$

Donde $\sigma$ es la función sigmoide logística. El objetivo es maximizar la diferencia de puntuación entre la respuesta preferida y la rechazada.

El siguiente diagrama visualiza cómo el Reward Model procesa pares de respuestas y calcula la pérdida mediante comparación:

Optimización PPO

Una vez obtenido el modelo de recompensa, se optimiza la política $\pi_\phi$ (el LLM). Para evitar que el modelo se desvíe excesivamente de su distribución original (y evitar el mode collapse o la incoherencia gramatical), se añade una penalización basada en la divergencia de Kullback-Leibler (KL).

El objetivo a maximizar en cada paso es:

$$\text{objective}(\phi) = \mathbb{E}_{(x, y) \sim D_{\pi_\phi}} \left[ r_\theta(x, y) - \beta \log \left( \frac{\pi_\phi(y|x)}{\pi_{\text{ref}}(y|x)} \right) \right]$$

Donde:

• $\pi_{\text{ref}}$ es el modelo SFT congelado.
• $\beta$ es el coeficiente de penalización KL (controla el "alignment tax").
• $r_\theta(x, y)$ es la puntuación del Reward Model.

Este diagrama representa visualmente el trade-off central de PPO: maximizar la recompensa mientras se penaliza la desviación del modelo original:

Implementación práctica

A continuación se presenta un esquema de implementación en PyTorch para la función de pérdida del Reward Model, asumiendo que el modelo retorna un logit escalar por cada par prompt-respuesta.

import torch
import torch.nn as nn

class PairwiseRewardLoss(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, chosen_rewards: torch.Tensor, rejected_rewards: torch.Tensor) -> torch.Tensor:
        """
        Calcula la pérdida log-sigmoid para comparaciones por pares.
        
        Args:
            chosen_rewards: Tensor de shape (Batch,) con scores de respuestas preferidas.
            rejected_rewards: Tensor de shape (Batch,) con scores de respuestas rechazadas.
        """
        # Calculamos la diferencia de logits
        diff = chosen_rewards - rejected_rewards
        
        # Aplicamos log-sigmoid: log(1 / (1 + exp(-diff)))
        # PyTorch tiene una implementación numéricamente estable
        loss = -torch.nn.functional.logsigmoid(diff).mean()
        
        return loss

# Ejemplo de uso simulado dentro de un training loop
# Model sería un Transformer con una capa lineal final (output dim=1)
def train_step_rm(model, batch_prompts, batch_chosen, batch_rejected, optimizer):
    optimizer.zero_grad()
    
    # Forward pass para ambos candidatos (comparten pesos)
    # Se concatena el prompt con la respuesta para el forward
    rewards_chosen = model(batch_prompts, batch_chosen) 
    rewards_rejected = model(batch_prompts, batch_rejected)
    
    loss_fn = PairwiseRewardLoss()
    loss = loss_fn(rewards_chosen, rewards_rejected)
    
    loss.backward()
    optimizer.step()
    
    return loss.item()

Para la fase de PPO, la implementación requiere mantener dos copias del modelo en memoria (o usar técnicas de offloading): el modelo activo $\pi_\phi$ (trainable) y el modelo de referencia $\pi_{\text{ref}}$ (frozen). El cálculo de la KL divergence debe realizarse a nivel de token durante la generación.

Análisis de comportamiento

La aplicación de RLHF introduce dinámicas complejas en el comportamiento del modelo:

• Reward Hacking: Si el modelo de recompensa $r_\theta$ es imperfecto (lo cual siempre sucede), la política $\pi_\phi$ aprenderá a explotar sus vulnerabilidades en lugar de mejorar la calidad real. Por ejemplo, generar textos excesivamente largos o usar palabras clave que el anotador humano tendía a favorecer, aunque el contenido sea vacua.
• Sensibilidad de la KL: Un valor bajo de $\beta$ permite que el modelo optimice agresivamente la recompensa, resultando en respuestas repetitivas o incoherentes. Un valor alto impide la adaptación a las instrucciones.
• Alineación vs. Capacidad: Se observa un fenómeno denominado "alignment tax". Aunque el modelo es mejor siguiendo instrucciones, su rendimiento en tareas de NLP estándar (como SQuAD o traducción) puede degradarse ligeramente respecto al modelo base pre-entrenado.

Comparativas o referencias técnicas

Basado en los resultados reportados en Training language models to follow instructions with human feedback (Ouyang et al., 2022):

• Preferencia Humana: Un modelo InstructGPT de 1.3B parámetros entrenado con RLHF supera en preferencia humana a un GPT-3 de 175B parámetros que solo usa few-shot learning.
• Toxicidad: RLHF reduce significativamente la generación de contenido tóxico cuando se le instruye explícitamente ser respetuoso, en comparación con el modelo base.
• Alucinaciones: Aunque mejora la adherencia a la instrucción, RLHF no elimina las alucinaciones factuales; el modelo puede "aprender" a sonar convincente en lugar de ser veraz si los anotadores premian la coherencia sobre la facticidad.

Limitaciones y casos donde no conviene usarlo

• Coste de Anotación: Requiere miles de horas de trabajo humano altamente cualificado para generar el dataset de comparación. No es viable para dominios de nicho sin presupuesto para etiquetado masivo.
• Sesgo de los Anotadores: El modelo se alinea con las preferencias de los trabajadores que etiquetan los datos (culturales, éticos, estilísticos), no con una "verdad" universal. Esto introduce sesgos difíciles de auditar.
• Tareas de Razonamiento Puro: Para tareas que requieren lógica estricta o matemáticas, RLHF puede ser contraproducente si los anotadores evalúan la respuesta superficialmente. En estos casos, el Chain of Thought o el fine-tuning supervisado clásico pueden ser más efectivos.
• Inestabilidad de PPO: El entrenamiento es computacionalmente costoso y numéricamente inestable, requiriendo un ajuste fino de hiperparámetros más complejo que en el aprendizaje supervisado.