M$^3$-Impute: Mask-guided Representation Learning for Missing Value Imputation

We appreciate the reviewer for the comments. Below we respond to the main concerns raised.

Q(1) First, the discussion of missing data mechanisms is incomplete.

• Thanks for the comment. We would like to point out that we indeed included the evaluation of M$^3$-Impute under all MCAR, MAR, and MNAR settings in the submitted manuscript. Please refer to Table 1, Table 16, and Table 19 for MCAR results, Table 7, Table 17, and Table 20 for MAR results, and Table 8, Table 18, and Table 21 for MNAR results.

Q(2) Second, the graph-based missing data method proposed by You et al. (2020) faces significant issues. In particular, it does not scale well to large datasets, as its efficiency depends heavily on both sample size and the number of variables.

• Thanks for the comment. With all due respect, we would like to clarify that graph-based imputation methods are highly scalable to large datasets. Regarding the training and inference efficiency on large datasets, this can be effectively managed using the graph mini-batch training/inference and the distributed computing techniques [1, 2]. Additionally, the number of model parameters remain independent of both the number of features and the number of samples in the dataset, making graph-based imputation methods efficient in terms of model size as well. In particular, we also demonstrated that M$^3$-Impute runs efficiently for inference on large datasets, such as Protein, Spam, Letter, and California-Housing, with a running time of less than one second to impute all the missing values in each dataset with 30% missingness when using a GPU, as shown in Table 10 of the submitted manuscript. We will soon also share the training time during this discussion period.

Q(3) Third, some concepts similar to the feature correlation unit (FCU) and sample correlation unit (SCU) have been previously explored in the following references [3-7]

• Thanks for the comment. Although there may be some relevance between our work and the previous studies [3,4,5,8] in capturing feature and sample correlations, our FCU and SCU are still substantially different from them. The primary advantage of FCU lies in its ability to dynamically adjust the importance of observable features when imputing different missing features for a sample (see Equation (5)). Similarly, the key benefit of SCU is its capability to adjust both the information passed from peer samples and the importance of observable features when imputing missing values across different samples. This ability to adaptively learn the importance of peers and observable features improves imputation accuracy, as shown in Table 2 for the ablation study in the submitted manuscript. None of the existing graph-based imputation methods and other previous studies [3-7] incorporate such dynamic capabilities introduced by FCU and SCU. Furthermore, we below differentiate our work from the previous studies [3-7] in more detail.

1. IGRM [3]: IGRM enables direct message passing between samples through extra edge creation. However, it incurred out-of-memory (OOM) errors for some of the datasets that we used with their official implementation. Thus we did not include it in our submitted manuscript.

2. HyperImpute [4]: Hyperimpute is an ensemble method that demonstrates highly competitive imputation accuracy and even outperforms the recent imputation methods [9, 10]. It selects a method (e.g., kNN, LR, NN) from its model pool to impute missing values. While these methods can model feature and sample correlations, they do not incorporate the dynamic adjustments introduced by the FCU and SCU. For example, the kNN method imputes missing values by identifying the k nearest neighbors and averaging their values. However, it fails to consider the importance or the weights of different features when calculating their similarities. In contrast, the FCU and SCU can adaptively learn the importance of each feature for imputing different values in different samples. We would also like to point out that M$^3$-Impute consistently outperforms HyperImpute for most of the 25 datasets. Please kindly refer to Table 1, Table 7, and Table 8 (as well as other tables and figures) of the submitted manuscript for more details.

3. GRAFENNE [5]: There are several key distinctions between M$^3$-Impute and GRAFENNE: (1) GRAFENNE does not impute missing values; it targets at completing specific downstream tasks only. (2) GRAFENNE requires labeled data while M$^3$-Impute does not, as M$^3$-Impute adopts a self-supervised learning approach. (3) M$^3$-Impute adaptively learns feature and sample correlations when imputing different values of different samples, whereas GRAFENNE cannot.

4. Lowrank [6]: It introduces prior information in estimating functional brain networks. There seems to be no connection between this paper and missing data imputation. We are not sure why it is mentioned here.

Q(1) You need to clearly describe how the MCAR, MAR, and MNAR settings are generated in the paper. This is currently missing. As a result, your claim regarding MNAR is inaccurate.

