Resumen Técnico: Recompensa Consciente de la Distribución para la Regresión con LLM

Enunciado del Problema

Los Modelos de Lenguaje Grande (LLM) han demostrado la capacidad de predecir cantidades de valor real a partir de entradas heterogéneas (texto, código, cadenas moleculares). Sin embargo, los objetivos de entrenamiento estándar para la regresión con LLM, como el Ajuste Fino Supervisado (SFT) o el Aprendizaje por Refuerzo (RL) con recompensas puntuales (por ejemplo, el Error Cuadrático Medio negativo), tratan cada número de punto flotante decodificado como una predicción independiente.

Este enfoque optimiza estimaciones puntuales individuales sin garantizar que el conjunto de predicciones muestreadas forme una distribución predictiva calibrada. En consecuencia, los modelos a menudo producen predicciones que son individualmente cercanas al objetivo pero que, colectivamente, están mal calibradas, son excesivamente estrechas o unilaterales. Este comportamiento de "búsqueda de la media" colapsa la distribución predictiva, fallando en capturar la incertidumbre necesaria o el ordenamiento relativo. Esta limitación obstaculiza aplicaciones que requieren la clasificación de candidatos, la estimación de incertidumbre o la toma de decisiones robusta bajo incertidumbre, donde la distribución de valores posibles importa tanto como la única mejor estimación.

Metodología: Recompensa Consciente de la Distribución (DAR)

Los autores proponen la Recompensa Consciente de la Distribución (DAR), un objetivo de aprendizaje por refuerzo en política diseñado para optimizar la calidad de toda la distribución predictiva formada por múltiples muestras decodificadas, en lugar de optimizar despliegues individuales de forma aislada.

Mecanismo Central

1. Distribución Predictiva Empírica: Para una entrada dada $x$, el modelo muestrea $K$ despliegues, produciendo un conjunto de predicciones $\{p_{ik}\}_{k=1}^K$. Este conjunto se trata como una distribución predictiva empírica $F_i$.
2. Puntuación mediante CRPS: La calidad de esta distribución se evalúa frente al objetivo de verdad fundamental $y_i$ utilizando la Puntuación de Probabilidad Jerarquizada Continua (CRPS). El CRPS es una regla de puntuación propia que recompensa distribuciones que son tanto precisas (concentradas cerca del objetivo) como apropiadamente dispersas (evitando la degeneración).
   • La recompensa para el conjunto completo se define como $S_i = -\text{CRPS}(F_i, y_i)$.
3. Asignación de Crédito Leave-One-Out: Dado que el RL requiere una recompensa escalar para cada despliegue individual para actualizar la política, DAR convierte la puntuación a nivel de distribución en retroalimentación a nivel de despliegue. Calcula la contribución marginal de cada despliegue $p_{ik}$ comparando el CRPS del conjunto completo ($S_i$) con el CRPS del conjunto con ese despliegue eliminado ($S_i^{(-k)}$).
   • La recompensa del despliegue es: $R_{ik}^{\text{DAR}} = S_i - S_i^{(-k)}$.
   • Si eliminar un despliegue degrada la distribución (reduce el CRPS negado), ese despliegue recibe una recompensa positiva, incluso si su error puntual individual es alto. Por el contrario, los despliegues redundantes que no mejoran la dispersión o la centración de la distribución reciben recompensas más bajas o negativas.

Distinción Teórica

El artículo contrasta esto con la pérdida cuadrática puntual (MSE), que es teóricamente de "búsqueda de la media" y alienta a la distribución predictiva a colapsar en un solo punto (la media condicional). En contraste, el CRPS es de "búsqueda de distribución", optimizando para toda la distribución condicional. DAR aprovecha esto utilizando la asignación de crédito leave-one-out para fomentar un conjunto de despliegues que estén bien centrados y útilmente dispersos.

Contribuciones Clave

1. Objetivo Novel de RL: Introducción de DAR, que cambia el objetivo de optimización de minimizar errores puntuales individuales a maximizar la calidad de la distribución predictiva empírica.
2. Asignación de Crédito Leave-One-Out: Un mecanismo específico para atribuir recompensas a despliegues individuales basándose en su contribución marginal a la calidad colectiva de la distribución, permitiendo al modelo recompensar la "dispersión útil" incluso cuando las muestras individuales son menos precisas.
3. Mitigación del Colapso de Diversidad: El método aborda la tendencia del ajuste fino con RL a reducir la diversidad de salida (entropía de la política) y colapsar las predicciones hacia la media, asegurando que los despliegues muestreados permanezcan lo suficientemente diversos para cubrir la verdad fundamental.

