Resumen Técnico: Modelado de Recompensas para el Alineamiento (en Tiempo de Inferencia) mediante la Perspectiva de un Juego de Stackelberg

1. Planteamiento del Problema

Los procesos actuales de alineamiento de Grandes Modelos de Lenguaje (LLM) suelen implicar el aprendizaje de un modelo de recompensa ($r_U$) a partir de datos de preferencias de usuario y su uso para optimizar la política del LLM. Esta optimización está constreñida por un regularizador de divergencia de Kullback-Leibler (KL) para mantener la política aprendida cerca de un modelo base ($\rho_{base}$), preservando la fluidez y la competencia lingüística.

Los autores identifican una limitación fundamental en este proceso estándar: optimizar directamente para la recompensa aprendida $r_U$ bajo una restricción KL es subóptimo para maximizar la utilidad del usuario. Cuando el modelo base presenta sesgos fuertes (por ejemplo, inclinación política) que entran en conflicto con las preferencias del usuario, la restricción KL impide que la política se desplace lo suficiente hacia la distribución preferida, incluso si el modelo de recompensa identifica correctamente la preferencia.

Por el contrario, amplificar simplemente las recompensas para forzar un desplazamiento conlleva el riesgo de reward hacking (ataque de recompensa), donde el modelo genera salidas incoherentes o degeneradas para maximizar la señal de recompensa mientras viola la intención de la restricción KL. El problema central abordado es cómo diseñar óptimamente (dar forma a) el modelo de recompensa utilizado para el alineamiento, con el fin de equilibrar la maximización de la utilidad del usuario frente al riesgo de reward hacking bajo la regularización KL.

2. Metodología: Formulación del Juego de Stackelberg

Los autores formalizan el problema del diseño de la recompensa como un juego de Stackelberg, un juego de dos jugadores secuencial:

• Líder (Proveedor del Modelo de Recompensa): Selecciona un modelo de recompensa $r$ para maximizar la utilidad esperada del usuario (basada en la función de utilidad real $r_U$), anticipando la respuesta del LLM. El líder no está obligado a usar $r_U$ directamente, sino que puede elegir cualquier función de recompensa $r: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}$ acotada por $B$.
• Seguidor (LLM): Observa el modelo de recompensa $r$ y optimiza su política $\rho_r$ para maximizar la recompensa esperada menos la divergencia KL de la política base:
  $$\rho_r(\cdot|x) = \arg\max_{\rho} \left( \mathbb{E}_{y \sim \rho}[r(x, y)] - \beta \cdot D_{KL}(\rho(\cdot|x) \parallel \rho_{base}(\cdot|x)) \right)$$

Caracterización Teórica

Los autores derivan una caracterización en forma cerrada del modelo de recompensa óptimo $r^*$:

1. Estructura de Umbral: El modelo de recompensa óptimo es un modelo de recompensa de umbral. Para un prompt dado $x$, las respuestas se particionan según si su utilidad real $r_U(x, y)$ está por encima o por debajo de un umbral específico $m^*(x)$.
   • Si $r_U(x, y) > m^*(x)$, la recompensa se establece en el límite superior $B$.
   • Si $r_U(x, y) < m^*(x)$, la recompensa se establece en $0$.
   • Si $r_U(x, y) = m^*(x)$, la recompensa está en $[0, B]$.
2. Condición de Autoconsistencia: El umbral óptimo $m^*(x)$ satisface la condición de que la utilidad esperada de la política resultante es igual al propio umbral:
   $$m^*(x) = \mathbb{E}_{y \sim \rho_{r^*}(\cdot|x)} [r_U(x, y)]$$
3. Computación: Dado que la esperanza exacta es intratable, los autores proponen calcular el umbral mediante muestreo de Monte Carlo desde la política base. Definen una función auxiliar $F_x(m)$ cuya raíz única corresponde al umbral óptimo. Esta raíz puede encontrarse eficientemente mediante un algoritmo de bisección.

Implementación Práctica: Suavizado de Umbral (SRS)

Aunque el umbral rígido es analíticamente óptimo, es discontinuo y sensible a pequeñas perturbaciones. Los autores introducen el Modelado de Recompensa de Stackelberg (SRS), una variante de umbral suave que utiliza una función sigmoide:
$$r_{\hat{m}^*, \alpha}(x, y) = B \cdot \sigma(\alpha \cdot (r_U(x, y) - \hat{m}^*(x)))$$
donde $\hat{m}^*(x)$ es el umbral estimado y $\alpha$ controla la nitidez de la transición. Esto permite que el método interpole entre un efecto de no alineamiento ($\alpha=0$) y el óptimo analítico ($\alpha \to \infty$).

Integración con el Alineamiento en Tiempo de Inferencia

