Imputation networks

Overview

Imputation networks are neural networks trained to fill in missing values in a table by predicting each missing cell from the cells that are observed. They learn, from the data itself, how columns depend on each other and how values are distributed, and then use that knowledge to produce a best guess for each gap. Good ones also report how sure they are, so the user can tell a confident fill from a rough one.

Missing data is one of the most common problems in real tables. Sensors fail, forms are left blank, systems are merged. Simple fixes such as filling with the column average are fast but ignore relations between columns and can distort results. Neural imputers use those relations. They sit at the heart of tabular foundation models, because the skill of predicting a hidden cell from its context is the same skill that masked cell modeling trains, and architectures like axial attention and Perceiver IO make it work across table shapes.

What it is

An imputation network is any neural model whose job is to estimate missing entries of a table. The word covers several designs.

• Autoencoders that compress a row and rebuild it, trained with corrupted inputs so they learn to repair gaps.
• Generative adversarial imputers in which a generator proposes values and a discriminator tries to tell real cells from imputed ones.
• Variational models that learn a distribution over rows and fill gaps by sampling or by taking the most likely value.
• Attention-based models that treat cells as tokens and predict hidden ones, as in masked pre-training.
• Graph and set-based models that treat rows or columns as nodes and pass messages between them.

What it takes in and gives out

Types of missingness

How a value went missing matters for which method is safe.

• Missing completely at random (MCAR). The gap has nothing to do with any value. Example: a random network drop.
• Missing at random (MAR). The gap depends on other observed values. Example: older people skip an online question more often.
• Missing not at random (MNAR). The gap depends on the missing value itself. Example: high earners leave the income field blank.

Imputation networks handle MCAR and MAR well. MNAR is harder because the information needed is exactly what is absent. A model can still help if it is given the missing pattern as an input.

Where it came from

Statisticians worked on missing data for decades. Key ideas included replacing values with averages, regression imputation, hot-deck imputation that borrows from similar rows, and multiple imputation, which creates several plausible completed tables to reflect uncertainty. Methods such as iterative chained-equation imputation and random-forest-based imputation became popular tools.

Neural approaches followed the rise of deep learning. Denoising autoencoders, introduced in the late 2000s, learned to restore corrupted inputs and were quickly applied to missing data. A method called MIDAS adapted overcomplete denoising autoencoders for multiple imputation. GAIN (2018), "Generative Adversarial Imputation Nets", used an adversarial setup in which a discriminator guessed which entries were real, with a hint mechanism. Variational autoencoder variants such as HI-VAE handled heterogeneous data types. Later, transformer-based models treated imputation as masked prediction, inheriting the self-supervised recipe from language models.

More recently, pre-trained tabular models have been used as imputers directly. Because they were trained to predict hidden cells across many tables, they can fill gaps in a new table with little or no additional training. See /model/research for that line of work.

Why neural imputation, given strong classical methods

• Classical methods fit one table at a time, while neural models can be pre-trained across tables.
• Neural models handle mixed types and nonlinear relations in one framework.
• They scale to many rows and columns more easily than some iterative methods.
• They can output full distributions rather than single values.

Classical methods remain valuable baselines and are often hard to beat on small tables. A fair study always includes them.

Families of imputation networks

Here is a closer look at the main designs, with what each one is good at.

Denoising autoencoders

• Corrupt the input by hiding cells, then rebuild the clean row.
• Simple, fast and easy to train.
• Often used with several runs to give multiple imputations.

Generative adversarial imputers

• A generator fills the gaps, and a discriminator guesses which cells are real and which are filled.
• The contest pushes the fills to look like real data.
• Training can be unstable, so careful tuning is needed.

Variational models

• Learn a latent code for each row and a way to decode it into a distribution over cells.
• Missing cells are handled by using only the observed ones to infer the code.
• Naturally give uncertainty and can sample.

Attention-based models

• Treat each cell as a token and let tokens attend to each other.
• Fit directly with masked pre-training.
• Share weights across table shapes.

Graph and set models

• Treat rows and columns as nodes in a bipartite graph and pass messages along the edges that have observed values.
• Missing cells are the edges to predict.
• Handle sparse tables well.

Everyday picture of the whole idea

Think of a class photo with a few faces smudged. A person who knows the families can guess the smudged faces from the siblings next to them, from the family resemblance, and from the typical look of the age group. An imputation network does the same: it uses neighbours in the row, in the column and in the table to guess the smudged cell.

