Missing Data Imputation by Reducing Mutual Information with Rectified Flows

TL;DR

This work tackles missing data imputation by reframing it as mutual information minimization between the imputed data and the missingness mask. It introduces Mutual Information Reducing Iterations (MIRI), an iterative framework that minimizes a KL-based objective between the imputed-data–mask joint and the previous iteration's product distribution, guaranteeing non-increasing mutual information. The optimal imputer is realized by solving an ODE whose velocity field is trained with a rectified-flow objective, tightly integrating flow-based generative modeling into imputation. Empirically, MIRI achieves strong distributional fidelity and competitive or superior performance on synthetic, tabular, and image benchmarks, with a public codebase for reproducibility.

Abstract

This paper introduces a novel iterative method for missing data imputation that sequentially reduces the mutual information between data and the corresponding missingness mask. Inspired by GAN-based approaches that train generators to decrease the predictability of missingness patterns, our method explicitly targets this reduction in mutual information. Specifically, our algorithm iteratively minimizes the KL divergence between the joint distribution of the imputed data and missingness mask, and the product of their marginals from the previous iteration. We show that the optimal imputation under this framework can be achieved by solving an ODE whose velocity field minimizes a rectified flow training objective. We further illustrate that some existing imputation techniques can be interpreted as approximate special cases of our mutual-information-reducing framework. Comprehensive experiments on synthetic and real-world datasets validate the efficacy of our proposed approach, demonstrating its superior imputation performance. Our implementation is available at https://github.com/yujhml/MIRI-Imputation.

Figures (8)

• Figure 1: Pairwise density plots of the imputed data and MMD/MI of MIRI.
• Figure 2: MMD on UCI datasets with 60% data missing. The lower the better.
• Figure 3: Comparison of imputed samples on CIFAR-10 and CelebA with 60% of missingness.
• Figure 4: Schematic of the sequential imputation algorithm. The process begins with an initial imputation, $\mathbf{X}^{(0)}$, and iteratively refines the missing values over $T$ steps.
• Figure 5: MCAR MMD on 10 UCI datasets (Above: 20% missingness, Middle: 40% missingness, Below: 60 % missingness). The lower the better.

Theorems & Definitions (12)

• Proposition 1
• Proposition 2
• Theorem 1
• proof
• proof
• Lemma 1
• Definition 1
• Lemma 2
• proof
• Lemma 3