TabSwift: An Efficient Tabular Foundation Model with Row-Wise Attention

Resumen Técnico: TABSWIFT

Planteamiento del Problema
Los modelos fundacionales tabulares (TFMs), como TabPFN, han demostrado un rendimiento competitivo en conjuntos de datos pequeños y medianos mediante el aprovechamiento del aprendizaje en contexto (ICL) sin actualizaciones de parámetros en tiempo de prueba. Sin embargo, los avances recientes en los TFMs suelen depender de arquitecturas cada vez más complejas (por ejemplo, atención alternante de filas/columnas) para capturar la heterogeneidad de las características. Estos diseños complejos aumentan significativamente los costos de inferencia y el uso de memoria, limitando su despliegue práctico en escenarios sensibles a la latencia. Aunque las arquitecturas ligeras de solo atención por filas (como TabPFN v1) son eficientes, corren el riesgo de alcanzar un techo de rendimiento en tareas altamente heterogéneas. La cuestión central que se aborda es si una arquitectura ligera de solo atención por filas puede lograr un rendimiento de modelo fundacional competitivo al ser mejorada con técnicas modernas de estabilización y preentrenamiento, ofreciendo así un compromiso superior entre precisión y eficiencia.

Metodología
Los autores proponen TABSWIFT, un modelo fundacional tabular ligero que conserva la estructura eficiente de solo atención por filas e introduce tres componentes clave:

1. Backbone por Filas Estabilizado:

   • Arquitectura: El modelo trata cada fila de la tabla como un token y aplica autoatención únicamente a través de las filas (instancias), evitando las costosas pasadas de atención por columnas presentes en modelos como TabPFN v2 o TabICL.
   • Atención con Compuerta (Gated Attention): Para mejorar la estabilidad de la optimización durante el preentrenamiento, los autores incorporan un mecanismo de atención con compuerta elemento a elemento. Este aplica una compuerta multiplicativa a la salida de la atención de producto escalar escalado (SDPA) antes de la proyección, introduciendo una dispersión y no linealidad dependientes de la entrada con un costo adicional insignificante.
   • Tokens de Registro Aprendibles (Learnable Register Tokens): Se antepone a la secuencia de entrada un pequeño conjunto de tokens de registro aprendibles. Estos tokens actúan como ranuras latentes para agregar el contexto global y la información a nivel de tarea, ayudando al modelo a mantener y refinar las representaciones del conjunto de datos a través de las capas.

2. Formulación de Preentrenamiento Unificada:

   • TABSWIFT se entrena con un objetivo unificado que soporta tanto clasificación como regresión dentro de un mismo modelo preentrenado.
   • El pipeline de preentrenamiento utiliza un gran grupo de tareas tabulares sintéticas (20,000 pasos, 512 tareas por paso). Para cada tarea sintética, los objetivos se procesan para proporcionar tanto un objetivo de clasificación discretizado como un objetivo de regresión normalizado, permitiendo la optimización simultánea de cabezales MLP separados para cada tipo de tarea.

3. Mecanismo de Salida Temprana Adaptativa por Capa (Adaptive Layer-Wise Early-Exit):

   • Para reducir aún más la latencia de inferencia, el modelo emplea una estrategia de salida temprana adaptativa por consulta.
   • Mecanismo: Se adjuntan cabezales de predicción y de "salida" ligeros a las últimas $E$ capas del transformador. El cabezal de salida toma la representación actual de la consulta concatenada con un resumen agrupado de los tokens de registro para predecir una puntuación de fiabilidad.
   • Política: Durante la inferencia, el modelo procesa las capas secuencialmente. Si la puntuación de fiabilidad predicha para una consulta específica supera un umbral $\tau$, el modelo realiza una salida temprana y entrega la predicción. Esto permite que las muestras "fáciles" sean procesadas con menos capas mientras que las muestras "difíciles" recorran toda la profundidad, permitiendo una inferencia de tipo "en cualquier momento" (anytime).

Contribuciones Clave

• Modelo TABSWIFT: Demuestra que un backbone de solo atención por filas, cuando se estabiliza mediante atención con compuerta y se aumenta con tokens de registro, logra un rendimiento competitivo con TFMs más pesados y complejos (por ejemplo, TabPFN v2, TabICL) manteniendo costos de inferencia significativamente menores.
• Arquitectura Unificada: Introduce un único modelo preentrenado capaz de manejar tareas de clasificación y regresión, simplificando el despliegue.
• Inferencia Adaptativa: Propone una estrategia de salida temprana basada en la predicción de fiabilidad que ajusta dinámicamente la profundidad de computación por muestra, reduciendo sustancialmente el tiempo promedio de inferencia con una pérdida mínima en el rendimiento predictivo.
• Código Abierto: El código se pone a disposición de la comunidad.

Resultados Experimentales
Los autores evaluaron TABSWIFT en el benchmark TALENT, una colección a gran escala de 300 conjuntos de datos de clasificación binaria, multiclase y regresión.

• Rendimiento: TABSWIFT logra un rango promedio comparable a los baselines más fuertes (TabPFN v2, TabICL, LimiX) tanto en tareas de clasificación (AUC) como de regresión (RMSE). Las pruebas estadísticas (Wilcoxon-Holm) indican que a menudo no es significativamente diferente de los métodos de alto rendimiento.
• Eficiencia: Las mediciones de tiempo de inferencia muestran que TABSWIFT opera con una latencia sustancialmente menor que los TFMs complejos. La Figura 1 ilustra que TABSWIFT alcanza un compromiso entre precisión y eficiencia notablemente mejor, posicionado más cerca del punto ideal de alto rendimiento y baja latencia.
• Estudios de Ablación: Los experimentos confirman que tanto la atención con compuerta como los tokens de registro contribuyen a las ganancias de rendimiento sobre un modelo base de atención por filas.
• Análisis de Salida Temprana (Early-Exit): Variar el umbral de salida $\tau$ revela una frontera de Pareto favorable. Los cabezales de salida aprendidos, condicionados por los tokens de registro, superan a los baselines de parada basados en entropía. Las visualizaciones del espacio de embeddings muestran que el modelo identifica correctamente las muestras ambiguas o superpuestas que requieren una computación más profunda, mientras que las muestras confiables salen temprano.

Significancia y Reivindicaciones
El artículo afirma que TABSWST permite un aprendizaje en contexto tabular eficiente y en cualquier momento (anytime), adecuado para despliegues prácticos donde los presupuestos de inferencia son estrictos. Al revisar el diseño minimalista por filas y aplicar técnicas modernas de estabilización, los autores argumentan que la comunidad no necesita necesariamente adoptar patrones de atención alternante computacionalmente costosos para lograr resultados de vanguardia. El trabajo sugiere que TABSWIFT sirve como un componente fuerte y eficiente para la selección de modelos o el ensamble, ofreciando una solución robusta para el servicio sensible a la latencia en aplicaciones tabulares del mundo real.