How it works

Here is the flow of a typical imputation network, with an everyday picture at each step.

Step 1: Record the mask

First, note which cells are observed and which are missing. This is a table of ones and zeros with the same shape as the data. Many networks take the mask as an extra input, so they know exactly which entries are real.

Analogy: before restoring a damaged painting, a conservator marks which parts are original and which are lost.

Step 2: Encode the table

Numbers are scaled, categories are turned into embeddings or one-hot vectors, and missing cells are set to a neutral placeholder plus the mask flag. The row, or the grid, is turned into vectors the network can use.

Step 3: Process with a network

The network mixes information. A denoising autoencoder compresses each row to a small code and expands it back. An attention model lets cells exchange information with their row and column neighbours. A latent model compresses the whole table into a summary and reads from it.

Analogy: the conservator studies the surviving parts of the painting to infer what the lost parts must have looked like.

Step 4: Predict values for missing cells

For each missing cell, a head produces a value, a class distribution, or the parameters of a distribution. Observed cells are kept as they are.

Step 5: Train with artificial gaps

The network is never given the true values of real gaps, because they do not exist. So training creates extra gaps by hiding observed cells, and scores the network on how well it recovers them. This is the same idea as masked cell modeling.

Step 6: Impute and optionally repeat

At use time, the real missing cells are filled. Some methods fill, then feed the filled table back in for another pass, refining the estimates. Others draw several samples to make multiple imputations.

Step 7: Report confidence

A good imputer returns more than a point guess. Options include a predicted spread, a set of samples, quantiles, or a calibrated probability for categories.

The shape of the loss

Let the observed mask be m, equal to 1 for observed cells. The training loss on artificially hidden cells h is computed only where h is 1 and m was 1 before hiding. Observed and still-visible cells are not scored.

loss = mean over hidden cells of  error(prediction, true value)

Where the error depends on type: squared error or a likelihood for numbers, cross-entropy for categories.

A compact code sketch

import torch
import torch.nn as nn
import torch.nn.functional as F

class DAEImputer(nn.Module):
    """A small denoising autoencoder that sees data and a mask."""
    def __init__(self, n_cols, hidden=64, code=16):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(2 * n_cols, hidden), nn.ReLU(),
                                 nn.Linear(hidden, code), nn.ReLU())
        self.dec = nn.Sequential(nn.Linear(code, hidden), nn.ReLU(),
                                 nn.Linear(hidden, n_cols))

    def forward(self, x, observed):
        x_in = torch.where(observed, x, torch.zeros_like(x))
        z = self.enc(torch.cat([x_in, observed.float()], dim=1))
        return self.dec(z)

# synthetic data: column 2 depends on columns 0 and 1
X = torch.randn(2000, 5)
X[:, 2] = 0.7 * X[:, 0] - 0.5 * X[:, 1] + 0.05 * torch.randn(2000)
observed = torch.ones_like(X, dtype=torch.bool)

model = DAEImputer(5)
opt = torch.optim.Adam(model.parameters(), 1e-3)
for step in range(500):
    hide = torch.rand_like(X) < 0.2                 # artificial gaps
    seen = observed & ~hide
    pred = model(X, seen)
    loss = F.mse_loss(pred[hide], X[hide])          # score hidden cells only
    opt.zero_grad(); loss.backward(); opt.step()
print(float(loss))

After training, the model fills column 2 well because it has learned the relation to columns 0 and 1.

Choices that shape an imputation network

Key concepts and terms

A worked example

Consider a table of five customers with age, income and spend. Some entries are missing.

The simple baseline

Fill with the column mean. The mean age of observed values is (25 + 38 + 52 + 31) / 4 = 36.5. Row 3 would get 36.5. The mean income is (40 + 85 + 80 + 60) / 4 = 66.25, so row 4 gets 66.25. The mean spend is (12 + 30 + 28 + 35) / 4 = 26.25, so row 5 gets 26.25.

These fills ignore that spend rises with income, and that row 3 looks like row 2.

What a trained imputer does

1. Row 3, age. Income is 80 and spend is 28. A model that has learned the relation between income, spend and age sees that row 2 (income 85, spend 30, age 38) is similar. It predicts an age around 38 to 40.
2. Row 4, income. Age is 52 and spend is 35. Spend is the highest in the table, which in this data goes with high income. The model predicts around 90.
3. Row 5, spend. Income is 60 and age is 31. Among the rows, income 60 sits between 40 (spend 12) and 80 (spend 28). The model predicts around 20.

