PaperlayerView exact revision on arXiv

Source-linked AI summary

Integrated Imputation-Classification for Supervised Learning with Missing Data

Yue Liu, Ben Liang, Ali Tizghadam, Ilijc Albanese

arXiv:2610.04273v1cs.LG

TL;DR

Missing features challenge supervised classification, where prediction should account for uncertainty in unobserved coordinates. IICN jointly trains an imputer and discriminator adversarially, and the paper reports strong performance across missingness settings alongside a Bayes-optimality result under MCAR.

  • Problem

    Missing features challenge supervised classification, which typically assumes fully observed inputs, and prediction under missingness must account for uncertainty in the missing coordinates.

    Source: §1 · pp. 1–2

  • Method

    IICN jointly trains an imputer and an (n+1)-class discriminator with a supervised adversarial objective that assigns imputed inputs an auxiliary class.

    Source: §1 · p. 2, §2.1.3 · p. 4

  • Results

    Across synthetic, image, and tabular settings, IICN outperforms impute-then-classify and other joint or generative baselines across missingness levels, including severe missingness.

    Source: §1 · p. 3, §4 · pp. 8–9

  • Takeaways & Limitations

    The reported results highlight the importance of task-aligned imputation for robust classification under missing data.

    Source: §4 · p. 9

    The Bayes-optimality argument is stated under the MCAR assumption.

    Source: §2.2.1 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Missing features complicate supervised classification, and imputation quality alone may not align with prediction. IICN jointly trains imputation and classification adversarially, with a Bayes-optimality guarantee and strong reported performance under missingness.

1 Introduction

  • The task is to classify complete and incomplete test data using training data containing both, while minimizing expected classification risk over uncertainty in missing features.The missing-feature objective motivates prediction-aligned learning rather than judging imputations only by reconstruction quality.Source: §1 · pp. 1–2
  • Impute-then-classify can be computationally challenging as missingness patterns grow with dimension, and improved imputations may bring only marginal classification gains or induce ambiguity or bias.Separating imputation from prediction also prevents direct optimization of the classification objective and complicates characterization of the resulting decision rule.Source: §1 · pp. 1–2
  • Existing alternatives include explicit search, marginalization-based predictors, and joint imputer-classifier training, each with stated computational, distributional, or theoretical limitations.Search can become infeasible as patterns grow; earlier marginalization methods have restricted assumptions, while joint methods leave gaps in guarantees or depend on generative-model fidelity.Source: §1 · p. 2
  • IICN is a unified, class-aware adversarial framework that couples imputation and classification in a single (n+1)-class objective.Its formulation is likelihood-free and expressive enough to capture the conditional distribution without parametric or variational restrictions.Source: §1 · p. 2
  • IICN’s imputer and discriminator implement Bayes-optimal marginalization at the adversarial equilibrium, with the imputer supplying imputations and the discriminator the class posterior.The framework uses an auxiliary class for imputed inputs and does not require a separate classifier.Source: §1 · pp. 2–3
  • IICN outperforms classical impute-then-classify pipelines and recent generative baselines on FashionMNIST, CIFAR-10, and tabular data, with robustness under severe missingness.The paper presents this as empirical validation of its theoretical prediction.Source: §1 · p. 3

2 IICN: Integrated Imputation and Classification Network

IICN jointly learns an imputer and discriminator through an (n+1)-class adversarial objective. Under MCAR, its optimal components align imputations with complete data and compose into a Bayes-optimal predictor.

2.1.1 Generator

  • The generator takes an incomplete sample, its mask, and latent noise, then fills missing coordinates while observed entries are retained in the completed sample.The completed sample is passed to the discriminator.Source: §2.1.1 · p. 3, Figure 1 · p. 3

2.1.2 Multi-class Discriminators

  • The discriminator adds an imputed class to the n semantic classes, allowing it to represent both class identity and completeness.Complete samples retain their semantic labels; incomplete, imputed samples receive the auxiliary label.Source: §2.1.2 · p. 3

2.1.3 Supervised Adversarial Loss

  • The supervised adversarial loss trains the discriminator to classify complete samples and detect imputations, while the generator tries to make imputations receive semantic labels.The objective extends the original label set with an auxiliary (n+1)th label.Source: §2.1.3 · p. 4

2.2.1 Optimal Discriminator

  • For a fixed generator, the optimal discriminator’s real-class outputs recover class posteriors on complete data under MCAR.Its real-class probabilities initially represent joint probabilities of completeness and class membership, then normalize to the complete-data posterior.Source: §2.2.1 · p. 4

2.2.2 Optimal Generator

  • Under MCAR, the optimal generator makes completed samples independent of completeness and distributionally aligned with complete data.The alignment is expressed as equality between the completed-sample and complete-data distributions.Source: §2.2.2 · p. 5

2.2.3 Predictive Behaviour at Optimality

  • At the optimum, G* supplies p(xmiss | xobs) and D* the complete-data posterior; averaging D* over imputations marginalizes missing coordinates to Bayes-optimal prediction.This composition does not require training a separate classifier.Source: §2.2 · p. 4, §2.2.3 · pp. 5–6

