Prior-Aligned Data Cleaning for Tabular Foundation Models

Laure Berti-Equille

Source-linked AI summary

Prior-Aligned Data Cleaning for Tabular Foundation Models

Laure Berti-Equille

arXiv:2604.25154v1Apr 28, 2026cs.LGcs.DB

TL;DR

Real-world tabular errors violate the approximately clean synthetic prior assumed by TFMs, harming prediction and confidence calibration. L2C2 uses deep reinforcement learning to sequence parameterized cleaning operators and align dirty data with the TFM prior. Across ten OpenML datasets, reward design strongly affected discovered behavior, while TFM-aware rewards, parameterized actions, and transfer learning produced the reported gains.

Problem
TFMs assume approximately clean inputs, but missing values, outliers, and duplicates create prior mismatch that degrades predictive accuracy and confidence calibration.
Method
L2C2 is a deep RL framework using a structured 9-dimensional observer, parameterized cleaning actions, and a TabPFN v2-based reward with a quadratic context-size penalty.
Results
TFMAwareReward selected different pipelines on 4 of 10 datasets and outperformed RF-reward on all four diverging cases, while parameterized actions improved reward on 9 of 10 datasets and transfer beat scratch training on all three held-out datasets at 2,000 steps.
Takeaways & Limitations
Prior alignment provides a supported data-preparation strategy for deploying TFMs on real-world tabular data, with reward engineering and reusable policies central to the framework’s reported performance.
Takeaways & Limitations
Small datasets can produce NaN policy logits when aggressive outlier removal leaves zero surviving rows and causes undefined TabPFN v2 inference.

Abstract

from arXiv · show

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.

1 Introduction

TFMs rely on approximately clean inputs, so real-world missing values, outliers, duplicates, and distributional shifts create a prior mismatch that can degrade prediction and calibration. L2C2 addresses this sequential cleaning problem with deep RL, while evaluating reward designs, TFM alignment, parameterized actions, and transfer across datasets.

Motivation: TFM prior mismatch arises when dirty empirical distributions violate the synthetic prior’s assumptions about approximately clean inputs.The mismatch is represented as a distributional gap between P_dirty(D) and P_synth.
Motivation: 30% row removal can shrink in-context learning context and increase predictive uncertainty because uncertainty scales as O(1/√n).This makes row preservation an important cleaning trade-off for small tabular datasets.
Why reinforcement learning: Cleaning is sequential because operator order and sub-parameters alter the distribution encountered by subsequent operators, defeating static preprocessing rules.The extended action space contains up to 302 valid sequences of up to three steps per dataset.
L2C2 framework: L2C2 uses deep policy networks, a 9-dimensional data-quality observer, parameterized cleaning actions, and TFMAwareReward evaluated directly against TabPFN v2.The action space includes imputers, outlier cleaners, and scalers with typed sub-parameters.
Reward design: Three of seven reward designs collapse to degenerate strategies, while TFMAwareReward selects structurally different pipelines from RF-reward on 4 of 10 datasets and is never outperformed.The reward comparison evaluates 112 pipelines across 10 datasets; only R3 provides stable non-trivial rankings among the initial candidates.
Evaluation findings: Parameterized actions improve best-found pipeline reward on 9 of 10 datasets, and pre-training on one dataset beats scratch training at 2,000 steps on all three held-out datasets.The transfer result supports reuse of prior-alignment knowledge across tabular domains.

2 Related Work

L2C2 extends prior work on automated pipeline search and model-aware cleaning by treating tabular cleaning as sequential, TFM-aware prior alignment with profiling and multi-objective rewards.

RL and Search for Data Pipeline Optimisation: L2C2 frames cleaning as a sequential decision problem where each operator changes the distribution seen by later steps, enabling credit assignment across actions.This distinguishes the approach from static preprocessing, greedy search, and conventional AutoML pipeline optimisation.
RL for Data Cleaning: L2C2 replaces shallow Q-learning with deep policy networks and evaluates cleaning against a TFM rather than a fixed downstream surrogate.The framework retains the sequential-cleaning MDP framing while adding seven rewards and cross-dataset policy learning.
RL for Data Cleaning: L2C2 combines TFM accuracy and calibration with explicit distributional-drift regularisation in its cleaning objective.This extends model-aware cleaning from classical models to TabPFN v2 and makes prior-distribution sensitivity part of reward design.
Data Profiling and Data Quality: The framework integrates automated data profiling before cleaning to capture missingness, duplicate patterns, and value-distribution signals.This connects profiling practice with the framework’s data-quality observer.
Data Profiling and Data Quality: L2C2 treats data quality as multidimensional because no single cleaning strategy dominates completeness, consistency, and accuracy simultaneously.This motivates the use of a multi-objective reward rather than a single cleaning criterion.
TFM Prior Alignment: Input-level cleaning addresses TFM mismatch directly, whereas post-hoc recalibration leaves the underlying corrupted inputs unchanged.The paper positions prior alignment as an input-level complement to model-level calibration methods.

