Speedrunning Tabular Foundation Model Pretraining

Resumo Técnico: Speedrunning de Pré-treinamento de Modelos de Fundação Tabulares

Problema
O pré-treinamento de modelos de fundação tabulares (TFMs) é um gargalo computacionalmente caro que retarda o ciclo de iteração para novas arquiteturas, priors e estratégias de otimização. Embora esforços significativos tenham sido feitos para acelerar o pré-treinamento de Grandes Modelos de Linguagem (LLMs), não existe um framework equivalente e orientado pela comunidade para TFMs. O alto custo do pré-treinamento limita a prototipagem rápida e o acúmulo de técnicas de aceleração na comunidade de aprendizado tabular.

Metodologia
Os autores introduzem uma competição pública de "speedrun" centrada no nanoTabPFN, uma reimplementação mínima e educacional do TabPFN v2. A competição estabelece um protocolo fixo onde os contribuidores modificam um script de treinamento de arquivo único para atingir um alvo específico de desempenho em downstream no menor tempo de relógio (wallclock time) possível.

• Alvo: O objetivo é pré-treinar uma rede neural que alcance um ROC AUC de validação médio superior ao de um baseline de Random Forest em conjuntos de dados TabArena subamostrados (38 tarefas de classificação).
• Restrições: Todas as execuções devem usar uma única GPU NVIDIA L40S. O tempo relatado é a mediana de múltiplas execuções de verificação para suprimir o ruído do cluster. Os registros não devem carregar pesos pré-treinados e devem usar um pipeline de avaliação fixo.
• Progressão: O speedrun é cumulativo; cada novo recorde constrói sobre as técnicas do anterior, isolando o impacto de otimizações específicas.

Contribuições Principais e Técnicas
O artigo não propõe uma nova família de modelos, mas sim um protocolo reproduzível e uma sequência de recordes demonstrando como várias técnicas são transferidas da modelagem de linguagem para dados tabulares. A progressão do baseline (74,32 minutos) para o melhor registro atual (0,92 minutos) envolveu as seguintes intervenções técnicas fundamentais:

1. Troca de Otimizador: Substituir o otimizador AdamW schedule-free pelo otimizador Muon para matrizes de peso 2D melhorou significativamente a eficiência de amostragem.
2. Arquitetura e Precisão:
   • Implementação de Scaled Dot-Product Attention (SDPA).
   • Mudança de post-norm para arquiteturas pre-norm.
   • Adoção de treinamento de precisão mista (bf16) e multiplicação de matrizes TF32.
   • Ajuste de hiperparâmetros (aumento da taxa de aprendizado para $10^{-3}$) e alargamento da dimensão de embedding (192 $\to$ 256) enquanto se reduz o número de cabeças de atenção.
3. Otimizações de Throughput: Agrupamento (batching) da iteração de Newton–Schulz para o Muon e compilação do forward pass do codificador transformer.
4. Modificações Arquiteturais:
   • Decaimento Residual (Residual Decay): Redução exponencial do peso das entradas para os blocos iniciais do transformer ($0.95^i$).
   • RMSNorm: Substituição do LayerNorm por RMSNorm de menor precisão.
   • Linhas de Pensamento (Thinking Rows): Prependagem de 16 (posteriormente 24) linhas aprendíveis de "pensamento" à entrada, para as quais as linhas de dados realizam atenção, o que impulsiona significativamente o desempenho.
   • Agrupamento Repetido de Características (Repeated Feature Grouping): Embedding de colunas conjuntamente com offsets específicos (ex: $j, j+1, j+3$) para evitar o colapso de representação em características com distribuições semelhantes.
5. Estratégias de Treinamento Avançadas:
   • Latest Weight Averaging (LAWA): Média dos últimos 10 checkpoints apenas para avaliação.
   • Decaimento de Peso (Weight Decay): Aplicação de decaimento de peso tanto para parâmetros AdamW quanto para Muon.
   • Modificação do Decoder: Alimentar o decoder com a média dos embeddings de características (em vez de embeddings de alvo).
   • Autoresearch: O registro final utilizou um loop de busca de hiperparâmetros e arquitetura dirigido por LLM (autoresearch) para descobrir a configuração ideal.

Resultos
O melhor registro atual atinge o desempenho alvo em 0,92 minutos em uma única GPU NVIDIA L40S. Isso representa:

• Um aceleração de 81× no tempo de relógio em comparação com o baseline de 74,32 minutos.
• Uma redução de 22× no número de conjuntos de dados sintéticos necessários (de 80.576 para 3.648).
• Quando treinado para o orçamento total do baseline (256.000 conjuntos de dados), o modelo otimizado atinge um ROC AUC médio de 0,82, aproximando-se do desempenho do TabPFN v2.5 (0,8286) sem pré-processamento ou ensembles.

Significância
O artigo afirma que a principal contribuição é o formato de speedrun em si, que fornece um protocolo simples e aberto para a comunidade adicionar, verificar e empilhar melhorias de pré-treinamento. Ao fixar o prior sintético e o pipeline de avaliação, o leaderboard isola quais otimizações práticas são eficazes para TFMs. Os autores observam que o melhor registro atual mantém o prior sintético fixo, sugerindo que o próprio prior permanece uma "direção intocada" para melhorias futuras. O trabalho demonstra que técnicas bem-sucedidas em modelagem de linguagem (como Muon, SDPA e ajustes arquiteturais específicos) podem ser efetivamente transferidas para modelos de fundação tabulares para reduzir drasticamente os custos de treinamento.