Thanks for the comment. We would like to point out that the steps for generating missingness under MCAR, MAR, and MNAR settings are detailed in Section A.2 of the manuscript, and this process is the same as the one in [1, 2].

1. For the MCAR setting, given a dataset containing $n$ samples and $m$ features, we generate an $n \times m$ matrix as a mask. Each entry value of the matrix is sampled from a uniform distribution. Entries with values no greater than the missingness ratio are treated as missing, i.e., masked out, and the remaining entries remain observable.
2. For the MAR setting, a subset of features is randomly selected to be fully observed. Thus, only the remaining features can have missing values. The values for the remaining features are removed according to a logistic model with random weights, using the fully observed feature values as input. The desired missingness ratio is achieved by adjusting the bias term.
3. For the MNAR setting, we first apply the MAR mechanism. Once its output is generated, we also remove additional feature values according to the MCAR mechanism.

[1] You, Jiaxuan, et al. "Handling missing data with graph representation learning." Advances in Neural Information Processing Systems (2020).

[2] Jarrett, Daniel, et al. "Hyperimpute: Generalized iterative imputation with automatic model selection." International Conference on Machine Learning. PMLR, 2022.

Dear reviewer pTfx,

We hope our responses have addressed your concerns. If there are any remaining questions, we would be happy to provide further clarification. May we kindly ask if our responses have resolved your concerns?

Sincerely,

Paper 5469 authors.

Q(6). The paper notes that M$^3$-Impute struggles with datasets where data points have low dependencies, as FCU and SCU rely on data similarity. This can lead to sparse or inadequate graph formation and poor model performance. Can we employ strategies such as enhancing data representation and refining graph construction techniques to tackle these challenges?

• Thanks for the comment. During graph formation, low data-point dependencies do not make the graph sparse and inadequate. Instead, the connectivity of the constructed bipartite graph is determined by the level of missingness in the tabular data. When missingness level is high, the bipartite graph exhibits low connectivity, as missing data is modeled as non-existent edges. We evaluated the robustness of M$^3$-Impute across different levels of missingness (up to 70%). As presented in Table 16 to Table 21. The results show that M$^3$-Impute is robust across different levels of missingness and consistently outperforms other methods. We appreciate the suggested strategies and will consider them as a future work. Thanks.

Q(7). Grape also engages in graph construction and employs graph representation learning for missing value imputation. However, the approach to graph construction varies between the two methods. Has the author conducted a comparative analysis? Is there a correlation between the efficiency of missing value imputation and the method of graph construction? Which approach proves to be more effective?

• Thanks for the comment. We would like to clarify that the graph construction in M$^3$-Impute is the same as the one in Grape. However, M$^3$-Impute improves the initialization of sample-node and feature-node embeddings in Grape. More specifically, Grape initializes all sample node embeddings as all-one vectors and feature node embeddings as one-hot vectors, which have a value 1 in the positions representing their respective features and 0's elsewhere. We observe, however, that such an initialization does not effectively utilize the information from the masked data matrix, which leads to inferior imputation accuracy. Thus, in M$^3$-Impute, we propose to initialize each sample node embedding based on its associated (initial) feature node embeddings instead of initializing them separately. While the feature node embeddings are randomly initialized as a learnable vector, the sample node embeddings are initialized in a way that reflects the embeddings of the features whose values are available in their corresponding samples, detailed in Section 3.2 of the submitted manuscript.

• Finally, we would like to point out that all the novel units proposed in M3-impute (including the initialization of node embeddings, the FCU, and the SCU) are independent of the bipartite graph modeling. Thus, all the novel units can be easily adapted to any graph that would be constructed with an improved modeling technique and can benefit from the improved graph model.

Q(4). The paper should delve into M$^3$-Impute's computational efficiency, including memory usage and optimization, to address resource limitations in practical applications.

• Thanks for the comment. M$^3$-Impute is efficient, with reasonable running time and memory usage. In Tables 4 and 10 of the submitted manuscript, we present a comparison of inference times between M$^3$-Impute and baseline methods. The results highlight that M$^3$-Impute is computationally efficient during inference, requiring less than one second to impute all missing values in each dataset with 30% missingness when using a GPU.

