End-to-End Compression for Tabular Foundation Models

AI summary

Overview

Research area: Tabular machine learning, specifically tabular foundation models (TFMs) that perform in-context learning, and model compression for transformer inference.

Technical level: Intermediate. The core ideas are conceptually simple (shrink the context, keep the accuracy), but the paper assumes familiarity with transformers, attention complexity, and tabular benchmarks.

Scope: The paper introduces TACO, a jointly trained compressor-plus-predictor architecture that reduces the training context of tabular foundation models an order of magnitude or more, and measures the resulting speed, memory, and accuracy trade-offs on TabArena and other tabular benchmarks.

What This Paper Is About

Tabular foundation models such as TabPFN, TabICL, and TabDPT predict on a dataset by conditioning on the entire training split in a single forward pass, but their attention mechanism scales quadratically with the number of training rows, making large tables slow and memory-hungry. The paper asks whether the training context can be compressed into a much smaller learned latent representation without losing predictive quality, and proposes an end-to-end method that learns this compression jointly with the predictor.

Key Contributions

1. TACO, an end-to-end context compression method for tabular in-context learning. A compressor transformer maps the training set into a smaller set of "mini-train" rows in latent space, and a predictor transformer then performs in-context prediction using only the compressed context.
2. Empirical speed and memory results. The authors report up to a 94x inference speedup and up to 98% memory savings (97.6% for the 1% compression setting in their memory table) with no statistically significant performance degradation at 1% of the original context size, evaluated on the TabArena benchmark.
3. A chunking-and-stitching strategy. Datasets are split into disjoint chunks, each compressed separately, and the per-chunk summaries concatenated into one global context, enabling ingestion of datasets with over one million rows (and, per the appendix, up to roughly 6M rows).
4. Ablations and baselines. Joint training of compressor and predictor is shown to be necessary relative to a frozen predictor, multi-rate training is shown not to hurt relative to rate-specific training, and TACO is compared against kNN-based and random subsampling compression.

Main Findings

• Large inference speedups at 1% compression. In the streaming-batch experiment (15,000 rows, 500 features, 100 sequential test batches of 50 rows each), TACO reaches up to about 94x speedup at a 1% compression rate without KV caching, and up to about 37x with KV caching. The paper reports that the gap in time between TACO and the predictor-only transformer (POT) increases with dataset size.
• Subsequent predictions become very cheap. In the inference time and memory table (no KV caching), POT takes 28.67s ± 0.05s for subsequent predictions, while TACO at 1% takes 306ms ± 18ms (93.6x), at 2% takes 382ms ± 18ms (75.2x), at 4% takes 544ms ± 17ms (52.7x), at 8% takes 943ms ± 17ms (30.4x), and at 16% takes 1.91s ± 0.01s (15x). First-prediction times are similar across all variants (29.68s for TACO 1% versus 28.55s for POT), since the compression cost is paid once.
• Memory reduction at prediction time. In the same no-KV-cache table, predict memory drops from 22.45GB for POT to 549MB for TACO 1% (−97.6%), 845MB for TACO 2% (−96.3%), 1.41GB for TACO 4% (−93.7%), 2.56GB for TACO 8% (−88.6%), and 4.89GB for TACO 16% (−78.2%). The paper also states TACO consumes 98% less memory during prediction in its insight heading, and the abstract says the method consumes "up to 97% less memory." Fit-time memory for TACO is 92MB versus 36MB for POT, which the authors describe as roughly 2.5x higher and not a practical concern in the no-KV-cache setting. With KV caching enabled, the paper reports 92% less memory than the predictor-only transformer.
• Accuracy is retained. On the TabArena classification tasks, mean ROC-AUC is 0.866 ± 0.103 for TabICL, 0.866 ± 0.103 for TabPFN V2.0, 0.862 ± 0.101 for POT, 0.855 ± 0.097 for TACO (r=1%), 0.857 ± 0.098 (r=2%), 0.857 ± 0.099 (r=4%), 0.858 ± 0.100 (r=8%), and 0.858 ± 0.101 (r=16%). The authors note POT and TACO were not trained as long as competing methods but reach comparable performance. Critical difference diagrams using Friedman–Nemenyi tests via autorank at α = 0.05 indicate no significant degradation up to a 1% compression rate.
• TACO enables KV caching where it otherwise fails. POT with KV caching ran out of memory on 32 out of 100 datasets in the synthetic grid (rows from 5,000 to 50,000, features from 100 to 1,000). The authors argue TACO makes cached inference feasible under tighter memory budgets, because the compressed context must be materialized anyway and the compressor is executed only once.
• Joint training matters. Comparing standard TACO to a variant where the predictor is initialized with TabPFN v2.0 weights and frozen (only the compressor trained), joint optimization consistently performs better across compression rates.
• Multi-rate training is not worse than rate-specific training. A single model trained with compression rates sampled from {1, 2, 4, 8, 16} shows performance differences centered near zero versus rate-specific models, with no statistically significant advantage for rate-specific training at the 95% confidence level.
• TACO beats simple compression baselines. Against kNN-based sampling and random sampling compression used with the predictor-only transformer, TACO is consistently and significantly better across compression rates from 1% to 16%. The gap narrows as compression rates increase.
• Chunking scales to large tables. On MetroPT-3 (approximately 1.5M rows, 15 features, 80%/20% time-based split, 0.1% compression rate), the training set of about 1.2M rows is compressed to 1,214 retained rows, and TACO (chunk) achieves precision 0.9886 ± 0.0089, recall 0.7925 ± 0.0711, and AUPRC 0.8955 ± 0.0084. For comparison, POT (Random) achieves AUPRC 0.7293 ± 0.1118, POT (kNN) 0.3715 ± 0.3437, TabPFNv2 (Random) 0.7439 ± 0.1180, TabPFNv2 (kNN) 0.2358 ± 0.2036, AutoGluon (extreme, 30m) 0.9656 ± 0.0022, and XGBoost 0.9593 ± 0.0072.

