Resumen Técnico: PFN-TS

Enunciado del Problema

Los banditos contextuales abordan la toma de decisiones secuencial bajo incertidumbre, donde un agente observa un contexto, selecciona una acción y recibe una recompensa estocástica. El objetivo es minimizar el arrepentimiento acumulado a lo largo del tiempo. Muestreo de Thompson (TS) es una estrategia estándar que mantiene una distribución posterior bayesiana sobre las funciones de recompensa y muestrea acciones de manera codiciosa a partir de esta posterior. Sin embargo, la eficacia de TS depende de la calidad del modelo de recompensa subyacente.

En aplicaciones tabulares (por ejemplo, salud móvil, recomendaciones), los modelos de recompensa lineales suelen estar mal especificados porque no capturan comportamientos humanos complejos como interacciones, efectos umbral y heterogeneidad. Aunque existen alternativas más ricas (métodos de kernel, redes neuronales, ensembles de árboles), adolecen de problemas de escalabilidad, alta sobrecarga computacional o suposiciones estructurales rígidas. Además, estos modelos suelen requerir hiperparámetros difíciles de ajustar en entornos en línea.

Redes Ajustadas a Datos Previos (PFN), como TabPFN y TabICL, ofrecen una alternativa prometedora. Estos son modelos basados en transformers entrenados para aproximar distribuciones predictivas posteriores (PPD) bayesianas a través de un amplio prior de procesos generadores de datos. Operan mediante aprendizaje en contexto (ICL), sin requerir actualizaciones de parámetros ni ajuste de hiperparámetros. Sin embargo, existe una brecha crítica: las PFN emiten PPDs (incertidumbre sobre recompensas futuras ruidosas), mientras que TS requiere incertidumbre sobre la función de recompensa media latente. Muestrear directamente de la PPD no equivale a muestrear de la posterior de la función media, lo que conduce a una exploración subóptima.

Metodología: PFN-TS

Los autores proponen PFN-TS, un algoritmo de Muestreo de Thompson que cierra esta brecha convirtiendo las predictivas posteriores de las PFN en muestras de la función de recompensa media latente. La metodología consta de tres componentes principales:

1. Teorema del Límite Central Predictivo Muestreado (SubCLT)

Para recuperar la distribución posterior de la media latente $f_0(x)$ a partir de la secuencia predictiva de la PFN, los autores aprovechan un Teorema del Límite Central (CLT) predictivo.

• Base Teórica: Bajo suposiciones específicas (varianza asintótica, condiciones de cuasi-martingala y momentos acotados), la cola de la secuencia de medias predictivas converge a una distribución gaussiana centrada en la media latente.
• Innovación en Eficiencia: Los enfoques tradicionales para estimar la varianza posterior requieren calcular la media predictiva sobre todos los $n$ prefijos del conjunto de datos, lo cual es computacionalmente prohibitivo para banditos en línea ($O(n)$ pasadas hacia adelante). PFN-TS introduce un estimador muestreado que calcula la varianza utilizando solo $O(\log n)$ prefijos espaciados geométricamente ($t_0, t_1, \dots, t_J$).
• Estimación de Varianza: La varianza se estima utilizando incrementos de bloque $D_j = m_{t_j}(x) - m_{t_{j-1}}(x)$ ponderados armónicamente ($w_j$). Los autores demuestran que este estimador muestreado es consistente ($\hat{V}_{s_n} \xrightarrow{p} V(x)$) y satisface el CLT, permitiendo una aproximación gaussiana tratable $N(m_{s_n}(x), \hat{V}_{s_n}(x)/s_n)$ para el muestreo de Thompson.

2. Caché y Optimización Computacional

Para reducir aún más el costo por ronda, PFN-TS explota la arquitectura de TabICL (la PFN específica utilizada):

