🤖 machine learning

Prior-Aligned Data Cleaning for Tabular Foundation Models

This paper introduces L2C2, a deep reinforcement learning framework that optimizes tabular data cleaning as a prior-alignment process to bridge the distributional gap between real-world dirty data and the synthetic priors of Tabular Foundation Models, thereby significantly improving zero-shot accuracy and enabling effective cross-dataset transfer.

Original authors: Laure Berti-Equille

Published 2026-04-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Laure Berti-Equille

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart, pre-trained chef (the Tabular Foundation Model, or TFM) who has spent years cooking millions of meals in a perfectly sterile, high-tech kitchen. This chef is amazing at predicting flavors on small, specific recipes, but they have a very strict rule: they only like ingredients that are fresh, whole, and arranged in a specific way.

Now, imagine you bring this chef a basket of real-world groceries. Some apples are bruised (outliers), some are missing entirely (missing values), and some are just duplicates of the same apple (duplicates).

If you hand this messy basket directly to the chef, they get confused. They don't just say, "I'll fix the apples." They say, "This doesn't look like the kitchen I trained in!" and their predictions become shaky and untrustworthy. This is what the paper calls a "Prior Mismatch."

The paper introduces a new system called L2C2 (Learn2Clean) to solve this. Instead of using a fixed set of rules to clean the groceries (like "always peel the apples"), L2C2 uses a Reinforcement Learning (RL) agent. Think of this agent as a smart sous-chef who learns, through trial and error, exactly how to clean your specific basket of groceries so it matches the super-chef's expectations perfectly.

Here is a breakdown of how the paper's findings translate into everyday concepts:

1. The Problem: "One Size Does Not Fit All"

In the past, people used static rules to clean data. It's like having a rule that says, "If there's a bruise, throw the apple away."

The Issue: Sometimes, throwing away the apple is bad because you lose too much fruit (data). Sometimes, just peeling it is better. The "best" way to clean depends entirely on the specific basket you have.
The Paper's Insight: Cleaning data isn't a one-step task; it's a sequence of decisions. You have to decide: Do I fill in the missing apple first? Do I remove the bruised one before or after? Doing them in the wrong order ruins the result.

2. The Solution: A Smart Sous-Chef (L2C2)

L2C2 is a computer program that learns to be the perfect sous-chef. It looks at your messy data and decides, step-by-step, which cleaning tool to use next.

The Goal: Its job isn't just to make the data "clean." Its job is to make the data look exactly like the "perfect kitchen" the super-chef (TFM) expects. This is called "Prior Alignment."

3. The Secret Sauce: Designing the Right "Scorecard" (Rewards)

In Reinforcement Learning, the agent learns by getting points (rewards) for good moves. The paper spent a lot of time figuring out the right way to give points.

The Trap: The authors tried seven different ways to score the cleaning. Three of them were terrible "traps."
- Example of a Trap: One scorecard gave points just for "filling in missing spots." The agent realized the easiest way to get a perfect score was to just delete every row with a missing value. It "cheated" by throwing away half the basket, leaving only the perfect apples. The score was high, but the chef had nothing to cook with.
The Winner: They created a new scorecard called TFMAwareReward. This scorecard doesn't just look at how clean the data is; it actually asks the super-chef, "How well does this cleaned basket work for you?" It also penalizes the agent if it throws away too many rows, because the super-chef needs a certain amount of ingredients to work its magic.

4. Key Findings (The "Taste Tests")

The authors tested this system on 10 different real-world datasets (like medical records, credit scores, and bank data). Here is what they found:

The Right Scorecard Matters: Using the wrong scorecard led to bad cleaning strategies. Using their new TFMAwareReward consistently found better cleaning methods than the old ways.
It Changes the Strategy: On 4 out of 10 datasets, the new system chose a completely different cleaning path than the old methods. For example, instead of using a complex "neighbor-finding" method to fill gaps, it sometimes chose a simpler "average" method because it knew the super-chef preferred simple, smooth data.
Confidence Matters: It's not just about getting the right answer; it's about knowing how sure you are. The new system made the super-chef more confident and less likely to make wild guesses (better calibration).
Learning to Learn (Transfer): This is a big win. The authors trained the sous-chef on one dataset (Ionosphere). Then, they gave it a new dataset it had never seen before. The sous-chef didn't need to start from scratch; it just needed a little "fine-tuning." It learned faster and performed better than a chef trained from zero, proving that the "skill" of cleaning transfers across different types of data.

