Phase 10 - Lesson 07

RLHF: Modelo de Recompensa + PPO

This lesson includes a graded coding exercise that runs in your browser, unlocked with lifetime access.

SFT enseña al modelo a seguir instrucciones. Pero no le enseña qué respuesta es MEJOR. Dos respuestas gramaticalmente correctas y factualmente precisas pueden diferir enormemente en utilidad. RLHF es cómo codificas el juicio humano en el comportamiento del modelo. Es lo que hace que Claude sea útil y GPT educado.

Tipo: Construcción Lenguajes: Python (con numpy) Prerrequisitos: Fase 10, Lección 06 (Ajuste de Instrucción / SFT) Tiempo: ~90 minutos

Objetivos de Aprendizaje

• Construir un modelo de recompensa que puntúe la calidad de la respuesta a partir de pares de preferencias humanas (elegida vs rechazada)
• Implementar el bucle de entrenamiento PPO que optimiza una política de modelo de lenguaje frente al modelo de recompensa con una penalidad KL
• Explicar por qué RLHF requiere tres modelos (SFT, recompensa, política) y cómo la restricción KL previene el "reward hacking" (hackeo de recompensa)
• Evaluar el efecto de RLHF comparando la calidad de la respuesta antes y después de la optimización de preferencias

El Problema

Pregunta a un modelo "Explica la computación cuántica" y podría producir:

Respuesta A: "La computación cuántica utiliza qubits que pueden existir en superposición, lo que significa que pueden ser 0, 1 o ambos simultáneamente. Esto permite a las computadoras cuánticas procesar ciertos cálculos exponencialmente más rápido que las computadoras clásicas. Los algoritmos clave incluyen el algoritmo de Shor para factorizar números grandes y el algoritmo de Grover para buscar en bases de datos no ordenadas."

Respuesta B: "La computación cuántica es un tipo de computación que utiliza fenómenos de la mecánica cuántica. Se propuso por primera vez en la década de 1980. Richard Feynman sugirió que los sistemas cuánticos podrían simularse con computadoras cuánticas. El campo ha crecido significativamente desde entonces. Muchas empresas están trabajando ahora en computadoras cuánticas. IBM, Google y otros han progresado. La supremacía cuántica fue reclamada por Google en 2019."

Ambas respuestas son factualmente correctas. Ambas son gramaticalmente correctas. Ambas siguen la instrucción. Pero la Respuesta A es claramente mejor. Es más concisa, más informativa y está mejor estructurada. Un ser humano elegiría la A siempre.

SFT no puede capturar esta distinción. Entrena al modelo en respuestas "correctas", pero no tiene ningún mecanismo para decir "esta respuesta es mejor que aquella". Trata cada ejemplo de entrenamiento como igualmente bueno. Si tanto A como B aparecieran en el dataset de SFT, el modelo aprendería de ambos por igual.

RLHF resuelve esto. Entrena un modelo de recompensa para prever qué respuesta preferiría un ser humano, y luego utiliza esa señal de recompensa para impulsar al modelo de lenguaje hacia salidas de mayor calidad. InstructGPT (el precursor de ChatGPT) utilizó RLHF para mejorar drásticamente la utilidad, veracidad e inocuidad de GPT-3. Los evaluadores internos de OpenAI prefirieron las salidas de InstructGPT sobre las de GPT-3 el 85% de las veces, a pesar de que InstructGPT era 135 veces más pequeño (1.3B vs 175B parámetros).

El Concepto

Las Tres Etapas

RLHF no es una única ejecución de entrenamiento. Es un pipeline de tres etapas secuenciales, cada una construida sobre la anterior.

Etapa 1: SFT. Entrena un modelo base en pares de instrucción-respuesta (Lección 06). Esto te da un modelo que puede seguir instrucciones pero no sabe qué respuestas son mejores que otras.

Etapa 2: Modelo de Recompensa. Recopila datos de preferencia humana: muestra a los anotadores dos respuestas al mismo prompt y pregunta "¿cuál es mejor?". Entrena un modelo para predecir estas preferencias. El modelo de recompensa toma (prompt, respuesta) como entrada y devuelve una puntuación escalar.

