CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

Resumen Técnico: CRUMB – Inferencia de Redes Ajustadas a Priores mediante el Agrupamiento de Contexto por Coincidencia Distribucional

Planteamiento del Problema
Las redes ajustadas a priores (PFN, por sus siglas en inglés), como TabPFN y TabICL, representan una clase de modelos fundacionales tabulares capaces de aprendizaje en contexto. Estas aceptan un conjunto de entrenamiento completamente etiquetado como contexto y generan predicciones para consultas de prueba en una única pasada hacia adelante (forward pass), superando a menudo a los métodos tradicionales de potenciación de gradiente (gradient-boosting) en conjuntos de datos pequeños a medianos. Sin embargo, un cuello de botella crítico limita su escalabilidad: el mecanismo de autoatención en las redes de transformadores basadas en PFN escala cuadráticamente ($O(N^2)$) con el número de muestras de entrenamiento ($N$). Esto hace que la inferencia sea computacionalmente prohibitiva para conjuntos de datos grandes ($N > 50,000$).

Las soluciones existentes enfrentan un compromiso entre precisión y eficiencia. El submuestreo uniforme mejora la velocidad al descartar aleatoriamente datos de entrenamiento, pero sacrifica información. Los métodos dependientes de la consulta, como los $k$-vecinos más cercanos ($k$NN), mejoran la precisión al adaptar el contexto a cada punto de prueba, pero impiden el procesamiento por lotes (batching), requiriendo $T$ pasadas hacia adelante separadas para $T$ puntos de prueba. Otros enfoques, como la Mezcla de Prompts en Contexto (MICP), intentan procesar consultas por lotes agrupando los datos de entrenamiento, pero pueden fallar al no alinear el contexto de entrenamiento con la distribución específica de los datos de prueba, particularmente bajo deriva de covariables (covariate drift).

Metodología
Los autores proponen CRUMB (Agrupamiento de Recuperación mediante el Minimizador de MMD para Procesamiento por Lotes), un envoltorio (wrapper) de inferencia agnóstico a la arquitectura que requiere tres etapas y no necesita el reentrenamiento del modelo. El método resuelve la tensión entre la calidad del contexto y la eficiencia del procesamiento por lotes mediante las siguientes etapas:

1. Agrupamiento del Lado de la Prueba: Las consultas de prueba ($D_{test}$) se particionan en $K$ grupos ($C_1, \dots, C_K$) utilizando $k$-means sobre las características de entrada estandarizadas. Esto permite al sistema agrupar puntos de prueba similares que pueden compartir un contexto de entrenamiento común.
2. Selección de Contexto con Coincidencia Distribucional: Para cada grupo de prueba $C_k$, se selecciona un pequeño subconjunto de entrenamiento $S_k$ (tamaño $n \ll N$) del conjunto completo de entrenamiento ($D_{train}$). A diferencia de los métodos de proximidad geométrica (por ejemplo, el vecino más cercano al centroide), CRUMB selecciona $S_k$ minimizando de forma codiciosa la Discrepancia de Media Máxima (MMD) entre la distribución empírica del grupo de prueba y el subconjunto de entrenamiento candidato.
   • Esto se logra mediante kernel herding, un algoritmo codicioso que selecciona iterativamente puntos de entrenamiento para minimizar el objetivo MMD. La función objetivo equilibra dos términos: recompensar la proximidad al grupo de prueba (relevancia) y penalizar la redundancia entre los puntos de entrenamiento seleccionados (diversidad).
   • Para acelerar este proceso, los autores emplean Características de Fourier Aleatorias (RFF) para aproximar el kernel RBF Gaussiano y utilizan una selección codiciosa por lotes.
3. Inferencia por Lotes: El modelo PFN ejecuta $K$ pasadas hacia adelante independientes. En cada pasada, el modelo recibe el subconjunto de entrenamiento seleccionado $S_k$ como contexto y el grupo completo $C_k$ como consultas. Esto reduce el número total de pasadas hacia adelante de $T$ (en el caso de $k$NN por consulta) a $K$, permitiendo un uso eficiente de la GPU mientras se mantiene la alineación distribucional.

Contribuciones Clave