Reading the uncertainty

With only five rows the model is not very sure. A good imputer shows this.

The row 4 income has a wide spread because the age and spend values are at the edge of what the table shows. A downstream analyst would treat that fill with more caution.

Validating on this table

To test the imputer, hide an observed cell, such as the income of row 2, run the model, and compare the fill with the true value of 85. Repeat for many cells and report the average error. This is the standard way to validate an imputer when true values of the real gaps are unknown.

Comparison on the example

The neural imputer uses all the observed columns at once, which is why it adapts better than a single-column rule.

Handling time and order

Some tables are time series where order matters. A sensor table with one row per minute is not a bag of independent rows. For such data, add time as a feature or use models that read sequences, so that a gap can be filled from the neighbours before and after. Axial attention can attend along the time axis in this way, and a time feature can be added to a Perceiver IO query.

Handling very sparse tables

When most cells are missing, as in a ratings matrix, an imputer is closer to a recommender than to a cleaner. Graph and low-rank methods are common there. The ideas still apply: learn from observed entries, hide some for practice, and predict the rest. See recommendation for the related topic.

Strengths and limitations

Strengths

• Uses relations between columns. Fills depend on context rather than on one column alone.
• Handles mixed types. Numeric, categorical and binary columns can share one model.
• Scales to wide and long tables. Training is based on batches and hardware-friendly operations.
• Can output uncertainty. Spreads, samples or intervals show how much to trust each fill.
• Can be pre-trained. A model trained across many tables can fill a new one with little extra training.
• Flexible missing patterns. Any gap pattern can be handled by changing the training masks.
• Integrates into pipelines. It can be a step before a downstream model, or built into it.

Limitations

• Needs enough data. With very few rows, simple methods are often as good or better.
• Mask mismatch. Training masks may not look like the real gaps, especially for MNAR.
• Overconfidence. Networks can report narrow intervals that are too optimistic.
• Sensitive to scaling and encoding. Poor preprocessing can ruin results.
• Hidden bias. Filling can smooth over minority patterns, shrinking variance and weakening real relations.
• Not a fix for bad data. If most of a column is missing, any fill is mostly guesswork.
• Opacity. It is harder to explain why a particular fill was chosen than with a simple rule.

Choosing a method in practice

1. If the table is small and clean, start with iterative regression imputation or a nearest neighbour method.
2. If you have many rows and nonlinear relations, try a denoising autoencoder.
3. If you need uncertainty and have the resources, use a variational or distribution-head model.
4. If you have many tables and want one tool for all, use a pre-trained tabular model.
5. If your final goal is prediction with a boosted tree, test native missing handling before adding imputation.

Cost and effort

How it compares with alternatives

Impute first or model with gaps

Some models, such as gradient boosted trees, handle missing values natively by learning which side of a split to send them to. In that case imputation may be unnecessary for prediction tasks. Imputation is still valuable when the completed table is itself the product, when a model cannot take gaps, or when you want consistent values for reporting.

Single versus multiple imputation

• Single imputation gives one completed table. It is simple, but it understates uncertainty.
• Multiple imputation gives several completed tables. Analysis is run on each, and the results are combined, which reflects the uncertainty from missing data.

Neural imputers can support multiple imputation by sampling from a predicted distribution or by running several networks.

Role in tabular foundation models

A tabular foundation model learns from raw tables with mixed types, missing values and often no column names. Imputation is not an add-on for such a model. It is one of the core skills.

Same skill, different name

Predicting a masked cell from the rest of the table is what masked cell modeling trains, and it is exactly imputation. A foundation model pre-trained this way is already an imputation network. At use time, real gaps take the place of artificial ones.

Raw rows and columns

Imputation networks in this setting use architectures that read the table as a grid. Axial attention lets each missing cell gather clues from its row and from its column. Perceiver IO lets a query for one missing cell read a global summary of the table. Both are shape-free, so the same weights fill tables of different sizes.

Mixed types and missing values

Because the model is trained on all types, it can fill numeric, categorical and date cells in one pass. Existing gaps are given a missing marker as input and are excluded from the loss, so they carry information rather than noise.

No column names

When headers are absent or meaningless, the model identifies columns by their values and relations. A column with values between 0 and 1 that rise with another column is understood without a name. If a header is present, its text adds a clue.

Pre-training on synthetic and real tables