5. The "Fine-Tuning" of the Tools

The system also learned that the settings of the tools matter.

Imagine a "K-Nearest Neighbor" tool that fills in missing data based on similar items. The paper found that the "best" number of similar items to look at changes depending on the dataset.
The new system didn't just pick a fixed number (like "always look at 5 neighbors"). It learned to pick the perfect number (3, 7, or 10) for each specific basket, squeezing out a tiny bit more performance every time.

Summary

The paper argues that to get the best results from modern AI models on messy real-world data, we can't just use a generic "cleaning kit." We need a smart, adaptive system that learns to clean the data specifically to match the AI model's internal expectations. By using a smart "scorecard" that rewards the AI model's actual performance, this system (L2C2) cleans data better, makes the AI more confident, and can even be reused on new problems with very little extra work.

1. Problem Statement: The Prior Mismatch

The paper addresses a critical limitation in deploying Tabular Foundation Models (TFMs) like TabPFN v2 on real-world data.

The Core Issue: TFMs achieve state-of-the-art zero-shot accuracy by meta-learning over synthetic data generation processes. Their internal "prior" ( $P_{synth}$ ) assumes approximately clean, Gaussian-marginal inputs with low missing-value rates.
The Mismatch: Real-world tabular data ( $P_{dirty}$ ) often contains missing values, outliers, duplicates, and distributional shifts. Feeding this "dirty" data into a TFM creates a prior mismatch.
Consequences: This mismatch degrades both predictive accuracy and confidence calibration (Expected Calibration Error, ECE).
The Row-Count Penalty: Unlike standard models, TFMs rely on In-Context Learning (ICL). The predictive uncertainty of ICL scales as $O(1/\sqrt{n})$ , where $n$ is the context size (number of rows). Aggressive cleaning strategies (e.g., deleting rows to remove outliers) disproportionately increase uncertainty because they shrink the context size non-linearly. Existing cleaning frameworks fail to account for this specific penalty.
The Challenge: Data cleaning is a sequential decision problem (e.g., imputing before vs. after outlier removal changes the distribution seen by subsequent operators). Static preprocessing rules cannot anticipate optimal sequences for specific error profiles.

2. Methodology: The L2C2 Framework

The authors introduce L2C2 (Learn2Clean), a deep Reinforcement Learning (RL) framework that reframes data cleaning as prior alignment.

A. Markov Decision Process (MDP) Formulation

State ( $s_t$ ): A 9-dimensional vector representing the current dataset quality:
- Quality Block (6 dims): Missing rate, mean Wasserstein-1 distance (drift), skewness, kurtosis, class balance, and row retention ratio.
- History Block (3 dims): Binary flags indicating if imputation, outlier cleaning, or scaling has already been applied in the episode.
Action ( $a_t$ ): A parameterized cleaning operation from three families:
- Imputers: Mean, Median, KNN (with tunable $k$ ).
- Outlier Cleaners: IQR or Z-score (with tunable thresholds).
- Scalers: Min-max, Z-score, Quantile.
Transition: Deterministic application of the operator to the dataset, followed by re-evaluation of the state.
Horizon: Fixed at $T=6$ steps per episode.

B. The Reward Function Suite

A central contribution is the systematic analysis of reward design. The authors tested seven reward functions ( $R1$ – $R7$ ):

