Speedrunning Tabular Foundation Model Pretraining

Resumen Técnico: Speedrunning del Preentrenamiento de Modelos Fundacionales Tabulares

Problema
El preentrenamiento de modelos fundacionales tabulares (TFM, por sus siglas en inglés) es un cuello de botella computacionalmente costoso que ralentiza el ciclo de iteración para nuevas arquitecturas, priors y estrategias de optimización. Mientras que se han realizado esfuerzos significativos para acelerar el preentrenamiento de los Modelos de Lenguaje Extensos (LLM), no existe un marco equivalente impulsado por la comunidad para los TFM. El alto costo del preentrenamiento limita el prototipado rápido y la acumulación de técnicas de aceleración dentro de la comunidad de aprendizaje tabular.

Metodología
Los autores presentan una competencia pública de "speedrun" centrada en nanoTabPFN, una reimplementación educativa y mínima de TabPFN v2. La competencia establece un protocolo fijo donde los contribuyentes modifican un script de entrenamiento de un solo archivo para alcanzar un objetivo específico de rendimiento downstream en el menor tiempo de reloj (wallclock time) posible.

• Objetivo: El objetivo es preentrenar una red neuronal que logre un ROC AUC de validación promedio superior al de un baseline de Random Forest en conjuntos de datos TabArena submuestreados (38 tareas de clasificación).
• Restricciones: Todas las ejecuciones deben utilizar una única GPU NVIDIA L40S. El tiempo reportado es la mediana de múltiples ejecuciones de verificación para suprimir el ruido del clúster. Los récords no deben cargar pesos preentrenados y deben utilizar un pipeline de evaluación fijo.
• Progresión: El speedrun es acumulativo; cada nuevo récord se construye sobre las técnicas del anterior, aislando el impacto de optimizaciones específicas.

Contribuciones Clave y Técnicas
El artículo no propone una nueva familia de modelos, sino un protocolo reproducible y una secuencia de récords que demuestran cómo diversas técnicas se transfieren del modelado de lenguaje a los datos tabulares. La progresión desde el baseline (74.32 minutos) hasta el mejor actual (0.92 minutos) involucró las siguientes intervenciones técnicas clave:

1. Cambio de Optimizador: Reemplazar el optimizador AdamW sin programación de decaimiento (schedule-free) por el optimizador Muon para matrices de pesos 2D, lo que mejoró significativamente la eficiencia de muestreo.
2. Arquitectura y Precisión:
   • Implementación de Atención de Producto Escalar Escalada (SDPA).
   • Cambio de una arquitectura post-norm a pre-norm.
   • Adopción de entrenamiento de precisión mixta (bf16) y multiplicación de matrices TF32.
   • Ajuste de hiperparámetros (aumento de la tasa de aprendizaje a $10^{-3}$) y ampliación de la dimensión de embedding (192 $\to$ 256) reduciendo simultáneamente los cabezales de atención.
3. Optimizaciones de Rendimiento (Throughput): Procesamiento por lotes (batching) de la iteración de Newton–Schulz para Muon y compilación del paso forward del codificador transformer.
4. Modificaciones Arquitectónicas:
   • Decaimiento Residual (Residual Decay): Pesado exponencial descendente de las entradas a los bloques iniciales del transformer ($0.95^i$).
   • RMSNorm: Reemplazo de LayerNorm por RMSNorm de menor precisión.
   • Filas de Pensamiento (Thinking Rows): Preceder 16 (luego 24) filas aprendibles de "pensamiento" a la entrada, a las cuales las filas de datos atienden, lo que aumenta significamente el rendimiento.
   • Agrupación de Características Repetida: Embedding de columnas conjuntamente con offsets específicos (por ejemplo, $j, j+1, j+3$) para prevenir el colapso de representación en características con distribuciones similares.
5. Estrategias de Entrenamiento Avanzadas:
   • Promedio de Pesos de Última Instancia (LAWA): Promediar los últimos 10 checkpoints únicamente para la evaluación.
   • Decaimiento de Pesos (Weight Decay): Aplicar decaimiento de pesos tanto a los parámetros de AdamW como a los de Muon.
   • Modificación del Decoder: Alimentar al decoder con la media de los embeddings de las características (en lugar de los embeddings de los objetivos).
   • Autobúsqueda (Autoresearch): El récord final utilizó un bucle de búsqueda de hiperparámetros y arquitectura impulsado por un LLM (autoresearch) para descubrir la configuración óptima.

Resultos
El mejor récord actual logra el objetivo de rendimiento downstream en 0.92 minutos en una sola GPU NVIDIA L40S. Esto representa:

• Una aceleración de 81× en el tiempo de reloj comparado con el baseline de 74.32 minutos.
• Una reducción de 22× en el número de conjuntos de datos sintéticos requeridos (de 80,576 a 3,648).
• Cuando se entrena para el presupuesto completo del baseline (256,000 datasets), el modelo optimizado alcanza un ROC AUC promedio de 0.82, acercándose al rendimiento de TabPFN v2.5 (0.8286) sin preprocesamiento ni ensambles.

Significancia
El artículo sostiene que la principal contribución es el formato de speedrun en sí mismo, que proporciona un protocolo abierto y sencillo para que la comunidad añada, verifique y acumule mejoras de preentrenamiento. Al fijar el prior sintético y el pipeline de evaluación, la tabla de clasificación aísla qué optimizaciones prácticas son efectivas para los TFM. Los autores señalan que el mejor récord actual mantiene el prior sintético fijo, lo que sugiere que el propio prior sigue siendo una "dirección no tocada" para futuras mejoras. El trabajo demuestra que las técnicas exitosas en el modelado de lenguaje (como Muon, SDPA y ajustes arquitectónicos específicos) pueden transferirse eficazmente a los modelos fundacionales tabulares para reducir drásticamente los costos de entrenamiento.