Etapa 3: PPO. Utiliza el modelo de recompensa para generar una señal de entrenamiento para el modelo de lenguaje. El modelo de lenguaje genera respuestas, el modelo de recompensa las puntúa y PPO actualiza el modelo de lenguaje para producir respuestas con puntuaciones más altas. Una penalidad de divergencia KL evita que el modelo de lenguaje se desvíe demasiado del checkpoint de SFT.

graph TD
    subgraph Stage1["Etapa 1: SFT"]
        B["Modelo Base"] --> S["Modelo SFT"]
        D["Datos de Instrucción\n(27K ejemplos)"] --> S
    end

    subgraph Stage2["Etapa 2: Modelo de Recompensa"]
        S --> |"Generar respuestas"| P["Pares de Preferencia\n(prompt, ganador, perdedor)"]
        H["Anotadores Humanos"] --> P
        P --> R["Modelo de Recompensa\nR(prompt, respuesta) → score"]
    end

    subgraph Stage3["Etapa 3: PPO"]
        S --> |"Inicializar política"| PI["Modelo de Política\n(siendo optimizado)"]
        S --> |"Congelar como referencia"| REF["Modelo de Referencia\n(SFT congelado)"]
        PI --> |"Generar"| RESP["Respuesta"]
        RESP --> R
        R --> |"Señal de recompensa"| PPO["Actualización PPO"]
        REF --> |"Penalidad KL"| PPO
        PPO --> |"Actualizar"| PI
    end

    style S fill:#1a1a2e,stroke:#51cf66,color:#fff
    style R fill:#1a1a2e,stroke:#e94560,color:#fff
    style PI fill:#1a1a2e,stroke:#0f3460,color:#fff
    style REF fill:#1a1a2e,stroke:#0f3460,color:#fff
    style PPO fill:#1a1a2e,stroke:#e94560,color:#fff

El Modelo de Recompensa

El modelo de recompensa es un modelo de lenguaje adaptado como evaluador. Toma el modelo SFT, reemplaza la cabeza de modelado de lenguaje (que genera una distribución sobre el vocabulario) por una cabeza escalar (que genera un solo número). La arquitectura es idéntica hasta la capa final.

Entrada: un prompt concatenado con una respuesta. Salida: una única puntuación escalar de recompensa.

Los datos de entrenamiento son pares de preferencias humanas. Para cada prompt, los anotadores ven dos respuestas y eligen la mejor. Esto crea tríos de entrenamiento: (prompt, respuesta_preferida, respuesta_rechazada).

La función de pérdida utiliza el modelo Bradley-Terry de preferencias pareadas:

loss = -log(sigmoid(reward(preferred) - reward(rejected)))

Esta es la ecuación clave. sigmoid(reward(A) - reward(B)) da la probabilidad de que la respuesta A sea preferida sobre la respuesta B. La pérdida empuja al modelo de recompensa a asignar una puntuación más alta a la respuesta preferida.

¿Por qué comparaciones pareadas en lugar de puntuaciones absolutas? Porque los humanos somos pésimos asignando puntuaciones de calidad absoluta ("¿Esta respuesta es un 7.3 o un 7.5 de 10?"), pero muy buenos en comparaciones relativas ("¿Es A mejor que B?"). El modelo Bradley-Terry convierte las comparaciones relativas en un sistema de puntuación absoluta consistente.

Números de InstructGPT: OpenAI recopiló 33,000 pares de comparación de 40 contratistas. Cada comparación tomó alrededor de 5 minutos. Eso representa 2,750 horas de trabajo humano para los datos de entrenamiento del modelo de recompensa.

PPO: Optimización de Política Próxima

PPO es un algoritmo de aprendizaje por refuerzo. En RLHF, el "entorno" es el modelo de recompensa, el "agente" es el modelo de lenguaje y la "acción" es generar un token.

El objetivo:

maximize: E[R(prompt, response)] - beta * KL(policy || reference)

El primer término empuja al modelo a generar respuestas de alta recompensa. El segundo término (penalidad de divergencia KL) evita que el modelo se desvíe demasiado del checkpoint de SFT.