Synthetic tables supply perfect ground truth and cover a huge variety of dependency structures. Real tables supply realistic quirks. Together they teach a broad imputation skill. See synthetic data generation.

In-context learning

For imputation, the observed part of the table serves as the context. The model reads it and fills gaps without any weight update. The more informative rows there are, the better the fill. This is the same mechanism behind in-context prediction in TabPFN and TabICL.

Continual learning

New tables keep arriving, and the imputer should improve without forgetting. Replay of earlier tables and regularization methods such as elastic weight consolidation and learning without forgetting protect old skills. See also experience replay.

Confidence scores

A foundation-style imputer can return an uncertainty with each fill. This lets downstream systems decide what to do: accept confident fills, flag uncertain ones, or request real data for them. Calibration should be checked on held-out tables.

Anomaly detection

If an observed value differs greatly from what the model would have predicted for that cell, it may be an error or a rare event. Treating each observed cell as if it were missing and comparing the model's guess with the real value yields an anomaly score. This links imputation to anomaly detection.

Sector and semantic classification

The representations learned while imputing carry the meaning of columns and tables. Light heads on top can classify a column's semantic type or a table's sector. See semantic and sector heads. Across tasks, the evaluation is tracked at /model/benchmarks, and the design is described at /model/architecture.

A simple workflow

1. Profile the table: missing counts, patterns, types.
2. Decide the goal: clean table, input to a model, or analysis with uncertainty.
3. Split the data before fitting anything.
4. Pick a baseline imputer and a neural imputer.
5. Hide observed cells to measure error for both.
6. Check calibration if you use intervals.
7. Run the downstream task with each imputer and compare.
8. Save the imputer with the rest of the pipeline.

Monitoring in production

• Track the share of imputed cells per column over time. A rise may signal a failing source.
• Track the average predicted uncertainty. A change suggests data drift. See schema drift detection.
• Keep a sample of rows where the true value later arrives, and measure real error.
• Flag rows with many imputed cells for review.

Practical guidance and best practices

• Understand the missingness first. Ask why values are missing. Count gaps by column and by row.
• Add missingness indicators. Keep a flag for each column that had gaps. The pattern can be informative.
• Scale numeric columns. Robust scaling avoids large columns dominating.
• Encode categories carefully. Use embeddings or one-hot vectors, and keep a category for unknown values.
• Train with realistic masks. Include random, column-wise and block-wise gaps, and match the real pattern if you know it.
• Validate by hiding observed cells. Report error per column and per type.
• Compare with simple baselines. Mean, median and iterative imputers set the floor.
• Check calibration. Intervals labeled 90 percent should cover about 90 percent of true values.
• Use multiple imputation for analysis. If you will draw statistical conclusions, keep the uncertainty.
• Do not impute the target in evaluation. Fit the imputer on training data only, then apply it to test data.
• Document the process. Record the method, settings and the share of cells imputed.
• Review rare groups. Check that filling does not erase patterns of small subgroups.

Fit on training data only

A frequent mistake in pipelines is fitting the imputer on the whole dataset before splitting. That leaks information from the test set into the training set. The right order is: split, fit the imputer on the training part, apply it to both parts.

Deciding how much missingness is too much

These ranges are rough guides, not strict thresholds.

Signs that your imputer is not working

• The imputed column has much smaller variance than the observed one.
• Correlations between columns are weaker in the filled data than in the observed data.
• Errors on hidden observed cells are no better than the mean fill.
• Intervals are far too narrow or far too wide.
• Results change a lot when the random seed changes.

Each of these points to a concrete fix: richer model, better masks, calibration, or more data.

Common mistakes

• Mean fill everywhere. It flattens variance and weakens correlations.
• Fitting on the full dataset before splitting. It leaks test information.
• Ignoring the mask. If missing values are replaced by zeros without a flag, zero looks like a real value.
• Using one global loss. Mixed types need type-aware losses.
• Trusting point fills blindly. Without uncertainty, a bad fill looks as good as a good one.
• Evaluating only on easy gaps. Random gaps are easier than structured ones. Test both.
• Imputing the label. Never fill the label in training data without a plan. It changes the problem.
• Overfitting the imputer. With small tables, a big network memorizes the observed cells.
• Forgetting MNAR. If missingness depends on the missing value, add the missing pattern as input and treat results with care.
• Skipping documentation. Later users cannot tell which values were real.

Frequently asked questions

What are imputation networks?

