Model-Aware Data Cleaning for Tabular Foundation Models

Abstract

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.