Expected Reward Prediction, with Applications to Model Routing

Título: Predicción de Recompensa Esperada con Aplicaciones a la Enrutación de Modelos (Expected Reward Prediction, with Applications to Model Routing)

1. El Problema

Los modelos de recompensa (Reward Models - RMs) son herramientas estándar en la alineación y evaluación de Modelos de Lenguaje Grande (LLM). Tradicionalmente, estos modelos asignan una puntuación a una respuesta específica ($y$) dada una instrucción ($x$), es decir, $r(x, y)$. Sin embargo, existe una brecha en la capacidad de predecir a priori qué tan bien se desempeñará un modelo específico en una instrucción dada, antes de generar la respuesta.

El desafío central es "elevar" la función de recompensa de nivel de respuesta a un nivel de modelo. Específicamente, se busca predecir la recompensa esperada de un modelo $\pi$ para una instrucción $x$, definida como el valor esperado de la recompensa sobre la distribución de respuestas del modelo:
$$ER_\pi(x) := \mathbb{E}_{Y \sim \pi(Y|x)}[r(x, Y)]$$

La pregunta de investigación es si es posible predecir este valor esperado basándose únicamente en la instrucción, sin necesidad de muestrear múltiples respuestas del modelo durante la inferencia. Esto es crucial para aplicaciones como la enrutación de modelos (Model Routing), donde se debe decidir dinámicamente qué modelo (de un conjunto con diferentes costos computacionales) es el más adecuado para una consulta específica, maximizando la calidad y minimizando el costo.

2. Metodología

Los autores proponen un enfoque empírico basado en la simplicidad y la eficiencia:

• Definición del Objetivo: Entrenar un modelo para predecir $ER_\pi(x)$ directamente desde la representación de la instrucción $x$.
• Generación de Datos de Entrenamiento:
  • Se utiliza el dataset open-perfectblend (4000 prompts divididos en categorías: chat general, seguimiento de instrucciones, razonamiento matemático y código).
  • Para cada prompt, se muestrean $K=32$ respuestas de cada modelo LLM en estudio.
  • Se calcula la recompensa media empírica de estas 32 respuestas para usarla como etiqueta de entrenamiento (el "ground truth" de la recompensa esperada).
• Modelo de Predicción (Linear Probe):
  • Se utiliza una representación de embeddings pre-entrenada de la instrucción (usando gte-large-en-v1.5).
  • Se entrena un modelo lineal con regularización Ridge para mapear el embedding de la instrucción a la recompensa esperada.
  • La función de pérdida minimiza el error cuadrático entre la predicción lineal y la recompensa media empírica.
• Configuración Experimental:
  • Modelos LLM: Llama 3.1-IT (8B, 70B), Gemma 2-IT (9B, 27B), Gemma 1-IT (7B).
  • Modelos de Recompensa (RMs): OpenAssistant-RM, GRM-2B-RM, InternLM-RM.
• Protocolo de Enrutamiento (ERP):
  • Se define una política de enrutamiento $\rho_\lambda(x)$ que selecciona el modelo $i$ que maximiza la recompensa esperada predicha menos un término de costo ponderado ($\lambda$):
    $$\rho_\lambda(x) = \arg\max_{i} \left( \widehat{ER}_{\pi_i}(x) - \lambda \cdot c(i) \right)$$
  • Donde $c(i)$ es el costo computacional (aproximado por el número de parámetros).

3. Contribuciones Clave

1. Predicción de Alta Fidelidad: Demuestran que la recompensa esperada de un modelo para un prompt específico es altamente predecible utilizando únicamente un modelo lineal simple sobre los embeddings del prompt.
2. Simplicidad y Escalabilidad: A diferencia de métodos anteriores que requieren modelos de preferencia complejos entrenados con datos de comparación par a par (que escalan cuadráticamente con el número de modelos), el enfoque ERP requiere entrenar un solo predictor por modelo. Esto hace que añadir nuevos modelos al "pool" sea trivial.
3. Validación en Enrutamiento: Muestran que estas predicciones son lo suficientemente precisas para discriminar entre modelos en prompts específicos, superando a estrategias de enrutamiento estáticas o basadas en promedios por categoría.
4. Explicación de Métodos Existentes: Sugieren que el éxito de protocolos de enrutamiento complejos (como Zooter) se debe fundamentalmente a la predictibilidad de la recompensa esperada, validando que una aproximación lineal directa es tan efectiva como redes neuronales más complejas.

4. Resultados

• Precisión Predictiva: Los modelos lineales logran valores de $R^2$ significativos (rango de 0.19 a 0.59 dependiendo del modelo y categoría), indicando que una gran parte de la varianza en la calidad de la respuesta puede explicarse solo por el prompt.
• Comparación de Enrutamiento:
  • La política basada en ERP supera a las líneas base que asignan prompts a un modelo fijo o al mejor modelo promedio por categoría.
  • En la frontera de Pareto (compromiso entre costo y arrepentimiento/regret), el ERP logra un rendimiento comparable o superior al método Zooter (que predice la distribución de ganadores) y a un oráculo que conoce las categorías de los prompts.
  • El ERP es particularmente efectivo cuando los costos de los modelos varían, permitiendo seleccionar modelos más pequeños para prompts donde su rendimiento esperado es suficiente, ahorrando recursos.
• Robustez: La predicción funciona bien tanto dentro de categorías específicas (ej. matemáticas) como entre ellas, demostrando que no es solo una detección de prompts "fáciles" o "difíciles", sino una medida de idoneidad del modelo para esa tarea específica.

5. Significado e Impacto

Este trabajo cambia el paradigma de cómo se evalúan y utilizan los LLMs en producción:

• Eficiencia Operativa: Permite la creación de sistemas de enrutamiento de inferencia altamente eficientes y escalables. Las empresas pueden gestionar pools de modelos heterogéneos (desde modelos pequeños y baratos hasta grandes y potentes) y dirigir el tráfico dinámicamente para maximizar la calidad por dólar gastado.
• Reducción de Complejidad: Elimina la necesidad de costosos procesos de entrenamiento de modelos de preferencia par a par para cada nueva combinación de modelos.
• Nueva Perspectiva Teórica: Sugiere que los modelos de recompensa modernos, al ser fine-tuneados sobre LLMs pre-entrenados, han internalizado una estructura que permite inferir propiedades distribucionales (como la recompensa esperada) a partir de la entrada, lo cual es una propiedad fundamental de estos sistemas.

En resumen, el artículo demuestra que predecir el rendimiento futuro de un modelo antes de ejecutarlo es posible, simple y altamente efectivo, abriendo la puerta a arquitecturas de inferencia más inteligentes y económicas.