Resultados Experimentales

Los autores evaluaron DAR en tres tareas distintas: una tarea de regresión sintética con mezcla gaussiana, predicción de rendimiento de código y predicción de propiedades moleculares.

1. Regresión Distribucional Sintética

• Tarea: Predicción de objetivos a partir de una mezcla gaussiana de dos componentes heterocedástica.
• Hallazgos: DAR logró la mayor correlación de Spearman entre los métodos basados en LLM, superando significativamente al SFT y a las recompensas MSE puntuales. Preservó mejor el ordenamiento global de la función objetivo y se adaptó bien a regiones de extrapolación donde otras líneas base neuronales tuvieron dificultades.

2. Predicción de Rendimiento de Código

• Tarea: Predicción de latencia del kernel (KBSS) y uso máximo de memoria (APPS) a partir de código fuente.
• Hallazgos: DAR logró la correlación de rango más fuerte (Spearman $\rho$) y un Error Absoluto Medio (MAE) competitivo o inferior en comparación con las líneas base de SFT y RL puntuales.
  • En KBSS, DAR mejoró la correlación de Spearman en 6 puntos sobre la línea base más fuerte.
  • En APPS, mejoró la correlación de Spearman en 4 puntos.
  • Esto indica una sensibilidad superior a la variación dependiente de la entrada y una clasificación de candidatos más fiable.

3. Predicción de Propiedades Moleculares (MoleculeNet)

• Tarea: Predicción de propiedades fisicoquímicas (FreeSolv, ESOL, Lipofilicidad) a partir de cadenas SMILES.
• Hallazgos: A pesar de usar solo cadenas SMILES (careciendo de sesgos inductivos explícitos de grafos o 3D utilizados por líneas base GNN fuertes), DAR permaneció competitiva.
  • Logró mejoras estadísticamente significativas en RMSE en FreeSolv y Lipofilicidad en comparación con otras estrategias de entrenamiento de LLM (SFT, MSE, Value Head).
  • Demostró que la optimización a nivel de distribución puede producir un comportamiento de regresión robusto incluso sin representaciones moleculares especializadas.

Diagnósticos de Incertidumbre

• Calibración: El análisis mostró una fuerte correlación log-Pearson (0.91) entre la desviación estándar de los despliegues y el error de predicción para DAR, indicando que las estimaciones de incertidumbre del modelo estaban bien calibradas con los errores reales.
• Diversidad: Durante el entrenamiento, DAR mantuvo una mayor entropía de despliegue y una mayor "tasa de paréntesis" (la fracción de ejemplos donde la verdad fundamental se encuentra entre las predicciones mínima y máxima de los despliegues) en comparación con el RL basado en MSE, confirmando la mitigación del colapso de diversidad.

Significado y Afirmaciones

El artículo afirma que optimizar directamente las distribuciones predictivas hace que la regresión basada en LLM sea más robusta y mejor calibrada. Al tratar múltiples muestras decodificadas como una distribución y recompensar su calidad colectiva, DAR mejora el rendimiento basado en rangos (crucial para la selección de candidatos) y la estimación de incertidumbre sin sacrificar la precisión puntual.

Los autores posicionan este trabajo como un paso hacia el uso de LLM no solo como generadores de texto o estimadores puntuales, sino como regresores probabilísticos flexibles para dominios científicos. Enfatizan que, si bien la regresión con LLM ofrece flexibilidad para manejar entradas heterogéneas (texto, código, cadenas), el enfoque consciente de la distribución propuesto es necesario para desbloquear el potencial completo de estos modelos para tareas que requieren cuantificación de incertidumbre y ordenamiento fiable, áreas donde el entrenamiento puntual tradicional falla. El trabajo no afirma reemplazar modelos de regresión especializados (como GP o GNN) en todos los dominios, pero demuestra que el RL consciente de la distribución puede hacer que los LLM de propósito general sean competitivos en configuraciones de regresión científica.