Integrated Imputation-Classification for Supervised Learning with Missing Data

Yue Liu; Ben Liang; Ali Tizghadam; Ilijc Albanese

cs.LGOct 3, 2026

Integrated Imputation-Classification for Supervised Learning with Missing Data

Authors: Yue Liu, Ben Liang, Ali Tizghadam, Ilijc Albanese

Organizations: University of Toronto · TELUS

Abstract

We study supervised classification problems with missing feature values. Existing approaches often decouple imputation from classification, producing imputations that may be plausible but uninformative for prediction. Instead, we propose the Integrated Imputation and Classification Network (IICN), which jointly trains an imputer and an ${(n{+}1)}$ -classdiscriminator adversarially with a single class supervised classification objective, where the discriminator learns to distinguish among the $n$ true classes and an additional ``imputed" class. We prove that at the global optimum, the imputer and discriminator together implement marginalization over missing coordinates and yield a Bayes-optimal classifier. We evaluate IICN on FashionMNIST, CIFAR-10, and tabular datasets with naturally occurring missingness. IICN outperforms classical impute-then-classify pipelines and recent generative baselines, showing strong robustness and accuracy in challenging settings.

Figures & tables

Figure 1: Schematic of the IICN network. The completed sample $\hat{\boldsymbol{x}}$ is made up of the output of the generator and the observed parts of an incomplete sample $\tilde{\boldsymbol{x}}$ . The completed sample $\hat{\boldsymbol{x}}$ is passed to a multi-class discriminator. The classification loss is applied to update both the generator and the discriminator.

Panel of Figure 2 — Figure 2: Results on the synthetic Burger dataset. Top: Data distribution, decision boundaries, and sample-level predictions for different imputation methods for an example with missing rate 0.9. The dataset is a two-class problem with $Y=0$ (blue) and $Y=1$ (orange), while black points indicate misclassified samples. Observed values of incomplete samples are shown as colored sticks at the bottom. Bottom: Classification accuracy on fully observed (left) and incomplete (right) datapoints under varying missing rates $\in\{0.1,0.3,0.5,0.7,0.9\}$ . The dashed line indicates the Bayes-optimal accuracy on incomplete data.

Missing Rate		Incomplete data						Fully observed data
$P_{\mathrm{obs}}$	$P_{\mathrm{miss}}$	Zero	Mean	GAIN	SupMIWAE	LAC	IICN	Zero	Mean	GAIN	SupMIWAE	LAC	IICN	CO †
$0.1$	$0.5$	$83.10\scriptstyle\pm.70$	$85.69\scriptstyle\pm.70$	$86.40\scriptstyle\pm.90$	$86.81\scriptstyle\pm.70$	$85.88\scriptstyle\pm.70$	$\textbf{88.01}\scriptstyle\pm.60$	$86.97\scriptstyle\pm.70$	$87.27\scriptstyle\pm.68$	$87.38\scriptstyle\pm.80$	$89.20\scriptstyle\pm 1.90$	$87.01\scriptstyle\pm.70$	$\textbf{89.13}\scriptstyle\pm.60$	$86.38\scriptstyle\pm.70$
$0.1$	$0.8$	$78.45\scriptstyle\pm.80$	$79.78\scriptstyle\pm.80$	$81.18\scriptstyle\pm.80$	$83.40\scriptstyle\pm.80$	$81.64\scriptstyle\pm.80$	$\textbf{83.06}\scriptstyle\pm.70$	$84.64\scriptstyle\pm.70$	$84.24\scriptstyle\pm.70$	$86.09\scriptstyle\pm.70$	$88.30\scriptstyle\pm 2.00$	$84.46\scriptstyle\pm.70$	$\textbf{86.06}\scriptstyle\pm.70$	$86.38\scriptstyle\pm.70$
$0.5$	$0.5$	$83.00\scriptstyle\pm.70$	$83.89\scriptstyle\pm.70$	$86.58\scriptstyle\pm.70$	$88.16\scriptstyle\pm.90$	$84.69\scriptstyle\pm.70$	$\textbf{88.59}\scriptstyle\pm.60$	$88.22\scriptstyle\pm.60$	$88.26\scriptstyle\pm.60$	$89.17\scriptstyle\pm.60$	$87.86\scriptstyle\pm.90$	$88.02\scriptstyle\pm.60$	$\textbf{90.01}\scriptstyle\pm.60$	$88.90\scriptstyle\pm.60$
$0.5$	$0.8$	$77.11\scriptstyle\pm.80$	$78.32\scriptstyle\pm.80$	$80.38\scriptstyle\pm.80$	$83.56\scriptstyle\pm 1.00$	$79.96\scriptstyle\pm.80$	$\textbf{83.69}\scriptstyle\pm.70$	$87.58\scriptstyle\pm.61$	$88.02\scriptstyle\pm.60$	$88.01\scriptstyle\pm.60$	$88.14\scriptstyle\pm.90$	$87.01\scriptstyle\pm.70$	$\textbf{89.23}\scriptstyle\pm.60$	$88.90\scriptstyle\pm.60$
$0.9$	$0.5$	$80.11\scriptstyle\pm.80$	$81.06\scriptstyle\pm.80$	$85.12\scriptstyle\pm.70$	$\textbf{88.10}\scriptstyle\pm 2.00$	$86.68\scriptstyle\pm.70$	${87.04}\scriptstyle\pm.70$	$88.38\scriptstyle\pm.60$	$87.79\scriptstyle\pm.60$	$89.29\scriptstyle\pm.60$	$88.32\scriptstyle\pm.70$	$89.41\scriptstyle\pm.60$	$\textbf{89.42}\scriptstyle\pm.60$