Methodology in Plain English

The authors take a tabular foundation model architecture (TabPFNv2, with 2D attention over rows and columns) and split it into two modules, each with roughly 7M parameters, 12 attention layers, 6 attention heads, and an embedding dimension of 192, with a two-layer residual MLP between them.

The compressor receives the full training table plus a small "dummy" table of K rows whose values are initialized from a random subset of the training data. The target column of the dummy rows is masked with a placeholder so the model can distinguish real from dummy cells. Through the transformer blocks, the compressor is trained to load the information from the full training set into the latent representation of those K dummy rows, producing a "mini-train" of K × (M+1) × L values.

The predictor then receives the compressed mini-train concatenated with the embedded test batch and produces class scores for the test points. Crucially, the compressor and predictor are meta-learned together (Equation 2 in the paper), so the compressor is directly rewarded for producing a compact representation that helps the predictor, rather than being trained to reconstruct the data.

Training follows the recipe of prior work: an initial synthetic prior of datasets generated by structural causal models (including tree-based causal mechanisms), using the synthetic prior released with TabICLv1. Sequence length is fixed at 1024, feature counts are sampled uniformly between 2 and 100, and training runs 80K optimization steps with a global batch size of 1024, yielding approximately 82M synthetic datasets. Continuation stages extend training to 81k steps with sequence lengths between 1k and 40k rows, then 81.025k steps with 40k to 60k rows, followed by a real-data adaptation stage (following TabDPT) up to 91k total steps with a maximum context of 20k rows, discretizing continuous targets into at most 10 bins. Optimization uses AdamW, mixed precision, gradient accumulation with a micro-batch size of 16, learning rate warmup to 1×10⁻⁴ with cosine annealing, gradient clipping at norm 1.0, and weight decay of 1×10⁻². Training was distributed over 8 NVIDIA H100 GPUs with 95GB of VRAM each, and pretraining took approximately 20 days.

