LLM REgression with a Latent Iterative State Head

1. Planteamiento del Problema

El artículo aborda la dificultad de adaptar los Grandes Modelos de Lenguaje (LLMs) a tareas de regresión (predicción de valores escalares continuos), en lugar de las tareas generativas de texto para las que fueron pre-entrenados.

• El Desajuste Paradigmático: El paradigma actual "texto-a-texto" es subóptimo para la regresión. Los métodos de decodificación autoregresiva tratan los números como tokens de texto, lo que ignora la proximidad numérica (ej. para un objetivo de 1.0, un modelo de lenguaje penaliza 0.1 y 0.9 por igual si son tokens distintos, mientras que la regresión prefiere 0.9).
• Limitaciones de los Enfoques Existentes:
  • Decodificación Autoregresiva: Simple pero ineficiente numéricamente y sensible al formato del prompt.
  • Inferencia Consciente de la Regresión (RAIL): Mejora la precisión mediante reglas de decisión bayesianas sobre la distribución de salida, pero requiere muestreo múltiple, lo que aumenta drásticamente el costo computacional en la inferencia.
  • Cabezas Predictivas (Predictive Heads): Evitan la generación de texto mapeando directamente las representaciones internas a un número. Sin embargo, los métodos actuales dependen de heurísticas de agrupación estáticas (como el mean pooling), lo que puede perder información fina necesaria para la regresión precisa.

2. Metodología: RELISH

Los autores proponen RELISH (REgression with a Latent Iterative State Head), una arquitectura ligera diseñada para extraer valores escalares directamente de las representaciones de un LLM congelado.

• Mecanismo Central: En lugar de agrupar estáticamente los tokens o generar texto, RELISH mantiene un estado latente aprendido que se refina iterativamente.
• Proceso de Arquitectura:
  1. Proyección: Las representaciones ocultas del LLM ($H(x)$) se proyectan a una dimensión de cabeza ($d_h$).
  2. Refinamiento Iterativo: Se inicia un estado latente $r^{(0)}$. Este estado se actualiza $L$ veces mediante bloques de refinamiento residual que utilizan atención cruzada (cross-attention). En cada paso, el estado latente actúa como consulta (query) sobre las representaciones de los tokens (keys y values) del LLM.
  3. Predicción: El estado final $r^{(L)}$ se mapea a un valor escalar mediante un regresor lineal simple.
• Entrenamiento: El LLM base permanece congelado. Solo se entrenan los parámetros de la cabeza (proyección, estado latente, bloques de atención y regresor lineal). Se utiliza la pérdida de Huber sobre objetivos normalizados para estabilizar el entrenamiento.
• Eficiencia: Es un método de "paso único" (single-pass) en la inferencia, evitando el muestreo múltiple de métodos como RAIL.

3. Contribuciones Clave

1. Nueva Arquitectura Predictiva: Introducción de un mecanismo de refinamiento iterativo de estado latente que supera las limitaciones del agrupamiento estático (pooling) en las cabezas predictivas.
2. Comparación Unificada: Un análisis exhaustivo que compara tres familias de métodos (decodificación autoregresiva, inferencia consciente de regresión y cabezas predictivas) a través de 5 conjuntos de datos, 4 backbones de LLM y diferentes regímenes de entrenamiento (congelado y LoRA).
3. Eficiencia de Parámetros: Demostración de que RELISH logra un rendimiento superior con una huella de entrenamiento mínima (solo ~3.4–3.7M parámetros), independientemente del tamaño del modelo base, a diferencia de métodos como LoRA cuyo costo escala con el tamaño del modelo.

4. Resultados Experimentales

El modelo fue evaluado en cuatro backbones (Llama 3.1 8B, Qwen3 8B/32B, Gemma 3 27B) y cinco datasets (STS-B, SICK-R, y tres tareas de estimación de calidad de traducción WMT).

• Rendimiento Predictivo: RELISH superó consistentemente a todos los baselines en todas las configuraciones.
  • Mejora respecto al segundo mejor método (RAFT): +6.7% en correlación de Pearson, +7.2% en Spearman y -19.9% en RMSE normalizado por rango.
  • Funciona mejor que la inferencia autoregresiva (incluso con few-shot) y que las cabezas predictivas estáticas (Lineal/MLP).
• Eficiencia de Parámetros:
  • RELISH requiere solo 0.01% – 0.04% de parámetros adicionales sobre el LLM base.
  • En comparación, LoRA/RAFT requiere 0.26% – 0.42% (aproximadamente 6 a 37 veces más parámetros dependiendo del tamaño del modelo).
  • A diferencia de LoRA, el costo de RELISH no escala con el número total de parámetros del LLM, sino principalmente con la dimensión oculta de la cabeza, manteniéndose casi constante al escalar el modelo base.
• Análisis de Ablación: Se confirmó que las ganancias provienen principalmente de la arquitectura iterativa (el refinamiento del estado latente) y no solo del uso de la pérdida de Huber. La profundidad de refinamiento ($L=3$) se identificó como un punto óptimo para la mayoría de los modelos.

5. Significado e Impacto

El trabajo de RELISH es significativo porque resuelve la tensión entre precisión y eficiencia en la regresión con LLMs:

• Superación del "Cuello de Botella de Resumen": Demuestra que las representaciones de los LLMs contienen señales de regresión ricas que no pueden ser capturadas por un simple promedio de tokens, sino que requieren un mecanismo de atención dinámica e iterativa para ser extraídas correctamente.
• Viabilidad Práctica: Al eliminar la necesidad de muestreo múltiple (costoso) y de ajustar miles de millones de parámetros (como en LoRA), RELISH hace que la regresión de alta precisión sea viable en entornos con recursos limitados o en aplicaciones en tiempo real.
• Aplicaciones Futuras: La arquitectura es prometedora para tareas como modelado de recompensas (Reward Modeling), LLMs como jueces (LLM-as-a-judge) y estimación de incertidumbre, donde se requieren predicciones numéricas precisas y eficientes sin sacrificar la capacidad del modelo base.

En resumen, RELISH establece un nuevo estado del arte en la regresión basada en LLMs al ofrecer un mecanismo de extracción de señales numéricas que es tanto arquitectónicamente superior como computacionalmente eficiente.