🤖 machine learning

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

This paper introduces a unified pre-training pipeline that generates pattern-specific transformer-based imputation specialists for tabular data, demonstrating that a single model trained exclusively on MCAR data (TabImpute) achieves robust performance across diverse missingness patterns while outperforming 14 established baselines on their target scenarios.

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

Published 2026-08-19

📖 6 min read🧠 Deep dive

CC BY 4.0

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

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

In the vast world of data, information is rarely perfect. Whether a doctor is reviewing patient records, an economist analyzing market trends, or an engineer monitoring factory sensors, the datasets they rely on often have holes. Numbers are missing, measurements were not taken, or entries were lost during transfer. This absence of data creates a significant problem: most statistical and machine learning tools cannot work with empty spaces. They require a complete picture to function. To fix this, scientists have long used a process called imputation, which means filling in the blanks with educated guesses. For decades, the field has been stuck in a difficult choice. Practitioners could use a general-purpose tool that tries to fill in any missing number but often performs poorly on specific, tricky cases. Or, they could wait for a researcher to design a highly specialized algorithm for their exact problem, a process that demands deep expertise and time. The reality is that missing data is rarely random; often, the reason a number is missing is directly related to the value itself, making the guessing game even harder.

A team of researchers from Columbia University and Cornell University has proposed a way to break this deadlock. They developed a flexible system that acts as a factory for creating custom data-filling experts without needing to redesign the machine every time. Instead of building a new algorithm for every type of missing data, they built a single pipeline that can produce a specialized "imputer" for any pattern of missingness. The core of their work is a method that treats every single missing number in a table as a unique puzzle piece, looking at the row and column it belongs to, as well as the values surrounding it, to make a prediction. By training this system on millions of synthetic tables—computer-generated data that mimics real-world scenarios—they created a model that can instantly adapt. If a user knows exactly why their data is missing, the system can generate a specialist for that specific pattern. If the user has no idea why the data is missing, the system offers a robust default model that works well across almost any situation.

The researchers tested this approach on a new benchmark they created, called MissBench, which includes 42 real-world datasets from fields like medicine, finance, and engineering. They simulated 11 different ways data can go missing, ranging from simple random gaps to complex scenarios where the missingness depends on the hidden values themselves. In these tests, their specialized models consistently outperformed fourteen established methods that were previously considered the best for specific tasks. Remarkably, the default model, which was trained only on the simplest type of missing data, proved to be strong enough to handle the most complex patterns better than many of the specialized tools designed for them. This suggests that the system learns the underlying structure of data so well that it does not need to be told the specific rules of the missingness to do a good job.

The process works by first generating a massive library of fake data tables. These tables are created using a mathematical framework that assumes real-world data often follows hidden patterns, similar to how a few underlying factors might influence many different measurements. The researchers then apply different "missingness modules" to these tables to simulate various ways data can disappear. One module might remove numbers completely at random, while another might hide values that are too high or too low, or hide data based on other values in the same row. The system then trains a transformer-based model—a type of artificial intelligence known for its ability to recognize patterns—to predict the missing values based on the context of the surrounding numbers. Crucially, the architecture of the model never changes. The only thing that changes is the missingness module used during the training phase. This means a user can swap out the training conditions to create a new specialist in a matter of hours, without needing to rewrite the code or re-engineer the model.

One of the most striking findings is that the system does not need to memorize the specific rules of missingness to succeed. When the researchers trained the model only on random missing data, it became a universal expert that performed well even when the missing data followed complex, non-random rules. This challenges the long-held belief that one must design a unique algorithm for every specific type of missing data problem. The system appears to learn a general understanding of how data relates to itself, allowing it to fill in gaps accurately regardless of the reason they are missing. For the rare cases where the reason for missing data is known and highly specific, the pipeline can produce a tailored model that beats even the best existing specialized methods.

The researchers also addressed the issue of speed. Traditional methods that try to fill in data often do so one column at a time, which is slow and can miss connections between different parts of the table. Their new approach looks at every entry in the table simultaneously, allowing it to use information from the entire dataset to make a single, fast prediction. This parallel processing makes the system significantly faster than previous methods, capable of handling small to medium-sized tables in a fraction of the time it takes older tools. While the current system is designed for tables that fit in standard computer memory, the researchers acknowledge that scaling it to massive datasets will require further engineering. However, for the vast majority of practical applications in science and business, the system offers a powerful new tool.

By making their code, models, and benchmark available to the public, the team hopes to change how practitioners approach missing data. Instead of waiting for a perfect algorithm or settling for a mediocre generalist, users can now generate a tool that fits their specific needs or rely on a strong default that works almost everywhere. The work demonstrates that with the right training strategy, a single system can master the art of filling in the blanks, turning the frustrating problem of missing data into a solvable puzzle.

Technical Summary: One Pipeline, Many Transformers

Problem Statement