El marco se integra en métodos existentes de alineamiento en tiempo de inferencia (específicamente Decodificación Controlada (CD) y Búsqueda Guiada por Recompensa de Alineamiento (ARGS)):

1. Fase Offline: Se muestrean $M$ respuestas del modelo base para un conjunto de prompts. Se calcula $r_U$ para estas muestras. Se estima el umbral $\hat{m}^*(x)$ mediante bisección. Se aplica la transformación SRS para generar un conjunto de datos de recompensa modelada. Se entrena un estimador de valor-Q sobre estos datos modelados.
2. Fase de Inferencia: Se utiliza el estimador de valor-Q entrenado (que captura la recompensa esperada bajo la recompensa modelada) para reponderar las probabilidades de tokens del modelo base durante la decodificación.

3. Contribuciones Clave

1. Suboptimalidad de los Procesos Estándar: Se demostró que el uso directo de un modelo de recompensa aprendido de las preferencias es subóptimo bajo objetivos de regularización KL, particularmente cuando los sesgos del modelo base entran en conflicto con las preferencias del usuario.
2. Formulación del Juego de Stackelberg: Se formuló la optimización del modelo de recompensa como un juego de Stackelberg, derivando una solución en forma cerrada que muestra que el modelo de recompensa óptimo sigue una estructura de umbral.
3. Aproximación Eficiente: Se desarrolló un algoritmo para aproximar el umbral óptimo utilizando el muestreo de Monte Carlo de la política base, permitiendo la integración en métodos existentes con un sobrecoste mínimo.
4. Validación Empírica: Se demostró que el método SRS propuesto mejora consistentemente la recompensa promedio y alcanza altas tasas de victoria-empate contra los modelos de referencia en entornos de tiempo de inferencia.

4. Resultados Experimentales

El método fue evaluado en los conjuntos de datos HH-RLHF y SHP utilizando arquitecturas Qwen3-8B y Llama3-8B.

• Maximización de la Recompensa: Los métodos integrados con SRS (SRS-CD y SRS-ARGS) alcanzaron consistentemente la mayor recompensa promedio en todos los entornos de evaluación en comparación con los modelos de referencia vanilla y los esquemas de modelado de recompensa heurísticos (Minmax, Meanstd).
• Métricas de Calidad: Los métodos mantuvieron la diversidad y coherencia comparables a los modelos de referencia, lo que indica que no hay degradación en la calidad del lenguaje.
• Evaluación con GPT-4: En comparaciones directas, los métodos SRS alcanzaron una tasa de Victoria-Empate superior al 66% contra todos los modelos de referencia, lo que sugiere que las mejoras no son producto de reward hacking.
• Robustez:
  • Generalización: SRS mostró un rendimiento superior cuando se entrenó en una distribución (SHP) y se evaluó en prompts fuera de la distribución (HH-RLHF y HarmfulQA).
  • Ablación: La estructura de umbral específica de SRS superó al simple recorte de recompensa y al umbral basado en la media, confirmando la importancia del umbral autoconsistente derivado de la formulación de la teoría de juegos.
• Sobrecoste: Para SRS-CD, el modelado de la recompensa es offline, con cero sobrecoste en tiempo de inferencia. Para SRS-ARGS, el sobrecoste es insignificante (~0.03 segundos por respuesta de 128 tokens).

5. Significado y Reivindicaciones

El artículo sostiene que el modelado de la recompensa no es meramente un heurístico, sino un problema de optimización que puede resolverse sistemáticamente. Al observar el proceso de alineamiento a través de la lente de la teoría de juegos, los autores demuestran que el modelo de recompensa utilizado para guiar al LLM no necesita ser (y no debe ser) un proxy directo de la utilidad del usuario. En su lugar, debe ser "modelado" para tener en cuenta las restricciones impuestas por el modelo base (divergencia KL).

La importancia radica en proporcionar un marco con principios para contrarrestar los sesgos del modelo base sin inducir reward hacking. Los autores enfatizan que su enfoque es ligero, se integra perfectamente en los métodos de alineamiento en tiempo de inferencia existentes y ofrece una solución robusta para adaptar modelos fijos a diversas preferencias de usuario.

Limitaciones señaladas por los autores:

• El marco asume un usuario benigno cuyas preferencias se alinean con los valores humanos. Si el modelo de recompensa refleja preferencias maliciosas, la política resultante podría desalinearse.
• El método depende de la disponibilidad de un modelo de recompensa y de muestras de Monte Carlo de la política base.
• El enfoque empírico actual se centra en el alineamiento en tiempo de inferencia, aunque los autores señalan el potencial de extenderlo a algoritmos de entrenamiento (por ejemplo, GRPO) como trabajo futuro.