¿Por qué la penalidad KL? Sin ella, el modelo encuentra soluciones degeneradas. El modelo de recompensa se entrena en un dataset finito de preferencias humanas. Tiene puntos ciegos. El modelo de lenguaje explotará esos puntos ciegos, encontrando salidas que obtienen puntuaciones altas en el modelo de recompensa pero que en realidad no tienen sentido. Ejemplos clásicos:

• Repetir "¡Soy tan servicial e inofensivo!" obtiene una puntuación alta en modelos de recompensa de utilidad/inocuidad
• Producir respuestas prolijas, de tono formal pero vacías que coinciden con el patrón de "alta calidad"
• Explotar frases específicas que casualmente se correlacionaron con una alta recompensa en los datos de entrenamiento

La penalidad KL dice: puedes mejorar, pero no puedes convertirte en un modelo completamente diferente. Mantente cerca de la versión SFT, que ya era razonable. Desvíate demasiado y el costo de KL dominará la recompensa.

Números de InstructGPT: El entrenamiento de PPO utilizó lr=1.5e-5, coeficiente KL beta=0.02, 256K episodios (pares de prompt-respuesta) y 4 épocas de PPO por lote. Todo el pipeline de RLHF tomó varios días en un clúster de GPUs.

graph LR
    subgraph PPO["Bucle de Entrenamiento PPO"]
        direction TB
        PROMPT["Muestra de prompt\ndel dataset"] --> GEN["La política genera\nla respuesta"]
        GEN --> SCORE["El modelo de recompensa\npuntúa la respuesta"]
        GEN --> KL["Calcular divergencia KL\nvs modelo de referencia"]
        SCORE --> OBJ["Objetivo:\nrecompensa - beta * KL"]
        KL --> OBJ
        OBJ --> UPDATE["Actualización de gradiente PPO\n(pérdida sustituta recortada)"]
        UPDATE --> |"repetir"| PROMPT
    end

    style PROMPT fill:#1a1a2e,stroke:#0f3460,color:#fff
    style SCORE fill:#1a1a2e,stroke:#51cf66,color:#fff
    style KL fill:#1a1a2e,stroke:#e94560,color:#fff
    style OBJ fill:#1a1a2e,stroke:#e94560,color:#fff

El Objetivo de PPO en Detalle

PPO utiliza un "objetivo sustituto recortado" (clipped surrogate objective) para evitar actualizaciones excesivamente grandes. La relación entre las probabilidades de la nueva política y la antigua política se recorta al rango [1 - epsilon, 1 + epsilon], donde epsilon es típicamente 0.2.

ratio = pi_new(action | state) / pi_old(action | state)
clipped_ratio = clip(ratio, 1 - epsilon, 1 + epsilon)
loss = -min(ratio * advantage, clipped_ratio * advantage)

La función de ventaja (advantage function) estima qué tan mejor es la respuesta actual en comparación con la calidad esperada. En RLHF:

advantage = reward(prompt, response) - baseline

El baseline suele ser el promedio de recompensa de las respuestas recientes. Una ventaja positiva significa que la respuesta fue mejor que el promedio; una ventaja negativa significa que fue peor. PPO aumenta la probabilidad de las respuestas superiores al promedio y disminuye la probabilidad de las inferiores al promedio.

El recorte (clipping) evita actualizaciones catastróficas. Si una sola respuesta obtiene una recompensa inusualmente alta, la relación sin recortar podría ser muy grande, lo que haría que el modelo cambie drásticamente hacia esa respuesta. El recorte limita la actualización, manteniendo la estabilidad del entrenamiento.

Reward Hacking

El lado oscuro de RLHF. El modelo de lenguaje se está optimizando frente al modelo de recompensa, que es un proxy imperfecto de las preferencias humanas. A medida que el modelo de lenguaje mejora en la maximización de la recompensa, comienza a explotar las debilidades del modelo de recompensa.

Modos de fallo comunes:

Estrategias de mitigación: penalidad KL más fuerte (evita que el modelo se desvíe lo suficiente como para explotar las debilidades), entrenar al modelo de recompensa con ejemplos adversarios (corregir modos de fallo conocidos) y usar múltiples modelos de recompensa con diferentes arquitecturas (más difícil de hackear todos simultáneamente).

Pipelines Reales de RLHF

