# Model-Aware Data Cleaning for Tabular Foundation Models

arXiv: https://arxiv.org/abs/2604.25154 · HTML: https://mlsanity.com/paper/2604.25154/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 33 references that are in the map.
- First posted on arXiv 28 Apr 2026.

## Builds on (most cited first, 11 of 11)

- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [AutoGluon-Tabular: Robust and Accurate AutoML for Structured Data](https://mlsanity.com/paper/2003.06505/) (2020)
- [TabICL: A Tabular Foundation Model for In-Context Learning on Large Data](https://mlsanity.com/paper/2502.05564/) (2025)
- [Data-centric Artificial Intelligence: A Survey](https://mlsanity.com/paper/2303.10158/) (2023)
- [OpenML Benchmarking Suites](https://mlsanity.com/paper/1708.03731/) (2017)
- [AlphaD3M: Machine Learning Pipeline Synthesis](https://mlsanity.com/paper/2111.02508/) (2021)
- [Automating biomedical data science through tree-based pipeline optimization](https://mlsanity.com/paper/1601.07925/) (2016)
- [DiffPrep: Differentiable Data Preprocessing Pipeline Search for Learning over Tabular Data](https://mlsanity.com/paper/2308.10915/) (2023)

## Abstract (arXiv)

Tabular Foundation Models (TFMs) achieve state-of-the-art zero-shot accuracy on small tabular datasets, but their in-context learning assumes approximately clean inputs: real- world missing values, outliers, and duplicates create a prior mismatch that degrades both accuracy and calibration. We study reinforcement learning for tabular data cleaning, a learned policy that sequences cleaning operators and introduce L2C-TFM with a model-aware reward (TFMAwareReward). We are explicit about what this reward optimizes: it regularizes the Wasserstein distance between the cleaned and the original (dirty) data, a distributional-stability term, which we measure as a diagnostic. Across six experiments on ten OpenML datasets: (i) three of seven reward designs collapse to degenerate strategies, so reward engineering is non-trivial; (ii) under an 8-seed repeated-holdout protocol the model-aware reward matches a random-forest-reward baseline on accuracy (p=0.38), with a benefit confined to minority-class macro-F1 under class imbalance that is partly a reward-agnostic calibrated-threshold effect; and (iii) a policy pre-trained on one dataset transfers to held-out datasets. A diagnostic analysis shows that the two distances are distinct objectives: cleaning tends to move data away from the prior, and prior-distance, not distance-to-dirty, is what tracks downstream quality. We therefore treat prior alignment as a motivating objective and a target for future work, not a property of the reward evaluated here. Code, datasets, and the nested evaluation harness are available at https://github.com/LaureBerti/Learn2Clean/tree/master/Learn2Clean_TFM.