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

Technische Samenvatting: TABSWIFT

Probleemstelling
Tabulaire fundamentele modellen (TFMs), zoals TabPFN, hebben competitieve prestaties op kleine tot middelgrote datasets getoond door gebruik te maken van in-context leren (ICL) zonder parameterupdates tijdens de testfase. Echter, recente ontwikkelingen in TFMs vertrouwen vaak op steeds complexere architecturen (bijv. alternerende rij/kolom-attention) om feature-wise heterogeniteit te vangen. Deze complexe ontwerpen verhogen de inferentiekosten en het geheugengebruik aanzienlijk, wat de praktische inzetbaarheid in latency-gevoelige scenario's beperkt. Hoewel lichtgewicht, uitsluitend rij-georiënteerde attention-backbones (zoals TabPFN v1) efficiënt zijn, lopen ze het risico een prestatieplafond te bereiken bij zeer heterogene taken. De centrale vraag die wordt geadresseerd is of een lichtgewicht, uitsluitend rij-georiënteerde attention-backbone competitieve fundamentele modelprestaties kan behalen wanneer deze wordt verbeterd met moderne stabilisatie- en pretrainingtechnieken, waardoor een superieur nauwkeurigheid-efficiëntie-evenwicht wordt geboden.

Methodologie
De auteurs stellen TABSWIFT voor, een lichtgewicht tabulair fundamenteel model dat de efficiënte rij-georiënteerde attention-only structuur behoudt, terwijl het drie belangrijke componenten introduceert:

1. Gestabiliseerde Rij-georiënteerde Backbone:

   • Architectuur: Het model behandelt elke tabelrij als een token en past alleen self-attention toe over de rijen (instanties), waardoor de dure kolom-georiënteerde attention-passes worden vermeden die aanwezig zijn in modellen zoals TabPFN v2 of TabICL.
   • Gated Attention: Om de optimalisatiestabiliteit tijdens de pretraining te verbeteren, voegen de auteurs een element-gebaseerd gated attention-mechanisme toe. Dit past een multiplicatieve gate toe op de scaled dot-product attention (SDPA) output vóór de projectie, wat input-afhankelijke sparsity en non-lineariteit introduceert met verwaarloosbare overhead.
   • Leerbare Register Tokens: Een kleine set leerbare register tokens wordt aan de inputsequentie toegevoegd. Deze tokens fungeren als latente slots om globale context en taak-niveau informatie te aggregeren, wat het model helpt om dataset-niveau representaties over verschillende lagen heen te behouden en te verfijnen.

2. Verenigde Pretraining Formulering:

   • TABSWIFT wordt getraind op een verenigde doelstelling die zowel classificatie als regressie ondersteunt binnen één enkel pretrained model.
   • De pretraining-pipeline maakt gebruik van een grote pool van synthetische tabulaire taken (20.000 stappen, 512 taken per stap). Voor elke synthetische taak worden de targets verwerkt om zowel een gediscretiseerde classificatie-target als een genormaliseerde regressie-target te bieden, waardoor simultane optimalisatie van afzonderlijke lichtgewicht MLP-heads voor elk type taak mogelijk is.

3. Adaptief Layer-Wise Early-Exit Mechanisme:

   • Om de inferentie-latency verder te verlagen, gebruikt het model een per-query adaptieve early-exit strategie.
   • Mechanisme: Lichtgewicht predictie-heads en "exit heads" worden bevestigd aan de laatste $E$ transformer-lagen. De exit head neemt de huidige query-representatie geconcateneerd met een gepoolde samenvatting van de register tokens om een betrouwbaarheidsscore te voorspellen.
   • Beleid: Tijdens de inferentie verwerkt het model de lagen sequentieel. Als de voorspelde betrouwbaarheidsscore voor een specifieke query een drempelwaarde $\tau$ overschrijdt, verlaat het model vroegtijdig de berekening en geeft de voorspelling uit. Dit stelt "makkelijke" samples in staat om met minder lagen te worden verwerkt, terwijl "moeilijke" samples de volledige diepte doorlopen, wat "anytime" inferentie mogelijk maakt.

Belangrijkste Bijdragen

• TABSWIFT Model: Demonstreert dat een rij-georiënteerde attention-only backbone, wanneer deze wordt gestabiliseerd door gated attention en wordt uitgebreid met register tokens, prestaties bereikt die competitief zijn met zwaardere, complexere TFMs (bijv. TabPFN v2, TabICL), terwijl het aanzienlijk lagere inferentiekosten behoudt.
• Verenigde Architectuur: Introduceert een enkel pretrained model dat zowel classificatie- als regressietaken kan afhandelen, wat de implementatie vereenvoudigt.
• Adaptieve Inferentie: Stelt een betrouwbaarheid-voorspellende early-exit strategie voor die de computationele diepte per sample dynamisch aanpast, wat de gemiddelde inferentietijd substantieel vermindert met minimaal verlies in voorspellende prestaties.
• Open Source: De code wordt beschikbaar gesteld aan de community.

Experimentele Resultaten
De auteurs hebben TABSWIFT geëvalueerd op de TALENT benchmark, een grootschalige collectie van 300 binaire, multi-class en regressie datasets.

• Prestaties: TABSWIFT bereikt een gemiddelde rang die vergelijkbaar is met de sterkste baselines (TabPFN v2, TabICL, LimiX) op zowel classificatie (AUC) als regressie (RMSE) taken. Statistische toetsing (Wilcoxon-Holm) geeft aan dat het vaak niet significant verschilt van de best presterende methoden.
• Efficiëntie: Inferentietijd-metingen laten zien dat TABSWIFT werkt met een aanzienlijk lagere latency dan complexe TFMs. Figuur 1 illustreert dat TABSWIFT een merkbaar beter nauwkeurigheid-efficiëntie-evenwicht bereikt, gepositioneerd dichter bij het ideale punt van hoge prestaties en lage latency.
• Ablatie-studies: Experimenten bevestigen dat zowel gated attention als register tokens bijdragen aan de prestatiewinst ten opzichte van een baseline rij-georiënteerde attention model.
• Early-Exit Analyse: Het variëren van de exit-drempel $\tau$ onthult een gunstig Pareto-front. De geleerde exit heads, geconditioneerd op de register tokens, presteren beter dan entropie-gebaseerde stopping baselines. Visualisaties van de embedding space laten zien dat het model correct identificeert welke ambigue of overlappende samples meer diepe berekening vereisen, terwijl zelfverzekerde samples vroegtijdig exciteren.

Betekenis en Claims
Het artikel claimt dat TABSWIFT efficiënte en anytime tabulaire in-context learning mogelijk maakt, geschikt voor praktische implementaties waar de inferentiebudgetten strikt zijn. Door terug te keren naar het minimalistische rij-georiënteerde ontwerp en moderne stabilisatietechnieken toe te passen, beargumenteren de auteurs dat de community niet noodzakelijkerwijs de computationeel dure alternerende attention-patronen hoeft te adopteren om state-of-the-art resultaten te behalen. Het werk suggereert dat TABSWIFT dient als een sterke, efficiënte component voor modelselectie of ensembling, en een robuuste oplossing biedt voor latency-gevoelige serving in real-world tabulaire toepassingen.