3 Problem Formulation

The paper formulates cleaning as constrained prior alignment: select a parameterized pipeline that reduces divergence from the TFM’s synthetic prior while meeting a minimum accuracy threshold, then solve the intractable search approximately with deep RL.

3.1 Prior Mismatch: A TFM is meta-learned on a synthetic prior, while a dirty dataset has an empirical feature distribution that may differ from that prior.The formulation separates the model’s training prior from the observed dataset distribution.
3.1 Prior Mismatch: Prior mismatch is defined as a distributional divergence between the dirty empirical distribution and the TFM’s synthetic prior.The divergence is instantiated using normalized mean column-wise Wasserstein-1 distance.
3.2 Cleaning as Prior Alignment: A cleaning pipeline is an ordered sequence of at most T deterministic actions with typed sub-parameters that maps a dirty dataset to a cleaned dataset.The pipeline space includes both action ordering and operator parameters.
3.2 Cleaning as Prior Alignment: The optimal prior-aligned pipeline minimizes prior mismatch subject to downstream TFM accuracy meeting a minimum acceptable threshold.The constraint preserves predictive performance while the objective targets distributional alignment.
3.2 Cleaning as Prior Alignment: Exact optimisation over the exponential pipeline space is intractable, so the constrained problem is scalarised into a reward and searched with deep RL.This converts the formal optimisation problem into a learnable sequential search procedure.
3.3 MDP Formulation: The cleaning problem is represented as a finite-horizon episodic Markov Decision Process with state, action, transition, reward, discount, and horizon components.The MDP provides the formal sequential-decision structure for pipeline construction.
3.3 MDP Formulation: Each state is a 9-dimensional quality descriptor of the full current dataset, combining quality statistics with binary action-family history.The state is not a windowed view of the data.
3.3 MDP Formulation: Actions are parameterized operations from imputer, outlier-cleaner, or scaler families, with deterministic transitions given the current state and action.The action representation makes operator choice and typed sub-parameter choice part of the search.

4 The L2C2 Framework

L2C2 combines dataset-wide quality observation, parameterized cleaning actions, and reward functions evaluated against TabPFN v2; its design addresses degenerate rewards, distributional faithfulness, and small-data training stability.

Data-Quality Observer: L2C2 maps each current dataset to a 9-dimensional observer state containing quality measures and action-family history.The observer recomputes the state after each cleaning step over the full dataset.
Data-Quality Observer: State construction costs O(d·n log n) per update because Wasserstein-1 drift requires sorting each numeric column.The paper reports that this cost is negligible relative to TabPFN v2 inference in its evaluation setting.
Reward Function Suite: L2C2 compares seven reward functions, including completeness retention, RandomForest accuracy, multi-objective accuracy, drift penalties, incremental gain, distributional faithfulness, and TFMAwareReward.All rewards are clipped to [−1, 1].
Reward Function Suite: Retention-only reward supplies no distributional-quality or downstream-performance signal, making it a lower-bound baseline.Its square-root form dampens the penalty for row deletion relative to a linear formulation.
Reward Function Suite: The multi-objective and drift-penalty rewards combine predictive accuracy with retention, quality, and Wasserstein-drift terms.The drift-penalty variant uses a substantially stronger Wasserstein penalty than the multi-objective baseline.
Reward Function Suite: TFMAwareReward combines TabPFN v2 accuracy, row retention, missingness and duplicate quality, and Wasserstein drift with fixed a priori weights.Its quadratic retention term penalizes row deletion nonlinearly, reflecting in-context uncertainty scaling with context size.
Parameterized Actions and Training: Parameterized actions expand the search from seven discrete actions to 17 actions by adding typed operator sub-parameters such as KNN counts and outlier thresholds.The framework includes imputers, outlier cleaners, and scalers, with PPO used by default for training.
Parameterized Actions and Training: Training episodes reset the dataset, observer reference distribution, reward baselines, action history, return, and replay buffer before applying policy-selected operations.Repeated-family actions receive a penalty and skip cleaning; PPO, DQN, and A2C are supported, while experiments use PPO.

