Technische Samenvatting: End-to-End Compressie voor Tabulaire Foundation Modellen

Probleemstelling

Tabulaire foundation modellen (TFMs), met name die gebruikmaken van in-context learning (ICL) met decoder-only transformer-architecturen, hebben recentelijk de dominantie van gradient-boosted decision trees uitgedaagd. Deze modellen kampen echter met een fundamentele schaalbaarheidsproblematiek: hun inferentiecomplexiteit schaalt kwadratisch ($O(N^2)$) of lineair ($O(N)$ met KV-caching) met het aantal trainingsinstanties $N$. Deze kwadratische afhankelijkheid van de conditionele contextgrootte creëert prohibitieve geheugen- en rekenoverhead, waardoor TFMs beperkt blijven tot kleine of middelgrote datasets. Bijgevolg verslechtert de prestatie vaak op grote, informatie-rijke datasets, en hebben de modellen moeite met real-world scenario's die massale contextvensters vereisen.

Methodologie: TACO

Om deze bottleneck aan te pakken, stellen de auteurs TACO voor (Tabular Compression), een end-to-end modulaire architectuur die ontworpen is om de trainingsdataset te comprimeren naar een latente ruimte voordat de voorspelling plaatsvindt. Het framework bestaat uit twee gezamenlijk getrainde componenten:

1. Compressor ($g$): Een transformer-gebaseerde module die de originele trainingsdataset $D_{train}$ (met $N$ rijen) mapt naar een compacte, lager-rangige representatie $D_{mini-train}$ bestaande uit $K$ rijen, waarbij $K \ll N$. De compressor is getraind om een gecondenseerde embedding te leren die prototypische patronen van de tabulaire data vastlegt. Het maakt gebruik van een "dummy" tabel die wordt geïnitialiseerd vanuit een subset van de data en die iteratief wordt verfijnd om de volledige dataset te representeren.
2. Predictor ($f$): Een standaard tabulaire transformer (gebaseerd op de TabPFNv2-architectuur) die in-context voorspellingen uitvoert. In plaats van te attentionen op de volledige $N$-rijen dataset, conditioneert de predictor op de gecomprimeerde context $D_{mini-train}$ en de testbatch-features.

Trainingsstrategie:
De compressor en de predictor worden gezamenlijk en end-to-end meta-geleerd. Het doel is om de voorspellingsverlies te minimaliseren over een prior van datasets, wat de compressor stimuleert om een representatie te produceren die de prestaties van de predictor maximaliseert. De trainingspipeline omvat:

• Synthetische Pretraining: Het gebruik van structurele causale modellen (SCMs) om diverse synthetische datasets te genereren.
• Curriculum Learning: Het progressief verhogen van sequentielengtes tijdens de training (van 1k naar 60k rijen).
• Aanpassing aan Real-Data: Fine-tuning op echte tabulaire datasets met behulp van self-supervised taken (bootstrapping van rijen/features).
• Multi-Rate Training: Het model wordt getraind om gelijktijdig verschillende compressieratio's ($r \in \{1, 2, 4, 8, 16\}\%$) te hanteren, waardoor één enkel model onder verschillende resource-beperkingen kan opereren.

Schaalbaarheidsmechanisme:
Voor datasets die de geheugencapaciteit van een enkele forward pass overschrijden (bijv. >1 miljoen rijen), gebruikt TACO een chunking-and-stitching strategie. De dataset wordt opgedeeld in disjuncte chunks, die elk onafhankelijk worden gecomprimeerd, waarna de resulterende samenvattingen worden samengevoegd om een globale context voor de predictor te vormen.

Belangrijkste Bijdragen

1. Nieuwe Architectuur: Introductie van TACO, de eerste end-to-end methode voor het comprimeren van de tabulaire context van foundation modellen, die de voorspellende kwaliteit ontkoppelt van de ruwe grootte van de conditionele context.
2. Efficiëntiewinst: Empirische demonstratie dat TACO een tot wel 94× versnelling bereikt in inferentietijd en 98% geheugenbesparing biedt vergeleken met state-of-the-art predictor-only transformers (POT), terwijl de prestaties behouden blijven bij een compressieratio van 1%.
3. Schaalbaarheid: Een gevalideerde strategie voor het verwerken van datasets met meer dan één miljoen rijen (bijv. de MetroPT-3 dataset met ~1,5M rijen) door ze te comprimeren tot enkele duizenden latente rijen, een taak die onhaalbaar is voor standaard TFMs vanwege geheugenbeperkingen.
4. Gezamenlijke Optimalisatie: Bewijs dat gezamenlijke training van de compressor en de predictor superieur is aan het bevriezen van de predictor, omdat dit de predictor in staat stelt zijn representatieruimte aan te passen aan de gecomprimeerde inputs.

Experimentele Resultaten

De auteurs evalueerden TACO op de TabArena benchmark en aanvullende grootschalige datasets (TabFSBench, TableShift, MetroPT-3):

• Prestatie vs. Snelheid: Bij een compressieratio van 1% behaalt TACO een 94× versnelling in daaropvolgende voorspellingen en vermindert het het geheugengebruik met 97,6% vergeleken met de ongecomprimeerde POT-baseline. De prestatieverslechtering is statistisch insignificant (verschil in ROC-AUC nabij nul) vergeleken met de full-context baseline.
• Geheugenefficiëntie: In scenario's met KV-caching krijgt POT vaak te maken met out-of-memory (OOM) fouten op datasets met meer dan 32.000+ rijen, terwijl TACO wel haalbaar blijft. TACO vermindert het VRAM-gebruik met tot wel 92% in KV-caching settings.
• Baseline Vergelijkingen: TACO presteert significant beter dan eenvoudige compressie-baselines zoals random sampling en k-Nearest Neighbors (kNN) subsampling over alle compressieratio's heen.
• Grootschalige Afhandeling: Op de MetroPT-3 dataset (~1,5M rijen) comprimeerde TACO de trainingsset tot ~1.200 rijen en behaalde een superieure precisie en AUPRC vergeleken met baselines en XGBoost, terwijl POT en TabPFNv2 niet konden draaien.
• Robuustheid: Op de TableShift benchmark opereerde TACO succesvol op 12 van de 15 datasets waar standaard TFMs faalden door geheugenlimieten, waarbij competitieve nauwkeurigheid werd behouden onder distributieverschuivingen.

Betekenis en Claims

Het artikel claimt dat TACO de tabulaire in-context inferentie fundamenteel herformuleert door datarepresentaties te leren die voorspellend blijven onder strikte inferentie-budgetten. De auteurs stellen dat hun methode een principes-gebaseerde nauwkeurigheid-latentie afweging mogelijk maakt, waardoor in-context tabulaire predictors kunnen schalen naar real-world, grootschalige datasets die voorheen ontoegankelijk waren voor transformer-gebaseerde benaderingen.

Het werk benadrukt dat de kwadratische complexiteit van attention-mechanismen geen onoverkomelijke barrière is als de context intelligent wordt gecomprimeerd. Door de effectieve contextgrootte te reduceren met een factor $N/K$, lineariseert TACO de schaling van geheugen- en rekenkosten, waardoor high-performance foundation modellen levensvatbaar worden voor industriële toepassingen met massieve tabulaire data. De auteurs concluderen dat hoewel TACO de schaalbaarheid vooruit helpt, toekomstig werk zich moet richten op diverse priors voor ontbrekende waarden, distributieverschuivingen en uitbreidingen naar regressie en tijdreeks-taken.