• In addition, the tables below provide a detailed analysis of M$^3$-Impute's computational complexity during the training phase. Table 1 reports the GPU runtime of M$^3$-Impute during training. As shown, M$^3$-Impute demonstrates high efficiency, requiring less than one second to complete the forward and backward passes for a single epoch across all examined datasets. Furthermore, Table 2 summarizes the GPU memory usage of M$^3$-Impute when employing a mini-batch training mechanism with a batch size of 256. The results highlight that M$^3$-Impute maintains a moderate memory footprint while demonstrating excellent scalability, even when applied to large-scale datasets such as Naval, Power, and Kin8nm.

  Table 1: GPU running time per epoch (in seconds) at the training stage.

  Method	Yacht	Wine	Concrete	Housing	Energy	Naval	Kin8nm	Power
  Init	0.02	0.02	0.02	0.02	0.02	0.08	0.04	0.03
  Init + FCU	0.02	0.02	0.02	0.02	0.02	0.10	0.05	0.04
  Init + SCU (Peer=5)	0.03	0.06	0.03	0.03	0.03	0.54	0.15	0.12
  M3-Impute (Peer=5)	0.03	0.06	0.04	0.04	0.04	0.54	0.16	0.14

  Table 2: Memory Usage at the training stage (batch_size=256).

  Method	Yacht	Wine	Concrete	Housing	Energy	Naval	Kin8nm	Power
  Init	1279MiB	1425MiB	1332MiB	1310MiB	1315MiB	3007MiB	1943MiB	1698MiB
  Init + FCU	1281MiB	1433MiB	1341MiB	1317MiB	1321MiB	3017MiB	1953MiB	1703MiB
  Init + SCU (peer=5)	1313MiB	1453MiB	1373MiB	1351MiB	1353MiB	3053MiB	1987MiB	1737MiB
  M3-Impute (Peer=5)	1317MiB	1455MiB	1379MiB	1359MiB	1353MiB	3063MiB	1992MiB	1743MiB

Q(5). The experiments, although extensive, are primarily conducted on a set of datasets that may not fully represent the diversity of real-world data. It would be beneficial to include more datasets with varying characteristics, such as different domains and levels of missingness (eg: a certain type of severe deficiency), to more comprehensively evaluate the robustness and generalizability of M$^3$-Impute.

• Thanks for the comment. We would like to point out that the 25 datasets used in the paper cover a wide range of real-world scenarios, including civil engineering (CONCRETE, ENERGY), physics and chemistry (YACHT), thermal dynamics (NAVAL), healthcare (BREAST cancer, DIABETES), and finance (California-Housing). These datasets also possess mixed data types with both continuous and discrete values. In addition, we would like to clarify that we indeed evaluated M$^3$-Impute against baseline methods across 4 different levels of missingness, ranging from missing 10% to missing 70%, and under three distinct missingness patterns, i.e, MCAR, MAR, and MNAR. Detailed results can be found in the submitted manuscript, where Table 16 and Table 19 present the MCAR results, Table 17 and Table 20 present the MAR results, and Table 18 and Table 21 present the MNAR results.

Q(1). To strengthen the paper, it would benefit from comparing M$^3$-Impute against the most recent state-of-the-art methods in missing value imputation, highlighting its unique advantages over the latest techniques.

• Thanks for the comment. In our manuscript, we included two highly competitive imputation methods. One is Hyperimpute, which is an ensemble method that demonstrates highly competitive imputation accuracy and even outperforms the state-of-the-art imputation methods [1, 2] in 2024. The other is Grape, a slightly older method that still delivers impressive results and occasionally outperforms HyperImpute. We believe the current baseline selection is fair and comprehensive, with a wide range of experiments with 10 baseline methods on 25 open datasets. We would also like to point out that M$^3$-Impute consistently outperforms both HyperImpute and Grape for most of the 25 datasets under all three missingness patterns. Please kindly refer to Table 1, Table 7, and Table 8 (as well as other tables and figures) of the submitted manuscript for more details.

Q(2). The paper should include a detailed examination of M$^3$-Impute's performance with large-scale datasets to assess its scalability and efficiency as data volumes expand.

• Thanks for the comment. We would like to point out that we indeed included some large-scale datasets, such as Protein, Spam, Letter, and California-Housing. M$^3$-Impute runs efficiently for inference on these large datasets and consistently achieves better imputation accuracy compared to baseline methods. The inference time on various datasets is presented in Table 10 of the submitted manuscript. Results show that M$^3$-impute takes less than one second to impute all the missing values in each dataset with 30% missingness when using a GPU. Please also kindly refer to Table 9 of the submitted manuscript for statistics of the large datasets and Table 19, Table 20, and Table 21 for imputation accuracy under various missingness conditions.

