GRPO paso a paso para ajustar LLMs

GRPO (Group Relative Policy Optimization) es un algoritmo de aprendizaje por refuerzo para afinar LLMs, creado por DeepSeek. Es clave en DeepSeek-R1. Es una variante de PPO (Proximal Policy Optimization) que elimina la necesidad de un modelo crítico (value network). 

Para cada prompt genera un grupo de respuestas, cada una recibe una recompensa (usando una función de recompensa que puede basarse en una recompensa automática), se calcula la ventaja relativa de cada respuesta, se actualiza la política (los parámetros del modelo) para aumentar la probabilidad de las respuestas con ventaja positiva y disminuir la de las negativas, usando una penalización KL contra el modelo de referencia.

Los pasos simplificados serían:

CONFIGURACIÓN Y GENERACIÓN DE RESPUESTAS

Primero, se carga un modelo (ej. Qwen/Qwen2-0.5B-Instruct) y su tokenizer. Para un prompt dado, se genera G (p.e. 4) respuestas diferentes usando una temperatura alta para fomentar la diversidad.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "Qwen/Qwen2-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)

prompt = "Prompt example...Solve the equation 2x + 3 = 7."
inputs = tokenizer(prompt, return_tensors="pt")

# Genera G=4 respuestas diferentes para el mismo prompt
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        do_sample=True,
        temperature=1.0,
        num_return_sequences=4, # 'G' en GRPO
        pad_token_id=tokenizer.eos_token_id
    )

# Decodifica las respuestas para inspeccionarlas
completions = [tokenizer.decode(out, skip_special_tokens=True) for out in outputs]
for i, comp in enumerate(completions):
    print(f"Completion {i+1}: {comp}")

CALCULAR LAS RECOMPENSAS

Se define una función de recompensa que verifique si la respuesta contiene el resultado correcto. Es la función de recompensa proxy que aproxima la función real. En un experimento real, se definen varias funciones de recompensa (formato, corrección, etc.)

def reward_function(completions, ground_truth="2", **kwargs):
    """
    Recompensa 1.0 si la respuesta contiene el resultado correcto, 0.0 en caso contrario.
    """
    rewards = []
    for completion in completions:
        # Extrae la parte de la respuesta después del prompt si es necesario
        # Aquí se asume que 'completion' es la respuesta completa.
        if ground_truth in completion:
            rewards.append(1.0)
        else:
            rewards.append(0.0)
    return rewards

# Calcula las recompensas para el grupo de completions
rewards = reward_function(completions)
print(f"Rewards: {rewards}")

CALCULAR LAS VENTAJAS

Este es el núcleo de GRPO. Las recompensas se normalizan dentro del grupo para obtener las ventajas para cada completion o respuesta i. La fórmula es:

$${\widehat{A}}_{i,t}=\frac{R_i-\text{mean}(\{R_j\})}{\text{std}(\{R_j\})}$$

Las ventajas indican cuánto mejor (valor positivo) o peor (valor negativo) es cada completion en relación con el promedio del grupo.

CALCULAR LA PÉRDIDA, EL GRADIENTE Y AJUSTAR LOS PESOS DEL MODELO

Para calcular la pérdida, necesitas las probabilidades de los tokens generados bajo la política actual ($P_{\theta}$​) y una política de referencia ($P_{ref}$​), que normalmente es una copia congelada del modelo antes de la actualización.

Como vamos a usar el descenso de gradiente tradicional y aquí nos interesa maximizar la recompensa, definimos la pérdida como el negativo de la recompensa. De esta forma el algoritmo buscará pérdidas menores que se corresponden con recompensas mayores.

La pérdida para un token generado en una posición determinada se define de la siguiente manera:

Donde Beta controla la penalización por desviarse del modelo de referencia. Normalmente, también se añade un ratio de clipping (ratio de probabilidades entre la política que estamos entrenando y la política justo antes del paso de optimización) que controla que los pesos no cambien de forma destructiva si una ventaja es demasiado grande.

De esta forma, el clipping controla que no haya gradientes muy grandes dentro de cada paso de optimización y la divergencia controla que el modelo no se desvíe mucho del modelo de referencia.

El gradiente de la pérdida, sin incluir el clipping, sería:

Que A sea mayor que Beta no garantiza que la probabilidad del token suba. A medida que el modelo se aleja de la referencia, la penalización KL aumenta de forma logarítmica hasta que su fuerza supera a la ventaja A, obligando al modelo a bajar la probabilidad del token para evitar el desvío.

El algoritmo detiene el crecimiento del token en el momento exacto en que:

Si las probabilidades del modelo actual y del modelo de referencia son muy cercanas, manda la comparación directa entre A y Beta. Si A es mayor que Beta, el token sube obligatoriamente. Como Beta menos A da un resultado negativo, al aplicar descenso de gradiente (que minimiza la pérdida restando el gradiente), el signo se invierte y obliga al modelo a incrementar la probabilidad del token.

Hay que tener en cuenta que los gradientes de cada token de la secuencia se calculan y se acumulan en los parámetros del modelo. El optimizador (optimizer.step()) toma el gradiente acumulado al final de la secuencia y ajusta los pesos del modelo una sola vez.

Esto tiene una ventaja, si un modelo genera una respuesta larga que empieza muy bien pero termina con una alucinación o un error lógico, la recompensa total del grupo bajará, haciendo que la ventaja A sea negativa. Gracias a que los gradientes se acumulan y se aplican al final, esa ventaja negativa afectará a los gradientes de todos los tokens de la secuencia por igual, enseñándole al modelo a corregir el rumbo completo de su razonamiento.