One Pipeline, Many Transformers: Pattern-Specific Imputation Specialists for Tabular Missing Data
cs.LG
Submitted: 2025-10-03
Updated: 2026-08-28
Code: https://github.com/PriorLabs/tabpfn-extensions
License: http://creativecommons.org/licenses/by/4.0/
The gist: Missing data in tabular datasets forces practitioners into a hard choice: deploy a general-purpose imputer that may perform poorly for the problem at hand, or wait for someone to design a specialized
Terminology
Abstract
Missing data in tabular datasets forces practitioners into a hard choice: deploy a general-purpose imputer that may perform poorly for the problem at hand, or wait for someone to design a specialized algorithm. This problem is worsened by the fact that real-world missingness rarely satisfies the textbook missing completely at random (MCAR) assumption, as entries are often missing not at random (MNAR), where the probability of being observed depends on the missing data itself. We collapse this trade-off into a single pre-training pipeline that builds transformer-based imputation specialists through three components: an entry-wise featurization that recasts imputation as supervised prediction over row--column context, a synthetic data generator with pluggable missingness modules, and prior-data fitting on millions of synthetic tables. Swapping only the missingness module, with no changes to architecture, loss, or training, yields a pattern-specific specialist that outperforms methods purpose-built for that pattern. We validate this on MissBench, a new benchmark of 42 OpenML datasets and 11 missingness patterns (including 9 MNAR variants) spanning medicine, finance, and engineering. Remarkably, training exclusively on MCAR yields a default model---TabImpute---robust across all tested patterns. In addition, the pattern-specific specialists produced by our pipeline outperform the 14 established baselines on their target patterns. We open-source the pipeline, models, and benchmark.
Sources
- N squared: A Unified Python Package and Test Bench for Nearest Neighbor-Based Matrix Completion
- ReMasker: Imputing Tabular Data with Masked Autoencoding
- CACTI: Leveraging Copy Masking and Contextual Information to Improve Tabular Data Imputation
- not-MIWAE: Deep Generative Modelling with Missing not at Random Data
- Decoupled Weight Decay Regularization
- TabDPT: Scaling Tabular Foundation Models on Real Data
- TabICL: A Tabular Foundation Model for In-Context Learning on Large Data
- A Closer Look at TabPFN v2: Understanding Its Strengths and Extending Its Capabilities
- TabFlex: Scaling Tabular Learning to Millions with Linear Attention
- DiffPuter: Empowering Diffusion Models for Missing Data Imputation
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks