Resumen Técnico: REAR: Realineación de Preferencias en Tiempo de Prueba mediante la Descomposición de Recompensas

Declaración del Problema

Alinear los Modelos de Lenguaje de Gran Escala (LLMs) con las diversas y personalizadas preferencias de los usuarios es un desafío crítico. Si bien los métodos de post-entrenamiento como el Aprendizaje por Refuerzo a partir de la Retroalimentación Humana (RLHF) son efectivos, requieren una costosa curación de datos y entrenamiento adicional, lo que los hace inflexibles para necesidades de usuario dinámicas o específicas. El Escalamiento en Tiempo de Prueba (TTS) ofrece una alternativa sin necesidad de entrenamiento, aprovechando la computación adicional durante la inferencia para mejorar la calidad de la salida. Sin embargo, los métodos de TTS existentes se han limitado en gran medida a dominios verificables (por ejemplo, matemáticas, programación) donde la corrección es fácilmente juzgable. Aplicar el TTS al alineamiento de preferencias es difícil porque la calidad de la respuesta es holística y no reducible a una respuesta verificable simple. Los enfoques actuales de alineamiento en tiempo de prueba suelen imponer sobrecargas computacionales fijas o requieren el entrenamiento de cabezales de valor, lo que limita su escalabilidad y eficiencia.

Metodología

El artículo introduce REAR (REAlignment Reward), un nuevo marco que plantea el alineamiento de preferencias en tiempo de prueba como un problema de realineación. La idea central es que el objetivo de entrenamiento de un modelo base puede no satisfacer de manera óptima la preferencia de un usuario específico en el momento de la inferencia. REAR aborda esto descomponiendo la función de recompensa subyacente en dos componentes y reescalando dinámicamente su importancia.

1. Descomposición de la Recompensa

Los autores formalizan la estructura de recompensa de un LLM preentrenado dentro del marco de RL de Máxima Entropía. Postulan que la recompensa $r(s \oplus x_p, a)$ para un estado $s$ (pregunta) y una preferencia $x_p$ puede descomponerse en:

• Recompensa relacionada con la pregunta ($r_0$): La recompensa implícita que el modelo base maximiza para responder a la pregunta correctamente.
• Recompensa relacionada con la preferencia ($r_p$): El componente de recompensa que captura qué tan bien la respuesta satisface la preferencia $x_p$.

La recompensa compuesta se modela como $r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$, donde $\alpha$ es un coeficiente de equilibrio inherente.

2. Recompensa de Realineación (REAR)

Para realinear el modelo para un usuario específico, REAR introduce un coeficiente controlable $\hat{\alpha}$ (o $\lambda = \hat{\alpha}/\alpha$) para reescalar el componente de preferencia durante la inferencia:
$$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$$

Crucialmente, los autores derivan una formulación computacionalmente eficiente para esta recompensa. Aprovechando la relación entre las políticas óptimas y las funciones de valor en MaxEnt RL, demuestran que $r_{REAR}$ puede expresarse como una combinación lineal de las log-probabilidades de los tokens de la política:
$$S_{REAR}(\tau) = \sum_{t=0}^{T} \gamma^t \left( (1-\lambda) \log \pi(a_t|s_t) + \lambda \log \pi(a_t|s_t \oplus x_p) \right)$$
Esta función de puntuación depende únicamente de las log-probabilidades del modelo base bajo dos contextos: el prompt de solo la pregunta y el prompt completo incluyendo las preferencias. No requiere entrenar un modelo de recompensa separado ni generar tokens adicionales para la puntuación.

3. Integración con el Escalamiento en Tiempo de Prueba (TTS)

La puntuación REAR se integra en dos algoritmos de TTS:

• Muestreo Best-of-N (BoN): Generar $N$ respuestas y seleccionar la que tenga la puntuación $S_{REAR}$ más alta.
• Búsqueda en Árbol de Verificador Diverso (DVTS): Un algoritmo de búsqueda en árbol que utiliza las puntuaciones REAR para guiar la expansión y selección de ramas, promoviendo la diversidad mientras se optimiza el alineamiento de preferencias.

Contribuciones Clave

