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

Riepilogo Tecnico: TABSWIFT

Definizione del Problema
I modelli fondazionali tabulari (TFM), come TabPFN, hanno dimostrato prestazioni competitive su dataset di piccole e medie dimensioni sfruttando l'apprendimento in-context (ICL) senza aggiornamenti dei parametri al momento del test. Tuttavia, i recenti progressi nei TFM spesso si affidano ad architetture sempre più complesse (ad esempio, l'attenzione alternata riga/colonna) per catturare l'eterogeneità delle feature. Questi design complessi aumentano significativamente i costi di inferenza e l'uso della memoria, limitandone l'implementazione pratica in scenari sensibili alla latenza. Sebbene le backbone leggere basate solo sull'attenzione riga per riga (come TabPFN v1) siano efficienti, rischiano di incontrare un limite prestazionale su compiti altamente eterogenei. La questione centrale affrontata è se una backbone leggera, basata solo sull'attenzione riga, possa raggiungere prestazioni da modello fondazionale competitive se potenziata con moderne tecniche di stabilizzazione e pre-addestramento, offrendo così un rapporto accuratezza-efficienza superiore.

Metodologia
Gli autori propongono TABSWIFT, una backbone fondazionale tabulare leggera che mantiene la struttura efficiente basata solo sull'attenzione riga introducendo tre componenti chiave:

1. Backbone a riga stabilizzata:

   • Architettura: Il modello tratta ogni riga della tabella come un token e applica l'auto-attenzione solo tra le righe (istanze), evitando le costose passaggi di attenzione colonna-per-colonna presenti in modelli come TabPFN v2 o TabICL.
   • Gated Attention (Attenzione a porta): Per migliorare la stabilità dell'ottimizzazione durante il pre-addestramento, gli autori incorporano un meccanismo di attenzione a porta elemento per elemento. Questo applica una porta moltiplicativa all'output dell'attenzione dot-product scalata (SDPA) prima della proiezione, introducendo una sparsità e una non-linearità dipendenti dall'input con un sovraccarico trascurabile.
   • Token di Registro Apprendibili: Un piccolo set di token di registro apprendibili viene preposto alla sequenza di input. Questi token agiscono come slot latenti per aggregare il contesto globale e le informazioni a livello di task, aiutando il modello a mantenere e raffinare le rappresentazioni del dataset attraverso i vari layer.

2. Formulazione di Pre-addestramento Unificata:

   • TABSWIFT è addestrato su un obiettivo unificato che supporta sia la classificazione che la regressione all'interno di un singolo modello pre-addestrato.
   • La pipeline di pre-addestramento utilizza un vasto pool di task tabulari sintetici (20.000 step, 512 task per step). Per ogni task sintetico, i target vengono elaborati per fornire sia un target di classificazione discretizzato che un target di regressione normalizzato, consentendo l'ottimizzazione simultanea di separati head MLP leggeri per ciascun tipo di task.

3. Meccanismo di Early-Exit (Uscita Anticipata) Adattivo per Layer:

   • Per ridurre ulteriormente la latenza di inferenza, il modello impiega una strategia di early-exit adattiva per query.
   • Meccanismo: Head di predizione e "exit head" leggeri sono attaccati agli ultimi $E$ layer del transformer. L'exit head prende la rappresentazione corrente della query concatenata con un riepilogo aggregato (pooled summary) dei token di registro per predire un punteggio di affidabilità.
   • Policy: Durante l'inferenza, il modello elabora i layer sequenzialmente. Se il punteggio di affidabilità predetto per una specifica query supera una soglia $\tau$, il modello esce anticipatamente e restituisce la predizione. Ciò consente ai campioni "facili" di essere elaborati con meno layer, mentre i campioni "difficili" attraversano l'intera profondità, permettendo un'inferenza "anytime".

Contributi Chiave

• Modello TABSWIFT: Dimostra che una backbone basata solo sull'attenzione riga, quando stabilizzata da gated attention e aumentata con token di registro, raggiunge prestazioni competitive con TFM più pesanti e complessi (ad esempio, TabPFN v2, TabICL), mantenendo costi di inferenza significativamente inferiori.
• Architettura Unificata: Introduce un singolo modello pre-addestrato capace di gestire sia task di classificazione che di regressione, semplificando l'implementazione.
• Inferenza Adattiva: Propone una strategia di early-exit basata sulla predizione dell'affidabilità che regola dinamicamente la profondità di calcolo per campione, riducendo sostanzialmente il tempo medio di inferenza con una perdita minima nelle prestazioni predittive.
• Open Source: Il codice è reso disponibile alla comunità.

Risultati Sperimentali
Gli autori hanno valutato TABSWIFT sul benchmark TALENT, una vasta collezione di 300 dataset di classificazione binaria, multiclasse e di regressione.

• Prestazioni: TABSWIFT raggiunge un ranking medio comparabile ai baseline più forti (TabPFN v2, TabICL, LimiX) sia nei task di classificazione (AUC) che di regressione (RMSE). I test statistici (Wilcoxon-Holm) indicano che spesso non è significativamente differente dai metodi con le migliori prestazioni.
• Efficienza: Le misurazioni del tempo di inferenza mostrano che TABSWIFT opera a una latenza sostanzialmente inferiore rispetto ai TFM complessi. La Figura 1 illustra come TABSWIFT raggiunga un rapporto accuratezza-efficienza marcatamente migliore, posizionandosi più vicino al punto ideale di alta prestazione e bassa latenza.
• Studi di Ablazione: Gli esperimenti confermano che sia la gated attention che i token di registro contribuiscono ai guadagni di prestazione rispetto a un modello baseline a attenzione riga.
• Analisi dell'Early-Exit: La variazione della soglia di uscita $\tau$ rivela una frontiera di Pareto favorevole. Gli exit head appresi, condizionati dai token di registro, superano i baseline di arresto basati sull'entropia. Le visualizzazioni dello spazio di embedding mostrano che il modello identifica correttamente i campioni ambigui o sovrapposti come richiedenti un calcolo più profondo, mentre i campioni sicuri escono anticipatamente.

Significatività e Rivendicazioni
Il paper sostiene che TABSWIFT consenta un apprendimento in-context tabulare efficiente e "anytime", adatto a implementazioni pratiche dove i budget di inferenza sono rigorosi. Revisitando il design minimalista riga per riga e applicando moderne tecniche di stabilizzazione, gli autori argomentano che la comunità non deve necessariamente adottare pattern di attenzione alternata computazionalmente costosi per raggiungere lo stato dell'arte. Il lavoro suggerisce che TABSWIFT funga da componente efficiente e robusto per la selezione del modello o l'ensembling, offrendo una soluzione solida per il serving a bassa latenza in applicazioni tabulari del mondo reale.