Pre-trained Tabular Foundation Models as Versatile Summary Networks for Neural Posterior Estimation

Resumen Técnico: Modelos Fundacionales Tabulares Preentrenados como Redes de Resumen Versátiles para la Estimación del Posterior Neural

Enunciado del Problema

La inferencia basada en simulación (SBI) aborda problemas de inferencia bayesiana donde la función de verosimilitud $p(x|\theta)$ es intratable o costosa de evaluar, pero se pueden generar muestras mediante un simulador $x \sim p(x|\theta)$. Los métodos estándar de SBI neuronal, como la Estimación del Posterior Neural (NPE), suelen entrenar estimadores de densidad condicional directamente sobre pares simulados de parámetros y observaciones. Sin embargo, la calidad de estos estimadores depende en gran medida del presupuesto de simulación y de la calidad de las estadísticas de resumen utilizadas para representar las observaciones. Mientras que las redes preajustadas (PFNs) como TabPFN ofrecen una ruta libre de entrenamiento para la inferencia amortizada adaptándose a nuevas tareas mediante aprendizaje en contexto, su aplicación a la SBI ha sido limitada. Trabajos anteriores (NPE-PFN) utilizaron TabPFN como un muestreador autoregresivo directo del posterior, limitado por su tamaño de contexto y su interfaz de objetivo escalar. Este artículo investiga un enfoque alternativo: aprovechar TabPFN no como un muestreador directo, sino como una red de resumen modular congelada para extraer representaciones informativas de las salidas del simulador para la inferencia aguas abajo.

Metodología: PFN-NPE

Los autores proponen PFN-NPE, una receta general que desacopla el aprendizaje de representaciones del muestreo del posterior. El método consta de tres componentes configurables:

1. Construcción en Contexto: Para cada coordenada de parámetro $d$ del vector de parámetros $\theta$, se construye un contexto de regresión escalar $C_d$ utilizando un subconjunto de $m$ simulaciones de entrenamiento. En este contexto, las salidas del simulador $x_i$ sirven como características tabulares, y la coordenada de parámetro escalar $\theta_{i,d}$ sirve como objetivo de regresión.
2. Extracción de Resumen Congelada: Un codificador TabPFN preentrenado (con pesos fijos) se consulta para cada simulación de entrenamiento, simulación de validación y observación de referencia $x_o$ utilizando el contexto $C_d$. Esto produce una incrustación de la última capa $h_d(x)$ para cada coordenada de parámetro. Estas incrustaciones se concatenan para formar un resumen crudo $h_{raw}(x)$.
3. Proyección e Inferencia: El resumen crudo se proyecta a una dimensión fija (predeterminada 64) utilizando una transformación lineal (por ejemplo, PCA) ajustada en el conjunto de entrenamiento. Luego, un estimador de densidad condicional (predeterminado: Flujo de Splines Neuronales) se entrena sobre pares de estos resúmenes fijos y los parámetros correspondientes $\{(s(x_i), \theta_i)\}$.

Crucialmente, TabPFN se consulta solo una vez durante la fase de precomputación del resumen. La cabeza de inferencia aguas abajo es específica de la tarea y se entrena sobre los resúmenes fijos, permitiendo que el modelo fundacional sirva como un componente reutilizable en diferentes tareas de SBI.

Contribuciones Clave

1. Formulación de PFN-NPE: El artículo formula un método que utiliza un TabPFN congelado como red de resumen para la Estimación del Posterior Neural, separando las capacidades de representación del modelo fundacional del mecanismo de muestreo del posterior.
2. Evaluación Empírica Exhaustiva: Los autores evalúan PFN-NPE en toda la suite de benchmarks SBIBM, incluyendo tareas estándar, variantes de coordenadas de estorbo (distractores) y simuladores de series temporales estructuradas, bajo presupuestos de simulación coincidentes (típicamente $n=10,000$).
3. Marco Diagnóstico: Más allá del rendimiento conjunto agregado, el estudio emplea C2ST marginal, C2ST en espacio de rangos y sondas de cuantiles para diseccionar dónde surgen las discrepancias del posterior. Esto incluye analizar si los resúmenes congelados preservan la ubicación marginal, la escala y las estructuras de dependencia conjunta.

Resultados

La comparación empírica revela un panorama de rendimiento dependiente de la tarea:

• Rendimiento: Con un presupuesto de 10,000 simulaciones, el método NPE-PFN autoregresivo generalmente logra el C2ST conjunto medio más bajo en tareas de referencia estándar (por ejemplo, GLM Bernoulli, SIR, Lotka-Volterra). Sin embargo, PFN-NPE supera tanto a NPE-PFN como a las líneas base de NPE con resumen aprendido en modelos de mezcla (por ejemplo, Mezcla Gaussiana) y tareas con coordenadas de estorbo (distractores).
• Información Marginal vs. Conjunta: Las sondas diagnósticas indican que los resúmenes derivados de TabPFN a menudo preservan información útil del posterior unidimensional. Las sondas lineales decodifican con éxito las medias del posterior y los cuantiles marginales de las incrustaciones congeladas en diversas tareas.
• Limitaciones en la Estructura Conjunta: A pesar de una buena recuperación marginal, PFN-NPE exhibe un C2ST conjunto y un C2ST en espacio de rangos elevados en tareas con estructuras de dependencia complejas (por ejemplo, SLCP, Lotka-Volterra, series temporales). La brecha "conjunto menos marginal" sugiere que, aunque los resúmenes capturan la ubicación y la escala marginales, luchan por representar la geometría completa del posterior conjunto necesaria para un muestreo multivariado preciso.
• Eficiencia: PFN-NPE incurre en un costo único para el entrenamiento de incrustaciones y flujo, pero ofrece un costo de muestreo insignificante para consultas repetidas, en contraste con NPE-PFN, que realiza inferencia en contexto para cada consulta.

Significado y Afirmaciones

El artículo afirma que TabPFN puede servir como una red de resumen efectiva y libre de entrenamiento para la SBI, particularmente en configuraciones que involucran distribuciones de mezcla o coordenadas de estorbo donde los resúmenes aprendidos estándar pueden tener dificultades. El significado radica en demostrar que los modelos fundacionales congelados pueden codificar información suficiente del posterior marginal para apoyar la estimación de densidad aguas abajo, incluso sin ajuste fino específico de la tarea.

Sin embargo, los autores son modestos respecto a la universalidad del método. Declaran explícitamente que el modo de fallo primario para este enfoque de resumen congelado es la representación de estructuras de dependencia multivariada. Los diagnósticos sugieren que, aunque el objetivo de preentrenamiento (predicción supervisada escalar en modelos causales estructurales) se alinea bien con tareas que requieren recuperación marginal, puede no capturar inherentemente las dependencias conjuntas complejas necesarias para un muestreo preciso del posterior en todos los regímenes de SBI. El artículo concluye que PFN-NPE es una alternativa versátil y modular que traslada la carga de la aproximación del posterior a la cabeza de inferencia, pero su eficacia depende de la alineación entre el prior de preentrenamiento del modelo fundacional y los requisitos estructurales específicos de la tarea de SBI.