El artículo de 2022 de Anthropic entrenó un modelo de recompensa de 52B en 22,000 comparaciones. Los modelos de recompensa más grandes producen señales más confiables, lo que hace que el entrenamiento con PPO sea más estable. Utilizar un modelo de recompensa pequeño para trenar un modelo de lenguaje grande es arriesgado: el modelo de recompensa no tiene la capacidad suficiente para capturar los matices de las respuestas buenas frente a las malas.

Constrúyelo

Paso 1: Datos Sintéticos de Preferencia

En producción, los anotadores humanos crean datos de preferencia. Crearemos pares sintéticos donde la respuesta "preferida" es objetivamente mejor (más concisa, más precisa, más útil).

import numpy as np

PREFERENCE_DATA = [
    {
        "prompt": "What is the capital of France?",
        "preferred": "The capital of France is Paris.",
        "rejected": "France is a country in Europe. It has many cities. The capital is Paris. Paris is known for the Eiffel Tower.",
    },
    {
        "prompt": "Explain gravity in one sentence.",
        "preferred": "Gravity is the force that attracts objects with mass toward each other.",
        "rejected": "Gravity is something that makes things fall down when you drop them.",
    },
    {
        "prompt": "What is 15 times 7?",
        "preferred": "15 times 7 is 105.",
        "rejected": "Let me think about this. 15 times 7. Well, 10 times 7 is 70, and 5 times 7 is 35, so the answer might be around 105.",
    },
    {
        "prompt": "Name three programming languages.",
        "preferred": "Python, Rust, and TypeScript.",
        "rejected": "There are many programming languages. Some popular ones include various languages like Python and others.",
    },
    {
        "prompt": "What year did World War II end?",
        "preferred": "World War II ended in 1945.",
        "rejected": "World War II was a major global conflict. It involved many countries. The war ended in the mid-1940s, specifically in 1945.",
    },
    {
        "prompt": "Define machine learning.",
        "preferred": "Machine learning is a field where algorithms learn patterns from data to make predictions without being explicitly programmed.",
        "rejected": "Machine learning is a type of AI. AI stands for artificial intelligence. Machine learning uses data to learn.",
    },
]

Las respuestas preferidas son concisas y directas. Las respuestas rechazadas exhiben modos de fallo comunes: relleno innecesario, evasivas, explicaciones redundantes e imprecisión. Este es exactamente el tipo de distinción que SFT no puede capturar pero RLHF sí.

Paso 2: Arquitectura del Modelo de Recompensa

El modelo de recompensa reutiliza la arquitectura transformer del mini GPT, pero reemplaza la cabeza de salida del tamaño del vocabulario por una única proyección escalar.

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "04-pre-training-mini-gpt", "code"))
from main import MiniGPT, LayerNorm, Embedding, TransformerBlock


class RewardModel:
    def __init__(self, vocab_size=256, embed_dim=128, num_heads=4,
                 num_layers=4, max_seq_len=128, ff_dim=512):
        self.embedding = Embedding(vocab_size, embed_dim, max_seq_len)
        self.blocks = [
            TransformerBlock(embed_dim, num_heads, ff_dim)
            for _ in range(num_layers)
        ]
        self.ln_f = LayerNorm(embed_dim)
        self.reward_head = np.random.randn(embed_dim) * 0.02

    def forward(self, token_ids):
        seq_len = token_ids.shape[-1]
        mask = np.triu(np.full((seq_len, seq_len), -1e9), k=1)

        x = self.embedding.forward(token_ids)
        for block in self.blocks:
            x = block.forward(x, mask)
        x = self.ln_f.forward(x)

        last_hidden = x[:, -1, :]
        reward = last_hidden @ self.reward_head

        return reward

El modelo de recompensa toma el estado oculto en la última posición del token y lo proyecta a un escalar. ¿Por qué el último token? Porque la máscara de atención causal significa que la última posición ha atendido a todos los tokens anteriores. Tiene la representación más completa de toda la secuencia (prompt, respuesta).

Paso 3: Pérdida de Bradley-Terry

Entrena el modelo de recompensa en pares de preferencia utilizando la pérdida pareada de Bradley-Terry.

