# In-context imputation for tabular data: filling missing values with a foundation model

In-context imputation fills the missing cells of a table by letting a pre-trained tabular model predict each incomplete column from the rows where that column is observed, without refitting the model to your data. Run in a loop over all columns, it competes with the strongest classical method, MICE, and it is most useful when it is checked by a validation layer that can fall back to simpler methods.

**In short.** Imputation quality for squared error depends on one thing: how well the imputed deviations from the column mean correlate with the true ones. Classical chained methods such as MICE raise that correlation by refitting a model per column. An in-context imputer does the same job with a pre-trained model that reads the observed rows as context. In my measurements a chained in-context imputer with a validation gate led MICE on a synthetic split and on 20 real datasets, led a published language-model study on average but lost two of its three missingness cells, and could not yet prove that the in-context part, rather than the gate, deserves the credit. The full paper is [An In-Context Imputation Head for Tabular Data](https://www.ozercemkelahmet.com/articles/in-context-imputation-head-paper).

## What is missing-value imputation?

Real tables have holes. A sensor dropped out, a form field was skipped, two systems were merged and one of them never recorded a value. Imputation means filling those blank cells with sensible estimates so that the rest of the pipeline, a report, a model, or a person, can use the row.

The task is measured by hiding cells you know, asking a method to reconstruct them, and scoring how close it gets. Two scores are common. NRMSE is a squared-error score normalised by the spread of the column, and MAE is the average absolute error. Both reward being close, not being right about the exact label.

## Why does the reason a value is missing matter?

Whether a method works depends on why the cells are missing, and benchmarks usually test three mechanisms.

MechanismMeaningEveryday example

MCAR, missing completely at randomWhich cells are missing has nothing to do with the data.A random batch of survey forms is lost in the post.
MAR, missing at randomMissingness depends on other columns you can see.Older patients are less likely to have a test recorded, and age is in the table.
MNAR, missing not at randomMissingness depends on the missing value itself.People with very high incomes skip the income question.

MNAR is the hardest, because the gap carries information that the observed data cannot show. A method that wins on average can still lose one mechanism, so results should be read per mechanism, not only as a single number.

## How do classical imputers work?

- **Column mean.** Replace every blank with the column's average. It is the baseline every method should beat.
- **kNN.** Find rows that look like the incomplete one and copy from them.
- **softImpute.** Fill the gaps with a low-rank approximation of the whole table.
- **MICE.** Multivariate imputation by chained equations. Fill every gap with a starting guess, then repeatedly predict each column from all the others using a regression, and update. Each pass makes the filled values more consistent with each other.
- **missForest.** The same chained idea with random forests as the column model.

The chained idea is the strongest classical paradigm, because predicting a column from all the others and repeating is a good way to use the correlation between columns. Single-shot reconstruction models, which fill everything in one pass, tend to shrink their answers toward the column mean.

## What is in-context imputation?

A [tabular foundation model](https://www.ozercemkelahmet.com/articles/what-is-a-tabular-foundation-model) such as [TabPFN](https://www.ozercemkelahmet.com/tools/tabpfn) or [TabICL](https://www.ozercemkelahmet.com/tools/tabicl) is pre-trained on many tables. At use time it predicts a column in a single forward pass by reading labelled rows as context, with no refitting. That makes it a natural column learner inside a chain: instead of fitting a regression for column j on every pass, hand the model the rows where j is observed as context and ask it to predict the blanks.

Think of an experienced analyst who has seen thousands of tables. You show them the filled rows of one column next to the other columns, and they read off a good estimate for the empty rows, without being retrained for your table.

## How a chained in-context imputer is built

The design in the paper has four parts, each answering a specific problem.

- **A chain.** Missing cells start at the column mean. Columns are visited from fewest to most missing values. For each column, the rows where it is observed are the context, the other columns (with their current completions and missing-value indicators) are the features, and the column itself is the target. Three sweeps are performed.
- **An in-context head.** A small module of 18.97 million parameters sits on top of a frozen pre-trained tabular encoder. It conditions on at most 384 support rows, 48 feature columns, and 16 classes in one forward pass. It is 0.87% of the parameters of the checkpoint that hosts it.
- **Bins and the expected value.** Because the head predicts classes, a continuous column is cut into ten quantile bins, and the answer is the probability-weighted average of the bin centres. For squared-error scores, the expected value beats the mode: replacing the mode plus rounding by the expected value recovered 75 to 84 percent of the categorical loss on five weak tables.
- **A validation gate.** A second chain uses BayesianRidge as the column model, and it reproduces MICE to numerical precision. Both chains run on the table with 15 percent of the observed cells hidden, and one candidate is chosen: the in-context chain, the ridge chain, a blend, a boosted tree, kNN, or the column mean. The column mean is always a candidate, so no column can end up worse than the column mean on the validation cells, and the ridge chain's presence means the result is no worse than MICE up to selection noise.

## What does the evidence show?

The paper evaluates the imputer on three boards. The figures below are from the paper, with its own caveats.

BoardMetricIn-context chainReference

Synthetic no-prior split, 81 cellsNRMSE0.1564MICE 0.3983, first in all 81 cells
Real-20, 20 datasets, 3 mechanisms, 3 ratesNRMSE0.3097MICE 0.3363, ahead in all nine cells
Gemini board, same 20 datasetsMAE, nine-cell mean0.2058Published 0.228

Real-20 overall NRMSE (lower is better)ICI chain (ours)ICI chain (ours): 0.30970.3097MICEMICE: 0.33630.3363missForestmissForest: 0.34740.3474kNNkNN: 0.35360.3536softImputesoftImpute: 0.35520.3552column meancolumn mean: 0.37250.3725Real-20 overall NRMSE (lower is better) Mean over 20 real datasets, 3 mechanisms, 3 missing rates.
Zero-shot foundation models used column by column did far worse in the same harness: TabPFN applied column-wise had an NRMSE of 0.8472 on MCAR, TabImpute 0.4828, against 0.2121 for the final chain. The lesson is that a foundation model is not an imputer by itself. The chain and the gate make it one.

Where it loses
On the Gemini board the advantage is carried entirely by MNAR (0.203 against 0.308). The imputer is behind on MCAR (0.197 against 0.150) and marginally behind on MAR (0.2175 against 0.216). That comparison also mixes a local run with a number published from another harness, so the headline deserves caution.

Gemini board: MAE by missingness mechanismMCARICI chain (ours): 0.19700.1970Gemini (published): 0.15000.1500MARICI chain (ours): 0.21750.2175Gemini (published): 0.21600.2160MNARICI chain (ours): 0.20300.2030Gemini (published): 0.30800.3080Nine-cell meanICI chain (ours): 0.20580.2058Gemini (published): 0.22800.2280Gemini board: MAE by missingness mechanism. Blue is ICI chain (ours), grey is Gemini (published). Lower is better.
Who gets the credit?
Because the gate contains a MICE-equivalent chain, a result no worse than MICE is guaranteed by construction. With the gate switched off, the in-context learner raised the correlation from 0.116 to 0.308 on one table and from 0.432 to 0.534 on another, but on densely correlated tables the gate picked the ridge chain. The in-context learner helps where the ridge chain is weak, yet the clean ablation, ridge chain alone against the full system, has not been run. That is the first experiment planned.

## What can go wrong when you evaluate an imputer?

The most useful part of the work was finding seven defects that each produced a convincing but false result. They are a good checklist for anyone benchmarking imputation.

- **The mask ignored the mechanism.** Every run used the hardest masker, so MCAR and MAR looked bad and almost identical. The giveaway was that the column-mean baseline got better as more data went missing.
- **A weak engine was called.** The harness used a reconstruction model that scored far worse than the real one. Two changes at once hid the cause.
- **A warm start silently failed.** A strict weight load broke on a later-added head, so the pre-trained variant was never actually tried.
- **A leaky A/B test.** A model trained on the complete table looked excellent when asked to fill a masked copy of it.
- **A candidate read the answer.** In the selection step, the hold-out cells were the true observed values, so the candidate won every comparison.
- **A rival was scored as NaN.** The ridge chain used the wrong mask, so it was eliminated before it could compete.
- **Mode instead of mean.** Squared error is minimised by the expected value, not the most likely class.

The general rule matches [A score is only as good as its report](https://www.ozercemkelahmet.com/articles/a-score-is-only-as-good-as-its-report): a result that looks too good, or a baseline that behaves in a physically implausible way, is a reason to audit the harness before believing it.

## Which imputer should you use?

SituationReasonable starting point

Need a quick, transparent baselineColumn mean, then MICE
Densely correlated numeric columnsMICE or a ridge chain, which the gate also tends to pick
Weakly correlated columns, mixed types, many tablesA chained in-context imputer with a validation gate
Missing not at randomTest explicitly per mechanism, since this is where the in-context chain gained most
Strict compute limitsClassical imputers, which take seconds. The chain averaged about 240 seconds per run on a GPU in the paper's setup
Categorical columns with more than 16 levelsNot supported directly by the head; use a classical fallback

Whatever you choose, compare on your own data with the mechanism you expect, hide cells you can check, and report the cells you lose.

## Imputation in the model documentation

Imputation is one of the outputs of the large tabular foundation model described on this site, alongside sector, column profile, cross-table map, and predictions. See the [capabilities](https://www.ozercemkelahmet.com/model/capabilities) page and the [imputation benchmark story](https://www.ozercemkelahmet.com/model/bench-imputation). A related story covers prediction under injected missingness, [missing-data robustness](https://www.ozercemkelahmet.com/model/bench-missing-data-robustness). The pre-training objective behind this skill is [masked cell modeling](https://www.ozercemkelahmet.com/tools/masked-cell-modeling).

## Frequently asked questions

What is in-context imputation?
Filling missing cells by letting a pre-trained tabular model predict each incomplete column from the rows where it is observed, read as context in a forward pass, without refitting the model to your table.

Is it better than MICE?
On the synthetic split and on 20 real datasets the chained in-context imputer had lower NRMSE than MICE in every reported cell. Because a MICE-equivalent chain is part of its selection set, the result is no worse than MICE by construction, and the ablation that would isolate the in-context part has not been run.

Why not just use TabPFN to fill each column?
Used zero-shot, column by column, it did much worse: an NRMSE of 0.8472 on MCAR in the same harness, against 0.2121 for the chained system. The chain, the bins, and the validation gate are what make a prediction model an imputer.

What is NRMSE?
Normalised root mean squared error: the squared-error distance between imputed and true values, scaled by the spread of the column. Lower is better. For squared-error metrics it depends on the correlation between the imputed and true deviations from the column mean.

Why use the expected value instead of the most likely value?
Squared error is minimised by the expected value. Replacing the mode plus rounding by the expected value recovered 75 to 84 percent of the categorical loss on five weak tables.

What is a validation gate?
A selection layer that hides some observed cells, tries several candidate imputers on them, and keeps the best one for each column, with the column mean always available as a safe fallback.

What are the limits of the method?
The head reads at most 384 support rows, 48 feature columns, and 16 classes, so columns with more than 16 categories need a fallback. It is slower than classical imputers, and the weights and code are not released.

Where is the paper?
Read it on this site as [An In-Context Imputation Head for Tabular Data](https://www.ozercemkelahmet.com/articles/in-context-imputation-head-paper), or download the [PDF](https://www.ozercemkelahmet.com/papers/in-context-imputation-head.pdf).