• Recuperación de Contexto Basada en MMD: El artículo introduce la minimización codiciosa de MMD (kernel herding) como un mecanismo para seleccionar subconjuntos de entrenamiento. Los autores demuestran que este enfoque de coincidencia distribucional supera a las estrategias de recuperación de vecino más cercano al centroide y de Voronoi-uniforme dentro del mismo marco de agrupamiento.
• Selección de Contexto por Lotes mediante Agrupamiento del Lado de la Prueba: Al agrupar las consultas de prueba y asignar un único contexto de entrenamiento, con coincidencia distribucional, a cada grupo, CRUMB permite un procesamiento por lotes eficiente sin sacrificar la relevancia de la consulta.
• Robustez ante la Deriva de Covariables: El método maneja naturalmente la deriva de covariables. Debido a que el contexto de entrenamiento se selecciona para coincidir con la distribución del lote de prueba actual (en lugar de una distribución de entrenamiento fija), CRUMB se adapta a los cambios en la distribución de los datos de prueba. En contraste, los métodos como MICP, que dependen de grupos de entrenamiento fijos, se degradan significamente cuando la distribución de prueba se desplaza de la distribución de entrenamiento.

Resultados Experimentales
Los autores evaluaron CRUMB en el benchmark TabArena (51 conjuntos de datos tabulares diversos) a través de tres arquitecturas PFN: TabPFNv2, TabICLv1 y TabICLv2.

• Rendimiento frente a Baselines: En el benchmark TabArena, CRUMB supera significativamente al submuestreo uniforme y a la técnica MICP con presupuestos de contexto idénticos. Aunque el rendimiento de CRUMB es estadísticamente cercano al de $k$NN por consulta, lo logra con un número fijo de $K$ pasadas hacia adelante (por ejemplo, $K=20$) en lugar de $T$ pasadas, lo que resulta en una reducción sustancial del tiempo de inferencia (por ejemplo, 35s frente a 361s en los experimentos reportados).
• Conjuntos de Datos Grandes: Los experimentos en conjuntos de datos con hasta 120,000 muestras (ej. GiveMeSomeCredit, Higgs) muestran que CRUMB con parada temprana (terminando la selección de contexto cuando la mejora de MMD se estanca) gana con mayor frecuencia contra MICP, particularmente para TabPFNv2 y TabICLv1.
• Deriva de Covariables: En experimentos de deriva controlada donde la distribución de prueba se desplaza a lo largo del primer componente principal, la ventaja de CRUMB sobre MICP crece monótonamente con la intensidad de la deriva. Con cero deriva, CRUMB lidera por +4.9%; con la deriva máxima (completamente fuera de la muestra), la brecha se amplía a +17.1%. Esto se atribuye a la capacidad de CRUMB para realinear el contexto de entrenamiento con la distribución desplazada de la prueba, mientras que los grupos de entrenamiento fijos de MICP quedan desalineados.

Significancia y Reivindicaciones
El artículo afirma que CRUMB hace que la inferencia de redes ajustadas a priores sea práctica para conjuntos de datos tabulares a gran escala sin requerir cambios arquitectónicos o el ajuste fino (fine-tuning) del modelo. Al combinar el agrupamiento del lado de la prueba con la recuperación de contexto por coincidencia distribucional (MMD), CRUMB alcanza un "punto ideal" entre la precisión de la recuperación por consulta y la eficiencia del procesamiento por lotes uniforme.

Los autores enfatizan que la innovación central reside en el diseño adaptativo a la prueba: agrupar puntos de prueba y seleccionar contextos basados en la alineación distribucional (MMD) en lugar de solo la proximidad geométrica. Este diseño proporciona un mecanismo de robustez natural contra la deriva de covariables, un régimen donde los métodos de selección de contexto existentes suelen fallar. El artículo concluye que, si bien CRUMB hereda el techo de rendimiento de los priores de preentrenamiento de la PFN subyacente, elimina efectivamente la barrera computacional que impide la aplicación de las PFN a grandes conjuntos de datos.

Limitaciones
Los autores reconocen que CRUMB requiere el conjunto completo de consultas de prueba de antemano para realizar el agrupamiento, lo que lo hace inadecuado para entornos de transmisión (streaming) en tiempo real sin modificaciones (por ejemplo, cachear centroides, lo que reduce efectivamente a un enrutamiento tipo MICP). Además, la dependencia de $k$-means euclidiano y kernels gaussianos isotrópicos puede limitar el rendimiento en espacios de características de alta dimensión o de tipos mixtos. Finalmente, como un envoltorio de inferencia puro, CRUMB no puede compensar un mal emparejamiento entre los priores de preentrenamiento de la PFN y el conjunto de datos específico.