1. Formalización de la Realineación: El artículo formaliza el alineamiento de preferencias en tiempo de prueba como un problema de realineación, proponiendo REAR como una recompensa computacionalmente eficiente derivada de un objetivo descompuesto.
2. Puntuación sin Entrenamiento: REAR se deriva únicamente de las probabilidades internas del modelo base, sin requerir entrenamiento adicional, entrenamiento de modelos de recompensa o pasos de generación especulativa.
3. Algoritmos Escalables: Los autores desarrollan dos algoritmos de tiempo de prueba escalables (BoN y DVTS) guiados por REAR que pueden mejorar el rendimiento de la generación con presupuestos de muestreo variables.
4. Generalizabilidad: Se muestra que el método se generaliza más allá del alineamiento de preferencias explícito hacia el razonamiento matemático y las tareas visuales mediante la inyección de texto de preferencia específico de la tarea.

Resultados Experimentales

Los autores evaluaron REAR en varios benchmarks utilizando modelos Qwen y Llama:

• Benchmarks de Alineamiento de Preferencias: En PrefEval (que cubre preferencias explícitas/implícitas y elecciones), Multifaceted Bench, y Ping-Pong (juego de roles), los métodos guiados por REAR (tanto BoN como DVTS) superaron consistentemente a los baselines, incluyendo el decodificación codiciosa (greedy decoding), Amulet, Alineamiento Lineal (LA) y Best-of-N con Modelos de Recompensa Generativos (GenRM). Por ejemplo, DVTS con REAR mejoró en promedio un 11.6% respecto a la decodificación codiciosa en todas las métricas.
• Robustez de Contexto Largo: REAR demostró una robustez superior en escenarios de contexto largo (hasta 16k tokens) en comparación con GenRM, el cual se degradó significamente a medida que aumentaba la longitud del contexto.
• Razonamiento Matemático: Al aplicarse a tareas matemáticas (MATH500, AIME, AMC) con el texto de preferencia apropiado (por ejemplo, "Actúa como un matemático riguroso"), la selección guiada por REAR superó al voto mayoritario (majority voting), mostrando ganancias de escalamiento consistentes.
• Tareas Visuales: En MMHal-Bench para la detección de alucinaciones visuales, REAR mejoró las puntuaciones generales y redujo las tasas de alucinación en comparación con la decodificación codiciosa y GenRM.
• Transferencia entre Modelos: El método se generalizó a Llama-3.1-8B-Instruct sin ajuste por modelo, confirmando que la descomposición de la recompensa no es específica de la familia Qwen.
• Eficiencia: Los métodos guiados por REAR fueron significativamente más rápidos que los baselines que requieren modelos de recompensa externos o optimización iterativa (como Amulet), ya que evitan la carga de modelos separados y dependen de un único paso hacia adelante (forward pass) para la puntuación.

Significado y Reivindicaciones

El artículo afirma que REAR proporciona una solución escalable y general para alinear LLMs con las preferencias del usuario en el tiempo de prueba sin la sobrecarga de entrenamiento o modelos de recompensa externos. Al descomponer la recompensa en componentes de pregunta y preferencia, REAR permite el control dinámico sobre el compromiso entre responder a la pregunta y adherirse a las preferencias.

Los autores enfatizan que REAR permite que el escalamiento en tiempo de prueba se extienda a dominios subjetivos y abiertos (como el alineamiento de preferencias y el juego de roles) donde los verificadores tradicionales fallan. Observan que, si bien un $\lambda = 20$ por defecto funciona de manera robusta en muchas tareas, el método permite ajustar la fuerza de la realineación. Este trabajo se posiciona como un paso hacia hacer que el alineamiento en tiempo de prueba sea más accesible y eficiente, cerrando la brecha entre el éxito del TTS en dominios verificables y las complejidades del alineamiento de preferencias subjetivas.

Limitaciones Reconocidas: Los autores señalan que REAR expone un hiperparámetro ajustable ($\lambda$) al usuario, lo que puede requerir pequeñas búsiones de validación para tareas con distribuciones de preferencia significativamente diferentes. Además, aunque el TTS es más ligero que el ajuste fino (fine-tuning), identificar el presupuesto computacional óptimo sigue siendo una dirección para el trabajo futuro.