2.3 IICN Algorithm

  • Two-phase discriminator updates alternate adversarial detection with true-label classification, countering collapse into an imputation detector; generator updates target incomplete samples while preserving observed entries.The schedule is described as discouraging artifact-based shortcuts and mitigating early discriminator domination.Source: §2.3 · pp. 6–7

3 Experiments

IICN is evaluated against classical imputation, generative, and joint-optimization baselines on controlled synthetic and image settings and naturally missing tabular data. Across these settings, it approaches the synthetic Bayes bound, mostly leads on incomplete images, and narrows the complete–imputed accuracy gap on UCI data.

3 Experiments

  • Baselines span simple zero and mean fills, downstream-task-independent GAIN, joint methods SupMIWAE and LAC, and complete-only training.SupMIWAE uses importance-weighted variational inference, LAC alternates a weighted prediction and reconstruction loss, and complete-only training discards incomplete datapoints.Source: §3 · p. 7

3.1 Experimental Setup

  • The experiments cover a controlled Burger data-generating process, artificially incomplete FashionMNIST and CIFAR-10, and UCI tabular datasets with natural missingness.The UCI datasets are Mammographic Mass, Diabetes, Hepatitis, and Thyroid.Source: §3.1 · p. 7

3.2 Numerical Results

  • The Burger experiment makes the second feature missing at random, leaving a multimodal conditional distribution for which deterministic single-point imputations collapse onto one mode.The Bayes-optimal incomplete-data classifier has P(Y = 1 | x1) = 2/3.Source: §3.2 · p. 8
  • At a 0.9 missing rate on Burger, IICN spans both modes and yields decision boundaries closest to the Bayes-optimal classifier.Zero and mean imputation collapse values to fixed points, GAIN has ambiguity near the class boundary, SupMIWAE has smoother regions but remains sensitive to missing-value uncertainty, and LAC concentrates on one mode.Source: §3.2 · p. 8
  • Across missing rates, IICN most consistently matches the 2/3 Bayes bound on incomplete Burger datapoints through marginalization, while single-imputation methods collapse toward a constant majority-class predictor.The single-imputation methods' rise to 67% incomplete-data accuracy coincides with the marginal class probability and accompanies a drop in complete-data accuracy at a 0.9 missing rate.Source: §3.2 · p. 8
  • IICN mostly achieves the highest accuracy on incomplete FashionMNIST and CIFAR-10 datapoints, with a larger performance gap as missingness increases.The evaluation controls the proportion of complete datapoints and independently masks incomplete datapoints at varying probabilities.Source: §3.2 · p. 8
  • On UCI datasets with naturally occurring missing values, IICN narrows the accuracy gap between complete and imputed inputs and can outperform complete-only training.The underlying missingness mechanism is unknown.Source: §3.2 · p. 8

4 Conclusion

The paper concludes that IICN jointly aligns imputation with downstream classification and reports gains across synthetic, image, and tabular evaluations. The supplied visuals identify the datasets and accuracy comparisons but do not include table cell values.

4 Conclusion

  • Figure 2 presents Burger data distributions, decision boundaries, and sample-level predictions at 0.9 missingness, plus accuracy on complete and incomplete points across five missing rates.Its dashed line marks Bayes-optimal accuracy on incomplete data.Source: Figure 2 · p. 9
  • Experiments on synthetic, image, and tabular datasets show IICN consistently outperforming impute-then-classify and other joint optimization baselines across a wide range of missingness levels.The conclusion identifies task-aligned imputation as important for robust classification under missing data.Source: §4 · pp. 8–9
  • Tables 1 and 2 report test accuracy for complete and incomplete FashionMNIST and CIFAR-10 images across missing probabilities.Source: Table 1 · p. 10, Table 2 · p. 10
  • Table 3 reports test accuracy on complete and imputed data across UCI datasets, using cross-validation ensembles.Source: Table 3 · p. 10

A Related work

Prior work spans staged imputation and classification, task-aligned joint training, generative imputation, semi-supervised adversarial models, and theory for prediction under missingness. IICN is positioned as a class-aware adversarial approach with a Bayes-optimality guarantee under MCAR.

Classification under missing data

  • Classical pipelines separate filling missing values from classifier training, while deep generative imputers can model complex feature distributions without ensuring that imputations aid classification.Joint approaches including GACN, SupMIWAE, and LAC address the mismatch; IICN adds a likelihood-free class-aware adversarial formulation [Ghanavi et al., 2021; Ipsen et al., 2022; Sugiyama & Uchida, 2025a,b].Source: A Related work · p. 13

Generative imputation models

  • Adversarial imputation methods such as GAIN and MisGAN primarily optimize distributional reconstruction, without considering downstream tasks or directly optimizing the downstream classifier’s decision boundary.IICN combines adversarial imputation with task-aware discrimination in a single (n + 1)-class objective; its extra class captures imputation artifacts.Source: A Related work · p. 13

Semi-supervised GANs

  • IICN adapts the semi-supervised GAN output structure so its extra class marks incomplete or improperly imputed inputs, while its generator completes features from observed values [Springenberg, 2015].This differs from using a generator to synthesize full samples from noise.Source: A Related work · p. 13