def tokenize_for_reward(prompt, response, vocab_size=256):
    prompt_tokens = [min(t, vocab_size - 1) for t in list(prompt.encode("utf-8"))]
    response_tokens = [min(t, vocab_size - 1) for t in list(response.encode("utf-8"))]
    return prompt_tokens + [0] + response_tokens


def sigmoid(x):
    return np.where(
        x >= 0,
        1.0 / (1.0 + np.exp(-x)),
        np.exp(x) / (1.0 + np.exp(x))
    )


def bradley_terry_loss(reward_preferred, reward_rejected):
    diff = reward_preferred - reward_rejected
    loss = -np.log(sigmoid(diff) + 1e-8)
    return loss


def train_reward_model(rm, preference_data, num_epochs=10, lr=1e-4, max_seq_len=128):
    print(f"Training Reward Model: {len(preference_data)} preference pairs, {num_epochs} epochs")
    print()

    losses = []
    accuracies = []

    for epoch in range(num_epochs):
        epoch_loss = 0.0
        epoch_correct = 0
        num_pairs = 0

        indices = np.random.permutation(len(preference_data))

        for idx in indices:
            pair = preference_data[idx]

            preferred_tokens = tokenize_for_reward(pair["prompt"], pair["preferred"])
            rejected_tokens = tokenize_for_reward(pair["prompt"], pair["rejected"])

            preferred_tokens = preferred_tokens[:max_seq_len]
            rejected_tokens = rejected_tokens[:max_seq_len]

            preferred_ids = np.array(preferred_tokens).reshape(1, -1)
            rejected_ids = np.array(rejected_tokens).reshape(1, -1)

            r_preferred = rm.forward(preferred_ids)[0]
            r_rejected = rm.forward(rejected_ids)[0]

            loss = bradley_terry_loss(r_preferred, r_rejected)

            if r_preferred > r_rejected:
                epoch_correct += 1

            diff = r_preferred - r_rejected
            grad = sigmoid(diff) - 1.0

            rm.reward_head -= lr * grad * rm.ln_f.forward(
                rm.embedding.forward(preferred_ids)
            )[:, -1, :].flatten()

            epoch_loss += loss
            num_pairs += 1

        avg_loss = epoch_loss / max(num_pairs, 1)
        accuracy = epoch_correct / max(num_pairs, 1)
        losses.append(avg_loss)
        accuracies.append(accuracy)

        if epoch % 2 == 0:
            print(f"  Epoch {epoch + 1:3d} | Loss: {avg_loss:.4f} | Accuracy: {accuracy:.1%}")

    return rm, losses, accuracies

La métrica de precisión es directa: ¿qué fracción de pares de preferencia clasifica correctamente el modelo de recompensa? Un modelo aleatorio obtiene un 50%. Un modelo de recompensa bien entrenado con datos limpios debería superar el 70%. El modelo de recompensa de InstructGPT logró aproximadamente un 72% de precisión en comparaciones de prueba retenidas, lo que parece bajo pero en realidad es bueno; muchos pares de preferencia son ambiguos incluso para los humanos (el acuerdo entre anotadores fue de aproximadamente un 73%).

Paso 4: Bucle PPO Simplificado

PPO completo es complejo. Esta implementación captura el mecanismo central: generar respuestas, puntuarlas, calcular la ventaja y actualizar la política con una penalidad KL.

def compute_kl_divergence(policy_logits, reference_logits):
    policy_probs = np.exp(policy_logits - policy_logits.max(axis=-1, keepdims=True))
    policy_probs = policy_probs / policy_probs.sum(axis=-1, keepdims=True)
    policy_probs = np.clip(policy_probs, 1e-10, 1.0)

    ref_probs = np.exp(reference_logits - reference_logits.max(axis=-1, keepdims=True))
    ref_probs = ref_probs / ref_probs.sum(axis=-1, keepdims=True)
    ref_probs = np.clip(ref_probs, 1e-10, 1.0)

    kl = np.sum(policy_probs * np.log(policy_probs / ref_probs), axis=-1)
    return kl.mean()