Q(3). Adding a section on the model's interpretability would enhance the paper, as it is vital to understand the model's predictions for trust and application in critical fields like healthcare and finance.

• Thanks for the comment. We appreciate your suggestion to add a section on model interpretability and incorporate an extra section (Section B) in the appendix of our revised manuscript. In the newly added Section B.1, we show that FCU utilizes the sample missingness information and dynamically adjusts the correlations between the target missing feature and observed features within each sample. In Section B.2, we demonstrate how SCU effectively measures sample correlations, resulting in improved imputation performance. To this end, we first show the limitations of using the popular cosine similarity measure to compute sample correlations, where it fails to consider the missingness pattern in a sample and the importance of different observed features to the target feature to impute. In contrast, SCU employs a novel mutual sample masking mechanism, integrating it with the FCU to jointly account for commonly observed features across samples and their importance for different imputation targets. This approach results in more robust modeling of sample correlations and significantly enhances imputation performance. Please kindly refer to Section B of the appendix (Pages 26–28) in the revised manuscript. Thanks again.

[1] Du, T, et al. "ReMasker: Imputing Tabular Data with Masked Autoencoding" ICLR 2024

[2] Zhang, H, et al. "Unleashing the Potential of Diffusion Models for Incomplete Data Imputation" arXiv preprint (2024).

Dear reviewer ykDf,

We hope our responses have addressed your concerns. If there are any remaining questions, we would be happy to provide further clarification. May we kindly ask if our responses have resolved your concerns?

Sincerely,

Paper 5469 authors.

Q(3) Is there any additional analysis regarding the computational complexity during the training phase?

• Thanks for the comment. Please kindly refer to the answer for Q(1).

Q(4) Is there any analysis when the proposed method significantly outperforms and when does not. In practicality, it is helpful to know when the proposed method is powerful since it entails more complexity. If the improvement is marginal, GRAPE can be a good alternative.

• Thanks for the comment. Please kindly refer to the answer for Q(2).

Q(5) How did you tune the proposed method? Information about baselines are attached in Appendix but I cannot see about the proposed method. Also, some ablation study only presents marginal changes, so it is difficult to get insights from them. Is there any different results about them?

• Thanks for the comment. (1) The settings of M$^3$-Impute can be found in Section 4.2 of the submitted manuscript. We also included the results of hyperparameter analysis in M$^3$-Impute in Table 12, Table 13, and Table 14 of the submitted manuscript. (2) We include additional ablation results on datasets with medium level of feature correlations in Table 4 below. The results show that M$^3$-Impute again outperforms GRAPE substantially. More importantly, all the three novel units in M$^3$-Impute contribute to the improvement.

   

  Table 4: Additional ablation studies. MAE scores under MCAR setting with 30% missingness are reported.

  Model	Ionosphere	Abalone	German	Steel
  Grape	1.08 $\pm$ .01	2.01 $\pm$ .02	2.02 $\pm$ .01	0.45 $\pm$ .00
  M$^3$-Impute	1.01 $\pm$ .01	1.84 $\pm$ .02	1.87 $\pm$ .02	0.39 $\pm$ .00
  Improv. Ratio	6.48%	8.46%	6.97%	13.33%
  Architecture
  Init Only	1.08 $\pm$ .01	2.00 $\pm$ .02	2.00 $\pm$ .01	0.45 $\pm$ .00
  Init + FCU	1.01 $\pm$ .03	1.89 $\pm$ .02	1.87 $\pm$ .02	0.42 $\pm$ .00
  Init + SCU	1.02 $\pm$ .00	1.84 $\pm$ .02	1.86 $\pm$ .03	0.42 $\pm$ .01

Q(2) Limitations in comparative analysis: While the authors conducted comprehensive experiments, but the improvements are mostly marginal, making it difficult to assess whether the improvements are significant or not. The authors should consider ways to present in a more compelling manner to highlight their significance.

• Thanks for the comment. To provide a clearer understanding of the conditions under which M$^3$-Impute significantly outperforms, we categorize real-world datasets based on the level of feature correlations (Pearson correlations) in each dataset.

