Tabular foundation models improve dataset representations for machine learning tasks
TabMF: Learning Dataset Representations with Tabular Foundation Models
Machine Learning
Summary
Machine learning models often benefit from understanding entire datasets as a whole to improve tasks like finding similar datasets or noticing when data changes. The authors created TabMF, a method that uses existing large pretrained models designed for tables to better represent entire datasets. They showed that adapting these models with special training on many synthetic datasets helps perform better on real data tasks like finding relevant datasets and detecting shifts in data. This method outperforms traditional ways of representing datasets and can help build better machine learning tools.
What this means in practice
- •For machine learning engineers: Create more accurate dataset similarity searches to find relevant data for new projects using pretrained tabular data models adapted with TabMF.
- •For data platform teams: Improve monitoring of data quality by assessing how severe distribution shifts are in tabular datasets with better dataset-level representations.
Authors
Rafiq Kamel, Nicolas Atienza, Aziz Fouché, Nicola Cancedda, Josif Grabocka
Abstract
Representations of supervised datasets support higher-level machine-learning tasks such as dataset retrieval and distribution-shift monitoring. However, existing approaches typically rely either on hand-crafted meta-features or on dataset encoders trained specifically for representation learning. We introduce TabMF, a framework that repurposes pretrained tabular foundation models (TFMs) as encoders of supervised datasets. TabMF equips a predictive TFM with an architecture-aware dataset-level readout and adapts the model end-to-end using contrastive post-training on stochastic views of synthetic supervised datasets. We instantiate TabMF on TabICLv2 and TabPFN-2.5 and evaluate the resulting representations across six tasks on unseen real datasets. Simple frozen TFM readouts yield competitive representations, indicating that predictive pretraining already induces substantial dataset-level structure. The two TabMF encoders obtain the strongest aggregate performance among the evaluated learned and classical dataset representations, with particularly strong results in dataset retrieval, compact benchmark construction, and distribution-shift severity assessment.