def generate_response(model, prompt_tokens, max_new_tokens=30, temperature=0.8, max_seq_len=128):
    tokens = list(prompt_tokens)

    for _ in range(max_new_tokens):
        context = np.array(tokens[-max_seq_len:]).reshape(1, -1)
        logits = model.forward(context)
        next_logits = logits[0, -1, :]

        next_logits = next_logits / max(temperature, 1e-8)
        probs = np.exp(next_logits - next_logits.max())
        probs = probs / probs.sum()
        probs = np.clip(probs, 1e-10, 1.0)
        probs = probs / probs.sum()

        next_token = np.random.choice(len(probs), p=probs)
        tokens.append(int(next_token))

    return tokens


def copy_model_weights(source, target):
    target.embedding.token_embed = source.embedding.token_embed.copy()
    target.embedding.pos_embed = source.embedding.pos_embed.copy()
    target.ln_f.gamma = source.ln_f.gamma.copy()
    target.ln_f.beta = source.ln_f.beta.copy()
    for s_block, t_block in zip(source.blocks, target.blocks):
        t_block.attn.W_q = s_block.attn.W_q.copy()
        t_block.attn.W_k = s_block.attn.W_k.copy()
        t_block.attn.W_v = s_block.attn.W_v.copy()
        t_block.attn.W_out = s_block.attn.W_out.copy()
        t_block.ffn.W1 = s_block.ffn.W1.copy()
        t_block.ffn.W2 = s_block.ffn.W2.copy()
        t_block.ffn.b1 = s_block.ffn.b1.copy()
        t_block.ffn.b2 = s_block.ffn.b2.copy()
        t_block.ln1.gamma = s_block.ln1.gamma.copy()
        t_block.ln1.beta = s_block.ln1.beta.copy()
        t_block.ln2.gamma = s_block.ln2.gamma.copy()
        t_block.ln2.beta = s_block.ln2.beta.copy()


def ppo_training(policy_model, reference_model, reward_model, prompts,
                 num_episodes=20, lr=1.5e-5, kl_coeff=0.02, max_seq_len=128):
    print(f"PPO Training: {num_episodes} episodes, lr={lr}, KL coeff={kl_coeff}")
    print()

    rewards_history = []
    kl_history = []

    for episode in range(num_episodes):
        prompt_text = prompts[episode % len(prompts)]
        prompt_tokens = [min(t, 252) for t in list(prompt_text.encode("utf-8"))]

        response_tokens = generate_response(
            policy_model, prompt_tokens,
            max_new_tokens=20, temperature=0.8, max_seq_len=max_seq_len
        )

        response_ids = np.array(response_tokens[:max_seq_len]).reshape(1, -1)
        reward = reward_model.forward(response_ids)[0]

        policy_logits = policy_model.forward(response_ids)
        ref_logits = reference_model.forward(response_ids)
        kl = compute_kl_divergence(policy_logits, ref_logits)

        total_reward = reward - kl_coeff * kl

        rewards_history.append(float(reward))
        kl_history.append(float(kl))

        for block in policy_model.blocks:
            update_scale = lr * total_reward
            block.ffn.W1 += update_scale * np.random.randn(*block.ffn.W1.shape) * 0.01
            block.ffn.W2 += update_scale * np.random.randn(*block.ffn.W2.shape) * 0.01

        if episode % 5 == 0:
            avg_reward = np.mean(rewards_history[-5:]) if rewards_history else 0
            avg_kl = np.mean(kl_history[-5:]) if kl_history else 0
            print(f"  Episode {episode:3d} | Reward: {reward:.4f} | KL: {kl:.4f} | "
                  f"Avg Reward: {avg_reward:.4f}")

    return policy_model, rewards_history, kl_history

El bucle principal: (1) muestrear un prompt, (2) generar una respuesta, (3) puntuarla con el modelo de recompensa, (4) calcular la divergencia KL frente a la referencia congelada, (5) calcular la recompensa ajustada (recompensa menos penalidad KL), (6) actualizar la política. La penalidad KL crece a medida que la política se desvía de la referencia, evitando automáticamente el reward hacking.

Paso 5: Comparación de Puntuaciones de Recompensa

Después de RLHF, las respuestas del modelo de política deberían obtener puntuaciones más altas en el modelo de recompensa que las respuestas del modelo SFT original.

