# An In-Context Imputation Head for Tabular Data: Chained, Validation-Gated, and Evaluated on Three Imputation Boards

This paper reports a chained imputer whose column learner is an in-context tabular classifier, the in-context imputation head (ICI head), and it reports the results on three imputation boards together with the head's failures. It is a draft dated 4 November 2026, written by Özer Cem Kelahmet.

**In short.** An 18.97M-parameter head on a frozen pre-trained tabular encoder fills missing cells column by column, in a chain, and a validation gate can always fall back to a MICE-equivalent chain or to the column mean. It leads MICE on a synthetic no-prior split (NRMSE 0.1564 against 0.3983) and on 20 real datasets (0.3097 against 0.3363), and is ahead of a published LLM-imputation study on its nine-cell mean (MAE 0.2058 against 0.228) but behind on MCAR and marginally behind on MAR. Because a MICE-equivalent chain is part of the selection set, the evidence that in-context learning itself helps is still indirect. The paper also catalogues seven defects in the harness and pipeline that each produced a convincing false result. [Download the PDF](https://www.ozercemkelahmet.com/papers/in-context-imputation-head.pdf). For a plain-language walk-through, read [In-context imputation for tabular data](https://www.ozercemkelahmet.com/articles/in-context-imputation-for-tabular-data).

## Abstract

Iterative imputation, in which each column is predicted from the others in turn, remains the strongest classical paradigm, while single-shot reconstruction models shrink toward the column mean. We replace the per-column learner of an iterative scheme by an in-context tabular classifier that conditions on a few hundred observed rows, discretises continuous targets into ten quantile bins and returns the expected value, and we gate the output with a validation layer that can always fall back to a MICE-equivalent chain or to the column mean.

We evaluate this in-context imputation head (ICI head, 18.97M parameters) on three boards. On a synthetic no-prior split it attains NRMSE 0.1564 against 0.3983 for MICE (81 of 81 cells; Wilcoxon p < 10−4). On 20 real datasets (3 mechanisms, 3 rates) it attains 0.3097 against 0.3363 for MICE and is ahead of MICE in all nine mechanism-rate cells, while TabPFN and TabImpute are far behind. On the MAE board of a published LLM-imputation study its nine-cell mean is 0.2058 against 0.228, but this rests entirely on MNAR (0.203 against 0.308): it is behind on MCAR (0.197 against 0.150) and marginally behind on MAR (0.2175 against 0.216).

Because a ridge chain that reproduces MICE is part of the selection set, a result no worse than MICE is guaranteed by construction, and we have no ablation isolating the in-context learner's own contribution. We also document seven defects in the evaluation harness and in our own pipeline, among them a mask generator that ignored the missingness mechanism and a selection step that let the in-context candidate read the true values, each of which produced a convincing but false result.

**Keywords:** missing data, imputation, in-context learning, tabular foundation models, MICE, benchmarking.

## 1. Introduction

For squared-error metrics, imputation quality is determined by one quantity: the correlation r between the imputed and true deviations from the column mean, since NRMSE ∝ σ√(1−r²), an identity verified on measured values to three digits. Iterative schemes such as MICE [8], missForest [7], and HyperImpute [4] raise r by repeatedly refitting a column model on the current completed table. In-context tabular foundation models [2, 3, 6] predict a column from observed rows without refitting, which makes them natural per-column learners in such a loop. TabImpute [1] builds an imputer on that idea.

The paper reports a chained imputer whose column learner is an ICI head, a 18.97M-parameter module on top of a frozen pre-trained tabular encoder, and reports its results on three boards with its failures. There are four contributions:

- the design of a chained in-context imputer with a validation gate that can fall back to classical components;
- results on a synthetic no-prior split, on 20 real datasets, and on the MAE board of a published LLM-imputation study;
- an explicit account of the cells it loses;
- a catalogue of harness and pipeline defects, each of which produced a plausible false conclusion at some point in development.

## 2. Related work

Chained-equation methods [8] and forest-based variants [7] iterate column-wise regressions; HyperImpute [4] adds per-column model selection. Pre-trained transformers are used either entry-wise [1] or as the column learner inside a chain; LimiX [9] reports imputation as one of several tasks. Mangussi et al. [5] compare large language models with classical and deep imputers on real and synthetic tables and show that the language models' advantage depends on world-knowledge priors and disappears on prior-free synthetic data, which is why a prior-free split is included here.

## 3. The model

The ICI head
The in-context imputation head has 18,972,865 parameters. It is a task block placed on top of a frozen pre-trained tabular encoder: the encoder maps each row to a fixed-length representation and is never modified, so the head is trained and evaluated without changing any other component. Given a support set of rows with a known target and a query row, the head predicts a distribution over class labels in a single forward pass, without gradient updates at test time, in the manner of TabPFN [2, 3] and TabICL [6].

Limits
The head conditions on at most 384 support rows, 48 feature columns, and 16 classes (a label embedding of fixed size). Continuous targets are therefore discretised into 10 quantile bins, and categorical columns with more than 16 levels are not supported directly. These limits, not the encoder, determine which columns the chain can model and are the reason a gate with classical fallbacks is needed.

Size
The head is 0.87% of the parameters of the checkpoint that hosts it; the cost of running the chain is dominated by repeated forward passes over the support set.

## 4. Method

Chain
Missing cells are initialised with column means. Columns are visited from fewest to most missing values; for column j the context is the set of rows where j is observed, the features are the current completions of all other columns plus missing-value indicators, and the target is column j. Three sweeps are performed.

In-context column learner
The learner is the ICI head of Section 3. Continuous targets are split into 10 quantile bins; the prediction is the probability-weighted average of the bin centres, the TabPFN bar-distribution device. For squared-error metrics the expected value, not the mode, minimises error; replacing the mode (plus rounding) by the expected value closed 75–84% of the loss on categorical columns of five weak tables. Real columns have priority over missing-value indicators when the 48-feature window is full; the opposite order drove the correlation negative on one table.

Validation gate
A second chain uses BayesianRidge as the column model; it reproduces MICE to numerical precision (maximum difference 2.6×10−15). Both chains run on the table with 15% of observed cells hidden, three times with different hidden sets, and one candidate (in-context chain, ridge chain, a blend, HistGradientBoosting, kNN, or the column mean) is chosen on the matrix as a whole, followed by a per-column shrinkage α ∈ [0, 1] toward the column mean. The column mean is always a candidate, so no column can be worse than the column mean on the validation cells, and the ridge chain's presence guarantees a result no worse than MICE up to selection noise.

Cost
Over 435 fold runs the mean was 240.5 s, the median 160.9 s, and the 95th percentile 733.7 s per run, on a Lambda GPU at $4.29 per hour (29.1 hours, $124.67 in total). Classical imputers take seconds.

## 5. Evaluation protocol

- **Synthetic no-prior split.** Nine synthetic tables without world-knowledge structure, 3 mechanisms × 3 rates, scored by NRMSE; rivals are MICE, missForest, and LimiX-16M.
- **Real-20.** Twenty real datasets from the harness of Mangussi et al., with MCAR, MAR, and MNAR masks at 5, 10, and 20%, 5 folds, scored by NRMSE; rivals are five classical imputers and several deep and foundation-model imputers run locally.
- **Gemini board.** The same 20 datasets scored by MAE against the numbers published for the Gemini model in that study.

The first two boards are computed in one harness with identical masks and folds; the third compares a local run with published numbers produced elsewhere, a weaker comparison. A contamination audit by dataset identifier found none of the 20 real datasets among the 2,205 identifiers of the pre-training corpus; near-duplicate content was not audited.

## 6. Results

6.1 Synthetic no-prior split
On the nine synthetic tables the system has mean NRMSE 0.1564 ± 0.0567. It is first in all 81 cells (27 of 27 for each mechanism). The mean difference to MICE is −0.2419 (95% interval [−0.2558, −0.2273], paired Wilcoxon p < 10−4). As a health check, the error rises with the missing rate as it should (MCAR 0.172, 0.190, 0.229 at 5, 10, 20%).

MethodNRMSE (mean)

ICI chain (ours)0.1564
LimiX-16M0.2258
MICE0.3983
missForest0.4022

Table 1. Synthetic no-prior split, 81 cells. Two published reference cells (synthetic MAR and MCAR with MICE 0.160 and 0.169) could not be reproduced locally (local MICE 0.240 for MAR); no claim is made on those published rows.

6.2 Twenty real datasets
Table 2 and Figure 1 give the real-data results of the 31 July full run. The system is ahead of MICE in all nine mechanism-rate cells (mean 0.3097 against 0.3363, about 8%) and ahead of missForest, kNN, softImpute, and the column mean. In an earlier configuration, with the mask defect fixed but before the chained design, the foundation-model and deep imputers were clearly behind (TabPFN 0.847, TabImpute 0.483, ReMasker 0.422, pMIVAE 0.408, SAEI 0.399 on MCAR in the harness logs), while all five classical imputers were ahead of the earlier system (mean NRMSE 0.422 against MICE 0.335).

MechanismICI chain (5 / 10 / 20%; mean)MICE (5 / 10 / 20%)

MCAR0.1893 / 0.2115 / 0.2354; 0.21210.2158 / 0.2422 / 0.2716
MAR0.2215 / 0.2354 / 0.2540; 0.23700.2386 / 0.2593 / 0.2883
MNAR0.4849 / 0.4908 / 0.4643; 0.48000.5044 / 0.5059 / 0.5002

Table 2. Real-20, NRMSE by mechanism and rate. Overall mean 0.3097 (MICE 0.3363, missForest 0.3474, kNN 0.3536, softImpute 0.3552, column mean 0.3725).

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.
Under MNAR the error falls as the missing rate rises (0.485 to 0.464). This is not a property of the model: the harness' thresholded MNAR mask changes the masked population with the rate, and MICE, kNN, and the column mean also improve (MICE 0.504 to 0.500). The ranking is unaffected.

6.3 Comparison with deep and foundation-model imputers
Table 3 compares all methods run in the real-20 harness. The deep and foundation-model values were recorded in the earlier configuration of the harness (mask defect fixed, before the chained design) for the MCAR mechanism, where the earlier single-shot head's mean was 0.3392; the last row gives the final chained system's MCAR mean (0.2121, from Table 2). Foundation models used zero-shot are the weakest group: TabPFN, applied column by column, errs by more than a factor of two, and TabImpute is behind the chain by a wide margin.

MethodClassNRMSE, MCAR

ICI chain (final)chained in-context0.2121
Earlier single-shot headreconstruction0.3392
SAEIdeep0.3987
pMIVAEdeep0.4077
ReMaskerdeep0.4220
HyperImpute (meta)ensemble of learners0.4610
HyperImpute (Sinkhorn)optimal transport0.4783
TabImputefoundation model0.4828
HyperImpute (EM)statistical0.5292
HyperImpute (MIWAE)deep0.5456
GAINdeep0.7563
TabPFN (column-wise)foundation model0.8472
HyperImpute (MIRACLE)deep2.8637

Table 3. Real-20, NRMSE for MCAR (lower is better). Rival values come from the earlier harness configuration; classical imputers are compared in Table 2.

6.4 The Gemini board
Table 4 and Figure 2 give the MAE board. The nine-cell mean is 0.2058 against the published 0.228, but this is carried by MNAR, where the margin is 0.105. The system is behind in the MCAR cell (0.197 against 0.150 for the cell target taken from the published study) and by 0.0015 on MAR. Forty rounds of a configuration search and six further configurations for MCAR did not close the gap; the tuning space was exhausted without reaching the published value. The system beats the classical imputers in 15 of 15 comparisons on this board, but the comparison with Gemini mixes a local run with a published number from another harness, so the headline is treated cautiously.

CellICI chain (5 / 10 / 20%)Cell meanGemini

MCAR0.1749 / 0.2005 / 0.21470.19670.150
MAR0.2001 / 0.2178 / 0.23450.21750.216
MNAR0.1793 / 0.2043 / 0.22620.20330.308
Nine-cell mean0.20580.228

Table 4. MAE board on the 20 real datasets.

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.

6.5 Signal of the in-context learner alone
With the validation gate switched off, the chained in-context learner raised the correlation r from 0.116 to 0.308 on one table (pima) and from 0.432 to 0.534 on another (wisconsin), against 0.94 for MICE on the latter. Single-column quality was comparable to MICE (pima: 0.464 against 0.467). On densely correlated tables the gate selects the ridge chain: on wisconsin, validation errors were 0.057 (ridge), 0.074 (blend), and 0.113 (in-context). The in-context learner therefore contributes where the ridge chain is weak, but the full ablation (ridge chain alone against the full system) that would quantify how much of the real-20 gain over MICE comes from it was not run.

## 7. Evaluation defects

Each defect below produced a plausible false result before it was found.

DefectEffect and evidence

Mask generator ignored the mechanismEvery run used the MNAR masker, the hardest case; MCAR and MAR looked bad and nearly identical (first-fold NRMSE 0.335, 0.350, 0.352; after the fix 0.304, 0.629, 0.382). Evidence: the column-mean baseline improved 17.5% as missingness grew.
Weak engine called by the harnessA single-shot reconstruction model of NRMSE about 0.62 was used; it also crashed on a missing-key error after a later code change. Switching engines alone looked harmful because the mask was fixed at the same time; the effects had to be separated.
Warm start silently offA strict state-dict load failed on a head added later, so the pre-trained variant was never tried. It wins on some tables (0.182 against 0.222) and loses on others.
Leaky A/B testA model trained on the complete table and asked to fill the masked copy gave values of 0.16–0.21 on one table (pima); the harness, which trains on masked data, gave 0.30 on the same table.
Candidate read the answerIn the gate, the in-context candidate's hold-out values were the true observed values because the chain had not masked those cells: it won every comparison.
Rival scored as NaNThe ridge chain used the original mask, not the hold-out mask, so its hold-out score was NaN and it was eliminated; the in-context candidate appeared to win 30 of 30 columns. This also invalidates earlier per-column win counts.
Mode instead of meanSquared error is minimised by the expected value; taking the mode plus rounding cost 75–84% of the categorical loss on five tables.

Table 5. Defects found in the harness and in the pipeline.

Under the defective harness, the change in error from 5% to 20% missing was +67.8% for the published MICE figure but −0.8% for local MICE, +0.6% for missForest, −7.3% for kNN, and −17.5% for the column mean. Local methods did not degrade with more missing data, which is physically implausible. A claim made before the fix (123 of 180 cells first, mean 0.3395 against 0.4127 for the best deep imputer) was an artefact of the broken mask and has been withdrawn. Two further lessons: validation scores on a mini-benchmark (NRMSE) did not predict the board metric (MAE), and a result averaged over mechanisms can hide a cell-level loss, as in Table 4.

## 8. Limitations

- The ridge chain in the selection set guarantees no worse than MICE, so "ahead of MICE" conflates the in-context learner with the gate; no ablation isolates it.
- The 0.3097 real-20 table is the 31 July overnight run; the categorical correction (expected value) was evaluated on paired masks and a re-run was launched, but its table is not in the record used for this draft.
- The Gemini comparison uses published numbers from a different harness.
- The deep and foundation-model rival values quoted for the earlier configuration come from the harness logs and were computed for the MCAR mechanism.
- The learner is limited to 384 support rows, 48 features, and 16 label classes.
- Contamination was audited by dataset identifier only.
- The per-mechanism configurations of the Gemini-board run (an MNAR offset found by a gate-seeker, a longer sweep, and finer bins for MCAR) were chosen on the evaluation datasets themselves, so those numbers include configuration selection.
- Weights and code are not released.

## 9. Conclusion

A chained in-context learner with a validation gate gives the lowest error among the tested imputers on a synthetic no-prior split and on 20 real datasets, and the best nine-cell mean on a published MAE board, but the last result rests on one mechanism and two of three cells are lost. Because the gate contains a MICE-equivalent chain, the evidence that in-context learning itself improves imputation is still indirect; the ablation that would settle it is the first experiment planned.

## References

- Feitelberg, Saha, Choi, Ahmad, Agarwal, Dwivedi. TabImpute: Accurate and fast zero-shot missing-data imputation with a pre-trained transformer. [arXiv:2510.02625](https://arxiv.org/abs/2510.02625), 2025.
- Hollmann, Müller, Eggensperger, Hutter. TabPFN: A transformer that solves small tabular classification problems in a second. ICLR, 2023.
- Hollmann, Müller, Purucker, Krishnakumar, Körfer, Hoo, Schirrmeister, Hutter. Accurate predictions on small data with a tabular foundation model. Nature, 637, 2025.
- Jarrett, Cebere, Liu, Curth, van der Schaar. HyperImpute: Generalized iterative imputation with automatic model selection. ICML, 2022.
- Mangussi, Pereira, Lorena, Abreu. Large language models for missing data imputation: Understanding behavior, hallucination effects, and control mechanisms. [arXiv:2603.22332](https://arxiv.org/abs/2603.22332), 2026.
- Qu, Holzmüller, Varoquaux, Le Morvan. TabICL: A tabular foundation model for in-context learning on large data. ICML, 2025.
- Stekhoven, Bühlmann. MissForest: non-parametric missing value imputation for mixed-type data. Bioinformatics, 28(1):112–118, 2012.
- van Buuren, Groothuis-Oudshoorn. mice: Multivariate imputation by chained equations in R. Journal of Statistical Software, 45(3):1–67, 2011.
- Zhang et al. LimiX: Unleashing structured-data modeling capability for generalist intelligence. [arXiv:2509.03505](https://arxiv.org/abs/2509.03505), 2025.

## How to cite

Kelahmet, Ö. C. An In-Context Imputation Head for Tabular Data: Chained, Validation-Gated, and Evaluated on Three Imputation Boards. Draft of 4 November 2026. [PDF](https://www.ozercemkelahmet.com/papers/in-context-imputation-head.pdf).

## Frequently asked questions

What is an in-context imputation head?
A small task block of 18,972,865 parameters on top of a frozen pre-trained tabular encoder. Given a support set of rows with a known target and a query row, it predicts a distribution over class labels in one forward pass, with no gradient updates at test time.

How does the chained imputer work?
Missing cells start at column means. Columns are visited from fewest to most missing values, and each column is predicted from the current completions of the others plus missing-value indicators. Three sweeps are performed.

Does it beat MICE?
On the synthetic no-prior split and on 20 real datasets it has lower NRMSE than MICE in every reported cell. A MICE-equivalent chain is part of the selection set, so a result no worse than MICE is guaranteed by construction, and no ablation isolates the in-context learner's own contribution.

Where does it lose?
On the Gemini MAE board it is behind on MCAR (0.197 against 0.150) and marginally behind on MAR (0.2175 against 0.216). Its nine-cell mean advantage comes entirely from MNAR.

Why does the paper list its own defects?
Each of seven harness and pipeline defects produced a convincing false result at some point, including a mask generator that ignored the missingness mechanism and a selection step that let the in-context candidate read the true values. Listing them shows how the final numbers were checked.

Are the weights and code available?
No. Weights and code are not released.