Theoretical analysis of classification under missingness

  • Prior theory characterizes the Bayes predictor as marginalization over missing-feature conditionals, while IICN claims exact recovery at adversarial equilibrium under MCAR without distributional assumptions or variational approximations [Le Morvan et al., 2021; Sugiyama & Uchida, 2025a,b].The cited LAC guarantees concern empirical loss rather than distance from the Bayes-optimal classifier.Source: A Related work · p. 13

B Proofs

The appendix establishes an optimization fact used in the proofs: a weighted sum of logarithms under a simplex constraint is maximized when each variable equals its weight. Its proof uses a Lagrangian stationary point and concavity.

B.1 Technical tools

  • For nonnegative weights a_i summing to one, the constrained objective sum_i a_i log x_i is maximized at x_i = a_i.Source: §B.1 · p. 14
  • The proof introduces a Lagrange multiplier and uses the equality constraint to identify the stationary point coordinatewise.Source: §B.1 · p. 14
  • Nonpositive second derivatives establish a maximum, and the maximum is globally and uniquely attained when all a_i are positive.Source: §B.1 · p. 14

B.2 Proof of Thm. 2.1

The proof maximizes an objective over the class discriminators and characterizes its unique optimum.

B.2 Proof of Thm. 2.1

  • The proof formulates a discriminator-only maximum with terms for true-class and imputed inputs.The displayed objective includes log-probability terms for both types of inputs.Source: §B.2 · p. 14
  • The proof identifies a unique maximizing solution for the class discriminators.The stated optimum relates the class terms to P(Y = i, R = 1 | ˆx).Source: §B.2 · p. 14

B.3 Proof of Thm. 2.2

The proof separates the generator objective into distributional and entropy terms, then uses its minimization to establish distribution alignment.

B.3 Proof of Thm. 2.2

  • The generator objective contains KL-divergence terms comparing the imputed distributions for R = 1 and R = 0 with the overall imputed distribution.The decomposition also includes class-conditional entropy and the binary entropy of R.Source: §B.3 · p. 15
  • The entropy terms are constant with respect to G, so the global minimum occurs when the imputed distributions for R = 1, R = 0, and the overall population align.The proof invokes the optimality criterion of KL-divergence.Source: §B.3 · p. 15

B.4 Proof of Thm. 2.5

The proof composes the optimal discriminator and generator: the generator supplies conditional imputations, and averaging the discriminator output yields the Bayes-optimal classifier under missingness.

B.4 Proof of Thm. 2.5

  • The imputation preserves observed coordinates exactly while replacing only the missing coordinates.The proof constructs the completed input by retaining observed entries and filling missing entries with G⋆’s output.Source: §B.4 · p. 16
  • The optimal generator’s imputed missing coordinates follow the true conditional distribution given the observed coordinates.Sampling z from pZ and transforming it with G⋆ produces samples from that conditional distribution.Source: §B.4 · p. 16
  • Taking the expectation of the complete-data class posterior over G⋆’s imputation distribution marginalizes missing coordinates and gives the Bayes-optimal classifier under missingness.The proof combines the optimality results for D⋆ and G⋆.Source: §B.4 · pp. 15–16

C Further Experimental details

The experimental details specify compute and optimizer choices, model architecture scope, and Monte Carlo marginalization at test time.

C Further Experimental details

  • Experiments used one NVIDIA H100 80GB GPU per job; Adam was used on real-world datasets and SGD on synthetic data.SGD was chosen for synthetic experiments to reduce overfitting.Source: C Further Experimental details · p. 17

C.1 Model archtecture

  • Table 4 specifies the generator architecture, with layer 1 receiving 2 input channels for FashionMNIST and 6 for CIFAR-10.The architecture table identifies dataset-specific input channel counts.Source: Table 4 · p. 17

C.2 Inference via Monte Carlo marginalization.

  • At test time, the method approximates marginalization by Monte Carlo averaging over K imputations.The passage identifies this procedure as an approximation to Thm. 2.5.Source: §C.2 · p. 17

C.3 Image reconstruction

Tables 6 and 7 compare qualitative reconstructions and missing-pixel MSE on FashionMNIST and CIFAR-10 across observation and missingness settings. The comparisons show that high reconstruction quality does not necessarily mean better classification performance.

C.3 Image reconstruction

  • High reconstruction quality does not necessarily correspond to superior classification performance.The tables compare reconstructions and MSE, but the section cautions against treating reconstruction quality as a proxy for classification performance.Source: §C.3 · p. 18
  • Tables 6 and 7 show representative reconstructions and corresponding MSE values across different observation and missingness settings.For each setting, MSE is reported alongside image blocks.Source: Table 6 · p. 18, Table 7 · p. 19
  • MSE is calculated per missing pixel, and lower values generally correlate with visually improved reconstruction quality.The captions define MSE as mean squared reconstruction error averaged over missing pixels only.Source: Table 6 · p. 18, §C.3 · p. 18, Table 7 · p. 19

Related papers

Loading 2610.04273v1…