Imputation networks are neural models that estimate missing values in a table. They learn from the observed cells how columns relate and how values are distributed, and then predict each missing cell from the cells around it. Many also report a confidence, such as a spread or an interval, so users can see which fills are reliable.

How are they trained without knowing the true missing values?

They create their own practice gaps. During training, some observed cells are hidden on purpose and the network is asked to recover them, with the loss computed only on those hidden cells. Because the true values are known for the practice gaps, this gives a clear learning signal, and the skill carries over to real gaps.

Are imputation networks better than filling with the mean?

Usually yes, when columns are related and there is enough data. A mean fill ignores the other columns and shrinks the spread of the data, which can distort later analysis. A neural imputer uses all observed values in the row and in the table, so its fills follow the real structure. On tiny tables, simple methods can match it.

What is the difference between single and multiple imputation?

Single imputation fills each gap with one value and treats it as if it were real, which understates uncertainty. Multiple imputation creates several plausible completed tables, runs the analysis on each, and combines the results, so the final uncertainty reflects the missing data. Neural imputers support it by sampling from predicted distributions.

How do imputation networks handle categorical data?

Categorical columns are turned into embeddings or one-hot vectors on the input side, and the output head produces a probability for each category. The network fills the gap with the most likely category or samples from the distribution. The loss for these cells is cross-entropy, separate from the loss used for numbers.

Can imputation networks deal with data that is missing not at random?

They can help, but not fully. When the chance of a value being missing depends on the value itself, the observed data is a biased sample. Giving the network the missingness pattern as an input lets it learn some of the dependence, and sensitivity checks with different assumptions are wise. No method can fully recover information that was never observed.

How do tabular foundation models relate to imputation networks?

A tabular foundation model pre-trained with masked cell prediction is itself an imputation network that works across many tables. It reads the observed part of a new table as context and fills the gaps without extra training, and it can return confidence scores. This makes imputation one of the natural uses of such models.

How can I check that the imputed values are good?

Hide a share of observed cells, impute them, and compare with the true values using errors scaled per column. For distribution outputs, check that stated intervals cover the right share of true values. Also check downstream results, such as the accuracy of a model trained on the imputed table, and compare with simple baselines.

Should I impute before training a prediction model?

It depends on the model. Gradient boosted trees can handle missing values natively, so imputation may not be needed. Many other models need a complete table. If you impute, fit the imputer on the training data only, include missingness flags, and keep the imputation step inside the same pipeline used at prediction time.

Can imputation help find errors in data?

Yes. If you treat an observed value as hidden and the model predicts something far from it, the value may be an error or a rare event. Large gaps between the prediction and the real value, measured relative to the model's uncertainty, make a useful anomaly score for cleaning and monitoring.

Glossary

• Adversarial imputation: training a generator and a discriminator against each other to make fills look real.
• Autoencoder: a network that compresses an input to a code and rebuilds it.
• Calibration: agreement between stated confidence and observed accuracy.
• Chained equations: an iterative method that predicts each column from the others in turn.
• Denoising autoencoder: an autoencoder trained to rebuild clean data from corrupted data.
• Hot-deck imputation: filling a gap with a value from a similar row.
• Imputation: replacing missing values with estimates.
• In-context learning: adapting to a task from examples placed in the input.
• Mask: an array that marks observed and missing cells.
• MCAR, MAR, MNAR: types of missingness by what the gap depends on.
• Missingness indicator: a flag that records whether a value was missing.
• Multiple imputation: creating several completed tables to reflect uncertainty.
• Prediction interval: a range expected to contain the true value with a stated probability.
• Single imputation: one filled value per gap.
• Tabular foundation model: a model pre-trained on many tables that serves new tables with little or no training.
• Variational autoencoder: a generative network that learns a distribution over data.

Summary

Imputation networks are neural models that fill missing table values by predicting each gap from the observed cells, often with a confidence attached. They learn from artificially hidden cells, handle mixed types, and capture relations between columns that simple fills ignore.

In tabular foundation models, imputation is a built-in skill. The same pre-training that teaches masked cell modeling produces a model that fills gaps in new tables, using row and column context from axial attention or a global summary from Perceiver IO. Good practice means understanding the missingness, using realistic masks, fitting only on training data, checking calibration and comparing with simple baselines. See /model/architecture, /model/research and /model/benchmarks for more.

Related

It belongs to Tabular foundation models in my stack, under Foundation models and applied ML. Others in the same group: