Prior-Aligned Data Cleaning for Tabular Foundation Models
| Authors | Laure Berti-Equille |
| Year | 2026 |
| HF Upvotes | 3 |
| arXiv | 2604.25154 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
Tabular Foundation Models (TFMs) achieve state-of-the-art zero-shot accuracy on small tabular datasets by meta-learning over synthetic data-generating processes -- making them highly attractive for practitioners who cannot afford large annotated corpora. However, their in-context learning mechanism assumes approximately clean inputs: missing values, outliers, and duplicates in the real-world data create a prior mismatch that degrades both accuracy and confidence calibration simultaneously. Correcting this mismatch requires sequential decisions over cleaning operators whose interactions no static preprocessing rule can anticipate -a natural fit for reinforcement learning~(RL). We introduce L2C2, the first deep RL framework framing tabular data cleaning as prior alignment: a learned policy sequences operators to minimize the distributional gap between dirty input and the TFM's synthetic prior. Six experiments on ten OpenML benchmark datasets establish: 1) three of seven reward designs collapse to degenerate trivial cleaning strategies -- principled reward engineering is scientifically non-trivial; 2) the novel TFMAwareReward reward we propose selects structurally distinct pipelines on 4/10 datasets and achieves higher TabPFN accuracy on those diverging cases (mean 0.851 vs. 0.843; Wilcoxon p=0.063, n=4) while never underperforming; 3) parameterized cleaning actions improve best-found pipeline reward on 9/10 datasets (Wilcoxon p=0.004); and 4) a policy pre-trained on one single source dataset exceeds scratch training at the 2,000-step fine-tuning checkpoint on all three held-out datasets (up to +28.8% after full fine-tuning) demonstrating cross-dataset transfer of prior-alignment knowledge. These findings establish that prior alignment is a principled data preparation strategy for TFM deployment on real-world tabular data.
Engineering Breakdown
Plain English
This paper addresses a critical gap in Tabular Foundation Models (TFMs), which achieve strong zero-shot performance on small tabular datasets through meta-learning but fail when real-world data contains missing values, outliers, and duplicates. The authors show that these data quality issues create a "prior mismatch" that simultaneously degrades accuracy and breaks confidence calibration — two problems that static preprocessing rules cannot solve because data cleaning operators interact in complex, unpredictable ways. They propose L2C2, a deep reinforcement learning framework that learns a policy to make sequential cleaning decisions tailored to each dataset, treating data cleaning as a prior alignment problem. The RL approach dynamically selects which cleaning operators to apply and in what order, adapting to the specific characteristics of the input data rather than following fixed rules.
Core Technical Contribution
The core novelty is framing tabular data cleaning as a prior alignment problem solvable by deep reinforcement learning, rather than as a static preprocessing step. L2C2 is the first framework to treat the sequence of cleaning operator decisions as a learned policy optimization problem, where the RL agent learns which cleaning transformations to apply and in what order to maximize downstream TFM performance. This is fundamentally different from prior work because it acknowledges that cleaning operator interactions are data-dependent and context-sensitive — a single best preprocessing pipeline doesn't exist. The RL formulation allows the policy to adapt to different datasets, data distributions, and TFM architectures, treating cleaning as an integral part of the learning pipeline rather than a separate preprocessing stage.
How It Works
L2C2 operates as a sequential decision-making process where an RL agent observes the current state of a tabular dataset and iteratively selects cleaning operators from a fixed set (e.g., imputation methods, outlier removal, deduplication). The state representation encodes data quality metrics — missing value patterns, statistical outliers, duplicate counts, and distributional properties — allowing the agent to condition cleaning decisions on the actual data characteristics. At each step, the agent selects an operator, applies it, and receives a reward signal based on the downstream performance of the TFM on a validation set, incentivizing cleaning decisions that improve both accuracy and confidence calibration. The policy is trained via deep RL (likely policy gradient methods) over a distribution of synthetic datasets generated with known data quality issues, enabling the policy to generalize to new, unseen tabular datasets. The final output is a cleaned dataset whose prior distribution aligns with the assumptions of the TFM's in-context learning mechanism, recovering performance lost to the mismatch.
Production Impact
In production systems using TFMs, this approach eliminates the need for hand-crafted, dataset-specific data cleaning pipelines that typically require domain expertise and extensive experimentation. Engineers would replace ad-hoc preprocessing scripts with a trained RL policy that automatically adapts cleaning strategies to each new dataset, reducing manual tuning time and improving consistency across different data sources. The framework simultaneously improves model accuracy and calibration — critical for high-stakes applications where confidence estimates guide downstream decision-making or uncertainty quantification. However, the approach introduces computational overhead: running RL inference adds latency before TFM inference, and retraining the RL policy on new data distributions requires access to TFM performance signals (validation sets), which may not always be available in truly zero-shot scenarios. Integration complexity increases because the system now depends on two learned components (RL policy + TFM) rather than one, and failure modes can cascade if the RL policy produces adversarial or unexpected cleaning sequences.
Limitations and When Not to Use This
The paper assumes access to validation data to compute reward signals for RL training, which contradicts the zero-shot setting that TFMs are designed for — in truly zero-shot scenarios, no labeled data exists to evaluate cleaning decisions. L2C2 is trained on synthetic data-generating processes, so its generalization to real-world data distributions with novel types of corruption, missing patterns, or domain-specific anomalies remains unclear. The approach also assumes a fixed set of cleaning operators; it cannot discover novel cleaning strategies or handle domain-specific data quality issues that weren't represented in the training distribution. Finally, the paper doesn't address computational efficiency: the RL inference cost may be prohibitive for real-time applications or when processing massive tabular datasets, and the trade-off between cleaning quality and inference latency is not explored.
Research Context
This work builds on the recent success of Tabular Foundation Models, which use meta-learning over synthetic data-generating processes to achieve strong zero-shot performance — a paradigm shift from traditional supervised learning on tabular data. It extends prior work on automated machine learning (AutoML) and hyperparameter optimization by moving beyond static pipeline selection to dynamic, instance-aware data cleaning. The paper contributes to a growing body of research applying reinforcement learning to data quality and preprocessing problems, though prior work has been limited to narrower settings (e.g., feature selection, imputation strategy selection). This opens a research direction toward end-to-end learned data pipelines where cleaning, feature engineering, and model inference are jointly optimized, moving away from the traditional separation between preprocessing and modeling.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::