Missing Data Imputation by Reducing Mutual Information with Rectified Flows
Authors: Jiahao Yu
Qizhen Ying
Leyang Wang
Ziyue Jiang
Song Liu
Venue: NeurIPS 2025 poster ()
Keywords: missing data
imputation
rectified flow
generative model
Primary area: general_machine_learning
TL;DR: We propose a missing data imputation method using rectified flow by following mutual information minimization criterion.

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}.

Missing Data Imputation by Reducing Mutual Information with Rectified Flows
Jiahao Yu" University of Cambridge
Qizhen Ying"
Leyang Wang"
University of Oxford
University College London
Ziyue Jiang
Song Liu University of Bristol
University of Bristol
Abstract
This paper introduces 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 decrease the predictability of missingness patterns, our method explicitly targets this reduction in mutual information. Specifically, our algorithm iteratively min- imizes 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 hensive experiments on synthetic and real-world datasets validate the efficacy of our proposed approach, demonstrating its superior imputation performance. Our imple- mentation is available at https: //github....... cww.comm//oooe mumm/ujjml. MMMIII//IIIIIIII IMIRIIIIIutttaattaaammm Imputation.
framework. Compre-
Introduction
The problem of missing data, referring to absent components (NaNs) in dataset, can arise in wide range of scenarios [31. 18]. For example, in a survey dataset, respondents might skip some questions accidentally intentionally, resulting in missing values. In fMRI data, there can be missing voxels due incomplete brain coverage and spatial variations images acquired across subjects [36]. Simple "one-shot" methods replace missing entries with summary statistics, such as the sample mean or median of the observed data. In contrast, approaches like MICE [37], MissForest [34], and Hyperlmpute [12] employ an iterative, dimension-wise scheme. In each iteration, they model one feature conditional on others and sample from that model to replace the feature's missing entries, cycling through all features repeatedly until the imputations stabilize. Because their updates are sequential and dimension-wise. these methods are difficult to parallelize and scale poorly to high-dimensional data such as images [3]. Their accuracy is also highly sensitive to the choice of per-variable models [37, 15].
In recent years, research on generative modeling has made huge progress. Since imputation is naturally a data-generating task, generative models have been increasingly applied to impute missing data. [30] proposes to maximize the likelihood of the imputed data using an EM algorithm, where the likelihood function modeled using a normalizing flow [29, 6). GAIN [38] leverages Generative Adversarial Nets (GANs) [8] to impute missing components. It trains a generator (a neural network) so that binary classifier cannot distinguish whether given component is imputed or not. However, both GANs and normalizing flows come with their own challenges. GANs rely on unstable adversarial
'Work partially completed at the University of Bristol. 'Correspondence to Song <song. bristol.ac.uks.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

training and may suffer from mode collapse [16, 40), while normalizing flows require specially designed architectures to ensure their invertibility. limiting their applications.
Currently, diffusion models and flow-based methods represent the state of the art in generative modeling techniques [33, 17, 21]. They first train velocity field according some criteria. Then, they sample from the target distribution by solving an ODE or SDE using the trained velocity field as the drift. The main challenge to adapt these methods for missing value imputation is that most flow- based methods are designed to transport reference samples to match the target distribution. However, it is unclear how to formulate the missing data imputation problem as a distribution transport problem. Nonetheless, progress has been made in this area, For example, MissDiff [27] learns a score model from partially observed samples, then uses an inpainter [23] impute missing values.
Recently, [20] observes that GAIN's criterion encourages the imputed dataset to be independent of the missingness pattern. In this paper, we further explore this idea within more rigorous framework. Our contributions are summarized as follows:
1. We introduce Mutual Information Reducing Iterations (MIRI), a novel framework that provably reduces mutual information between imputed data and their missingness mask when using an optimal imputer;
2. We show that optimal MIRI imputer be obtained by solving ODE whose velocity field is trained by rectified flow objective, naturally integrating generative modeling;
3. We reveal that several existing imputation methods can be viewed as approximate special cases of the MIRI framework;
4. We demonstrate that our proposed approach achieves promising empirical performance on both tabular and image datasets.
Background
We now briefly review the missing data imputation problem, Generative Adversarial Imputation Nets (GAIN) [38], and Rectified Flow [21], which form the foundation of our proposed method.
Notations. We denote scalars with lowercase letters (e.g., xi y) and vectors with bold lowercase letters (e.g.. y). Random variables are denoted by uppercase letters (e.g., X, Y), and random vectors by bold uppercase letters (e.g., X, Y). The superscript E {1,-.., 7) denotes the iteration count of the algorithm. The subscripts 0, 1, and r denote the continuous time index in an ODE. Given fixed missingness mask m and vector x of the same dimension, we write Xm as the subvector of x whose elements correspond to entries where m. = 1, and Xi-m as the subvector whose elements correspond to entries where m; O. Following convention, we refer Xm as the "non-missing components" of x and X1 xm-m as the "missing components" of x. The notation A B means that the random variables A and B are equal in distribution.
2.1 Sequential Imputation
Now, we formally define the missing data imputation problem. The true data vector X*
[Xi X31 is a random vector taking values on X' C Rd, while the missingness mask M
[M,.... Ma) is a random vector in (0, 1]d. In the Missing Completely at Random (MCAR) [31] setting, assume that the true data is independent of the missingness mask i.e., X 1 M (See Section F for discussions on Missing At Random (MAR) setting). The observed data vector
Mj = 1
X = (x,,..., Xa) is defined as X, = NaN M, = 0
In words, the variable j in the observation
X is missing when M, = 0 and is observed when M, = 1. Equivalently,
M + (1 M) @ NaN,
(1)
where O represents element- wise product. In this paper, we want to construct an imputation vector
X(g) = M x + (1- M) o X, M),
(2)

---