def compare_models(sft_model, rlhf_model, reward_model, prompts, max_seq_len=128):
    print("Model Comparison (reward scores)")
    print("-" * 60)
    print(f"  {'Prompt':<35} {'SFT':>10} {'RLHF':>10}")
    print("  " + "-" * 55)

    sft_total = 0.0
    rlhf_total = 0.0

    for prompt in prompts:
        prompt_tokens = [min(t, 252) for t in list(prompt.encode("utf-8"))]

        sft_response = generate_response(
            sft_model, prompt_tokens,
            max_new_tokens=20, temperature=0.6, max_seq_len=max_seq_len
        )
        rlhf_response = generate_response(
            rlhf_model, prompt_tokens,
            max_new_tokens=20, temperature=0.6, max_seq_len=max_seq_len
        )

        sft_ids = np.array(sft_response[:max_seq_len]).reshape(1, -1)
        rlhf_ids = np.array(rlhf_response[:max_seq_len]).reshape(1, -1)

        sft_reward = reward_model.forward(sft_ids)[0]
        rlhf_reward = reward_model.forward(rlhf_ids)[0]

        sft_total += sft_reward
        rlhf_total += rlhf_reward

        truncated_prompt = prompt[:33] + ".." if len(prompt) > 35 else prompt
        print(f"  {truncated_prompt:<35} {sft_reward:>10.4f} {rlhf_reward:>10.4f}")

    n = len(prompts)
    print("  " + "-" * 55)
    print(f"  {'Average':<35} {sft_total/n:>10.4f} {rlhf_total/n:>10.4f}")

    return sft_total / n, rlhf_total / n

Úsalo

Demostración del Pipeline Completo de RLHF

if __name__ == "__main__":
    np.random.seed(42)

    print("=" * 70)
    print("RLHF PIPELINE: REWARD MODEL + PPO")
    print("=" * 70)
    print()

    print("STAGE 1: SFT Model (from Lesson 06)")
    print("-" * 40)
    sft_model = MiniGPT(
        vocab_size=256, embed_dim=128, num_heads=4,
        num_layers=4, max_seq_len=128, ff_dim=512
    )
    print(f"  Parameters: {sft_model.count_parameters():,}")
    print()

    print("STAGE 2: Train Reward Model")
    print("-" * 40)
    rm = RewardModel(
        vocab_size=256, embed_dim=128, num_heads=4,
        num_layers=4, max_seq_len=128, ff_dim=512
    )

    rm, rm_losses, rm_accuracies = train_reward_model(rm, PREFERENCE_DATA, num_epochs=10, lr=1e-4)
    print()

    print("Reward Model Evaluation:")
    print("-" * 40)
    correct = 0
    for pair in PREFERENCE_DATA:
        pref_tokens = tokenize_for_reward(pair["prompt"], pair["preferred"])[:128]
        rej_tokens = tokenize_for_reward(pair["prompt"], pair["rejected"])[:128]

        r_pref = rm.forward(np.array(pref_tokens).reshape(1, -1))[0]
        r_rej = rm.forward(np.array(rej_tokens).reshape(1, -1))[0]

        if r_pref > r_rej:
            correct += 1
        print(f"  Preferred: {r_pref:+.4f} | Rejected: {r_rej:+.4f} | {'Correct' if r_pref > r_rej else 'Wrong'}")

    print(f"\n  Accuracy: {correct}/{len(PREFERENCE_DATA)} = {correct/len(PREFERENCE_DATA):.1%}")
    print()

    print("STAGE 3: PPO Training")
    print("-" * 40)

    policy_model = MiniGPT(
        vocab_size=256, embed_dim=128, num_heads=4,
        num_layers=4, max_seq_len=128, ff_dim=512
    )
    reference_model = MiniGPT(
        vocab_size=256, embed_dim=128, num_heads=4,
        num_layers=4, max_seq_len=128, ff_dim=512
    )

    copy_model_weights(sft_model, policy_model)
    copy_model_weights(sft_model, reference_model)

    train_prompts = [pair["prompt"] for pair in PREFERENCE_DATA]

    policy_model, rewards, kls = ppo_training(
        policy_model, reference_model, rm,
        train_prompts, num_episodes=20, lr=1.5e-5, kl_coeff=0.02
    )
    print()

    print("=" * 70)
    print("COMPARISON: SFT vs RLHF")
    print("=" * 70)
    print()

    eval_prompts = [
        "What is the capital of France?",
        "Explain gravity.",
        "Name three programming languages.",
    ]

    sft_avg, rlhf_avg = compare_models(sft_model, policy_model, rm, eval_prompts)
    print()

    print("=" * 70)
    print("KL DIVERGENCE ANALYSIS")
    print("=" * 70)
    print()

    if kls:
        print(f"  Initial KL: {kls[0]:.4f}")
        print(f"  Final KL:   {kls[-1]:.4f}")
        print(f"  Max KL:     {max(kls):.4f}")
        kl_threshold = 0.1
        print(f"  KL > {kl_threshold}: {'Yes (model drifted significantly)' if max(kls) > kl_threshold else 'No (model stayed close to reference)'}")