Table 1: Table shows the test accuracy (

\%

) of different methods on complete images and incomplete images with different missing probabilities on FashionMNIST.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Generator (shared; $n_{\text{ch}}=1$ for FashionMNIST, $n_{\text{ch}}=3$ for CIFAR-10)
Layer	Operation	Kernel / Stride / Padding	Out Ch.
1	Conv2D	$3\times 3\ /\ 1\ /\ 1$	128
	BatchNorm	–	128
	ReLU	–	128
2	Conv2D	$3\times 3\ /\ 1\ /\ 1$	64
	BatchNorm	–	64

Appendix

Table 4: Generator architecture. Input to layer 1 has

2n_{\text{ch}}

channels (FashionMNIST: 2, CIFAR-10: 6).

Panel of Table 6 — Table 6: Qualitative reconstruction results and corresponding MSE values on FashionMNIST under different observation and missing probabilities. MSE denotes the mean squared reconstruction error averaged over missing pixels only. For each setting ( $P_{\mathrm{obs}},P_{\mathrm{miss}}$ ), we visualize representative samples together with their mean squared error (MSE), reported on the right of each image block. Lower MSE values generally correlate with visually improved reconstruction quality. This figure complements the main paper by providing additional qualitative comparisons.

Panel of Table 7 — Table 7: Qualitative reconstruction results and corresponding MSE values on CIFAR-10 under different observation and missing probabilities. MSE denotes the mean squared reconstruction error averaged over missing pixels only. For each setting ( $P_{\mathrm{obs}},P_{\mathrm{miss}}$ ), we visualize representative samples together with their mean squared error (MSE), reported on the right of each image block. Lower MSE values generally correlate with visually improved reconstruction quality. This figure complements the main paper by providing additional qualitative comparisons.

Explore similar work

When Does Task-Aware Imputation Matter? An Analysis of Robustness Under Train-Test Mismatch
Oct 14, 2026Maria V. Salvador, Ana F. Sequeira, Jaime S. CardosoIncomplete Data Imputation Learning with Missing Data
INESC TEC University of Porto
Learning What Not to Impute: An Uncertainty-Aware Diffusion Framework for Meaningful Missingness
Jun 3, 2026Lixing Zhang, Yidong Ouyang, Weifu Li +3Incomplete Data Imputation Learning with Missing Data
University of Minnesota · University of California, Los Angeles · Carnegie Mellon University
FILLER: Feature Imputation via Latent Location Exploration and Retrieval
Jul 25, 2026Santu Mondal, Chayan Maitra, Rajat K. DeIncomplete Data Imputation Learning with Missing Data
Machine Intelligence Unit Indian Statistical Institute 203 Barrackpore Trunk Road Kolkata 700108

Oct 14, 2026J/K move · Enter open · S save

When Does Task-Aware Imputation Matter? An Analysis of Robustness Under Train-Test Mismatch

Maria V. Salvador, Ana F. Sequeira, Jaime S. Cardoso

INESC TEC University of Porto

Incomplete Data Imputation Learning with Missing Data Tabular ML

Open paper AskPDF arXiv

Missing data is a pervasive challenge in machine learning, often degrading predictive performance and complicating model deployment. Most existing approaches treat imputation as a preprocessing step independent of the downstream task, a separation that becomes problematic under train-test mismatches in missingness mechanisms or rates. We introduce a unified framework that treats missing entries as variables in a task-conditioned optimization problem. We instantiate this framework in two settings: an amortized approach for scenarios with missing data at both training and test time, and a latent optimization approach for settings where only training data is incomplete. The evaluation is conducted on 98 tabular datasets, and robustness under train-test mismatches is assessed on 36 datasets. Results reveal that reconstruction error constitutes a poor proxy for predictive utility (Spearman rho approximately 0.16 to 0.32), and that task-aware imputation becomes critical in scenarios where classical methods fail to recover complete-data performance, achieving win rates of up to 80% in challenging train-test mismatch scenarios.

Abstract

Explore similar work

When Does Task-Aware Imputation Matter? An Analysis of Robustness Under Train-Test Mismatch

Learning What Not to Impute: An Uncertainty-Aware Diffusion Framework for Meaningful Missingness

FILLER: Feature Imputation via Latent Location Exploration and Retrieval