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

Resumo Técnico: TABSWIFT

Definição do Problema
Modelos de fundação tabulares (TFMs), como o TabPFN, demonstraram um desempenho competitivo em conjuntos de dados de pequeno a médio porte ao alavancar o aprendizado em contexto (ICL) sem atualizações de parâmetros em tempo de teste. No entanto, avanços recentes em TFMs frequentemente dependem de arquiteturas cada vez mais complexas (por exemplo, atenção alternada entre linhas e colunas) para capturar a heterogeneidade das características. Esses designs complexos aumentam significativamente os custos de inferência e o uso de memória, limitando sua implementação prática em cenários sensíveis à latência. Embora backbones leves de apenas atenção por linha (como o TabPFN v1) sejam eficientes, eles correm o risco de atingir um teto de desempenho em tarefas altamente heterogêneas. A questão central abordada é se um backbone leve, de apenas atenção por linha, pode alcançar um desempenho de modelo de fundação competitivo quando aprimorado com técnicas modernas de estabilização e pré-treinamento, oferecendo assim uma relação superior entre precisão e eficiência.

Metodologia
Os autores propõem o TABSWIFT, um modelo de fundação tabular leve que mantém a estrutura eficiente de apenas atenção por linha, introduzindo três componentes principais:

1. Backbone por Linha Estabilizado:

   • Arquitetura: O modelo trata cada linha da tabela como um token e aplica autoatenção apenas entre as linhas (instâncias), evitando as passagens dispendiosas de atenção por coluna encontradas em modelos como o TabPFN v2 ou TabICL.
   • Atenção com Portão (Gated Attention): Para melhorar a estabilidade da otimização durante o pré-treinamento, os autores incorporam um mecanismo de atenção com portão elemento a elemento. Isso aplica um portão multiplicativo à saída da atenção de produto escalar escalado (SDPA) antes da projeção, introduzindo esparsidade e não-linearidade dependentes da entrada com um overhead negligenciável.
   • Tokens de Registro Aprendíveis: Um pequeno conjunto de tokens de registro aprendíveis é prefixado à sequência de entrada. Esses tokens atuam como slots latentes para agregar o contexto global e informações ao nível da tarefa, ajudando o modelo a manter e refinar as representações do conjunto de dados através das camadas.

2. Formulação de Pré-treinamento Unificada:

   • O TABSWIFT é treinado em um objetivo unificado que suporta tanto classificação quanto regressão dentro de um único modelo pré-treinado.
   • O pipeline de pré-treinamento utiliza um grande pool de tarefas tabulares sintéticas (20.000 etapas, 512 tarefas por etapa). Para cada tarefa sintética, os alvos são processados para fornecer tanto um alvo de classificação discretizado quanto um alvo de regressão normalizado, permitindo a otimização simultânea de cabeças MLP separadas e leves para cada tipo de tarefa.

3. Mecanismo de Saída Antecipada Adaptativa por Camada (Adaptive Layer-Wise Early-Exit):

   • Para reduzir ainda mais a latência de inferência, o modelo emprega uma estratégia de saída antecipada adaptativa por consulta.
   • Mecanismo: Cabeças de predição leves e "cabeças de saída" são anexadas às últimas $E$ camadas do transformer. A cabeça de saída recebe a representação atual da consulta concatenada com um resumo agrupado dos tokens de registro para prever uma pontuação de confiabilidade.
   • Política: Durante a inferência, o modelo processa as camadas sequencialmente. Se a pontuação de confiabilidade prevista para uma consulta específica exceder um limiar $\tau$, o modelo realiza a saída antecipada e fornece a predição. Isso permite que amostras "fáceis" sejam processadas com menos camadas, enquanto amostras "difíceis" percorrem toda a profundidade, possibilitando uma inferência "a qualquer momento" (anytime).

Principais Contribuições

• Modelo TABSWIFT: Demonstra que um backbone de apenas atenção por linha, quando estabilizado por atenção com portão e aumentado com tokens de registro, alcança um desempenho competitivo com TFMs mais pesados e complexos (ex: TabPFN v2, TabICL), mantendo custos de inferência significativamente menores.
• Arquitetura Unificada: Introduz um único modelo pré-treinado capaz de lidar com tarefas de classificação e regressão, simplificando a implementação.
• Inferência Adaptativa: Propõe uma estratégia de saída antecipada baseada em predição de confiabilidade que ajusta dinamicamente a profundidade de computação por amostra, reduzindo substancialmente o tempo médio de inferência com perda mínima de desempenho preditivo.
• Código Aberto: O código é disponibilizado para a comunidade.

Resultados Experimentais
Os autores avaliaram o TABSWIFT no benchmark TALENT, uma coleção de grande escala de 300 conjuntos de dados de classificação binária, multiclasse e regressão.

• Desempenho: O TABSWIFT atinge um ranking médio comparável aos baselines mais fortes (TabPFN v2, TabICL, LimiX) em tarefas de classificação (AUC) e regressão (RMSE). Testes estatísticos (Wilcoxon-Holm) indicam que ele frequentemente não é significativamente diferente dos métodos de alto desempenho.
• Eficiência: Medições de tempo de inferência mostram que o TABSWIFT opera com uma latência substancialmente menor do que TFMs complexos. A Figura 1 ilustra que o TABSWIFT atinge uma relação entre precisão e eficiência marcadamente melhor, posicionando-se mais próximo do ponto ideal de alto desempenho e baixa latência.
• Estudos de Ablação: Experimentos confirmam que tanto a atenção com portão quanto os tokens de registro contribuem para os ganhos de desempenho sobre um modelo de base de apenas atenção por linha.
• Análise de Saída Antecipada (Early-Exit): A variação do limiar de saída $\tau$ revela uma fronteira de Pareto favorável. As cabeças de saída aprendidas, condicionadas aos tokens de registro, superam os baselines de interrupção baseados em entropia. Visualizações do espaço de embedding mostram que o modelo identifica corretamente amostras ambíguas ou sobrepostas como exigindo computação mais profunda, enquanto amostras confiantes saem antecipadamente.

Significância e Alegações
O artigo afirma que o TABSWIFT possibilita um aprendizado em contexto tabular eficiente e "anytime", adequado para implementações práticas onde os orçamentos de inferência são rigorosos. Ao revisitar o design minimalista por linha e aplicar técnicas modernas de estabilização, os autores argumentam que a comunidade não precisa necessariamente adotar padrões de atenção alternados computacionalmente caros para alcançar resultados de estado da arte. O trabalho sugere que o TABSWIFT serve como um componente forte e eficiente para seleção de modelos ou ensembling, oferecendo uma solução robusta para o atendimento (serving) sensível à latência em aplicações tabulares do mundo real.