Missing data in tabular datasets presents a significant challenge for practitioners, who must choose between deploying general-purpose imputation methods that may underperform on specific problems or waiting for researchers to design specialized algorithms for specific missingness patterns. This trade-off is exacerbated by the reality that real-world missingness rarely satisfies the Missing Completely At Random (MCAR) assumption. Instead, data is often Missing Not At Random (MNAR), where the probability of an entry being missing depends on the unobserved value itself. Distinguishing between Missing At Random (MAR) and MNAR is often impossible using observed data alone, particularly in small datasets where statistical tests lack power. Consequently, practitioners face two scenarios: either they lack knowledge of the missingness mechanism and require a robust default, or they possess domain knowledge and seek to exploit it for higher accuracy. Existing methods typically address only one of these scenarios, failing to provide a unified framework that is both robust across diverse patterns and adaptable to specific mechanisms without redesigning the algorithm.

Methodology

The authors propose TabImpute, a flexible pre-training pipeline that generates transformer-based imputation specialists. The pipeline eliminates the need to modify model architecture or training algorithms to adapt to different missingness patterns. It consists of three core components:

1. Entry-Wise Featurization (EWF)

The pipeline recasts imputation as a supervised prediction problem over row-column context. Unlike previous column-wise approaches (e.g., iterative TabPFN), EWF transforms every entry in a table into a feature vector comprising:

Row index and column index.
Row context (the rest of the row).
Column context (the rest of the column).
The target value (the entry itself).

This transformation creates a feature matrix with $m \times n$ rows (where $m$ and $n$ are the number of rows and columns in the original table). During training, the model predicts the target values for all entries. Crucially, the model handles missing inputs by replacing NaNs with column means and appending a binary mask, while the true values of missing entries are only used in the loss function, not as input. This approach enables parallel prediction of all missing entries, offering a significant speedup over column-wise methods.

2. Synthetic Data Generation with Interchangeable Modules

The models are pre-trained on millions of synthetic tables generated via a two-step process:

Data Generation: A Linear Factor Model (LFM) generates the underlying complete data matrix $Y = UV^T$ , where latent factors are sampled from various distributions (Gaussian, Laplace, Student's-t).
Missingness Injection: A pluggable missingness module applies a mask $M$ to the data. The pipeline supports a flexible family of patterns, including MCAR, MAR (specifically Col-MAR), and nine distinct MNAR variants (e.g., Self-MNAR, Censor-MNAR, Block-MNAR, NN-MNAR).

The key innovation is that the missingness module is interchangeable. By swapping only this module, the pipeline can generate training data for any specific pattern without altering the model architecture or the training algorithm.

3. Architecture and Training

The model utilizes a Transformer architecture similar to TabPFN v2, featuring a two-way attention mechanism (inter-feature and inter-sample attention). A critical modification for imputation is the removal of the attention mask that typically restricts information flow between query points (missing entries). In imputation, missing values are often correlated through shared latent structures; allowing query points to exchange information better reflects the problem structure.

The models are trained using Negative Log Likelihood (NLL) loss on approximately 1.2 million synthetic datasets per model. The training process involves a single forward pass (in-context learning) where the network approximates the posterior predictive distribution.

Key Contributions

Unified Pre-training Pipeline: A framework that addresses both practical scenarios (unknown vs. known missingness mechanisms) by producing either a robust default model or a pattern-specific specialist simply by changing the synthetic data generator's missingness module.
Entry-Wise Featurization: A novel featurization strategy that enables parallel prediction of missing entries, resulting in a 100× speedup over prior column-wise TabPFN-based imputation.
MissBench Benchmark: A comprehensive benchmark comprising 42 real-world OpenML datasets across diverse domains (medicine, finance, engineering, etc.) and 11 missingness patterns (including 9 MNAR variants).
Specialized Specialists: The demonstration that models trained on specific patterns (TabImpute-P) outperform methods purpose-built for those patterns, while the default MCAR-trained model (TabImpute) remains robust across all tested patterns.

Results

Evaluated on MissBench against 14 established baselines (including MissForest, MICE, GAIN, MIWAE, and TabPFN):

Robustness: The default TabImpute model, trained exclusively on MCAR data, achieved the lowest overall Root Mean Squared Error (RMSE) across all 11 patterns and ranked in the top-2 for 9 out of 11 individual patterns.
Specialization: Pattern-specific TabImpute models (e.g., TabImpute-Self-MNAR) achieved the lowest RMSE on their respective target patterns, outperforming specialized baselines like not-MIWAE.
Efficiency: TabImpute demonstrated over 100× faster runtime per entry compared to column-wise TabPFN and competitive or superior performance to other methods on small-to-medium tables.
Generalization: The pipeline successfully handled complex MNAR scenarios (e.g., censoring, value-dependent masking) where traditional methods often struggle or require manual specification of propensity models.

Significance and Claims

The paper claims that TabImpute resolves the trade-off between general-purpose robustness and pattern-specific accuracy. By leveraging prior-data fitting on synthetic data, the pipeline allows practitioners to:

Deploy a strong, robust default imputer when the missingness mechanism is unknown.
Rapidly train specialized imputers for known mechanisms without designing new algorithms or modifying model architectures.

The authors emphasize that this approach eliminates the need for practitioners to wait for specialized algorithm development or settle for underperforming general methods. While the current implementation is optimized for small and medium-sized tables (due to the $O(mn)$ expansion in EWF), the pipeline establishes a foundation for scalable, pattern-agnostic imputation that can be extended to larger datasets and more complex data types in future work. The code, models, and benchmark are open-sourced to facilitate adoption and further research.