Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation

Resumen Técnico: Los Transformers como Experimentadores Bayesianos en Contexto

Formulación del Problema
El artículo aborda el desafío de diseñar experimentos adaptativos para estimar el Efecto de Tratamiento Promedio (ATE, por sus siglas en inglés) en entornos secuenciales de dos brazos con covariables. La estrategia óptima para tales experimentos es la asignación de Neyman "oráculo", que asigna probabilidades de tratamiento basándose en las varianzas residuales desconocidas condicionales al brazo para minimizar la varianza del estimador del ATE. Implementar esto de manera óptima requiere un experimentador bayesiano que actualice secuencialmente las creencias no paramétricas sobre los momentos de los resultados potenciales (medias y segundos momentos) y traduzca estas incertidumbres posteriores en probabilidades de asignación de tratamiento.

Un cuello de botella crítico en espacios de covariables de alta dimensión es la suavidad desconocida de las funciones de resultados potenciales. Los priors excesivamente suaves inducen sesgo, mientras que los excesivamente rugosos convergen lentamente. Además, desplegar un proceso completo de actualización y diseño bayesiano en tiempo real es computacionalmente intensivo, lo que requiere ingeniería por cada experimento. Los autores investigan si este complejo proceso de inferencia y diseño secuencial puede ser "amortizado" mediante el aprendizaje en contexto, específicamente a través de arquitecturas Transformer, para crear una política que mapee directamente las historias experimentales hacia probabilidades de asignación óptimas sin necesidad de realizar una computación bayesiana explícita en cada paso.

Metodología
Los autores proponen un marco centrado en un "experimentador bayesiano en contexto", implementado como una política Transformer entrenada para imitar a un "maestro de Neyman con posterior bayesiana".

1. El Maestro Bayesiano:

   • Modelo: El maestro mantiene creencias bayesianas no paramétricas sobre las medias de los brazos ($\mu_w$) y los segundos momentos ($q_w$) utilizando priors de series gaussianas. Estos priors están indexados por un parámetro de suavidad $\beta$.
   • Adaptación Jerárquica: Para manejar la suavidad desconocida, el maestro emplea una extensión bayesiana jerárquica donde $\beta$ es tratado como una variable aleatoria con un prior sobre un intervalo compacto. Esto permite que la posterior se concentre en la clase de suavidad apropiada a medida que se acumulan los datos.
   • Asignación: El maestro actualiza las creencias utilizando una cuasi-verosimilitud gaussiana simplificada (pérdida cuadrática) y asigna la siguiente unidad de tratamiento según una propensión de Neyman posterior: $e_t(x) \propto \sqrt{\bar{v}_{1,t-1}(x)} / (\sqrt{\bar{v}_{0,t-1}(x)} + \sqrt{\bar{v}_{1,t-1}(x)})$, donde $\bar{v}$ es la media posterior de la varianza residual.
   • Garantías Teóricas: El artículo demuestra que, bajo condiciones de regularidad leves, la posterior del maestro se contrae a la tasa de convergencia no paramétrica minimax (adaptada a la suavidad real) y que el diseño resultante converge a la asignación de Neyman oráculo, asegurando la eficiencia semiparamétrica para el estimador final del ATE (vía Ponderación Inversa de Probabilidad Aumentada, AIPW).

