Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context

1. El Problema

El artículo aborda la tarea de regresión de texto a distribución utilizando Grandes Modelos de Lenguaje (LLMs). A diferencia de la regresión tradicional que predice un único valor puntual (estimación de punto), muchas aplicaciones del mundo real (como la predicción de precios, la previsión de demanda o la evaluación de riesgos) requieren predecir una distribución de probabilidad completa para capturar la incertidumbre, la dispersión y el comportamiento de las colas.

Los enfoques actuales basados en LLMs presentan tres limitaciones clave:

1. Cuello de botella de representación: Los métodos existentes (como el de Vedula et al., 2025) utilizan un estado oculto final compartido del que se derivan múltiples cabezas lineales para predecir diferentes cuantiles. Esto crea una conexión indirecta entre las características de entrada y cada cuantil específico.
2. Falta de anclaje local: Los modelos predicen distribuciones basándose únicamente en el texto de la consulta, sin aprovechar la evidencia empírica de instancias similares (vecinos) que podrían informar mejor sobre la dispersión y la forma de la distribución.
3. Supervisión limitada en métodos de recuperación: Los métodos de recuperación aumentada (RAG) anteriores suelen adjuntar solo una etiqueta de punto a los vecinos recuperados, perdiendo la riqueza de la información distribucional que estos vecinos poseen.

2. Metodología Propuesta

Los autores proponen una arquitectura novedosa llamada Regresión de Tokens de Cuantil (Quantile Token Regression) aumentada con contexto de vecinos recuperados.

A. Tokens de Cuantil (Quantile Tokens)

En lugar de usar cabezas lineales separadas sobre una representación compartida, el modelo inserta tokens dedicados de cuantil ($\langle Q_{\tau_1} \rangle, \dots, \langle Q_{\tau_Q} \rangle$) directamente en la secuencia de entrada del transformador.

• Mecanismo: Cada token de cuantil participa en el mecanismo de autoatención a través de todas las capas del transformador.
• Ventaja: Esto permite que cada cuantil acumule información específica de las partes más relevantes de la entrada. Por ejemplo, el token para el cuantil 10 ($Q_{10}$) puede atender a señales de popularidad (prediciendo respuestas rápidas), mientras que el token para el cuantil 90 ($Q_{90}$) atiende a indicadores de complejidad (prediciendo respuestas lentas). Esto crea una vía directa de entrada-salida para cada nivel de cuantil.

B. Recuperación Aumentada de Distribuciones (Retrieval-Augmented Distribution Estimation)

El modelo recupera $K$ instancias vecinas semánticamente similares de un conjunto de datos de entrenamiento.

• Innovación: A diferencia de métodos previos que solo traen una etiqueta de valor, aquí cada vecino recuperado se presenta junto con su distribución empírica completa (representada como un conjunto de cuantiles observados).
• Objetivo: Proporcionar al modelo evidencia local sobre la dispersión y las colas de la distribución objetivo, mejorando la calibración.

C. Funciones de Pérdida y Análisis Teórico

El artículo realiza un análisis teórico de las funciones de pérdida para la regresión de cuantiles bajo supervisión empírica:

• Pérdidas $\ell_1$ y $\ell_2$ (Matching de Wasserstein): Se identifican como las más adecuadas. Minimizan la distancia entre los cuantiles empíricos interpolados y los predichos, ofreciendo estimaciones Fisher-consistentes y robustas.
• Pérdida Pinball (Pinball-Q y Pinball-Med): Se demuestra que aplicar la pérdida pinball estándar a los cuantiles empíricos o solo a la mediana introduce sesgos sistemáticos (inflación/deflación) y pierde información sobre la forma de la distribución, especialmente en las colas.

3. Contribuciones Clave

1. Arquitectura de Tokens de Cuantil: Primera propuesta de insertar tokens aprendibles específicos para cada cuantil en la entrada, permitiendo vías de atención directas y representaciones más coherentes.
2. Recuperación de Distribuciones Completas: Extensión de la recuperación aumentada (RAG) para incluir distribuciones empíricas completas de los vecinos, no solo puntos de datos, mejorando la estimación de la incertidumbre.
3. Análisis Teórico de Pérdidas: Clarificación de qué objetivos de distribución optimiza cada función de pérdida bajo supervisión de cuantiles empíricos, recomendando el uso de distancias de Wasserstein ($\ell_1$/$\ell_2$) sobre el pinball tradicional.
4. Evaluación Exhaustiva: Pruebas en dos conjuntos de datos diversos (Airbnb y StackSample) con modelos LLM de escala de 1.7B a 14B parámetros.

4. Resultados Experimentales

Los experimentos se realizaron en los conjuntos de datos Inside Airbnb (predicción de precios) y StackSample (predicción de tiempo de respuesta en preguntas de Stack Overflow), utilizando modelos Qwen3.

• Rendimiento Superior: La combinación de Tokens de Cuantil y Recuperación de Vecinos (QT + Neighbors) superó consistentemente a los baselines (Regresión de Cuantiles estándar sin vecinos y con vecinos).
• Precisión y Agudeza:
  • Reducción del MAPE promedio en un ~4% en comparación con los baselines.
  • Intervalos de predicción 2 veces más estrechos (RCIW reducido significativamente), lo que indica distribuciones más precisas y menos conservadoras.
  • En el conjunto de datos más pequeño y desafiante (StackSample), la mejora fue drástica: reducción del MAPE del 68% y reducción de la anchura del intervalo de confianza del 99% en un factor de 131x en comparación con el baseline sin vecinos.
• Impacto de la Recuperación: La recuperación de vecinos fue crucial, especialmente en conjuntos de datos más pequeños o con alta incertidumbre, donde la evidencia local de instancias similares mejoró drásticamente la calibración.
• Pérdidas: El uso de la pérdida $\ell_1$ de Wasserstein demostró ser la mejor opción, logrando el mejor equilibrio entre precisión y agudeza (intervalos estrechos).

5. Significado e Implicaciones

Este trabajo establece un nuevo estándar para la regresión de texto a distribución con LLMs.

• Superación de Limitaciones de Representación: Demuestra que la arquitectura de "cabeza compartida" es subóptima para tareas distribucionales y que la atención directa mediante tokens dedicados es superior.
• Valor del Contexto Local: Valida empíricamente que las distribuciones de resultados de instancias similares son una fuente de información crítica para estimar la incertidumbre, algo que los modelos puramente generativos a menudo pasan por alto.
• Aplicabilidad Práctica: La metodología es escalable y efectiva en modelos de diferentes tamaños, ofreciendo soluciones robustas para problemas donde la gestión del riesgo y la incertidumbre son tan importantes como la predicción del valor central.

En resumen, el paper propone un marco robusto que combina arquitectura de atención especializada (tokens de cuantil) con evidencia empírica local (vecinos con distribuciones completas) y optimización teórica adecuada (Wasserstein loss) para lograr predicciones de distribución de alta calidad a partir de texto no estructurado.