Zero-shot tabular imputation

TabImpute

One Pipeline, Many Transformers: Pattern-Specific Imputation Specialists for Tabular Missing Data

A single pre-training pipeline that turns synthetic data into transformer imputers. The default TabImpute model is robust across every missingness pattern we tested. Swap one module, and the same pipeline builds a specialist for the pattern you expect.

Jacob Feitelberg, Dwaipayan Saha, Kyuseong Choi, Zaid Ahmad, Anish Agarwal, Raaz Dwivedi

Paper Code Weights PyPI
Entry-wise featurization · click a missing cell
Every entry becomes one row of features: its row index, column index, its whole row and its whole column. All missing entries are predicted in parallel. Illustrative values.

Zero-shot

Pre-trained once on about 1.2 million synthetic tables. At test time it imputes a new table in a forward pass, with no fitting or tuning.

One pipeline, many specialists

Only the missingness module of the data generator changes. The architecture, loss and training stay fixed, and each pattern gets its own specialist.

Over 100× faster

Entry-wise featurization predicts all missing entries in parallel, over 100× faster per entry than column-wise imputation with TabPFN.

Results on MissBench

42 real OpenML datasets, each under 11 missingness patterns at 40% missingness, against 14 established baselines. Scores are column-normalized RMSE, averaged over datasets: lower is better. The default TabImpute model, trained on MCAR alone, has the lowest overall error and is best on 8 of the 11 patterns.

Whiskers are 95% confidence intervals across the 42 datasets. Rows are sorted by overall RMSE, and that order is kept for every pattern.

Show data table

Head-to-head with MissForest

MissForest is the strongest baseline from prior work. Across all 462 dataset–pattern pairs, TabImpute has the lower RMSE on 326 (70.6%; Wilcoxon p < 10−4).

Each comparison is one dataset under one missingness pattern.

Pattern-specific specialists

When you know how your data went missing, swap that pattern into the synthetic generator and train a specialist. Nothing else changes. Each column below is a model trained on one pattern; each row is the pattern it was tested on.

Cells show RMSE. Shading compares each cell with the best in its row: row best2× the best or worse trained on the pattern it is tested on Bold = top two in the row

Self-MNAR, Pol-MNAR and Censor-MNAR specialists are the best model on their own pattern, and most specialists rank in the top two on theirs. NN-MNAR and Latent-MNAR are the exceptions, which the paper attributes to an untuned, uniform training setup. The MCAR-trained default (first column) stays a strong all-rounder.

Show data table

How it works

TabImpute follows the prior-data fitted network (PFN) recipe behind TabPFN: generate millions of synthetic problems, and train one transformer to solve them in context. Three pieces make it work for imputation.

1 · Entry-wise featurization

Imputation is recast as supervised prediction. For a table with m rows and n columns, every entry (i, j) becomes one row of a new feature table:

(i, j) ↦ i ⊕ j ⊕ Xi,: ⊕ X:,j  →  target Xij

That gives mn rows with m + n + 2 features. Observed entries are the training points and missing entries are the queries. Features are built from the incomplete table, never the true values.

  • Parallel: all missing entries are predicted in one pass on a GPU.
  • Missing inputs: NaNs in the features are replaced with the column mean, plus a binary indicator of which values were missing.
  • Trade-off: the table grows to mn rows, which suits small and medium tables. Our follow-up, MatrixFormer, works on the matrix directly to scale further.

2 · A TabPFN-style transformer without the query mask

The architecture follows TabPFN v2: attention alternates between features within a row and samples within a feature. One change matters for imputation. TabPFN masks attention between query points, encoding the assumption that test points are independent. Missing entries are not independent, especially when missingness is not at random, so TabImpute removes that mask and lets queries share information. The output is a Riemann (bar) distribution over values, trained with negative log-likelihood.

3 · Synthetic data with an interchangeable missingness module

Data: complete tables come from linear factor models, Y = UVᵀ, with low rank k and latent vectors drawn from Gaussian, Laplace or Student-t distributions. A nonlinear factor model scored almost as well in ablations.

Missingness: a mask is then applied by a pluggable module: MCAR, column-wise MAR, or any of nine MNAR generators (see MissBench). This module is the only part that changes between specialists.

Training recipe per model
Training data
~1.2 million synthetic tables
Loss
Negative log-likelihood (bar distribution)
Optimizer
AdamW
Learning rate
3 × 10⁻⁴
Weight decay
10⁻²
Batch size
16
Compute
~10–15 hours on one H200
Why train the default on MCAR only? Training on patterns one after another caused catastrophic forgetting, and mixing MCAR, Col-MAR and Self-MNAR made the model fixate on Self-MNAR. Pure MCAR exposes the model to a huge variety of masks that overlap with those of structured patterns. Because the model learns the data distribution given what is observed, not the missingness mechanism itself, this makes a robust general-purpose imputer.

MissBench

A benchmark of 42 small real-world OpenML datasets from medicine, finance, engineering, education and more, each masked with 11 missingness patterns. Nine of the patterns are missing not at random (MNAR), where whether a value is missing depends on unobserved values. Illustrations below are schematic.

Quickstart

Install from PyPI. The default TabImpute checkpoint downloads from Hugging Face on first use.

bash
pip install tabimpute
python
import numpy as np
from tabimpute.interface import ImputePFN

imputer = ImputePFN(device="cpu")  # or "cuda"

X = np.random.rand(20, 5)
X[np.random.rand(*X.shape) < 0.2] = np.nan

X_imputed = imputer.impute(X.copy())
  • Input: a 2D NumPy array with np.nan for missing values.
  • Best fit: small and medium tables. Each entry becomes a feature row, so cost grows with rows × columns.
  • GPU recommended. On CPU, a table can take a few minutes.
  • Reproduce the paper: MissBench and every figure and table can be regenerated from the command line; see the benchmark README.

Working with larger matrices? MatrixFormer is our follow-up foundation model for matrix completion. It applies axial attention directly to the matrix instead of featurizing every entry.

Citation

If you use TabImpute or MissBench, please cite the paper.

bibtex
@article{feitelberg2025tabimpute,
  title={One Pipeline, Many Transformers: Pattern-Specific Imputation Specialists for Tabular Missing Data},
  author={Feitelberg, Jacob and Saha, Dwaipayan and Choi, Kyuseong and Ahmad, Zaid and Agarwal, Anish and Dwivedi, Raaz},
  journal={arXiv preprint arXiv:2510.02625},
  year={2025}
}