2. Expresividad y Construcción del Transformer:

   • Arquitectura: Los autores demuestran constructivamente que un Transformer puede representar el mapeo del maestro.
     • La Atención como Estadísticos Suficientes: Los bloques de auto-atención agregan los tokens históricos para calcular estadísticas de respuesta empíricas y productos de vectores de covarianza, actuando efectivamente como estadísticos suficientes para el modelo de series gaussianas.
     • Las FFN como Descenso de Gradiente: Las Redes de Alimentación hacia Adelante (FFN) por token implementan pasos de descenso de gradiente proyectado para resolver el problema de regresión de Ridge con máscara de rango finito correspondiente a la actualización de la media de la posterior bayesiana.
     • Enmascaramiento Dinámico: La arquitectura utiliza una máscara de dimensión efectiva dinámica que activa un número de coordenadas de base ($J$) dependiente del tamaño de la muestra ($N$) y del prior de suavidad. Esto impone el compromiso entre sesgo y varianza: $J$ se elige de tal manera que la matriz de Gram empírica esté bien condicionada y el prior de precisión no domine.
     • Mezcla de Expertos (MoE): Una compuerta de MoE aproxima la posterior jerárquica sobre las clases de suavidad. Esta pondera diferentes "redes expertas" (cada una correspondiente a un nivel de suavidad $\beta_k$) basándose en puntuaciones de verosimilitud marginal aproximadas derivadas de las estadísticas de atención.
   • Generalización vía ERM: El artículo analiza el aprendizaje de esta política mediante la Minimización del Riesgo Empírico Supervisado (ERM). El Transformer es entrenado en trayectorias de preentrenamiento finitas para imitar el comportamiento de entrada-salida del maestro (el mapeo de historia a propensión). La teoría garantiza que el modelo entrenado puede aproximar el comportamiento del maestro hasta los errores de aproximación y estimación, sin necesidad de recuperar explícitamente el estado interno de la posterior o ejecutar los pasos de gradiente durante la inferencia.

Contribuciones Clave

1. Diseño Bayesiano y Adaptación de Suavidad: Los autores definen un maestro de Neyman con posterior bayesiana que se adapta a la suavidad de los resultados desconocida mediante un prior jerárquico. Demuestran que este diseño alcanza tasas de contracción cercanas al oráculo y una eficiencia de ATE óptima.
2. Expresividad del Transformer: Proporcionan una prueba constructiva de que los Transformers pueden representar el mapa de actualización y diseño bayesiano. Específicamente, muestran cómo los mecanismos de atención y los bloques de descenso de gradiente proyectado pueden realizar una regresión de Ridge con máscara de rango finito y cómo una compuerta MoE puede implementar la adaptación de suavidad.
3. Teoría de la Generalización: El artículo establece un límite de generalización de ERM, demostrando que el preentrenamiento supervisado sobre trayectorias finitas permite al Transformer aprender el mapa de experimentación adaptativo, imitando el comportamiento de historia-a-propensión del maestro.
4. Validación Empírica: Experimentos con datos sintéticos muestran que:
   • Un único Transformer, entrenado solo con etiquetas de la posterior de mezcla bayesiana (sin una cabeza MoE explícita), logra tasas minimax adaptativas a la suavidad en niveles de suavidad no vistos.
   • El Transformer de diseño reproduce el comportamiento de asignación de Neyman bayesiano, respondiendo direccionalmente a las varianzas residuales específicas de cada brazo, a pesar de que no aprende explícitamente los momentos de los resultados en sus estados latentes.
   • La política aprendida mejora la precisión de la estimación del ATE (reduciendo el MSE) en comparación con la aleatorización uniforme y se acerca al rendimiento de la asignación de Neyman oráculo, incluso bajo relaciones de varianza fuera de la distribución.

Resultados y Significado
El artículo sostiene que los Transformers pueden servir como "experimentadores bayesianos en contexto amortizados". Al aprender el mapeo de la historia experimental directamente hacia la propensión de tratamiento, estos modelos evitan la necesidad de resolver un problema de diseño bayesiano completo en el momento del despliegue.

La significancia radica en la demostración de que:

• La amortización es factible: La compleja inferencia bayesiana secuencial y el diseño óptimo pueden destilarse en una política de red neuronal fija.
• La adaptatividad emerge: El modelo puede adaptarse a la suavidad de las funciones y a las estructuras de varianza desconocidas sin necesidad de ingeniería explícita del parámetro de suavidad o de tuberías de estimación de momentos.
• La eficiencia se preserva: Las políticas aprendidas logran una eficiencia estadística comparable a la asignación de Neymax oráculo, validando el uso de Transformers para el diseño experimental basado en principios estadísticos.

Los autores señalan limitaciones, indicando que el estudio actual está restringido a experimentos de dos brazos y entornos sintéticos, y que extender este enfoque a espacios de acción más ricos, resultados con retraso y registros del mundo real sigue siendo una dirección futura. El impacto más amplio se enmarca en la mejora de la eficiencia de los datos en dominios como la salud digital y los servicios en línea.