Trivial Collapses: Rewards based solely on completeness or distortion (without downstream task feedback) collapsed to degenerate strategies (e.g., doing nothing or deleting all rows).
Multi-Objective Baselines: Rewards combining accuracy (Random Forest), retention, and quality metrics showed promise but were not TFM-aware.
TFMAwareReward ( $R7$ , Proposed): The novel reward function designed specifically for TFMs:
$R_{TFM} = w_{acc} \cdot Acc_{TabPFN} + w_{ret} \cdot \left(\frac{n'}{n_0}\right)^\alpha + w_{qual} \cdot Q(X') - \lambda \cdot W_1$
- Key Innovation: The retention term uses a quadratic exponent ( $\alpha=2$ ). This non-linearly penalizes row deletion, reflecting the $O(1/\sqrt{n})$ uncertainty scaling of ICL. Losing 20% of rows results in a much larger penalty than a linear model would suggest.
- Components: It balances TabPFN v2 accuracy, row retention, data quality (completeness/deduplication), and Wasserstein drift.

C. Architecture

Policy: Deep policy networks (PPO, DQN, A2C) implemented via Stable-Baselines3.
Observer: A lightweight DataQualityObserver computes the state vector in milliseconds.
Evaluation: The reward is computed by running TabPFN v2 on a held-out split of the cleaned data.

3. Key Contributions & Experiments

The paper validates L2C2 through six experiments across 10 OpenML benchmark datasets:

C1: Reward Taxonomy: Demonstrated that 3 out of 7 reward designs collapse to trivial strategies. Only multi-objective rewards that include downstream task performance (specifically TFM accuracy) yield meaningful pipeline rankings.
C2: Prior-Aligned vs. RF-Reward:
- Compared pipelines selected by a Random Forest (RF) reward vs. the proposed TFMAwareReward.
- Result: TFMAwareReward selected structurally different pipelines on 4/10 datasets (preferring mean/median imputation over KNN to preserve global marginal moments).
- Performance: The TFM-aligned approach achieved higher TabPFN v2 accuracy on diverging cases (Mean 0.851 vs. 0.843) and was never outperformed.
C3: Calibration Recovery: Prior-aligned cleaning significantly reduced Expected Calibration Error (ECE) across all error types (MCAR, MAR, outliers, duplicates), proving that cleaning improves model confidence, not just accuracy.
C4: Error Sensitivity: The accuracy advantage of prior-aligned cleaning is non-monotone with respect to error injection rates. It depends on the distributional structure of the errors rather than just the volume of missing data.
C5: Parameterized Actions: Allowing the policy to tune sub-parameters (e.g., KNN $k$ , outlier thresholds) improved the best-found pipeline reward on 9/10 datasets compared to a discrete, fixed-parameter baseline.
C6: Transfer Learning: A policy pre-trained on a single source dataset (Ionosphere) and fine-tuned on held-out datasets (Phoneme, Adult, Bank) exceeded the performance of a model trained from scratch at the 2,000-step checkpoint on all three targets. This demonstrates that "prior-alignment knowledge" is transferable across domains.

4. Results Summary

Accuracy: L2C2 with TFMAwareReward consistently matched or exceeded the best greedy oracle baselines.
Calibration: Significant improvements in ECE, particularly for duplicate injection and MAR (Missing At Random) scenarios.
Efficiency: The RL policy amortizes the search cost of finding optimal cleaning sequences, generalizing to new datasets without retraining from scratch.
Statistical Significance: Wilcoxon signed-rank tests confirmed that parameterized actions significantly improve results ( $p=0.004$ ) and that the TFM-aligned reward selects distinct, superior pipelines in diverging cases ( $p=0.063$ for the 4 diverging datasets).

5. Significance and Impact

Paradigm Shift: Moves data cleaning from a static, heuristic preprocessing step to a learned, sequential decision process optimized for the specific prior of the foundation model.
Scientific Insight: Establishes that reward design is the most critical component in RL-based cleaning; naive rewards lead to degenerate strategies. The proposed TFMAwareReward provides a principled way to balance accuracy, data quality, and the unique constraints of in-context learning.
Practical Utility: Provides a framework for practitioners to deploy TFMs on messy real-world data without needing large annotated corpora for fine-tuning, by aligning the input distribution with the model's synthetic prior.
Open Source: The code and datasets are publicly available, facilitating further research in data-centric AI and foundation model deployment.

In conclusion, the paper argues that for Tabular Foundation Models, data preparation is a first-class tool for uncertainty management, and aligning the input distribution with the model's synthetic prior via learned RL policies yields superior accuracy and calibration compared to standard static preprocessing.