Evaluation uses the classification subset of TabArena. Most experiments use the same 26 datasets as the TabPFN v2.0 evaluation (20 binary, 6 multiclass); the large-scale chunking experiment uses all 36 TabArena classification datasets with fewer than 10 classes (8 multiclass). Metrics are ROC-AUC for binary tasks and ROC-AUC (one-vs-one) for multiclass tasks, with predefined cross-validation folds and the TabPFN v2.0 preprocessing pipeline. Statistical comparisons use Friedman tests with Nemenyi post-hoc tests via autorank at a 0.05 significance level. The compression rate is defined as r = K / N_train, and "POT" denotes a predictor-only transformer trained under the same setup without any compressor.

For very large tables, TACO avoids loading the whole dataset at once: a dataset with N = 10⁶ rows is partitioned into N_C = N/C = 100 disjoint chunks of size C = 10⁴, each compressed (for example at 1%, producing 100 rows per chunk), and the per-chunk summaries are concatenated into one global context for the predictor. Code is released at https://github.com/machinelearningnuremberg/TACO.

Why This Matters

The paper targets the central scaling obstacle for tabular foundation models: attention cost that grows with the conditional context. By showing that the context can be compressed to as little as 1% (and to 0.1% in the large-table case) with no statistically significant accuracy loss, it suggests that in-context tabular prediction can be made practical under realistic inference budgets and also enabled where it previously could not run at all due to memory limits.

Real-world applications named or implied by the paper:

• Real-time fraud detection, where predictions must be produced continuously for incoming transactions.
• Ad click-through prediction on streaming data.
• Predictive maintenance, including the MetroPT-3 compressor-failure dataset used in the experiments.
• Large relational or industrial tables with millions of rows, where the chunking-and-stitching procedure allows inference that a full-context model could not fit in memory.

Industry relevance: The reported speedups (up to about 94x without KV caching) and memory reductions (down to 549MB from 22.45GB in the 1% setting) directly affect serving costs. The paper also positions TACO as an enabler of KV caching, which was infeasible for the uncompressed predictor on 32 out of 100 grid datasets due to out-of-memory errors. The authors openly release code and note that competing SOTA numbers come from open-weight checkpoints.

Future Directions

• Stronger and more diverse priors. The authors state pretraining could benefit from priors that better capture real-world datasets, explicitly including those with a higher degree of missing values.
• Distribution shift. They identify datasets reflecting distribution shifts as important for real-world deployment, and point to prior work (Helli et al., 2024) on temporal distribution shift within the TFM regime as an approach that could be adapted to TACO.
• Beyond classification. The evaluation is limited to classification tasks; extending to regression and time-series settings is described as an important direction for future work.
• Open question on the accuracy-latency frontier. The paper frames TACO as providing a principled accuracy-latency trade-off via the compression rate, which invites further work on tuning that trade-off under deployment constraints and on how far compression can be pushed for specific domains.

Target Audience

Researchers and engineers working on tabular machine learning, tabular foundation models, and efficient transformer inference will benefit most. It is also relevant to practitioners deploying in-context tabular predictors under latency or memory constraints, and to readers interested in learned context compression as an alternative to row or feature subsampling. Some familiarity with attention complexity, in-context learning, and tabular benchmarking conventions is helpful, but the central argument (compress the context, keep the accuracy) is accessible to a broader machine learning audience.

Authors’ abstract

The long-standing dominance of gradient-boosted decision trees for tabular data has recently been challenged by in-context learning tabular foundation models. In-context learning methods fit and predict in one forward pass without parameter updates by leveraging the training data as context for predicting on query test points. While recent tabular foundation models achieve state-of-the-art performance, their transformer architecture based on the attention mechanism has quadratic complexity regarding dataset size, which in turn increases the overhead on training and inference time, and limits the capacity of the models to handle large-scale datasets. In this work, we propose TACO, an end-to-end tabular compression model that compresses the training dataset in a latent space. We test our method on the TabArena benchmark, where our proposed method is up to 94x faster in inference time, while consuming up to 97\% less memory compared to the state-of-the-art tabular transformer architecture, all while retaining performance without significant degradation. Lastly, our method not only scales better with increased dataset sizes, but it also achieves better performance compared to other baselines.