where g R x (RR UNaN) x {0. - R is called an imputer. Zo is random seed of the imputer. Different random seeds Zo can lead different imputations.
In many cases, imputation is an iterative process: Instead of computing X in a single step, the imputation is improved through iterative updates. This sequential imputation involves the following procedure: for each iteration = I, we construct the imputation
g.) = M x (- + + M) + (Z) X ((-1) M),
where (+) indicates the imputed data vector at iteration t. The initial vector, X is initialized with any standard imputation technique (e.g., Gaussian noise, mean or median imputation). schematic of this sequential imputation algorithm is provided in Figure 4 of Appendix A.
The central task of this iterative process is to learn the imputer g: at each iteration f. In the following sections, we show how a popular missing data imputation scheme can be formulated within this framework and how g: can be trained using GAN-type objective.
2.2 Generative Adversarial Imputation Nets (GAIN) [38]
GAIN arises from natural intuition: if imputation perfect, the imputed entries should be indistinguishable from the originally observed ones. Thus, GAIN trains a generator so that a classifier cannot predict the missingness mask M accurately.
Let fmy, x, m 5) E [o, 1] be a probabilistic classifier modeling the conditional probability P M2 , which is the probability that the j-th component observed given the imputed data from the previous step, (t-1), and the rest of the missingness mask, M-j.
For each iteration = 1,..., T, GAIN performs the following steps:
1. Classifier fej arg miny -E [log I (M,,X(1-) ((-1), M M. for all j 1...d.
2. Imputer g: arg maxg" - E, E [log Sej (M,. X (g). M -)] (reconstruuctio error).
3. Update (1) impute
(g:), t
Here, (g) is the imputed vector from the imputer as defined in (3). At iteration t, we first train classifiers fiy by minimizing the cross-entropy loss. Subsequently, we train the imputer g parameterized as neural network, to reduce the classifiers' predictive accuracy by maximizing the sum of cross-entropy losses across all j. while simultaneously minimizing the reconstruction error.
In practice, the minimization and maximization steps are performed via alternating gradient updates, leading adversarial training. We observe that GAIN works well in practice. However, the adversarial training is known to be difficult and may suffer from mode collapses [16, 40J. Moreover, balancing the reconstruction error and the cross entropy loss requires careful tuning of the hyperparameter X.
Recently, it was observed that the generator training of GAIN can be viewed as a process of breaking the dependency between and (g) [20]. Indeed, if (1)) predictive of M at all, then the cross-entropy loss is maximized. In this paper, rigorously explore this idea imputation by dependency reduction. Guided by this principle, our algorithm explicitly aims to reduce the dependency between M and (g)
M, meaning (() is not
2.3 Flow-based Sampling and Rectified Flow
As an alternative to GANs, flow-based generative models have attracted significant interest recent years. These algorithms first train "velocity field" over time E [0, according to some loss function, and then generate samples by solving the ODE or SDE using the learned velocity field as the drift. In some cases, the loss function is simple least-squares loss, leading to a training routine simpler and more stable than GAN. Methods in this category include Score-based generative models [32, 33), Rectified Flow [21], Flow Matching [17], Diffusion Schrödinger Bridge [5]. among others.
Rectified flow is one of the simplest flow-based generative models. Given two random vectors X, (from a reference distribution) and X, (from target distribution), the velocity field v" is trained by

---

Algorithm 1 MIRI with Rectified Flow (Single Imputation)
Require: Paired, i..d. observations and masks {(X, M)), Maximum Iterations T,, Maximum SGD steps N, Batch Size B, Optimizer SGD _GD_update update and an ODE solver ODE_Solven 1: (X()) M) + initial_impute(X M).
2: for t= to T do
3: Initialize neural network model v.
4: for to N do
Sample batch { {(x) MD ) ~ 1x
6: Sample batch 1B 11=1 ~ ~P x and time indices {T₃ [T,b=1 Uniform(0, 1)
7:
Vj. X + (1 X; +Tzx(( Y
x(x)
x
V: f SGD uGd_te
r [v
v(x 3 M
v
@ x
3 @ x M +3)
9: end for
10: ( - - + M) ODE ODE_Soller((1 - M) O (
11: end for
12:
return (X
minimizing the following objective:
V arg min
E X1 Xo- - -v(X, 7)1? dr.
(4)
where X, the linear interpolation between Xo and X,, defined as X, = T + (1 T) Xo for E [0, 1]. Samples are then generated by solving the ODE:
dz, = v- (Z,,T))
(5)
dr
One can prove that, if Zo Xn. then Z, _ X, for all T E [0,1]. This prreseriin property" guarantees that solving the ODE from 0 transports samples from the reference distribution to the target distribution.
3 Reducing Dependency with Mutual Information
In this section, we introduce the MIRT imputation framework, obtain the optimal imputer under this framework using rectified flow, and theoretically justify its validity under the Missing Completely at Random (MCAR) setting. The validity for the Missing at Random (MAR) scenario is in Appendix
3.1 Mutual Information Reducing Iterations (MIRI)
We study the problem of training an imputer g to reduce the dependency between the imputed sample X(g) and missingness mask M. The approach adopted by [20] is to minimize the mutual information between them, i.e., the optimal g is given as:
g E arg min (gX M) arg min MlP & PM])
(6)
8
where D[P[[Q] is the Kullback-Leibler (KL) divergence of probability distributions Q from IP.
However, the mutual information is not directly computable as we do not have access to the true distribution Px(g).M and Px(g) X PM. A potential solution is to transform this optimization problem into bi-level min-max adversarial optimization problem: Fixing g. it is possible to estimate the mutual information using samples from I'x ( .M and Pxxgg) & IPM with Mutual Information Neural Estimation (MINE) [1]. which maximizes the Donsker-Varte Varadhan lower bound [7] to approximate the KL divergence. After that, we can minimize the estimated mutual information with respect to g. This process is repeated until convergence. If one solves both optimization problems gradient step at time, this algorithm is the classic adversarial training that GAN is known for.
Contrary to the suggestion in [20], the mutual information in (6) cannot be minimized simply by directly applying forward or reverse KI. Wasserstein Gradient Flow (WGF). While WGF can minimize KL divergence D[P[|[Q] when the distribution to be optimized is either IP Q, the

---

distribution of interest in our setting (P) appears simultaneously in both I and Q. As a result, neither the forward nor reverse KL formulation of WGF directly applicable. Therefore, while the WGF procedure proposed in [20] is conceptually appealing, it does not correctly minimize the mutual information objective as formulated (6).
Motivated by the limitation of both adversarial training and WGF, we propose a sequential imputation algorithm that reduces the mutual information, called Mutual Information Reducing Iterations (MIRI). Instead of targeting the unknown marginal distribution P (g) we use the imputed data distribution from the previous iteration, ((-1),, stable target. Let us denote the joint probability of the pair (X (+) g). M) as (x) and the product measure of X (f-1) and M as IP 1x @ PM.
For each iteration 1... T, the algorithm performs the following steps:
1. Find the Optimal Imputer: g: E arg ming D ([ 8 PM].
2. Update Data: (
imputee X (g.), + 1.
We can confirm that MIRI does indeed reduce the mutual information: Proposition 1. The mutual information between ( and M is non-increasing after each iteration.
See Appendix B for the proof. In Section 4.1, we show that the GAIN algorithm can be viewed as an approximate implementation of the above MIRI algorithm.
Moreover, we show that minimizing the KL divergence in the first step of MIRI has sufficient and necessary condition:
Lel PR(x, m) be the density of (X(0 (g:), M) and q(x) be the density of X (1-1)) g arg ming D) (x) x 69 PM]] if and only if
Proposition 2.
Pg (x Xm m) ((11 Xm). Vx, m.
The proof can be found the Appendix E. This result shows that the optimal g should sample from the conditional distribution of the missing components given the non-missing components according to the marginal distribution of x (t-1), the imputed data in the previous iteration. (7) inspires us to
construct a flow-based generative model with target distribution q(x1-ml Xm))
As we discuss in Section 4.2 and 4.3, (7) is also the key design principle behind many classic and contemporary data imputation algorithms.
3.2 Imputation by Rectified Flow
In this section, we show how to construct an optimal imputer for MIRI using an imputation ODE obtained via rectified flow training. The optimality condition of Proposition 2 requires an imputer to sample from target distribution q Xm). The standard diffusion models are restrictive, as they are constrained to fixed Gaussian reference. We therefore use rectified flow to transport samples from the current imputation (reference) to the target distribution.
We now zoom in on Step of the MIRT algorithm at iteration t.
Let (X1,M) M,) be pair of random vectors drawn from the product measure P Pxit-n) PM and (Xo, Mo) from the joint distribution ( The pairs are drawn such that X, 1 (Xo, Mo). X, is defined as the linear interpolation of X, and X,, following the same construction in Section 2.3.
Consider the following rectified flow training objective function:
v = arg min / LE [[XI - Xo v(X,, Xo.Mo XI, Mo. +]] dr,
(8)
where Xo,M, := M, @ X, and others are defined similarly.
We can define an imputation process by using the optimal velocity field v*. Given partially observed vector where the missing entries are padded with zeros, i.e., [0, xm]. the imputation process dynamics are governed by the following ODE:
dzz
= (1
m
C v (Zr. xm,,
1 [0, Xm.) m. T),
(9)
dr
5

---

with initial condition
Zo IP X0,1-m (0,,ml Manm
(10)
We define the imputer g" as the solution to the imputation ODE (9) evaluated at terminal time r -= 1:
(Zo. Xm, m) Z1.
(11)
We show that, g" is an optimal imputer to the MIRI algorithm that introduced in Section 3.1. Theorem 1. g" desined in (11) is an optimal imputer in the sense that
g" E arg min D (M) ( @ PM]
(12)
The main technical challenge of the proof is show that the ODE (9) can indeed establish pressrrvnng process from M to Px(--)). This is different from the classic proof since both our training objective and ODE are different from the regular rectified flow. See Appendix C for the proof. With an imputer g", we can proceed to the second step of the MIRI algorithm.
3.3 Initial Condition
To run the imputation process, we require least a sample of X0,1 (Xo.m Xm: M m}.
(t-1) By the definition of Xg. this is equivalent to sampling from X|--b| [x((1-) [[_
Xm. M m} This poses no difficulty, since at each iteration t, we maintain access to the joint sample (x M), from which we can extract the relevant conditional sample by slicing. However, in the first iteration,
we need samples from x (0)
- ! Xm. M m) kickstart the MIRI algorithm.
One can sample from any distribution for those initial samples and in our experiment, we simply draw independent samples from the normal N(0,) or uniform (00, 1) to fill out each missing component.
Alternatvelly, one can use another imputation algorithm to suggest initial samples.
3.4 Practical Implementation
Algorithm 1 presents the full MIRI procedure on a finite sample set (X, M). Note that the superscript (t) represents iteration count while (j) is the sample index in a batch. We model v, with neural network trained via stochastic gradient descent. The function initial_ inipte replaces NaNs with initial guesses.
According the requirement in Section 3.2, we need to ensure (Xo, Ma) 1 X,. In practice, propose to sample a batch (Xaj, Mo.; }2, from the paired dataset and sample a batch (X1,) from shuffled set of (Xo,) to weaken their dependency. Although, in this case, X, is still not independent of Xo and Mo. we observe that this strategy works well in practice.
MIRI and Other Imputation Algorithms
In this section, we show how MIRT is related to other existing imputation algorithms.
4.1 GAIN
Let q(x, m) the density function of the joint distribution of (X(++1), M). Suppose the discrimi- nator is well-trained in the first step in the GAIN algorithm, i.e., [(m; x, m-,) RS q(m, x, m. then the generator training the second step (without the reconstruction error) is
g: = arg min E
Σ log al (M, M = arg min
log II
q (8) M ,).
(13)
Using the pseudo-likelihood approximation [2), II, q(M, X M,j) q x (t) )
thus:
( (())
g: RJ arg min E log q( [M]( [)) arg min E g
log xxxx
(14)
R
(3)
)p()))pM)))

---

By Gibbs' inequality,
log (M, x x(xg (P(
Px (M )
- E log
[ ( Xa @ PM].
=
(15)
((x (M)
thus, we can see that GAIN learns its imputer gr by approximately minimizing lower bound of D ( (( @ PM].
There approximations made in the above argument: the pseudo beeudo-likeihoodd approximation and the Gibbs' inequality. First, since GAIN only performs one gradient step update to the imputer g: at each iteration, the gap between D (g) M x) FM] and its lower bound may not be large. Second, the pseudo-likelihood approximation can be a good approximation, for example, when the correlation among Mj is not too strong. The discriminator training of GAIN resembles dimension- wise strategy that has been widely used in Ising model estimation [28, 25]. Note that the optimization in (14) is fundamentally intractable due to the high-dimensional density { (g))) m is a binary variable defined on {0, 1J4, thus, the density does not have a tractable normalizing constant. This further restricts to the dimension-win
sreudo-likelihook approach in GAIN algorithm.
4.2 Round-Robin Approaches: MICE, HyperImpute, etc.
Recall that Proposition 2 states that minimizing the KL. divergence (g) & PM] equivalent to choosing a g such that Pg Pg(X: Px(Xxm||xm) Xm, m) = g(xi-m|Xx where Px(x, m) is the density of (x (g), M) and ((x) is the density of (t-1), One can simply build an im- puter that samples from ( Xm) to impute the missing components. Then, by construction, Pg( Xm. m) = (xm))
This algorithm has two issues: First, m changes for every sample, so we need to learn different im- puter for every sample. Second, sampling from the conditional probability distribution q(x1-m Xm) is not trivial as the dimension of X can still be high.
In principle, we can leam the joint probability g then sample from the conditional distributions ((11mm xm) ) for any fixed m (see the next section). However, learning a high dimensional joint probability density function q not trivial. Thus, we adopt the pseudo-likelihodd approximation again: factorize q(x) over simension-wise conditional probabilities and learn ( x-ك] for every j. This dimension- wise conditional probability learning scheme not only makes the learning of the joint distribution easier but also solves the second issue we can perform Gibbs sampling, one missing component time. This pseudo- likelihod approximation Gibbs sampling strategy gives rise to the round-robin algorithms such as MICE and Hyperlmpute.
Let ʃg be the density model of q(zg) x x-j), then the "MIRI MICE" performs the following algorithm.
For = 1, I,...,T: T:
1. Vj. learn fj to approximate (x
2. Let g: be Gibbs sampler for g(x1 mlx)) using (fz)-
3. (() ingute (g)) ++
In practice, MICE interlaces first step and the later two steps, meaning Gibbs sampling step and imputation is immediately performed after learning each fy. However, the sequential nature of Gibbs sampling prevent us from parallelizing this procedure for high-dimensional samples.
4.3 Contemporary Approaches: Learn Jointly, Sample/Inpaint Conditionally
Traditionally, learning a high-dimensional joint distribution q(x) intractable due its normalising constant. Recently, it was realized that one can learn a joint model g(x) without dealing with the normalising constant. [35] first estimates the joint model q via Score Matching [11] or Noise Contrastive Estimation [10], then sample from 9 via importance sampling. This method works well when Xi-m is in relatively low-dimensional space. [4] noticed that given a joint model g. this conditional sampling problem could be solved using Stein Variational Gradient Descent (SVGD) [19]. However, SVGD utilizes properties of RKHS function, and require kernel computations. Therefore, does not scale to high-dimensional datasets.
7

---

(a) MIRI
(b) Hyperlmpute
(c) KnewImp
(d) MMD and MI
Figure 1: Pairwise density plots of the imputed data and MMD/MI of MIRI.
MIRI
HyperImpute
Knewlmp
TabCSDI
2.4
0.00
bike
energy
forest
gas
housing
parkinsons pumadyn
solar
stock
wine
Figure 2: MMD on UCT datasets with 60% data missing. The lower the better.
In recent years, it is also realized that one can "inpaint" samples given diffusion model trained on the joint samples [23]. This idea gives rise DiffPuter [39], an imputer that trains joint diffusion model to sample from g. Then use an inpainter to impute missing values given the observed vector Xm- However, inpainter is not standard backward process of diffusion dynamics, thus in general, is not guaranteed to generate samples from g(x1-m) xm).
5 Experiments
This section evaluates MIRI's performance variety of imputation tasks. Details experimental settings and implementations can be found in Appendix G. Additional results can be found in Appendix H. To ensure fair comparisons, match model architectures across all methods wherever possible. For example, if our method uses CNN, we use the same architectures for other neural- network-based methods well.
Imputation methods aim to either predict the conditional mean E[ Xm). evaluated by metrics like MAE/RMSE, or learn sample from the conditional distribution IP Px, -',-m/XX. evaluated by the distributional fidelity. Our method, MIRI, is designed for the latter. Thus, we mostly evaluate the performance by using metrics that measure the distributional fidelity (such MMD [9]).
5.1 Toy Example: Assessing Imputation Quality Using MMD and Mutual Information
In this section, showcase the performance of the proposed MIRI imputer on toy example that a mixture of 2-dimensional Gaussian: X ~ ([5N([[2, - 22,, 0.521) + (,5./ 2], 0.5²1). The missingness mask is generated as M ~ II j== Bemouuli (0.7), i.e., 30% of the data is Missing Completely at Random. We draw 6000 i.i.d. samples from Fx and PM to create missing obser- vations. The imputed samples are visualized in the pairwise density plots in Figure 1. We choose Hyperlmpute [12] and Knewlmp [4] comparisons. It can be seen that MIRI not only captures the bimodal structure of the true data but also preserves the proportion of the two modes. On the other hand, Hyperlmpute and KnewImp fail to recover the correct proportion of the two modes. We also show performance metrics (MAE, RMSE, MMD) with various other imputation algorithms (such as MissDifff GAIN, TabCSDI [42]) in Table 5 in the Appendix H. Finally, in Figure ((d), we show that both MIRT's MMD and MI decrease over iterations, proving that MI between the imputed data and the missingness mask is good criterion for assessing how well the imputed data recovers Fx..

---

Table 2: Quantitative results on C-FAR-10. Methods are evaluated at three levels of missingness (20%, 40%, 60%) using FID, PSNR, and SSIM. The best results are highlighted bold.
20% Missingness
40% Missingness
60% Missingness
Method
FID (
PSNR (1)
SSIM (1)
FID (U)
PSNR (t)
SSIM (1)
FII (L)
PSNR (t)
SSIM (1)
GAIN [38]
164.11
21.21
0.7803
281.62
16.20
0.5576
285.53
11.99
0.2933
KnewImp [4]
153.09
18.84
0.6463
193.68
15.81
0.4740
264.40
14.04
0.3317
MissDiff' [27]
90.51
22.29
0.7702
129.84
19.65
0.6648
197.91
16.78
0.4989
Hyperimpute [12]
8.92
34.09
0.9750
65.01
23.22
0.7931
130.36
20.17
0.6533
MIRI (Ours*)
6.01
0.9736
27.53
27.14
0.9126
68.58
23.22
0.8063
5.2 Real-World Tabular Imputation: UCI Regression Benchmarks
We evaluate MIRI on 10 standard UCI regression bench- marks [13]. We select these datasets for their diversity in sample size and dimensionality (see Table in Appendix G.1). Following [26], we simulate missingness under MCAR/MAR/MNA mechanisms. We simulate 20%, 40%, and 60% missingness for MCAR and MNAR, and 40% and 80% for MAR. For the MAR mechanism, the missingness mask is conditioned on a randomly selected 50% of the variables. In the main paper, we compare method with Hyperlmpute, KnewImp, TabCSDI, GAIN, MIWAE [24], and TDM [41].
Table 1: Aggregated MMD rankings for UCI (40% missingness) across 10 datasets. Lower is better.
Method
MCAR
MAR
MNAR
TabCSDI [42]
6.4
6.4
6.1
GAIN [38]
6.2
6.2
5.9
TDM [41]
4.5
4,4
4.2
Knewlmp [4]
4.1
4.2
3.8
MIWAE [24]
3.0
2.9
2.8
Hyperlmpute [12]
1.6
1.5
MIRI (Ours*)
1.4
2.0
1.3
Figure reports the average MMD over trials (error bars denote one standard deviation). As shown, even in the 60% missing data cases, MIRI outperforms performs comparably all competitors. We aggregate performance using mean rank across 10 datasets. As shown in Table for the 40% missingness setting, MIRI attains the best mean rank under MCAR and MNAR settings and is highly competitive under MAR setting.
Comprehensive results, including performance across all specified missingness, levels and comparisons with an extended set of baseline methods, are deferred Appendix H.2.
5.3 Real-World Image Imputation: CIFAR-10 and CelebA Benchmarks
We assess MIRI on high-dimensional image datasets. First, on CIFAR-IO [14] (32x32 RGB), we train us- ing only 5000 samples (<10% of the full set) with 60% of pixels randomly removed (all RGB channels jointly). Figure 3(a) compares MIRI's ' 60% missingness reconstructions against generative approaches (GAIN, KnewImp, MissDiff). Table reports FID, PSNR, and SSIM at 20%, 40%, and 60% missingness, showing MIRI strongly outperforms the best generative alterna- tives across this range, even in challenging cases.
Table 3: Accuracy of 10-class CIFAR-10 classification on imputed data at varying missingness. Higher is better.
Method
20%
40%
60%
KnewImp [4]
0.267
0.137
0.106
GAIN [38]
0.337
0.133
0.096
MissDiff [27)
0.486
0.269
0.152
Hyperlmpute [12]
0.804
0.405
0.212
MIRI (Ours*)
0.812
0.525
0.364
See Appendix H.3 for additional results, and Table 3 for downstream classification results.
Second, we test MIRI on higher-resolution CelebA [22] (64x64 RGB) with 5 500 samples, using a mechanism that masks each RGB channel independently. Figure 3(b) again substantiates MIRI's strong imputation performance.
6 Limitations and Future Works
Despite the competitive results, MIRI has several limitations that open avenues for future work.

---

Ground Truth
Partial Observations
MIRI (Ours*)
GAIN
Knewlmp
Miss
(a) 15 uncurated 32 32 CIFAR-I0 10 images and their imputations. Pixels are removed from all RGB channels.
Ground Truth
Partial Observations
MIRI (Ours*)
GAIN
Knewimp
MissDiff
(b) 15 uncurated 64x64 CelebA images and their imputations. Pixels are removed from each RGB channel independently.
Figure 3: Comparison of imputed samples CIFAR-10 and CelebA with 60% of missingness.
primary limitation computational intensity. Each MIRI iteration requires running a full rectified flow algorithm, including training velocity field and solving the corresponding ODE. One interesting future work is to incorporate recent distillation methods to accelerate sample generation.
Our Theorem is population-level result, meaning it assumes the optimal velocity field v* is obtained and the imputation ODE is solved exactly. does not consider errors due to finite sample approximation, optimisation and inaccurate ODE solver. Extending our theoretical results to the finite-sample optimal v and the approximate imputation ODE solver (like the Fuler method) another avenue for future works.
more fundamental limitation is that our objective, which minimizes the mutual information between the data and the mask, is not designed for Missing Not at Random (MNAR) scenarios where the missingness mechanism depends on the true values. This is a common challenge for many modern imputation frameworks, and extending our method to handle such cases is crucial direction.
Finally, MIRI could be extended artificially masked, fully-observed sequences, we could learn a general-purpose imputation vector field. The key advantage here is flexibility: because MIRI's vector fields agnostic to missingness patterns, this approach could yield single, pre-trained foundation model for robust, general-purpose time-series imputation and forecasting.
self-supervised framework for time-series. By training on
7 Conclusion
In this work, we introduced Mutual Information Reducing Iterations (MIRI), framework that reduces data-mask dependency by iteratively minimizing D[Px( MP 1 & P.... We proved that MIRI decreases mutual information I(X; M). The optimal imputer is constructed by solving an ODE trained with a rectified flow objective. We also showed that existing imputation methods such as GAIN are approximate special cases of this approach. Empirically, MIRI also demonstrated strong distributional fidelity across synthetic, tabular, and image benchmarks.
10

---

Acknowledgments
We thank the four anonymous reviewers and the area chair for their insightful comments and suggestions. SL thanks Josh Givens, Sam Power, Katarzyna Reluga and other colleagues at Bristol Mathematics for helpful discussions. Partial funding for this work was provided by the University of Bristol School of Mathematics Summer Research Bursary 2024; JY was funded by the Alumni Foundation, and LW by the Heilbronn Institute for Mathematical Research. JY also acknowledges support from the Scholar Award and the ED Davies Travel Fund, provided by the Neural Information Processing Foundation and the Fitzwilliam College, University of Cambridge, respectively. Partial computational resources were provided by the Advanced Computing Research Centre at the University of Bristol (https //wp://www..oo bristol ac uk/acrc/).
References
[1] M. I. Belghazi, A. Baratin, S. Rajeshwar, S. Ozair, Y. Bengio, A. Courville, and D. Hjelm. Mutual information neural estimation. In Proceedings of International Conference on Machine Learning, 2018.
[2] J. Besag. Statistical analysis of non-lattice data. The Statisticiam, pages 179-195, 1975.
[3] A. Brini and R. van den Heuvel. Missing data imputation with high digh-dimensional data. The American Statistician, 78(2)2400 252, 2024.
[4] Z. Chen, Li, F. Wang, O. Zhang, H. Xu, Jiang, Z. Song, and H. Wang. Rethinking the diffusion models for missing data imputation: A gradient flow perspective. In Advances in Neural Information Processing Systems, 2024.
[5] V. De Bortoli, J. Thornton, J. Heng, and A. Doucet. Diffusion Schrödinger bridge with applications to score- score-based generative modeling. In Advances in Neural Information Processing Systems, 2021.
[6] L. Dinh, Krueger, and Y. Bengio. NICE: Non-linear independent components estimation. In 3rd International Conference on Learning Representations, ICLR 2015, Workshop Track Proceedings, 2015.
[7] D. Donsker and S. R. Varadhan. Asymptotic evaluation certain markov process expectations for large time, i. Communications on Pure and Applied Mathematics, 28(1):1-47, 1975.
[8] Goodfellow, J. Pouget- Abadie, M. Mirza, Xu. D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, 2014.
[9] A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf, and A. Smola. kemell two-sample test. Journal of Machine Learning Research, 13(1):723- 773, 2012.
[10] Gutmann and A. Hyvärinen. Noise-contrastive estimation: A new estimation principle unnormalized statistical models. In Proceedings of International Conference on Artificial Intelligence and Statistics, 2010.
[11] A. Hyvärinen. Estimation of non nonmomzal statistical models by score matching. Journal of Machine Learning Research, 6:695-709, 2005.
[12] D. Jarrett, B. Cebere, T. Liu, A. Curth, and M. van der Schaar. Hyperlmpute: Generalized iterative imputation with automatic model selection. In Proceedings of International Conference on Machine Learning, 2022.
[13] M. Kelly, R. Longjohn, and K. Nottingham. UCI machine learning repository. Available at: https: //archive iic.. iws.iii.a.................... edu/mi, 2017. (Accessed: 16 May 2025).
[14] A. Krizhevsky, V. Nair, and G. Hinton. CIFAR-10 and CIFAR-100 datasets. Available at: https: //www. cs. toronto. edu/"kriz/ kriz/cifarr html, 2009, (Accessed: 16 May 2025).
[15] H. S. Laqueur, A. B. Shev, and R. M. C. Kagawa. SuperMICE: An ensemble machine learning approach to multiple imputation by chained equations. American Journal of Epidemiology, 191(3):516-35) 1-5525, Mar. 2022.
[16] J. Li, A. Madry. J. Peebles, and L. Schmidt. On the limitations of first-order approximation in GAN dynamics. In Proceedings of International Conference Machine Learning, 2018.
[17] Y. Lipman, R. T. Q. Chen, H. Ben-Hamu M. Nickel, and M. Le. Flow matching for generative modeling. In International Conference Learning Representations, 2023.
11

---

[18] R. J. A. Little and D. B. Rubin. Statisticul analysis with missing data. John Wiley & Sons, 2019.
[19] Q. Liu and D. Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. In Advances in Neural Information Processing Systems, 2016.
[20] S. Liu, J. Yu, J. Simons, M. Yi, and M. Beaumont. Minimizing fdiieergences by interpolating velocity fields. In Proceedings of International Conference on Machine Learning, 2024.
[21] X. Liu, C. Gong, and Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. In International Conference on Learning Representations, 2023.
[22] Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attrihutes the wild. In Proceedings of International Conference on Computer Vision (ICCV), 2015.
[23] A. Lugmayr, M. Danelljan, A. Romero, F. Yu, R. Timofte, and I.. Van Gool. RePaint: Inpainting using denoising diffusion prohabilistic models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2022.
[24] P.-A. Mattei and J. Frellsen. Miwae: Deep generative modelling and imputation of incomplete data sets, In Proceedings of International Conference on Machine Learning, 2019.
[25] N. Meinshausen and P. Bühlmann. High dimensional graphs and variable selection with the Lasso, Annals of Statistics, 34(3):1436 1462. June 2006.
[26] B. Muzellec, J. Josse, C. Boyer, and M. Cuturi. Missing data imputation using optimal transport. In Proceedings of International Conference on Machine Learning, 2020.
[27] Y. Ouyang, L. Xie, C. Li, and G. Cheng. MissDiff: Training diffusion models tabular data with missing values. In ICML 2023 Workshop on Structured Probabilistic Inference & Generative Modeling, 2023.
[28] P. Ravikumar, M. J. Wainwright, and J. D. Lafferty. High-dimensional Ising model selection using & Ei-regularized logistic regression. Annals of Statistics, 38(3):1287-1311 2010.
[29] D. J. Rezende and S. Mohamed. Variational inference with normalizing flows. In Proceedings of International Conference on Machine Learning, 2015.
[30] T. W. Richardson, W. Wu, Lin, B. Xu, and E, A. Bernal. McFlow: Monte Carlo flow models for data imputation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
[31] D. B. Rubin. Inference and missing data, Biometrika, 63(3):581-55) 5992, 1976.
[32] Y. Song and Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 2019.
[33] Y. Song, J. Sohl-Dickstein D. Kingma, K. A., S. Ermon, and B. Poolc. Score- generative modcling through stochastic differential equations. In International Conference on Learning Representations, 1021.
[34] D. J. Stekhoven and P. Bühlmann. MissForest--nn--paaaaeee non-parametrie missing value imputation for mixed-type data Bioinformatics, 28(1):1 2012.
[35] M. Uehara, T. Matsuda, and J. K. Kim. Imputation estimators for unnormalized models with missing data. In Proceedings of International Conference on Artificial Intelligence and Statistics, 2020.
[36] K. L. Vaden, M. Gebregziabher.. S. E. Kuchinsky, and M. A. Eckert. Multiple imputation of missing IMRI data in whole brain analysis. Neuroimage, 6033):843-18855 2012.
[37] S. van Buuren and K. Proothuix-Oudssoone MICE: Multivariate imputation by chained equations in R. Journal of Statistical Software, 45(3):1667, 2011.
[38] J. Yoon, J. Jordon, and M. Schaar. GAIN: Missing data imputation using generative adversarial nets. In Proceedings of International Conference on Machine Learning, 2018.
[39] H. Zhang, I.. Fang, Q. Wu, and P. S. Yu. DiffPuter: An EM-driven diffusion model for missing data imputation. In International Conference on Learning Representations, 2025.
[40] Z. Zhang, M. Li, and J. Yu. On the convergence and mode collapse of GAN. In SIGGRAPH Asia 2018 Technical Briefs, SA 11, pages 1-4, New York, NY, USA. 2018. Association for Computing Machinery.
[41] H. Zhao, Sun, A. Dezfouli, and E. V. Bonilla. Transformed distribution matching for missing value imputation. In Proceedings of International Conference on Machine Learning, 2023.
[42] S. Zheng and N. Charoenphaklee. Diffusion models for missing value imputation in tabular data. In NewrIPS Table Representation Learning (TRL) Workshop, 2022.
12

---

A Illustration of Sequential Imputation
Fort 1..
O
M
(g))
M
X
[
impute
[gc
[gr]
[X 1z
train
[x(c-1)) J2
[x 12
(x(t-11,, M)
ft
[X
1
[x(t-10]s la
[x Js
[X]
[0.]
[Ørls
[g.]s
LØrls_
M aberrad
Figure 4: Schematic of the sequential imputation algorithm. The process begins with initial imputation, (0), and iteratively refines the missing values over I steps.
Figure 4 illustrates the sequential imputation framework, where a model is iteratively trained to refine estimates for missing data.
B Proof of Proposition 1
Proof. Consider the difference between the mutual information iteration and
D
(
(t) X
PM
D X D.M
Px Pxa-u
60 [P]]
X.M
D
p())
X.M
FM
D
(p(0) X ( & PM +D
(p()) X.M
I ( 8 PM DD [P & PM]
(16)
Since ( M () @ PM]] minimized at g: and ( (1) (g:) by definition,
D
p(t) X.M 15 2 PM] < [P ( P P @ [M Vg.
(17)
Moreover, D [P MMIlP PM] = D [PP | P & PM] Identiiy, so B is non- positive. To show A is non-positive, we write out the KI.-divergences in A. Let p(x, m) be the density of (X(1), M) and q(x) be the density of (t-1)
p(t) X.M P & FM
D
X.M
X
&
PM
(18)
(X
M) (M)
p (( M)
=E log
E log
(19)
p
((X()))pM))
E
log ()))
log (0) s
(20)
where the last inequality is due to the Gibbs inequality. Thus, A B desired.
C Proof of Theorem 1
Proof. First, we introduce the following lemma, based on the marginal preserving property of rectified flow.
13

---

Lemma 1. The imputation ODE with the initial condition given in (10), i.e.,
Zo X0.1-m/ (Xo.m Xm, Mo m)
(21)
has solution the terminal time 1, Z1 X1,1-m/ {X1.m xm) for every fixed x, m.
See the Appendix D for the proof.
Now derive the main result. Lemma shows that under the initial condition, the solution Z is
Zi
d X 1, 1,1-ml
{X1,m Xm)
X (x(-1)) Xm
(22)
Moreover, Z is also our imputer gr, so by construction (9),
Z m [ xi Mo m
(23)
Combining (23) and (22), we obtain the following equality
10)
X
==xmm M=m
x 1 u 1) ==="
(24)
The imputation ODE does not change observed components where m, 1, x have
x We
P (1))
[x Xm
IP
(
(25)
X
(24) and (25) imply that Pg- (- (mx m) = ((((xm)) According to Proposition 2, we have shown that g" is an optimal imputer that minimizes D (g) x @ PM]
D Proof of Lemma 1
Definition 1. For fixed values y, u, v, m. the event defined as
E(u, w. m) (Xu.m = u, X1m = w, My = m).
(26)
and the optimal solution of (8) is
u, : E X.. m X X,,1 = E,u, w. w.
(27)
We denote the density of random variable X w, m) as pr (ylu, w, m).
Lemma 2. poly|u, w, m) = Po (y u, m) and p1 (y u, w, m) = (y)
Proof. We can see that due to the independence X, Mo).
X X X1m M, X Xo.m. Mp.
(28)
hence the first equality holds. What's more,
X1, Xo,m. X X1,,, Mg X1,1 1
(29)
hence the second equality holds.
Lemma 3. The density function Py (y|u, u, w, m) satisfies the continuity equation
Qypr(ylu, w. m) + Vy (p-(ylu, w, m) v" w, m,,,)) =
(30)
14

---

Proof. Consider the following expectation:
w, Jga, w,)
(31)
whose time derivative
1mm) w, w.m] p w,
(32)
On the other hand, that X, is deterministic function of X. and X1, thus using the law of unconscious statistician:
0, E. h(( ,, E(u m)
(33)
= Xx. Xo.X. X1 (x (. (u, w, m)]
(34)
=Ex,,x, x, ( [E, w, m)]
(35)
=E (Xr,1- E w, m))
(36)
Expanding the previous expression using the chain rule,
E h X 1 w. m)
(37)
=R 1- X -- E(, w, m)
(38)
=E Vh X
X
E(u, w. m)
(39)
=E E V
X 1-
X - w, m)]
E(u, w. m)
(40)
=E VM(X, (-) B X( - Eu,, w, m)]
E(u, w, w))
(41)
= / VM(y)T u, u, w. m)
(42)
=- f Myy v. (v u. w,, w.) dy
(43)
where the last equality by performing integration by parts the right- hand side (with respect to y).
Now we equate the two representations ((32) and (43)) for 3 E,E[(, 1.r1-m)) E(u w, m)]:
[ h(y) O.p, (ylu, w. m) dy = -D h(y) V.
(v V m. p. dy. (44)
Since this equality holds for every smooth, compactly supported test function h(y), the fundamental lemma of the calculus of variations implies that,
0,p-(ylu, w, m) + V. (y" u, w, w. w.m) 0.
(45)
We can see that v satisfies the continuity equation for a time-varying density function Pr.
We can see that with the initial condition
Zo = Xo,,-m/E(u, w, m),
Lemma 3 says
Z, X1,,,,,mEE(,, w. m).
Using Lemma 2, we simplify Zo and Z
Zg = X [Xo. = u, M, = m), Z, = X1,1 (X1, w)
Letting
Xm- we obtain the desired results in Lemma 1.
15

---

E Proof of Proposition 2
Proof. Since the sequential imputer (3) never changes the observed part X nor the missing mask M, we have
Pg (x, m) Pg (x Xm: m ()) m):
(46)
where the marginal density p(xm: m) does not depend on g. Now we show that the KL divergence only depends on g through the KL divergence DIPx(-Ix1-mi
D (F
x @ [M
(47)
log Pg (x(x, m) ((x)p(m)
= (x.m)~(X M)
(48)
Px ( m Xm, m)
=E (x.m)- M)
log
(49)
q(x1 1-x1-m= Xm ) m
Pg [X] m Xm, m)
log p(xm: m)
=IE (x.m)-(XX M)
log
E (x.,)-(x M)
(50)
g [1 1x1mmi Nm
q( g(xm)p(m)
Pg (x1 Xm, m)
Σ
p(m)E
log
const.
(51)
x M
q(11 (m))
(0.10.)
Σ p(m)
p(xm m) D)Dg (x1 m Xm m) (
Xm [x)]]xx dxx const.
(52)
(0.14))
where const. does not depend g. The constant in (51) is due to the factorization (46):
log p(Xm m)
p(Xm: m)
(xxm) MM)
Ep(xm m) E Bpe (xi -m Mm)
log ( ((m))((m
const.
g(Xm) ((m)
(53)
Assuming p(Xm m) is positive everywhere, Ag is minimized if and only if
D(r& ( Xm, m)]g(x1 - [xm)]
(54)
The KL divergence D(pg (xi ( Xm: m)] ((x) xm)] zero if and only
Ps Px(x1-m Xm m) (xm))
(55)
Thus, if p(m) is strictly positive, D ( PM] is minimized if and only if (x(xi-mlo Xm. m) ((x1-
F Validity of MIRI under MAR Setting
Recall that we denote the observed components of vector X as XM and the missing components as X1-M- Under the Missing Random (MAR) assumption, the goal minimize the conditional mutual information:
Pg (x III Xmm
- M = -(X.m))(XX.M)) log [g((x i-m Xm Xm)p(mml (pm (m)mm)m)))))))))))))))
(56)
Following similar argument to the proof of Proposition we can define iterative algorithm to reduce this conditional mutual information. At each iteration t, we find the optimal imputer g: by solving:
(x.))>(x()) .M) log PR (x i xm)
E arg min E, g
(57)
q(x1-m m Xm )g( (m xm))
where Pg is the density corresponding the distribution of (x(²) M), and y the density of ((-1)
from the previous iteration.
16

---

We now show that Pg Pg((1mml) m. xm) = q(X1- Xm) is sufficient condition for g to be optimal. Analogous to the proof of Proposition 2, we can rewrite the objective function. Since the missingness mechanism is fixed and does not depend on the imputer g. we have p(m[xm) = q( (- The objective then simplifies:
Px (X] mm; m Xmm)
log
(58)
(x .M)
q (X1-mi Xm )g( m Xm
log Pg (11mm m, Xm m Xm
(59)
(x.m)~(X .M)
( Xm m Xm
PR Px m Xm)
log
(60)
(x. M))
Xm)
log Pg X1 mm m. Xm)
Ex Pal mm))
(61)
q( [ ( mxm)
=I [(m,,))) D [g( [,, Xm [lg(-)>m)
(62)
The inner expectation the KL divergence between the conditional distributions (x m, Xm) and g(x1-m Xm). The KL divergence is non-negative and is minimized (to zero) if and only if the two distributions are equal:
Pg (1-m m m, xm) q Xm))
(63)
This demonstrates that the same optimality condition derived for the MCAR setting in Proposition also holds under the MAR assumption.
G Experimental Setup
This appendix provides all information necessary for reproducibility: datasets, evaluation protocol, and computational resources. The hyperparameter settings and its sensitivity studies can be found in our supplemental material. Hyperparameters were selected through a sensitivity analysis. The synthetic datasets used in this sensitivity study consist of 1000 samples and 20 features. A missing rate 20% was applied to each dataset. The data types include Gaussian, Uniform Correlated, and Mixed (Gaussian and Uniform). The results indicate that MIRI is not sensitive to the hyperparameters selections. See our supplementary material for more details.
G.1 Datasets
Synthetic Data. We generate N 6 6.00 samples in IR2 by drawing two equally-sized clusters of 71 = 3.00 points each from isotropic Gaussians. One cluster is centered at (-2, 22) and the other at (2, 2), both with standard deviation 0.5.
UCI Regression Benchmarks. Table 4 lists the ten UCI datasets used. For each, we report sample size and feature dimensionality.
Table 4: UCI datasets used in study.
Dataset # Samples
# Features
wine
599
11
energy
768
8
parkinsons
875
20
stock
536
11
pumadyn32nm
192
32
housing
506
13
forest
517
12
bike
17379
17
solar
1066
10
gas
2565
128
17

---

We randomly sample 5 000 x2 RGB images and apply pixel-level MCAR masks at varying rates. All three colour channels each missing pixel are masked.
CIFAR-10.
CelebA. We randomly sample 5 000 64 RGB images and apply channel-level MCAR masks at varying rates. All three colour channels each missing pixel are masked independently.
G.2
Evaluation Protocol
For each dataset, we generate independent MCAR masks and rerun the full training and imputation pipeline. Reported results are mean + standard deviation over these runs.
G.3
Hyperparameter Selection
Values were chosen via sensitivity analysis. Complete settings and search ranges are provided in our supplemental material. MIRI exhibits low sensitivity to these choices.
G.4
Computational Resources
Tabular Experiments. NVIDIA P100 GPU (16 GB), Intel Xeon E5-2680 v4 CPU (8 cores, 2.4 GHz), 24 GB RAM.
Image Experiments. GHz), 90 GB RAM.
NVIDIA RTX 3090 GPU (24 GB), Intel Xeon Gold 6330 CPU (14 cores, 2.0
G.5 Baseline Implementations
We evaluate Hyperlmpute, MICE, MIWAE, and Sinkhorn using the official Hyperlmpute reposi- tory with default settings."
We adapt the official implementations of KnewImp", TabCSDF GAIN', and TDM from their respective GitHHub repositories. MissDiff is reimplemented based the algorithm described in the original publication [27]. All baselines use default byperparameters unless otherwise stated.
H
Additional Experimental Results
H.1 Performance on Synthetic Data
We evaluate each imputation method using three criteria, computed only on entries originally masked under the MCAR mechanism: (1) Root Mean Square Error (RMSE); (2) Mean Absolute Error (MAE); (3) Maximum Mean Discrepancy (MMD). Table reports the mean = standard deviation over independent MCAR masks (30% missing).
Overall, Hyperlmpute is optimal when minimizing per-entry error, achieving roughly 8% lower RMSE and 14% lower MAE than MIRI. In contrast, MIRI reduces MMD by more than 60% relative Hyperlmpute, thereby better preserving the underlying data distribution despite modest increase in wise error. MissDiff offers a balanced trade-off, whereas GAIN, TabCSDI, and KnewImp underperform both fronts.
H.2 Additional UCI Regression Experiments
We comprehensively evaluate imputation performance on ten UCI regression benchmarks under MCAR, MAR, and MNAR settings (Figures 5, 6, and 7). In some high-dimensional or high- missingness scenarios, certain baselines failed; for example, Hyperlmpute produced runtime errors
hhttp:///s/ ht//g//thhb..cmm/ssooo00p2////1
hhttps://C httpp:::///tthbbccoo/. http::////ghhth.o//./eeaaa
18

---

Table 5: Performance metrics at 30% missingness. Metrics computed over 10 runs; values denote mean standard deviation. The best values are highlighted in bold.
Method
RMSE (L)
MAE (L)
MMD (L)
GAIN
1.128 + 0.004
0.600 = 0.002
0.342 + 0.015
TabCSDI
1.128 = 0.004
= 1600=00000
0.337 = 0.010
KnewImp
1.130 = 0.004
0.600 =6000
0.335 0.009
MissDiff
0,951 = 0.008
0.437 = 0.004
0.189 = 0.009
HyperImpute
0.862 = 0.020
0.278 = 0.006
0.094 - 0.018
MIRI (Ours*)
0.938 + 0.022
0.325 = 0.009
0.036 - 0.007
hypertmpuse
Knawtmp
DoCSS
GAIN
MisiO
Sinchom
MINAE
2:00
5.30
3.06
9.02
bike
energy
forest
gas
housing parkinsonspumadyn
solar
stock
wine
hypertmmute
Knewwrpp
GAIN
IisDoff
Sirichm
MINAE
DScrepancyy
0:20
9.08-
bike
energy
forest
gas
housing parkinsonspumadyn
solar
stock
wine
Hypertmpute
Knewimp
MIsS
0.60
Dhorrppacy
0:30
9.02
bike
energy
forest
gas
housing parkinsonspumadyn
solar
stock
wine
Figure 5: MCAR MMD on 10 UCI datasets (Above: 20% missingness, Middle: 40% missingness, Below: 60 % missingness). The lower the better.
and MIWAE encountered out-of-memory (OOM) errors. The results confirm that MIRI consistently delivers strong distributional fidelity, remains robust under high missing rates, and scales effectively to high-dimensional datasets.
H.3 Additional CIFAR-10 Experiments
We evaluate the imputation quality of MIRI, GAIN, KnewImp, and HyperImpute on 15 randomly selected C1FAR-10 images corrupted under pixel-level MCAR with missing rates of 20%, 40% and 60% (see Figures 8(a), 8(b) and 8(c)). Across all missing-data regimes, MIRI consistently produces accurate, visually coherent reconstructions, in agreement with the quantitative metrics in Table 2. Although Hyperlmpute achieves competitive performance at 20% missingness, MIRI outperforms at both 40% and 60% missingness.
19

---

MIRI
Hyperfmpue
Krewimp
TOOCSEI
GAIN
NICE
Hiso
Sinihom
TDM
MINAE
0:00
0.20
0.03
6:92
bike
energy
forest
gas
housing parkinsonspumadyn solar
stock
wine
NIRI
Knewimp
TobCEDI
GAIN
NICE
Nissomr
Smithom
TDM
Hyperimpute
MINAE
0.50
2.20
2.00
0.02
bike
energy
forest
gas
housing parkinsonspumadyn
solar
stock
wine
Figure 6: MAR MMD on 10 UCI datasets (Above: 40% missingness, Below: 80 % missingness). The lower the better.
Hypertmpute
0:60-
2:30
3.00
3.92
bike
energy
forest
gas
housing parkinsonspumadyn
solar
stock
wine
Knewimo
TabCSDI
Smithom
Hypertmpute
3.5
2:29
2.02
bike
energy
forest
gas
housing parkinsonspumadyn solar
stock
wine
Knewimm
TabcSDI
GAIN
HissoW
Hysertmpute
3.40
2.00
bike
energy
forest
gas
housing parkinsonspumadyn
solar
stock
wine
Figure 7: MNAR MMD on 10 UCT datasets (Above: 20% missingness, Middle: 40% missingness, Below: 60 % missingness). The lower the better.
20

---

Observations
(Ourst)
lmp
Hyperlmpute
(a) 15 uncurated 32x32 CIFAR-10 images and their imputations. 20% of pixels randomly removed from all RGB channels.
Ground Truth
Partial Observatiors
MIRI (Ourst)
GAIN
Knewimp
Miss
Hyperimpute
fife
(b) 15 uncurated 32x32 CIFAR- 10 images and their imputations. 40% of pixels randomly removed from all RGB channels.
Ground
Observations
MIRI (Ourst)
Knewimp
MissDiff
Hyperimpute
(e) 15 uncurated 32x32 CIFAR-10 images and their imputations. 60% of pixels randomly removed from all RGB channels.
I Computational Time
Computational cost matters high dimensions. MIRI is end-to-end vectorized ("vector in, vector out") and exploits parallel hardware common in image/tabular pipelines, whereas round-robin methods such as MICE and Hyperlmpute [12, 37] are inherently sequential and scale poorly with dimensionality.
Wall-clock runtimes. Table reports mean std hours to impute 1000 samples across dimension- alities. We include diffusion-style baselines (DifiPuter [39]. TabCSDI [42]), round-robin (HyperIm- pute), distributionnmatching (MOT [26]. TDM [41]) and VAE (MIWAE [24]). MIRI is on par with diffusion- based methods and substantially faster than round-robin in high dimensions. On CIFAR-10 10 (5k images, 60% missing). Hyperlmpute did not finish within a 24h budget on our PyTorchCUUDA setup, while MIRI finished in s=3.5h the same hardware.
J Euler ODE Solver
In this section, we provide the Euler solver used in our experiments.
21

---

Table Average computation time (hours) for 1000 samples across dimensions (mean-tstd),
Method
50d
200d
500d
1000d
2000d
5000d
Hyperlmpute
0.08 + 0.03
0.23 1 0.02
1.18 + 0.51
3.84 + 0.99
11.65 + 3.34
57.54 + 0.00
TabCSDI
0.00
0.011
0.26 1 0.00
0.00
1 0.01
1.10 0.01
DiflPuter
2.30
0.18
5.06 1 0.01
7.70
12.40 1: 0.21
29.10 -: 0.00
MIWAE
0.00
0.30
0.90 + 0.08
2.02
00M
OOM
MOT
0.31
0.30
0.30
0.38
0.49 0.02
0.56 + 0.00
TDM
0.32
0.55
1.30
3.30
13.90 0.83
OOM
MIRT (Ours*)
0.17
0.21
0.28 0.02
0.37 1: 0.02
0.59 + 0.02
1.30 + 0.01
Algorithm 2 ODE Solver with Euler Method
Require: Velocity field Initial condition {(X,, M.). Number of Euler steps N. N do
for
k
2:
N
1
3:
Vi, X, + X, + N
v(X,,,,7)
4: end for
5: return
{X,}
22

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction list concrete contributions (MIRI, rectified- floowwwwwwwwwoo imputer, relationships existing methods, empirical validation) and clearly state scope and assumptions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The "Conclusion and Future Works" section explicitly covers computational cost, -istiiionn---i-ieee-eeeeeeeeeeee gaps, and challenges under MNAR.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
23

---

Answer: [Yes]
Justification: Propositions/Therren are stated with assumptions; full proof' are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Datasets, masking schemes (MCAR/MAR/MNAR) metrics, protocols (10 runs; is)), baselines and their sources, and compute details are documented. Code link provided.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: The repository URL is given in the abstract; appendix outlines datasets and baselines.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https:/nnips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. cc/ppbblc/guu http//gudeee/eee/ee (CodeSubmissionll
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Appendix "Experimental Setup" gives dataset sizes, masking rates, protocol, hardware; sensitivity study noted. (Optimizens/leaming rates are to be in the repo.)
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Results are reported as me over runs; bar plots include error bars; variability sources are the random masks/seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
25

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Dedicated subsection lists GPUs/CPU//RAM: a runtime table vs, dimensional- ity is included.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: h/pepr/p////r/..../e/ hcc/public/Ethicsuiieeee
Answer: [Yes]
Justification: No human subjects; standard public datasets; foreseeable harmful use beyond typical generative modeling caveats.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Broader impacts are discussed the "Introduction" section.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
26

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: No high-risk models or sensitive datasets are released; scope is standard imputation on public benchmarks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We use only public datasets (UCI, C1FAR-10, and CelebA), which have appropriate licenses/terms of use; citations are provided in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
27

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com comddataeess has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The codebase is announced; include README with environment, commands, and license in the repo
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human subjects or crowdsourcing were used.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects research.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
28

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
The technical contributions do not rely on L.Mss.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
29


Reviews:
[{"review_id": "xJssuok4iH", "reviewer": "Reviewer_iwML", "summary": "This 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.", "questions": "1. 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)?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors have addressed the concerns regarding the motivation for using rectified flow and the long running time. One remaining issue is the moderate performance on entry-wise metrics such as RMSE. Although the authors claim that the method is primarily designed for distributional metrics. This method can be adapted to entry-wise metrics through repeated sampling and expectation estimation. However, the paper lacks experiments to discuss this matter. Therefore, I am increasing my score to borderline accept.", "rebuttal": "# Computational Cost\n> ... 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\nWe appreciate that computational cost is crucial for high-dimensional data imputation. However, we would like to highlight that it is decisively faster than dimension-wise method in high dimensional settings. Hyperimpute or MICE does not scale well in high-dimensional settings, due to their round-robin algorithm. Moreover, the “sequential” nature of these algorithms, making parallelization difficult.  \n\nIn comparison, the proposed method is \"vector in, vector out\", which can fully utilize a variety of hardware which are specialized in dealing with high-dimensional structured data (e.g., images). \n\nOn CIFAR dataset, Hyperimpute failed to finish imputation in a preset time-limit (24 hours), so is not included in the final result, while MIRI finishes in around 3.5 hours on the same environment (Pytorch with CUDA). \n\nWe would also like to highlight that, DiffPuter [3], a state-of-the-art imputation algorithm also adopts a similar algorithm (training and sampling from a flow-based model repeated), which will have a very similar computational cost as our method. \n\nBelow, we compare the average computation time (hours) of 1000 samples with various dimensions. We include methods based on diffusion (DiffPuter, TabCSDI), round-robin (HyperImpute), distribution matching (MOT [1], TDM [2]) and VAE (MIWAE [4]). It can be seen that MIRI is on par with all other diffusion-based methods, if not faster. It is way faster than HyperImpute in high-dimensional space. \n\n| **Methods**| **50d**| **200d**| **500d**| **1000d**| **2000d**| **5000d**|\n|-|-|-|-|-|-|-|\n| HyperImpute| 0.08±0.03| 0.23±0.02|1.18±0.51|3.84±0.99| 11.65±3.34| 57.54±0.00|\n| TabCSDI| 0.18±0.00| 0.20±0.01| 0.26±0.00|0.36±0.00| 0.53±0.01| 1.10±0.01|\n| DiffPuter| 2.30±0.01| 3.22±0.18| 5.06±0.01|7.70±0.05| 12.40±0.21| 29.10±0.00|\n| MIWAE| 0.11±0.00| 0.30±0.00|0.90±0.03|2.02±0.01| OOM| OOM|\n| MOT | 0.31±0.00| 0.33±0.00|0.36±0.00|0.38±0.00| 0.49±0.02| 0.56±0.00|\n| TDM | 0.32±0.00| 0.55±0.00|1.30±0.02|3.30±0.04| 13.90±0.83| OOM|\n| **MIRI (Ours\\*)** | 0.17±0.00|0.21±0.00|0.28±0.02| 0.37±0.02| 0.59±0.02| 1.30±0.01|\n\n[1] Muzellec et al., Missing data imputation using optimal transport. ICML 2020.\n\n[2] Zhao et al.,transformed distribution matching for missing value imputation. ICML 2023.\n\n[3] Zhang et al., DiffPuter: Empowering Diffusion Models for Missing Data Imputation. ICLR 2025.\n\n[4] Mattei and Frellsen, MIWAE: Deep Generative Modelling and Imputation of Incomplete Data Sets. ICML 2019.\n\n---\n\n# Justification of Rectified Flow\n> the paper does not clearly explain why it is preferable to alternatives like diffusion models ... A deeper justification—either theoretical or experimental—would strengthen the case for this design decision. \n\nThanks for raising this point. We would like to clarify that diffusion model **cannot** be used, as the initial distribution of the backward process must be a Gaussian distribution. However, Proposition 2 requires us to transport from a specific conditional distribution $p$ to $q$, (eq. 5), where $p$ is almost always not a Gaussian distribution. Rectified flow is an example of the Flow Matching (Lipman et al., 2023), which are designed to transport between two distributions. We choose rectified flow for its simplicity both in theory and practice. \n\nIn short, diffusion model is designed to \"sample\", while rectified flow is designed to \"transport\". \n\nMore specifically, we also utilized a mathematical property of Rectified Flow, called Marginal Preserving Property (Theorem 3.3 in Liu et al., 2023) to prove Theorem 1 (Lemma 1) where we have shown that sampling using the vector field in eq. 6 can indeed transport samples from $p$ of $q$. We will clarify this point in the revision and discuss the rationale of our choice. \n\nLipman et al., Flow Matching for Generative Modeling, ICLR, 2023\n\n---\n\n# Underperformance on Entry-Wise Metrics \n> For 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\nThanks for pointing it out! There are two types of imputation: Mean imputation (HyperImpute, DiffPuter) and distributional imputation (MIRI, GAIN). \n\nThe former imputes data such that  $X_\\mathrm{imp} = E[X_{\\mathrm{miss}} \\mid X_{\\mathrm{obs}}]$,  while the latter imputes data such that  $X_\\mathrm{imp} \\sim P(X_{\\mathrm{miss}} \\mid X_\\mathrm{obs})$.  Distributional imputation is harder than mean imputation, as we can draw multiple imputations from  $P(X_{\\mathrm{miss}} \\mid X_\\mathrm{obs})$  and compute  $E[X_{\\mathrm{miss}} \\mid X_{\\mathrm{obs}}]$,  but not vice versa. Our method is inspired by GAIN, and therefore also focuses on **distributional imputation**.\n\nThe key in distributional imputation is to test if the impute data looks like it is drawn from the probability distribution that generated the ground truth. For this reason, we test the statistical discrepancy between imputed data and the ground truth using a divergence measure such as MMD.  \n\nMoreover, we can see that Mean impute is optimized for the means squared error metric as $E[X_{\\mathrm{miss}}|X_{\\mathrm{obs}}] = \\arg\\min_V E[\\|X_{\\mathrm{miss}} - V\\|^2|X_{\\mathrm{obs}}]$, so it is not surprising that methods like Hyperimpute achieves a slightly better MSE. \n\nWe also want to highlight that, although our method is not designed to minimize MSE, it also performs well on entry-wise metrics based on MSE such as PSNR in the CIFAR experiment (Table 1):\n\\begin{equation}\n\\text{PSNR} = 10 \\cdot \\log_{10} \\left( \\frac{L^2}{\\text{MSE}} \\right),\n\\end{equation}\nThis indicates though MIRI can still achieve superior entry-wise performance on **high dimensional image dataset**, especially when the missing rate is high. The PSNR metric is widely used in inpainting (which is similar to our setting), like in Elharrouss et al. 2020. \n\nWe will add this explanation to the revision.\n\nWe also tested MIRI under multiple imputations in synthetic data used in Section 6.1 ($p_{\\mathrm{miss}}=0.6$), which can effectively reduce RMSE:\n\n|#Mult. Imp.|MCAR|MAR|MNAR |\n|---|---|---|---|\n|1|1.816|1.047|2.167|\n|5|1.562|0.898|1.700|\n|**10**|**1.513**|**0.866**|**1.666**|\n\nElharrouss et al., Image Inpainting: A Review. Neural Process. Lett., 2020.\n\n---\n\n# MAR, MNAR\n> What will be the performance for other missing patterns like Missing at random (MAR) and Missing Not At Random (MNAR)?\n\nWe have done theoretical analysis and empirical evaluation on MAR data.  \n    \nTheocratically, our method is also justified for the MAR setting. Our method is motivated by the key assumption of MCAR, that true data is independent of the mask.\n\nIn the MAR setting, the independent assumption becomes $X_{\\mathrm{miss}} \\perp M | X_{\\mathrm{obs}}$. Following the same rationale, we can minimize the **conditional mutual information** $\\mathrm{I}[X_\\mathrm{miss}(g); M | X_\\mathrm{obs}]$ given the observed variables. \n\nWe prove that the condition in Proposition 2 is also the sufficient condition of minimizing the conditional mutual information. In other words, running the unmodified MIRI on MAR data will reduce the conditional mutual information $\\mathrm{I}$ iteratively. \n    \nWe provide a proof at the end. \n\nIn addition to these new theoretical results, we also provide empirical proofs that our method works on MAR data. Please see our response to the Reviewer k1XP. \n\nWe will add new empirical results and theoretical results to the revised paper. \n\n---\n\n# Experiment Settings, Missing Baseline\n\nWe perform 10 iterations of the MIRI algorithm, which we observe is long enough for the mutual information to converge. We will clarify this in the revision. \n\nPlease see the Appendix E and supplementary materials (`hyperparameters.md`) for model configuration and training details. \n\nWe also included two methods MOT and TDM as baselines in MAR and MNAR experiments, please see our response to reviewer k1XP. \n\n---\n\n# Proof for MAR Setting\n\nDefine the notation:  $X_\\mathrm{obs} = X_M$ and $X_\\mathrm{miss} = X_{1-M}$, and consider the goal of minimizing the **conditional mutual information**:\n\n$$\\mathrm{I}[X_\\mathrm{miss}(g); M | X_\\mathrm{obs}] = E_{(x,m)\\sim (X,M)}\\log \\frac{p_g(x_{1-m}, m |x_m )}{p(x_{1-m}| x_m)p(m|x_m)}$$\n\nFollowing the a very similar proof in proposition 1, we have the iterative algorithms to reduce the mutual information. At iteration $t$: \n\n- $$\n  g^* = \\arg\\min_{g} E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m) q(m \\mid x_m)}\\right],\n  $$  \n  where $p$ is the density function of $P_{X^{(t)}, M}$ and $q$ is the density function of $P_{X^{(t-1)}}$, the probability of $X$ in the previous iteration.\n\n- Impute using $g^*$ and repeat.\n\nWe now show that  $p(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)$  is a sufficient condition for $g$ to be optimal in the first step.\n\nAnalogous to Proposition 2, we can rewrite the objective function:\n\n\\\\[\n\\begin{aligned}\nE_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m) p(m \\mid x_m)}\\right] =&E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right] \\\\\\\\\n=&E_{(m, x_m) \\sim (M, X_M)} E_{x_{1-m} \\sim X_{1-M} \\mid M, X_M} \\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right]\n\\end{aligned}\n\\\\]\n\nNote that $p(m \\mid x_m)$ does not change across iterations, so we assume $p(m \\mid x_m) = q(m \\mid x_m)$.\nThe inner expectation is a KL divergence between $p_g(x_{1-m} \\mid m, x_m)$ and $q(x_{1-m} \\mid x_m)$,  which is minimized (to zero) when the two distributions are equal:\n\n\\\\[\np(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)\n\\\\]\n\nThis is exactly the sufficient condition stated in Proposition 2."}, {"review_id": "G9GHw3gcn0", "reviewer": "Reviewer_Dwu5", "summary": "This 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.", "questions": "- 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.", "limitations": "The 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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## 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.", "was_revised": true, "final_justification": "The authors have provided a detailed and thoughtful rebuttal that effectively addresses the main concerns raised in my initial review. I particularly appreciate the expanded comparison with DiffPuter, the discussion on the limitations under MNAR, and the inclusion of MIWAE as a strong baseline, which enhances both the theoretical positioning and the empirical rigor of the paper. While the clarity and motivation in the Introduction could still benefit from revision, I trust the authors will incorporate the reviewers’ suggestions in the final version. Overall, I find the contribution original and well-justified, and I have updated my score to reflect my strengthened confidence in the work.", "rebuttal": "# Comparison with DiffPuter\n>The proposed method seems closely related to DiffPuter, ... 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\nDiffPuter [1] and MIRI are motivated from different objectives: MIRI minimizes the mutual information between the mask and imputed data while DiffPuter or MICE [2] performs a pseudo likelihood maximization repeatedly where the missing data is replaced with mean-impute (see Section 3.2 in DiffPuter). The ingenuity of DiffPuter is that instead of maximizing the joint likelihood explicitly over both missing and observed variables, it trains a generative model (or implicit likelihood model) on the imputed dataset to produce the mean impute used in the next iteration. This overcomes the difficulties of training and sampling from a high-dimensional likelihood model, which is often intractable and has to be reduced to a dimension-wise regression problem (like MICE). \n\nBoth methods train flow-based generative models but use them differently: We use flow model to \"match distributions\" (as required by Proposition 2) while DiffPuter learns a joint score model to sample conditionally. These two schools of thoughts leads to different choices of flow models. \n\nMIRI chooses rectified flow as it efficiently constructs a flow between two probability distributions. DiffPuter chooses diffusion model as it can use the joint score model to \"inpaint\" missing variables. \n\nComparing to DiffPuter, MIRI has two clear advantages: \n- MIRI has a clear information theoretic objective. However, DiffPuter does not maximize a clear utility function (e.g., the marginal likelihood of the observed variables). \n- Inpainting is not a proper conditional sampling procedure which is described in the algorithm in Section 3.2 of [1]. The inpainting procedure undermines the forward-backward process that is the key to the diffusion model. In fact, conditional sampling using diffusion model is non-trivial, and often requires training a dedicated conditional score [3]. However, training a conditional score is difficult in missing data setting as the unobserved variables vary from sample to sample. In our algorithm, the transport between two conditional probabilities using rectified flow is justified through Theorem 1. \n\n[1] Zhang et al., DiffPuter: Empowering Diffusion Models for Missing Data Imputation. ICLR 2025.\n\n[2] van Buuren and Groothuis-Oudshoorn, mice: Multivariate imputation by chained equations in R. J. Stat. Softw., 2011.\n\n[3] Sharrock et al., Sequential Neural Score Estimation: Likelihood-Free Inference with Conditional Score Based Diffusion Models, ICML 2024.\n\n---\n\n# GAIN Experiment Setup\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? \n\nWe thank the reviewer for raising this point. For the CIFAR-10 experiments, we adapted the official GAIN implementation [4]. We used a CNN-based generator–discriminator identical to MIRI’s network architecture. To ensure a fair comparison, we applied the same learning rate, batch size, and optimizer settings as used in our method.\nWe tested our setting using the MNIST dataset, and our selection of parameters produces convincing handwriting samples which are comparable to Figure 4 in [4].\n\nOur implementation of GAIN can be found at code/competitors/gain_imp.py in the supplementary material. \n\nAs we briefly mentioned in Section 2.2, GAN is more unstable than training a normal deep neural network model due to the adversarial training. The parameters like learning rates and number of epochs may have significant impact to the final outcome. \nHowever, the proposed method requires only optimizes a least squares objective using SGD, thus, is significantly more stable and robust against different choices of tuning parameters (see our ablation study, sensitivity experiments in the **supplementary material**). \n\n[4] Yoon et al., GAIN: Missing Data Imputation using Generative Adversarial Nets. ICML 2018.\n\n---\n\n# MAR, MNAR\n> 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. ... \n\nWe thank reviewer for raising this excellent point and we will clearly state the limitations of missingness assumption in the revision. Indeed, our method may not work well under MNAR settings as the missing mask may be strongly correlated with the ground truth, making any algorithm that minimizes the mutual information unlikely to work. Interestingly, Section 6 highlights this limitation across a broad range of existing imputation methods. The methods discussed there also implicitly minimize mutual information, making them similarly unsuitable for MNAR scenarios. We will discuss this matter in more detail in the revision. \n\nHowever, we would also like to highlight that our method not only works with MCAR, but also MAR as a way of enforcing the conditional independence $X_\\\\mathrm{miss} \\\\perp M | X_\\\\mathrm{obs}$. \n\nPlease see a proof at the end of the reply. \n\n---\n\n# VAE Baseline & Additional Results\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\nThank you for raising this point! We have now included MIWAE [5] as an additional baseline. \n\nWe compare our method with MIWAE under MCAR, MAR, and MNAR settings. Below table shows the MMD comparisons. We will include this result in the revision. \n\n|Mechanism|$p_{\\mathrm{miss}}$|Method|forest|gas|housing|parkinsons|solar|stock|\n|-|-|-|-|-|-|-|-|-|\n|MCAR |0.2|**MIWAE**|0.220 ± 0.042|OOM|0.164 ± 0.010|0.057 ± 0.016|0.157 ± 0.011|0.085 ± 0.007|\n|||**MIRI (Ours\\*)**|**0.066 ± 0.014**|**0.045 ± 0.008**|**0.083 ± 0.010**|**0.031 ± 0.007**|**0.076 ± 0.007**|**0.033 ± 0.007**|\n||0.4|**MIWAE**|0.247 ± 0.018|OOM|0.166 ± 0.016|0.058 ± 0.010|0.165 ± 0.015|0.086 ± 0.014|\n|||**MIRI (Ours\\*)**|**0.069 ± 0.010**|**0.044 ± 0.010**|**0.083 ± 0.004**|**0.031 ± 0.007**|**0.080 ± 0.007**|**0.030 ± 0.006**|\n||0.6|**MIWAE**|0.245 ± 0.018|OOM|0.154 ± 0.014|0.061 ± 0.012|0.159 ± 0.019|0.089 ± 0.011|\n|||**MIRI (Ours\\*)**|**0.062 ± 0.012**|**0.043 ± 0.006**|**0.089 ± 0.009**|**0.029 ± 0.009**|**0.077 ± 0.007**|**0.026 ± 0.003**|\n|MAR|0.4†|**MIWAE**|0.313 ± 0.030|0.216 ± 0.006|0.136 ± 0.014|0.091 ± 0.014|0.166 ± 0.048|0.108 ± 0.030|\n|||**MIRI (Ours\\*)**|**0.084 ± 0.020**|**0.045 ± 0.008**|**0.121 ± 0.003**|**0.037 ± 0.012**|**0.092 ± 0.021**|**0.037 ± 0.007**|\n|MNAR|0.4|HyperImpute|0.178 ± 0.008|0.049 ± 0.005|**0.087 ± 0.017**|Error|Error|0.069 ± 0.010|\n|||GAIN|0.704 ± 0.005    |0.805 ± 0.002|0.750 ± 0.006|0.197 ± 0.011|0.325 ± 0.004|0.179 ± 0.008|\n|||TabCSDI|0.704 ± 0.006|0.810 ± 0.013|0.739 ± 0.010|0.195 ± 0.013|0.327 ± 0.006|0.190 ± 0.004|\n|||MissDiff|0.085 ± 0.028|0.076 ± 0.015|0.116 ± 0.007|0.049 ± 0.005|0.093 ± 0.014|0.057 ± 0.004|\n|||Sinkhorn|0.259 ± 0.036|0.307 ± 0.019|0.262 ± 0.017|0.089 ± 0.018|0.179 ± 0.027|0.098 ± 0.005|\n|||**MIWAE**|0.220 ± 0.042|OOM|0.186 ± 0.012|0.071 ± 0.017|0.146 ± 0.027|0.087 ± 0.002|\n|||TDM|0.284 ± 0.056|0.176 ± 0.006|0.209 ± 0.007|0.098 ± 0.020|0.171 ± 0.042|0.111 ± 0.001|\n|||**MIRI (Ours\\*)**|**0.066 ± 0.014**|**0.045 ± 0.008**|0.097 ± 0.002|**0.043 ± 0.007**|**0.093 ± 0.003**|**0.052 ± 0.002**|\n\n†MAR: 50% of observed variables are used to generate mask.\n\nWe also provide additional results on MAR and MNAR data. Please see our response to the Reviewer k1XP.\n\n[5] Mattei and Frellsen, MIWAE: Deep Generative Modelling and Imputation of Incomplete Data Sets. ICML 2019.\n\n---\n\n# Proof for MAR Setting\n\nDefine the notation:  $X_\\mathrm{obs} = X_M$ and $X_\\mathrm{miss} = X_{1-M}$, and consider the goal of minimizing the **conditional mutual information**:\n\n$$\\mathrm{I}[X_\\mathrm{miss}(g); M | X_\\mathrm{obs}] = E_{(x,m)\\sim (X,M)}\\log \\frac{p_g(x_{1-m}, m |x_m )}{p(x_{1-m}| x_m)p(m|x_m)}$$\n\nFollowing a very similar proof in proposition 1, we have the iterative algorithms to reduce the mutual information. At iteration $t$: \n\n- $$\n  g^* = \\arg\\min_{g} E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m) q(m \\mid x_m)}\\right],\n  $$  \n  where $p$ is the density function of $P_{X^{(t)}, M}$ and $q$ is the density function of $P_{X^{(t-1)}}$, the probability of $X$ in the previous iteration.\n\n- Impute using $g^*$ and repeat.\n\nWe now show that  $p(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)$  is a sufficient condition for $g$ to be optimal in the first step.\n\nAnalogous to Proposition 2, we can rewrite the objective function:\n\n\\\\[\n\\begin{aligned}\nE_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m) p(m \\mid x_m)}\\right] =&E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right] \\\\\\\\\n=&E_{(m, x_m) \\sim (M, X_M)} E_{x_{1-m} \\sim X_{1-M} \\mid M, X_M} \\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right]\n\\end{aligned}\n\\\\]\n\nNote that $p(m \\mid x_m)$ does not change across iterations, so we assume $p(m \\mid x_m) = q(m \\mid x_m)$.\nThe inner expectation is a KL divergence between $p_g(x_{1-m} \\mid m, x_m)$ and $q(x_{1-m} \\mid x_m)$,  which is minimized (to zero) when the two distributions are equal:\n\n\\\\[\np(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)\n\\\\]\n\nThis is exactly the sufficient condition stated in Proposition 2."}, {"review_id": "I0HQcbW3i4", "reviewer": "Reviewer_yRf4", "summary": "The 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.", "questions": "Can the authors think more about how their method situates among MAR vs. MNAR imputation methods?", "limitations": "The 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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "I 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.", "was_revised": true, "final_justification": "I think the authors have addressed the limitations of this work, namely that it is restricted to MAR problems in their responses, edits to the paper, and acknowledgments that this is a future area for research. As such, I recommend the paper for a borderline accept.", "rebuttal": "# MAR, MNAR\n\n> This method is restricted to MAR problems... Can the authors think more about how their method situates among MAR vs. MNAR imputation methods?\n\nThanks for pointing it out! We will make sure to include a paragraph in the introduction to clarify missingness assumption of the proposed method works. \n\nOur method is motivated by the key assumption of MCAR, that true data is independent of the mask. For MCAR data, the proposed method reduces the mutual information. For MAR data, the proposed method reduces the conditional mutual information, given the observed variables (see proof at the end of the reply). \n\nHowever, for MNAR dataset, the ground truth could be _highly correlated_ with the missing mask, thus minimizing any type of mutual information is unlikely to work. To impute MNAR datasets, one must have some knowledge about the missingness pattern and _learn the probability of missing_. Our method does not incorporate such an assumption, thus will produce a biased result. However, extending MIRI to MNAR dataset is certainly an interesting future direction. \n\nIt is also worth noting that, we can view other imputation algorithms as MIRI under various modeling and optimization approximations (Section 5), meaning these popular imputation methods are also unlikely to work under the MNAR setting as they are also implicitly minimizing the mutual information between $X$ and $M$. We will include these discussions in the revision. \n\nWe include some additional empirical results on MNAR data below. \n\n>The 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). \n\nThank you for raising this point. We will definitely clarify our assumption in the revision! \n\n---\n\n# MNAR Experiments\n\nAs empirical validation, we also performed additional experiments under MNAR settings. As shown in the table below, although our proposed method is _not_ designed for MNAR, it still consistently delivers low MMD across the baselines. Note that we also added a few other baselines that other reviewers requested and please see these reviews for references. \n\n**MMD of MNAR imputations on UCI datasets ($p_{\\mathrm{mis}}=0.4$). All variables are used to generate mask.**\n| **Method**        | **forest**          | **gas**             | **housing**          | **parkinsons**       | **solar**           | **stock**           |\n|-------------------|---------------------|---------------------|----------------------|----------------------|---------------------|---------------------|\n| HyperImpute       | 0.178 ± 0.008       | 0.049 ± 0.005       | **0.087 ± 0.017**    | Error                | Error               | 0.069 ± 0.010       |\n| GAIN              | 0.704 ± 0.005       | 0.805 ± 0.002       | 0.750 ± 0.006        | 0.197 ± 0.011        | 0.325 ± 0.004       | 0.179 ± 0.008       |\n| TabCSDI           | 0.704 ± 0.006       | 0.810 ± 0.013       | 0.739 ± 0.010        | 0.195 ± 0.013        | 0.327 ± 0.006       | 0.190 ± 0.004       |\n| MissDiff          | 0.085 ± 0.028       | 0.076 ± 0.015       | 0.116 ± 0.007        | 0.049 ± 0.005        | 0.093 ± 0.014       | 0.057 ± 0.004       |\n| MOT          | 0.259 ± 0.036       | 0.307 ± 0.019       | 0.262 ± 0.017        | 0.089 ± 0.018        | 0.179 ± 0.027       | 0.098 ± 0.005       |\n| MIWAE             | 0.220 ± 0.042       | OOM                 | 0.186 ± 0.012        | 0.071 ± 0.017        | 0.146 ± 0.027       | 0.087 ± 0.002       |\n| TDM               | 0.284 ± 0.056       | 0.176 ± 0.006       | 0.209 ± 0.007        | 0.098 ± 0.020        | 0.171 ± 0.042       | 0.111 ± 0.001       |\n| **MIRI (Ours\\*)** | **0.066 ± 0.014**   | **0.045 ± 0.008**   | 0.097 ± 0.002        | **0.043 ± 0.007**    | **0.093 ± 0.003**   | **0.052 ± 0.002**   |\n\nWe will add these additional empirical results in the revision.\n\n---\n\n# Proof for MAR Setting\n\nDefine the notation:  $X_\\mathrm{obs} = X_M$ and $X_\\mathrm{miss} = X_{1-M}$, and consider the goal of minimizing the **conditional mutual information**:\n\n$$\\mathrm{I}[X_\\mathrm{miss}(g); M | X_\\mathrm{obs}] = E_{(x,m)\\sim (X,M)}\\log \\frac{p_g(x_{1-m}, m |x_m )}{p(x_{1-m}| x_m)p(m|x_m)}$$\n\nFollowing the a very similar proof in proposition 1, we have the iterative algorithms to reduce the mutual information. At iteration $t$: \n\n- $$\n  g^* = \\arg\\min_{g} E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m) q(m \\mid x_m)}\\right],\n  $$  \n  where $p$ is the density function of $P_{X^{(t)}, M}$ and $q$ is the density function of $P_{X^{(t-1)}}$, the probability of $X$ in the previous iteration.\n\n- Impute using $g^*$ and repeat.\n\nWe now show that  $p(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)$  is a sufficient condition for $g$ to be optimal in the first step.\n\nAnalogous to Proposition 2, we can rewrite the objective function:\n\n\\\\[\n\\begin{aligned}\nE_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m)\\, p(m \\mid x_m)}\\right] =&E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right] \\\\\\\\\n=&E_{(m, x_m) \\sim (M, X_M)} E_{x_{1-m} \\sim X_{1-M} \\mid M, X_M} \\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right]\n\\end{aligned}\n\\\\]\n\nNote that $p(m \\mid x_m)$ does not change across iterations, so $p(m \\mid x_m) = q(m \\mid x_m)$.\nThe inner expectation is a KL divergence between $p(x_{1-m} \\mid m, x_m)$ and $q(x_{1-m} \\mid x_m)$, which is minimized (to zero) when the two distributions are equal:\n\n\\\\[\np(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)\n\\\\]\n\nThis is exactly the sufficient condition stated in Proposition 2."}, {"review_id": "CINrKtCV2Z", "reviewer": "Reviewer_k1XP", "summary": "The 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.", "questions": "Following 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?", "limitations": "The 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).", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## 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).", "was_revised": true, "final_justification": "Most concerns were resolved during rebuttal so moved my recommendation to the acceptance side.", "rebuttal": "# MAR, MNAR\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 ...\n\nOur method works on **MAR and MCAR** data and we provide both theoretical and empirical justification. \n    \nIn the MAR setting, the independent assumption becomes $X_{\\mathrm{miss}} \\perp M | X_{\\mathrm{obs}}$. Following the same rationale of MIRI, we can minimize the **conditional mutual information** given the observed variables. \nWe can show the condition in Proposition 2 is also the sufficient condition to minimize the conditional mutual information. In other words, running the unmodified MIRI on MAR data will reduce the conditional mutual information $\\mathrm{I}[X_{\\mathrm{miss}}(g); M | X_{\\mathrm{obs}}]$, justifying MIRI's capability on MAR datasets. \n\nWe provide a proof at the end of this reply. \n\nWe also conduct a series of experiments to evaluate MIRI's performance under **MAR and MNAR** settings, following settings in [1]. The results demonstrate the robustness and effectiveness of our method. \n\nAt 40\\% missingness, MIRI demonstrates highly competitive performance. As shown in the table below, it achieves the best (lowest) MMD on three over six datasets, (gas, parkinsons, solar) and is a close second on the rest.  \n\n**MMD of MAR imputations on UCI datasets ($p_{\\mathrm{miss}}=0.4$). 50% of observed variables are used to generate mask.**\n|**Method**|**forest**|**gas**|**housing**|**parkinsons**|**solar**|**stock**|\n|-|-|-|-|-|-|-|\n|HyperImpute|0.243 ± 0.059|0.049 ± 0.008|**0.107 ± 0.021**|0.067 ± 0.006|0.109 ± 0.028|0.127 ± 0.018|\n|GAIN|0.628 ± 0.038|0.807 ± 0.014|0.624 ± 0.082|0.316 ± 0.057|0.313 ± 0.086|0.249 ± 0.071|\n|TabCSDI|0.641 ± 0.046|0.801 ± 0.019|0.626 ± 0.064|0.327 ± 0.065|0.311 ± 0.081|0.248 ± 0.076|\n|MissDiff|**0.072 ± 0.008**|0.081 ± 0.011|0.124 ± 0.007|0.042 ± 0.016|0.098 ± 0.029|**0.028 ± 0.005**|\n|MOT|0.408 ± 0.045|0.302 ± 0.003|0.181 ± 0.014|0.228 ± 0.036|0.267 ± 0.082|0.178 ± 0.059|\n|MIWAE|0.313 ± 0.030|0.216 ± 0.006|0.136 ± 0.014|0.091 ± 0.014|0.166 ± 0.048|0.108 ± 0.030|\n|TDM [1]|0.471 ± 0.028|0.301 ± 0.015|0.249 ± 0.027|0.188 ± 0.052|0.270 ± 0.059|0.201 ± 0.065|\n|**MIRI (Ours\\*)**|0.084 ± 0.020|**0.045 ± 0.008** |0.121 ± 0.003|**0.037 ± 0.012**|**0.092 ± 0.021**|0.037 ± 0.007|\n\nMIRI also demonstrates highly competitive performance on MNAR data. It achieves the best (lowest) MMD on five over six datasets and is a close second on the rest.  \n\n**MMD of MNAR imputations on UCI datasets ($p_{\\mathrm{miss}}=0.4$). All variables are used to generate mask.**\n|**Method**|**forest**|**gas**|**housing**|**parkinsons**|**solar**|**stock**|\n|-|-|-|-|-|-|-|\n|HyperImpute|0.178 ± 0.008|0.049 ± 0.005|**0.087 ± 0.017**|Error|Error|0.069 ± 0.010|\n|GAIN|0.704 ± 0.005|0.805 ± 0.002|0.750 ± 0.006|0.197 ± 0.011|0.325 ± 0.004|0.179 ± 0.008|\n|TabCSDI|0.704 ± 0.006|0.810 ± 0.013|0.739 ± 0.010|0.195 ± 0.013|0.327 ± 0.006|0.190 ± 0.004|\n|MissDiff|0.085 ± 0.028|0.076 ± 0.015|0.116 ± 0.007|0.049 ± 0.005|0.093 ± 0.014|0.057 ± 0.004|\n|MOT|0.259 ± 0.036|0.307 ± 0.019|0.262 ± 0.017|0.089 ± 0.018|0.179 ± 0.027|0.098 ± 0.005|\n|MIWAE|0.220 ± 0.042|OOM|0.186 ± 0.012|0.071 ± 0.017|0.146 ± 0.027|0.087 ± 0.002|\n|TDM [1]|0.284 ± 0.056|0.176 ± 0.006|0.209 ± 0.007|0.098 ± 0.020|0.171 ± 0.042|0.111 ± 0.001|\n|**MIRI (Ours\\*)**|**0.066 ± 0.014**|**0.045 ± 0.008**|0.097 ± 0.002|**0.043 ± 0.007**|**0.093 ± 0.003**|**0.052 ± 0.002**|\n\n[1] Zhao, He, et al. \"Transformed distribution matching for missing value imputation.\" International Conference on Machine Learning. PMLR, 2023.\n\nWe will add these theoretical and empirical results in the revision. \n\n---\n\n# Time-series \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\nYes, MIRI works on time-series. We can consider a **self-supervised version** of MIRI. \n\n1. We delete all missing values from the time-series and obtain an zero-padded dataset $X_0$. \n2. Similar to CSDI [2], we randomly sample many masks to form pairs of samples $(X_0, M)$ and run MIRI on it. \n3. We impute the original time-series using vector fields obtained from MIRI. \n\nThis approach is similar to the self-supervised scheme used by CSDI: delete missing values first, if the time-series is not missing too many values, the imputer trained on the remaining time-series should also be a good imputer for the original time-series. \n\nComparing to the existing diffusion-based method (CSDI), our method can impute future time-series with unknown missing lengths while **the imputation space of CSDI is rigid**. CSDI learns to sample from $\\mathbb{P}(X_\\mathrm{miss} | X_\\mathrm{obs})$ using a conditional diffusion model. Once a diffusion model is trained, the dimension of $X_\\mathrm{miss}$ is fixed. Thus, the diffusion model trained on one time-series cannot be used to impute other time-series if the lengths/sizes of the missing patterns are different. However, in our algorithm, the algorithm is agnostic to the space of $X_\\mathrm{miss}$: MIRI does not learn $\\mathbb{P}(X_\\mathrm{miss} | X_\\mathrm{obs})$ but only to finds $X_\\mathrm{imp}$ that satisfies some independence condition. MIRI vector fields trained on one time-series can be easily used to impute other time-series with different missing lengths/sizes. \n\nDue to this flexibility, we can consider a large foundation model, trained by MIRI on many fully observed or partially observed sequences, then use the pretrained vector field to impute/predict any similar unseen time-series. \n\nWe have run MIRI using this algorithm on real-world time-series and obtain promising results as in the below table. For the baselines, we adopted the PyPOTS [5] library with fair experimental settings. The datasets were loaded from the statsmodels library.\n\n**RMSE of imputations on time‐series datasets (MCAR, $p_\\mathrm{mis}=0.3$).**\n|Method|Airpassengers|CO2|Nile|Sunspots|Toy (Sine)|\n|---|---|---|---|---|---|\n|CSDI [2]|21.658±13.127|13.954±14.451|33.806±30.647|23.128±7.211|7.105±4.396|\n|TimeMixer++ [3] |0.989±0.070|0.944±0.024|**1.056±0.109**|0.964±0.115|0.958±0.036|\n|TSLANet [4] |1.010±0.062|0.956±0.045|1.060±0.106|0.966±0.084|0.949±0.050|\n|**MIRI (Ours\\*)**|**0.310±0.050**|**0.342±0.163**|1.107±0.159|**0.518±0.117**|**0.460±0.106**|\n\nWe will discuss this extension to time-series in the revision of our paper and include these empirical results. \n\n[2] Tashiro et al., CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation. NeurIPS 2021.\n\n[3] Wang et al., TimeMixer++: A General Time Series Pattern Machine for Universal Predictive Analysis. ICLR 2025.\n\n[4] Eldele et al., TSLANet: Rethinking Transformers for Time Series Representation Learning. ICML 2024.\n\n[5] Du. PyPOTS: a Python toolbox for data mining on Partially-Observed Time Series, 2023.\n\n---\n\n# Downstream Performance\n> How would MIRI benefit downstream classification/regression tasks using the imputed data? Could authors quantify and show the improvements in experiments?\n\nWe thank the reviewer for raising this issue. We have performed a 10-class classification experiment on imputed CIFAR-10 imputations and the results are shown below. \n\n| **Method**         | **20 %** | **40 %** | **60 %** |\n|--------------------|--------------:|--------------:|--------------:|\n| MissDiff           | 0.486        | 0.269        | 0.152        |\n| KnewImp            | 0.267        | 0.137        | 0.106        |\n| GAIN               | 0.337        | 0.133        | 0.096        |\n| HyperImpute        | 0.804        | 0.405        | 0.212        |\n| **MIRI (Ours\\*)**  | **0.812**    | **0.525**    | **0.364**    |\n\nIt can be seen that the proposed method achieves the highest overall accuracy among all compared imputation methods and we will add this result to the revision. \n\n---\n\n# Proof for MAR Setting\n\nDefine the notation:  $X_\\mathrm{obs} = X_M$ and $X_\\mathrm{miss} = X_{1-M}$, and consider the goal of minimizing the **conditional mutual information**:\n\n$$\\mathrm{I}[X_\\mathrm{miss}(g); M | X_\\mathrm{obs}] = E_{(x,m)\\sim (X,M)}\\log \\frac{p_g(x_{1-m}, m |x_m )}{p(x_{1-m}| x_m)p(m|x_m)}$$\n\nFollowing the a very similar proof in proposition 1, we have the iterative algorithms to reduce the mutual information. At iteration $t$: \n\n- $$\n  g^* = \\arg\\min_{g} E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m) q(m \\mid x_m)}\\right],\n  $$  \n  where $p$ is the density function of $P_{X^{(t)}, M}$ and $q$ is the density function of $P_{X^{(t-1)}}$, the probability of $X$ in the previous iteration.\n\n- Impute using $g^*$ and repeat.\n\nWe now show that  $p(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)$  is a sufficient condition for $g$ to be optimal in the first step.\n\nAnalogous to Proposition 2, we can rewrite the objective function:\n\n\\\\[\n\\begin{aligned}\nE_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m}, m \\mid x_m )}{q(x_{1-m} \\mid x_m)\\, p(m \\mid x_m)}\\right] =&E_{(x,m) \\sim (X,M)}\\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right] \\\\\\\\\n=&E_{(m, x_m) \\sim (M, X_M)} E_{x_{1-m} \\sim X_{1-M} \\mid M, X_M} \\left[\\log \\frac{p(x_{1-m} \\mid m, x_m )}{q(x_{1-m} \\mid x_m)}\\right]\n\\end{aligned}\n\\\\]\n\nNote that $p(m \\mid x_m)$ does not change across iterations, so $p(m \\mid x_m) = q(m \\mid x_m)$.\nThe inner expectation is a KL divergence between $p(x_{1-m} \\mid m, x_m)$ and $q(x_{1-m} \\mid x_m)$, which is minimized (to zero) when the two distributions are equal:\n\n\\\\[\np(x_{1-m} \\mid m, x_m) = q(x_{1-m} \\mid x_m)\n\\\\]\n\nThis is exactly the sufficient condition stated in Proposition 2."}]

Author rebuttal:

Meta-review: The authors introduced novel method for data imputation by minimizing the mutual information between the imputed data and the missing mask. This proposed technique has solid theoretical foundation. Furthermore, the empirical evidence was substantially enhanced throughout the rebuttal period. One limitation is that the proposed approach lacks a theoretical guarantee in MNAR scenarios; however, the authors still present empirical evidence for these settings. I highly recommend that the authors incorporate all the discussions during the rebuttal into the final camera-ready version of their manuscript.
Decision: Accept (poster)