5 Experiments

Across ten OpenML datasets, the experiments show that reward design determines whether cleaning pipelines are informative, while TFMAwareReward aligns reward with TabPFN performance in divergent cases. Parameterized actions improve pipeline rewards, but comparisons remain constrained by a shared 20-pipeline budget and calibration trade-offs.

C2: Prior-Aligned Cleaning vs. RF-Reward Cleaning: TFMAwareReward achieves the highest mean reward and TabPFN v2 accuracy, while its selected pipelines differ from RF-reward on 4/10 datasets.The divergence cases are D2, D4, D5, and D8, and all involve imputer choice.
C1: Reward Function Taxonomy: Three rewards collapse to degenerate strategies: R1 fills missing cells, while R6a and R6b select no-op pipelines across all 10 datasets.R1 scores 1.0000±0.0000; R6a scores 0.9997±0.0009 and R6b 1.0000±0.0000.
C1: Reward Function Taxonomy: R4 selects no-op on 7/10 datasets because its Wasserstein drift penalty can outweigh accuracy gains.Its mean reward is 0.772 ± 0.167.
C1: Reward Function Taxonomy: R3 provides the strongest non-trivial baseline, with mean reward 0.984 ± 0.036 and statistically higher scores than R2 and R4.R3 selects KNN imputation on 9/10 datasets and jointly optimizes accuracy, row retention, and data quality.
C2: Prior-Aligned Cleaning vs. RF-Reward Cleaning: B-greedy-TFM reaches mean accuracy 0.8513 versus 0.8428 for B-greedy-RF, wins on 4 datasets, ties on 6, and is never exceeded.The difference is not conventionally significant (Wilcoxon p=0.063, n=4 diverging datasets), and the ≥7/10 hypothesis threshold is unmet.

6 Discussion

L2C2 aligns cleaning with the TFM prior by complementing TabPFN v2’s fixed preprocessing and penalizing context loss. Its evaluation remains bounded by synthetic benchmark errors, classification scope, scalability uncertainty, inference cost, and model-specific reward calibration.

Prior alignment: Prior-aligned cleaning complements TabPFN v2 by restoring outlier-distorted dynamic range and reducing structured missingness beyond its fixed preprocessing.TabPFN v2 applies z-normalization, a power transform, and uniform missing-value flags unconditionally.
Prior alignment: A quadratic retention penalty values row preservation because in-context uncertainty scales as O(1/√n), especially at small context sizes.Dropping from 100 to 80 rows yields a retention score of 0.64 rather than 0.80 under α=2.
Positioning: L2C2 avoids known functional-dependency assumptions and uses the TFM synthetic prior as its reference distribution, while non-monotone rewards invalidate monotone-pruning arguments.Aggressive outlier removal can decrease accuracy by distorting alignment, and learned policies are intended to generalize across datasets.
Limitations: L2C2 is evaluated on ten OpenML datasets with synthetic errors, while natural errors, multi-table schemas, regression, and scalability beyond 834 sequences remain unevaluated.The framework targets classification, and regression would require a different calibration objective and reward recalibration.
Limitations: Approximately 0.3-second per-step TabPFN v2 inference rules out online settings, while D9/D10 subsampling may introduce sampling bias.These constraints affect deployment latency and the representativeness of the large-dataset evaluation.
Limitations: Reward weights are calibrated to a specific TabPFN v2 version and must be recalibrated when the underlying model changes.The reward is therefore tied to the evaluated TFM configuration rather than being model-version invariant.

7 Conclusion

L2C2 reframes tabular cleaning as TFM prior alignment and shows that reward design affects the pipelines selected and their calibration behavior. The authors identify future extensions across TFMs, benchmarks, profiling signals, and multi-agent policies.

Conclusion: Three of seven reward candidates collapse to degenerate strategies, while TFMAwareReward selects different pipelines from RF-reward on 4 of 10 datasets.TFMAwareReward outperforms RF-reward on all four diverging cases, with Wilcoxon p=0.063 and n=4, and ties on the remaining six.
Conclusion: Prior-aligned cleaning improves TabPFN v2 ECE across all four error types on five representative datasets, with the advantage over RF-reward confined to duplicate injection.The improvement over standard preprocessing holds under MAR, outlier, and duplicate injection, but not MCAR.
Future work: Future work includes adapting TFMAwareReward to TabICL and CARTE, evaluating CleanML, expanding profiling signals, and exploring specialized multi-agent cleaning policies.The proposed extensions target model transfer, benchmark comparison, proactive action masking, and operator specialization.