Title: Missing Data Imputation by Reducing Mutual Information with Rectified Flows
Venue: NeurIPS 2025
Accepted: yes
Keywords: missing data

imputation

rectified flow

generative model
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 \url{https://github.com/yujhml/MIRI-Imputation}.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduced a missing data imputation method that follows along findings around GAIN's strengths, i.e., by enforcing the orthogonality between missing data and data missingness patterns. Specifically, the authors introduces MIRI that is built on top of rectified flows where the objectives can be formulated as ODEs. Experiments were ran on three data set with missing-completely-at-random (MCAR) missingness setup.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n* Motivation and methodology are clearly written.\n* Theoretical analyses are provided.\n\n## Weaknesses\n\n* Unclarity of the problem setup and the actual use-cases/assumptions of MIRI. The approach sits on the basis that imputed data can be orthogonal to the missingness patterns, which is basically the MCAR setting. However, it was not mentioned specifically if MIRI can only work work MCAR or it can be extended to MAR which covers many real-world use cases in healthcare data and other time-series extensively. Also looks like all the experiments are ran based on the MCAR setup.\n* Moreover, imputation for time-series [see 1-14 below as a non-exhaustive list] is an important aspect of missing data imputation and has been extensively studied. While it does not seem the author claimed if MIRI would work for time-series imputation, the experiments used datasets that are not time-seires. \n* Missing discussion and comparison to related works. Recently there are a few works that emphasized if perfect imputation are needed for  downstream tasks [13-16], i.e., [15, 16] argued that mean imputation could be sufficient for downstream tasks, [13] showed that missingness patterns themselves could provide useful information for downstream tasks, with both [13] and [14] introduced methods that can directly achieve downstream tasks without the need for imputation. \n* Missing experimental setup. A majority of existing imputation approaches (i.e., listed above) also focused on how their imputed data could benefit the performance of downstream classification/regression tasks, as part of experiments. Such experiments are missing in this paper.\n* Missing baselines and related work in using diffusion models for missing data imputation. Upon some preliminary search there also exists a number of papers that use diffusion to impute missing data [4-7]. It would be helpful to discuss them in the paper and compare with them in experiments if possible.\n\n[1] Luo, Yuan, Peter Szolovits, Anand S. Dighe, and Jason M. Baron. \"3D-MICE: integration of cross-sectional and longitudinal imputation for multi-analyte longitudinal clinical data.\" Journal of the American Medical Informatics Association 25, no. 6 (2018): 645-653.\n\n[2] Fortuin, Vincent, Dmitry Baranchuk, Gunnar Rätsch, and Stephan Mandt. \"Gp-vae: Deep probabilistic time series imputation.\" In International conference on artificial intelligence and statistics, pp. 1651-1661. PMLR, 2020.\n\n[3] Xinyu Yang, Yu Sun, Xiaojie Yuan, and Xinyang Chen. Frequency-aware generative models for multivariate time series imputation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024\n\n[4] Lopez Alcaraz, Juan Miguel, and Nils Strodthoff. \"Diffusion-based time series imputation and forecasting with structured atate apace models.\" Transactions on machine learning research (2023): 1-36.\n\n[5] Jianping Zhou, Junhao Li, Guanjie Zheng, Xinbing Wang, and Chenghu Zhou. Mtsci: A conditional diffusion model for multivariate time series consistent imputation. In CIKM, pages 3474–3483, 2024.\n\n[6] Xu Wang, Hongbo Zhang, Pengkun Wang, Yudong Zhang, Binwu Wang, Zhengyang Zhou, and Yang Wang. An observed value consistent diffusion model for imputing missing values in multivariate time series. In SIGKDD, 2023.\n\n[7] Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: Conditional score-based diffusion models for probabilistic time series imputation. In NeurIPS, 2021.\n\n[8] Shuai Liu, Xiucheng Li, Gao Cong, Yile Chen, and Yue Jiang. Multivariate time-series imputation with disentangled temporal representations. 2022.\n\n[9] Yonghong Luo, Ying Zhang, Xiangrui Cai, and Xiaojie Yuan. E2GAN: End-to-end generative adversarial network for multivariate time series imputation. In IJCAI, 2019.\n\n[10] Guojun Liang, Prayag Tiwari, Sławomir Nowaczyk, and Stefan Byttner. Higher-order spatio-temporal physics-incorporated graph neural network for multivariate time series imputation. In CIKM, 2024.\n\n[11] Andrea Cini, Ivan Marisca, and Cesare Alippi. Filling the g ap s: Multivariate time series imputation by graph neural networks. In ICLR, 2022.\n\n[12] Zhengping Che, Sanjay Purushotham, Kyunghyun Cho, David Sontag, and Yan Liu. Recurrent neural networks for multivariate time series with missing values. Scientific Reports, 8(1), Apr 2018.\n\n[13] Gao, Qitong, Dong Wang, Joshua David Amason, Siyang Yuan, Chenyang Tao, Ricardo Henao, Majda Hadziahmetovic, Lawrence Carin, and Miroslav Pajic. \"Gradient Importance Learning for Incomplete Observations.\" In International Conference on Learning Representations.\n\n[14] Cao, Wei, Dong Wang, Jian Li, Hao Zhou, Lei Li, and Yitan Li. \"Brits: Bidirectional recurrent imputation for time series.\" Advances in neural information processing systems 31 (2018).\n\nQUESTIONS:\nFollowing from the weaknesses above, the reviewer would appreciate additional clarifications around, i.e.,\n\n* Would MIRI require data to be MCAR? If so, what would make a difference when imputing MAR data and how does MIRI perform on MAR datasets (e.g., MIMIC, physionet or synthetic MAR masking).\n* Would MIRI also work for imputing time-series? If so, the reviewer is curious in how it compares to other diffusion-based methods in time-series imputation.\n* How would MIRI benefit downstream classification/regression tasks using the imputed data? Could authors quantify and show the improvements in experiments?\n\nLIMITATIONS:\nThe authors acknowledged limitations in computing resources and for part of their theories. However, to the reviewer's point, the more important discussion is missing, i.e., if the approach can work with MAR data and for time-series (see the weaknesses section above).", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors develop multivariate imputation engines, constructed using a so-called imputation ODE obtained through Rectified Flow training. The author's imputation engine is optimal due to the adversarial nature of their training: they optimize their imputation engines to minimize the KL divergence between the joint distribution of the missingness masks and study variables and a product of the marginal distributions of two. Iterative estimation algorithms are introduced which successively decrease the mutual information between X and M, so that that it becomes difficult to distinguish imputed from observed values.  The authors also connect their approach (MIRI) to popular iterative imputation methods like MICE. The effectiveness of their methods are demonstrated in two simulation studies.\n\nSTRENGTHS AND WEAKNESSES:\nI enjoyed this fresh perspective on missing data imputation. The idea to treat imputation as an adversarial learning problem is quite interesting. The author's formulation is quite interesting -- to build a generative imputation engine that is optimized such that joint imputations are essentially indistinguishable from the observed data based on their mutual information with the missingness mask.\n\nWeakness: This method is restricted to MAR problems. With informative missingness (MNAR). The missing data is dependent on the missingness process, so this formulation would not work. The authors should be sure to clarify this limitation in their manuscript more clearly.\n\nQUESTIONS:\nCan the authors think more about how their method situates among MAR vs. MNAR imputation methods?\n\nLIMITATIONS:\nThe authors should be more transparent about the fact that their methods are specifically designed for the case where Y_{j}^{mis} (The missing portion of Y_j) is conditionally independent of M given Y^{obs} (observed data). In this setting we want to preserve the distribution of the observed data among the imputed values, which the proposed methods are optimized to do. Unless, I am missing something, it seems that the methods presented are almost tailored to the MAR, and not MNAR case.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces MIRI, a novel method for missing data imputation based on rectified flows. The key idea is to iteratively reduce the mutual information between the imputed data and the missingness mask, which corresponds to minimizing the KL divergence between their joint distribution and the product of their marginals. The authors provide theoretical justification for this approach and show that several existing methods can be cast within this general framework. They present a practical iterative algorithm based on rectified flows to implement this objective and validate its performance across a variety of datasets. Extensive experiments show that MIRI consistently achieves strong results. The paper also discusses MIRI’s computational limitations in an informative manner.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- The paper proposes a novel and theoretically grounded approach to missing data imputation using rectified flows. The method provides a clear and principled objective that unifies and generalizes several prior works.\n- Theoretical results are solid and insightful. The authors rigorously show how existing approaches can be reinterpreted within the MIRI framework.\n- Empirical results across a variety of datasets support the effectiveness of the proposed method. The paper also transparently discusses MIRI’s computational limitations, which adds to its credibility.\n\n## Weaknesses\n### Main concerns\n- The motivation and presentation in the introduction could be clearer. The writing flow is somewhat disjointed, making it difficult to understand what specific challenges MIRI addresses and how it stands out from prior work. Important limitations of existing approaches (e.g., GAIN) that are only discussed later should be summarized earlier in the paper.\n- The literature review underrepresents recent progress in deep generative modelling approaches for missing data imputation, particularly VAE-based imputation methods (e.g. [1-4]), which are strong one-shot alternatives with lower computational cost. Including such baselines would provide a fairer evaluation of MIRI’s benefits.\n- The method is specifically designed for settings where the missingness is independent of the missing values (MCAR). However, MNAR settings—more common in practice—are not addressed. The authors should explicitly acknowledge early in the paper that MIRI is unlikely to perform well under MNAR conditions, and clarify that its benefits are tied to the MCAR assumption.\n\n### Minor concerns\n- Although the Gaussian Mixture with uniform weights can be easily drawn from the context, the equation in L270 is not a valid density without including the weights because it’s unnormalized. I suggest either including the “proportional to” sign or explicitly including the weights.\n\n## References\n\n[1] Nazabal, Alfredo, et al. \"Handling incomplete heterogeneous data using vaes.\" Pattern Recognition 107 (2020): 107501.\n\n[2] Mattei, Pierre-Alexandre, and Jes Frellsen. \"MIWAE: Deep generative modelling and imputation of incomplete data sets.\" International conference on machine learning. PMLR, 2019.\n\n[3] Ma, Chao, et al. \"Vaem: a deep generative model for heterogeneous mixed type data.\" Advances in Neural Information Processing Systems 33 (2020): 11237-11247.\n\n[4] Peis, Ignacio, Chao Ma, and José Miguel Hernández-Lobato. \"Missing data imputation and acquisition with deep hierarchical models and hamiltonian monte carlo.\" Advances in Neural Information Processing Systems 35 (2022): 35839-35851.\n\nQUESTIONS:\n- The proposed method seems closely related to DiffPuter, which also applies iterative imputation via a generative model using an EM-like approach. In particular, the EM updates in Diffuser resemble SGD and the imputation steps of Algorithm 1. Could the authors elaborate further on the similarities and differences beyond the brief discussion in lines L258–262? A more in-depth comparison could clarify the novelty of MIRI and affect the originality rating.\n\n- GAIN performs surprisingly poorly in the CIFAR MCAR experiments. Could the authors provide insight into this result? Was the model properly tuned or trained under comparable conditions? If GAIN is underperforming due to experimental choices, clarifying this could impact the validity of the empirical comparisons.\n\nLIMITATIONS:\nThe authors acknowledge the computational limitations of their approach, which is appreciated. However, a more explicit discussion is needed regarding the assumptions behind their method—particularly the reliance on MCAR missingness mechanisms. MIRI is unlikely to perform well in MNAR settings, which are common in real-world data, and this should be clearly stated in the paper. Including a short discussion of this limitation in the introduction or a dedicated paragraph in Section 7 would help set clearer expectations about the scope of the method.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new iterative approach for missing data imputation that progressively reduces the statistical dependence between imputed values and the observed mask. At each iteration, a rectified flow model is trained to transform the current imputation distribution toward a reference distribution where imputed values are statistically independent of the mask.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The core idea of iteratively reducing dependence between the imputation and the mask with a generative model is novel.\n\n2. The paper is in general well-written and structured. \n\n3. The method is theoretically grounded.  \n\n4. The method shows better performance on distribution level metric MMD compared to previous SOTA methods like Hyperimpute.\n\n**Weakness**\n\n1. High computational cost:\nThe proposed method requires training a new rectified flow model and performing sampling at each iteration, which can be computationally expensive. However, the paper does not report imputation time or compare runtime performance with other baselines. A systematic evaluation of efficiency would be helpful for understanding the method’s practicality.\n\n2. Motivation for Choosing Rectified Flow Not Justified:\nWhile rectified flow is used as the generative backbone at each step, the paper does not clearly explain why it is preferable to alternatives like diffusion models. Nor is there any ablation to validate this choice empirically. A deeper justification—either theoretical or experimental—would strengthen the case for this design decision.\n\n3. Underperformance on Entry-Wise Metrics:\nFor standard entry-wise imputation metrics such as RMSE and MAE, the evaluation is only conducted on synthetic data. On this task, the proposed method performs worse than existing baselines like HyperImpute, raising concerns about its overall performance.\n\n4. Incomplete Baseline Comparisons:\nThe experimental section lacks several related baselines, e.g. MOT [1], TDM [2], which are both based on distribution matching.\n\n[1]  Muzellec, Boris, et al. \"Missing data imputation using optimal transport.\" International Conference on Machine Learning. PMLR, 2020.\n\n[2]  Zhao, He, et al. \"Transformed distribution matching for missing value imputation.\" International Conference on Machine Learning. PMLR, 2023.\n\nQUESTIONS:\n1. How many iterations are needed for the proposed method to converge? \n\n2. What are the model configuration and training details of the rectified flow model?\n\n3. What will be the performance for other missing patterns like Missing at random(MAR) and Missing Not At Random (MNAR)?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}