Entrégalo

Esta lección produce outputs/prompt-reward-model-designer.md, un prompt para diseñar pipelines de entrenamiento de modelos de recompensa. Dado un comportamiento objetivo (utilidad, capacidad de programación, seguridad), produce un protocolo de recopilación de datos, pautas para los anotadores y criterios de evaluación del modelo de recompensa.

Ejercicios

1. Modifica el modelo de recompensa para usar la media de todos los estados ocultos en lugar de solo la última posición. Compara la precisión. El enfoque de pooling de media le da a cada token el mismo peso, mientras que el enfoque de la última posición depende de la atención causal para agregar información. Realiza pruebas en los 6 pares de preferencia e informa qué enfoque obtiene una mayor precisión.

2. Implementa la calibración del modelo de recompensa. Después del entrenamiento, pasa todos los pares de preferencia por el modelo de recompensa y calcula: (a) la recompensa promedio para las respuestas preferidas, (b) la recompensa promedio para las respuestas rechazadas, (c) el margen (preferido menos rechazado). Un modelo bien calibrado debería tener un margen claro. Luego agrega 4 nuevos pares de preferencia y verifica si el margen se mantiene en datos no vistos.

3. Simula el reward hacking. Crea un modelo de recompensa que asigne puntuaciones altas a respuestas largas (recompensa = len(respuesta) / 100). Ejecuta PPO con este modelo de recompensa defectuoso y observa cómo el modelo de política genera salidas cada vez más largas y repetitivas. Luego agrega una penalidad KL de 0.1 y demuestra que previene el comportamiento degenerado.

4. Implementa una recompensa multiobjetivo. Entrena dos modelos de recompensa: uno para la utilidad (helpfulness) y otro para la concisión. Combínalos como R = 0.7 * R_helpful + 0.3 * R_concise. Demuestra que el objetivo combinado produce respuestas que son tanto útiles como concisas, evitando la trampa de la prolijidad de una sola recompensa de utilidad.

5. Compara diferentes coeficientes KL. Ejecuta PPO con beta=0.001 (demasiado bajo, reward hacking), beta=0.02 (estándar) y beta=0.5 (demasiado alto, sin aprendizaje). Grafica la curva de recompensa y la curva KL para cada uno. La ejecución con beta=0.02 debería mostrar una mejora constante de la recompensa con una KL acotada.

Términos Clave

Lecturas Adicionales

• Ouyang et al., 2022 -- "Training language models to follow instructions with human feedback" (InstructGPT) — el artículo que hizo que RLHF fuera práctico para grandes modelos de lenguaje
• Schulman et al., 2017 -- "Proximal Policy Optimization Algorithms" — el artículo original de PPO de OpenAI
• Bai et al., 2022 -- "Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback" — el artículo de RLHF de Anthropic con un análisis detallado del reward hacking y la penalidad KL
• Stiennon et al., 2020 -- "Learning to summarize with human feedback" — RLHF aplicado a la resumición, demostrando que los modelos de recompensa pueden capturar juicios de calidad matizados
• Christiano et al., 2017 -- "Deep reinforcement learning from human preferences" — el trabajo fundacional sobre el aprendizaje de funciones de recompensa a partir de comparaciones humanas