1. Low feature correlations: This refers to the datasets where each feature is weakly correlated or even independent of most or all the other features. Imputation in such datasets is particularly challenging, since knowing any of the non-missing features offers little help in imputing the missing ones due to their independence. An example dataset in our submitted manuscript is "Kin8nm", where most features are independent of each other. Nonetheless, as shown in Table 1, M$^3$-Impute achieves the second-best result and does remain competitive, with a difference of less than 0.02 in MAE compared to the top performer.

2. High feature correlations: This refers to the datasets where each feature is highly correlated with most or even all the rest features. For such datasets, imputation is relatively easy, since knowing any of the non-missing features makes the inference of missing values easier. Such datasets in our manuscript are "Naval", "Power", "Breast", "Diabetes", and "Protein". As shown in the Table 2 below, the improvements brought by M$^3$-Impute over SOTA method Grape are marginal (for the comprehensive results, please refer to Table 1, 16-21 in the submitted manuscript). Nonetheless, M$^3$-Impute consistently outperforms GRAPE.

    

   Table 2: Results for datasets with high feature correlations. MAE scores under MCAR setting with 30% missingness are reported.

   Model	Naval	Power	Breast	Diabetes	Protein
   Grape	0.07	1	0.37	1.31	0.25
   M$^3$-Impute	0.06	0.99	0.36	1.29	0.24
   Improv. Ratio	14.29%	1.00%	2.70%	1.53%	4.00%

3. Medium level of feature correlations: This refers to the rest of the datasets (other than the ones exhibiting low and high feature correlations), which is the most common case in real-world applications. Specifically, a dataset has a medium level of feature correlations when each feature is highly correlated with a small subset of features and weakly correlated or nearly independent of the others. For such datasets, imputation is also challenging because we need to consider feature-feature, sample-sample, and sample-feature correlations. Given a subset of observable features and a target feature value to impute in a sample, an ideal model should identify which features are highly or weakly correlated with the target, i.e., sample-feature and feature-feature correlations. In addition, the model should also consider which similar samples to use in imputation, i.e., sample-sample correlations altogether to achieve good imputation results. Our proposed feature correlation unit (FCU) and sample correlation unit (SCU) are designed to capture these three types of correlations. In our submitted manuscript, datasets having a medium level of feature correlations are: "Yacht", "Concrete", "Energy", "Ionosphere", "Abalone", "German", and "Steel". As shown in the Table 3 below, M$^3$-Impute outperforms GRAPE substantially.

    

   Table 3: Results for datasets with medium level of feature correlations. MAE scores under MCAR setting with 30% missingness are reported.

   Model	Yacht	Concrete	Energy	Ionosphere	Abalone	German	Steel
   Grape	1.46	0.75	1.36	1.08	2.01	2.02	0.45
   M$^3$-Impute	1.33	0.71	1.31	1.01	1.84	1.87	0.39
   Improv. Ratio	8.90%	5.33%	3.68%	6.48%	8.46%	7.43%	13.33%

Q(1) Complexity issue: The proposed method is much more complex than other methods, including its predecessor GRAPE. Although the authors already provided the computing time at test phase, but I am curious about how long the training phase takes.

• Thanks for the comment. M$^3$-Impute is also efficient during training. Table 1 presents the GPU running time of M$^3$-Impute components during training. The time taken to impute all the missing values in a training epoch for any dataset we tested is less than one second under the setting of MCAR with 30% missingness.

   

  Table 1: GPU running time per epoch (in seconds) at the training stage.

  Method	Yacht	Wine	Concrete	Housing	Energy	Naval	Kin8nm	Power
  Init	0.02	0.02	0.02	0.02	0.02	0.08	0.04	0.03
  Init + FCU	0.02	0.02	0.02	0.02	0.02	0.10	0.05	0.04
  Init + SCU (Peer=5)	0.03	0.06	0.03	0.03	0.03	0.54	0.15	0.12
  M3-Impute (Peer=5)	0.03	0.06	0.04	0.04	0.04	0.54	0.16	0.14

Dear reviewer 7xNf,

We hope our responses have addressed your concerns. If there are any remaining questions, we would be happy to provide further clarification. May we kindly ask if our responses have resolved your concerns?

Sincerely,

Paper 5469 authors.