• Reutilización de Caché KV: TabICL almacena en caché matrices de atención clave-valor para un conjunto de entrenamiento fijo. PFN-TS reutiliza estas representaciones en caché a lo largo de las rondas.
• Puntos de Actualización: El algoritmo solo realiza nuevas pasadas hacia adelante en los "puntos de actualización" geométricos ($t_j$) definidos por la cuadrícula SubCLT. Entre los puntos de actualización, utiliza la aproximación gaussiana de instantánea derivada del estado en caché más reciente. Esto reduce la complejidad de $O(n)$ a $O(\log n)$ pasadas hacia adelante por brazo por ronda.

3. Estrategia de Codificación Adaptativa

El artículo aborda la elección de diseño sobre cómo representar la identidad del brazo en la entrada:

• Codificación Disjunta: Mantiene modelos independientes por brazo.
• Codificación Conjunta: Añade la identidad del brazo como un vector one-hot al contexto, permitiendo el intercambio de información entre brazos.
• Selección Adaptativa: Dado que ninguna codificación domina universalmente, PFN-TS ejecuta ambas en paralelo durante una fase de "codificación dual". Selecciona la codificación activa basándose en la Puntuación de Probabilidad Clasificada Continua (CRPS), una regla de puntuación estrictamente propia que mide la calibración. La codificación con menor CRPS acumulativo se retiene y el retador se descarta después de un tiempo de cambio preespecificado.

Contribuciones Clave

1. Puente Algorítmico: PFN-TS es el primer método que adapta con éxito las PFN (que emiten PPDs) para el Muestreo de Thompson derivando un estimador consistente para la varianza posterior de la media latente mediante un CLT muestreado.
2. Garantías Teóricas:
   • Consistencia: Demuestra la consistencia del estimador de varianza muestreado bajo suposiciones leves.
   • Límite de Arrepentimiento: Deriva un límite de arrepentimiento bayesiano para PFN-TS. El límite se descompone en el arrepentimiento del Muestreo de Thompson exacto bajo el prior de la PFN más un término de error dependiente de la calidad de aproximación de la PFN y del proceso de muestreo.
3. Eficiencia: Demuestra que $O(\log n)$ pasadas hacia adelante son suficientes para una estimación de varianza consistente, haciendo viables los banditos basados en PFN para su despliegue en línea.
4. Rendimiento Empírico:
   • Benchmarks Sintéticos: PFN-TS logra el mejor rango promedio en tareas sintéticas no lineales (variantes de Friedman) y se mantiene competitivo en tareas lineales.
   • Benchmarks OpenML: Supera a las líneas base (LinTS, NeuralTS, BFTS, XGBoostTS) en seis de ocho conjuntos de datos de clasificación a bandito.
   • Aplicación del Mundo Real: En una evaluación fuera de línea del ensayo de salud móvil Drink Less, PFN-TS alcanza el mayor valor de política estimado (SNIPS) en el horizonte final, superando a BFTS y otras líneas base.

Significado y Afirmaciones

El artículo afirma que PFN-TS demuestra que el aprendizaje en contexto tabular puede servir como un modelo de recompensa efectivo para la toma de decisiones secuencial una vez que se implementa una extracción eficiente de incertidumbre y una codificación adaptativa.

• Flexibilidad del Modelo: A diferencia de los modelos tradicionales, PFN-TS no requiere ajuste de hiperparámetros y se adapta automáticamente a la estructura de los datos (esparcidad, suavidad, no linealidad) inherente al amplio prior de la PFN.
• Viabilidad Práctica: Al reducir el cuello de botella computacional de la estimación de incertidumbre de lineal a logarítmico en el tamaño del conjunto de datos, el método hace práctico el despliegue de modelos fundacionales en entornos de banditos en línea.
• Robustez: Se demuestra que el método es robusto a través de diversos procesos generadores de datos, funcionando bien tanto en tareas sintéticas no lineales como en datos reales de salud móvil.

Los autores mantienen modestia respecto a las limitaciones, señalando que el tiempo de ejecución sigue siendo una restricción debido a las evaluaciones repetidas del modelo, y que la calibración en muestras finitas puede ser pobre cuando las historias de los brazos son cortas o poco informativas. Sugieren trabajo futuro en estimación de varianza regularizada y límites de arrepentimiento más precisos para los priores específicos de TabPFN/TabICL.