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

We thank the reviewer for recognizing our method as a novel approach with innovative masking schemes and that our experiments comprehensively demonstrate superior performance of our method to baseline methods. We also appreciate the constructive comments. Below we provide our response to the concerns raised.

Q1. Detailed sensitivity analysis for other critical hyperparameters.

Thanks for the comment. The results of a detailed sensitivity analysis for other critical hyperparameters are presented in Tables 5-7 in the rebuttal PDF. We will incorporate them into our final manuscript.

Q2. Differentiation with GRAPE and IGRM.

Thanks for pointing this out. GRAPE is a pioneering work that models tabular data with missing values as a bipartite graph. IGRM builds on the bipartite graph modeling and further computes sample correlations to connect similar samples by creating new edges for learning. However, they consider each sample or feature as a whole in computing correlations, whereas in M$^3$-Impute, we treat each entry in the tabular dataset individually to compute its correlation with other entries. This is achieved by our novel soft masking design with Feature Correlation Unit (FCU) and Sample Correlation Unit (SCU). We will clarify this point in the final manuscript. Thanks.

Q3. RMSE comparison between M$^3$-Impute and baseline methods

Thanks for the suggestion. The RMSE results are now provided in Tab.1 below, which shows that our method outperforms the baselines as well. We will incorporate them into the final manuscript.

Tab.1: RMSE under MCAR setting with 30% missingness. RMSE scores are enlarged by 10 times for better clarity.

Q4. Misclassification of MIDA and GAIN.

Thanks for the correction. We will restructure our introduction as suggested to make our arguments more precise in the final manuscript.

Q5. When discussing statistical methods (line 70), the authors should mention that FCS approaches such as MICE are flexible in imputing different types of variables.

Thanks for the comment. We will restructure the related work section and discuss such approaches correctly in the final manuscript.

Q6. More precise statement in Sec. 4.2.

Thanks for the detailed comment. We will make the statements more precise in the final manuscript.

Q7. How can interpretability techniques be incorporated into M3-Impute to help users understand the imputation decisions?

Thanks for the suggestion. We plan to add a plot in the final manuscript to show the intermediate correlation learning process. This plot will illustrate the dependency on other entries when imputing the value of a missing entry in the table.

Q8. Summarize limitations in the conclusion.

Thanks for the suggestion. We will summarize the limitations of our methods, e.g., presence of the datasets where our method is not the top performer, in the conclusion of the final manuscript.

We thank the reviewer for recognizing our method as an advanced and novel approach that achieves significant performance improvements, and that our contribution is excellent. We also appreciate the constructive comments. Below we respond to the main concerns raised.

Q1. The size of the bipartite graph may drastically increase with data size and dimensionality.

Yes, the size of the bipartite graph would increase with larger datasets. However, as demonstrated in Table 4 in the rebuttal PDF, the computation cost would not increase drastically as the mini-batch neighborhood aggregation is efficient in generating node embeddings, which use a reasonable amount of computational resources.

Q2. No investigation on various missingness scenarios and missing rates.

Thanks for the comment. We would like to point out that we indeed carried out this study, and the results are available in Tables 5, 6, 8, and 9 in the appendix of the submitted manuscript.

Q3. Computational and space complexity analysis is needed. Can the proposed method be used for a very large training dataset with high dimensionality?

Yes, we have conducted experiments on new large datasets, including Protein, Spam (with 57 features), Letter, and California Housing, which have a large number of samples and features. The results confirm that our method is feasible and efficient in handling such large datasets. The dataset statistics and running times can be found in Table 1 and Table 4 in the rebuttal PDF, respectively. Thanks.

Q4. How does the proposed method perform under lower and higher missing rates? The inference phase is unclear. Please elaborate on how the proposed method makes imputations for a query instance containing missing values?

Thanks for the comment. The results of different missing ratios are available in Table 8 in the submitted manuscript.

In the inference phase, the pipeline for imputing the missing values of a given sample is detailed in Algorithm 1 in the manuscript. For instance, given a sample with 5 features, where 2 features are missing, M$^3$-Impute first introduces new sample node and establishes edges between the sample node and the three non-missing feature nodes, which already exist in the original graph. The sample node embedding is initialized as per Equation (2). The embeddings of feature nodes and the sample node are then updated through the neighborhood aggregation process using the node embeddings learned from training.

Next, these updated node embeddings are processed through Feature Correlation Unit (FCU) and Sample Correlation Unit (SCU) to learn the feature-wise and sample-wise correlations. Suppose we are to impute the first missing feature, denoted by $f$. FCU computes the correlations between $f$ and all three non-missing features, resulting in the vector $c^{f}_s$. In addition, SCU first measures the similarity between the query sample and a set of peer samples when imputing the missing feature $f$ according to Equation (6), and then fuses the information from the peer samples as in Equation (9), resulting in the vector $z^{f}_s$. Finally, the missing feature $f$ is imputed using Equation (11).

Q5: While the authors reviewed other recent methods that leveraged graph neural networks ([40], [52], [54]), only the method in [52] was compared in the experiments. Why were the methods in [40] and [54] not compared?

Thanks for the comment. The code of [54] was unavailable at the submission time of our paper. While their code was recently available, we have not been able to complete running their code on all the 25 datasets without out of memory errors on our experiment platform (NVIDIA A100 GPU with 80GB memory). Hence, we don't report the results here. Nonetheless, we have been able to run the code of [40] and report the results in Tab.1.

Tab. 1: Imputation accuracy in MAE under MCAR setting with 30% missingness

We thank the reviewer for recognizing that our FCU and SCU are particularly compelling, our method is methodologically sound, and the experiments are extensive. We also appreciate the reviewer for the constructive comments. Below we provide our response to the concerns raised.

Q1. M3-Impute demonstrates strong performance across many datasets but shows limitations in handling datasets with highly independent features or strong linear correlations, such as the KIN8NM and NAVAL datasets. The model's robustness in general scenarios may degrade when confronted with datasets containing extreme correlation structures.

Thanks for the comment. While our method does not show a significant improvement compared to state-of-the-art methods in such extreme cases, it does remain competitive. For instance, in the KIN8NM and NAVAL datasets, M$^3$-Impute achieves second-best results, with a difference of less than 0.02 in MAE compared to the top performers. We have considered 10 new open datasets in our experiments, in addition to the original 15 datasets, leading to a total of 25 open datasets. Our method consistently outperforms the baselines, demonstrating its general applicability.

Q2. Future improvements could concentrate on enhancing M3-Impute's adaptability to these challenging cases to broaden its applicability and robustness across diverse datasets.

Thanks for the comment. We will explore challenging cases where data have extremely strong or weak correlations in future work.

Q3. In ablation study (4.3), Could you provide more details on the refined initialization process of feature-node and sample-node embeddings? How does it differ specifically from the initialization used in Grape?

We appreciate the opportunity to clarify our approach regarding node embedding initialization.

Feature-node Embeddings: In Grape, all feature-node embeddings are initialized as one-hot vectors. Since one-hot vectors are orthogonal, this implicitly assumes that all features are independent. However, in reality, features are often correlated, and initializing feature nodes as one-hot vectors can hinder the modeling of these correlations. To address this issue, we have refined the initialization of feature-node embeddings by using a learnable vector for each feature node. This approach enables the feature node embeddings to be learned during training, allowing correlated features to potentially have similar embeddings from the early stage of learning, which better captures the relationships between features.

Sample-node Embeddings: In Grape, sample-node embeddings are initially set to all-one vectors, which results in identical embeddings for all sample nodes at the beginning. This approach does not capture the unique missingness information of each sample. In contrast, we propose using Equation (2) for initializing sample-node embeddings, where observable feature values of a sample are kept and a small value $\epsilon$ is attached to unobserved features. This method ensures that the initial embeddings reflect the missingness information that is specific to each sample.

We hope this explanation clarifies the rationales behind our refinements and their importance in improving imputation accuracy.

Q4. Can you provide more insights into the computational complexity and runtime performance of M3-Impute compared to the baselines, especially for large datasets?

Thanks for the suggestion. The running time of each method is now shown in Table 4 in the rebuttal PDF. The results show that our method is both accurate and time-efficient. Among the datasets, PRotein, SPam, and LEtter are large datasets, where our method takes a little more time to achieve higher accuracy, which we believe is worthwhile.

Q5. In missing ratio test, why does M3-Impute perform similarly to Grape on the KIN8NM dataset, and what characteristics of KIN8NM contribute to this result?

As can be seen from Figure 4 in the appendix of the submitted manuscript, the features in KIN8NM are highly independent. Almost all the imputation methods perform similarly to Grape on this dataset. Please also refer to our response to Q1. Thanks.

We sincerely thank you for recognizing that our paper is well-written and our empirical results are extensive. Below we provide our response to the concerns raised.

Q1. The paper does not support categorical features. This is a big weakness compared to other imputation methods that can handle categorical features such as iterative approaches like hyperimpute.

We appreciate your feedback and would like to clarify that our method is indeed capable of handling mixed features, including categorical ones.

Specifically, during the forward computation, in M$^3$-Impute, we first convert categorical feature values into numerical values (e.g., categories 1, 2, and 3 are converted to real numbers 1, 2, and 3) and incorporate them alongside numerical features in the initialization stage. Subsequently, M$^3$-Impute performs GNN computations to obtain node embeddings. These embeddings are then processed through Feature Correlation Unit (FCU) and Sample Correlation Unit (SCU) to compute the corresponding context vectors. Finally, M$^3$-Impute uses an MLP with a ReLU activation function for imputing numerical features and an MLP with a softmax activation function for imputing categorical features. This is feasible because we know the data type of the target feature to be imputed and can switch the activation function accordingly. For parameter updates, we use the L1 error for numerical features and the cross-entropy loss for categorical features.

In our submitted manuscript, we already reported the results on the performance of our method on the datasets with mixed feature types, such as the housing and airfoil datasets. To further demonstrate the efficacy of our method in handling categorical features, we have evaluated M$^3$-Impute on four new datasets with mixed feature types. The dataset statistics are presented in Tab.1 and the results are shown in Tab.2.

Tab.1: 4 Additional datasets that contain categorical features.

Tab.2: Imputation accuracy in MAE under MCAR setting with 30% missingness.

Q2. The paper does not discuss the impact of missing value imputation on downstream tasks. Imputation is usually a preprocessing step, and thus assessing the impact on possible downstream tasks is paramount. For example, in supervised learning, some recent evidence suggests that mean/zero imputation is as good as more complex imputations.

Thanks for your constructive comments. We have conducted new experiments on the impact of imputation methods on downstream tasks and report the results in Tab.3. The results show that M$^3$-impute outperforms the baselines, including the mean imputation, for most cases. We will include the new results in the final manuscript.

Tab.3: Averaged RMSE of label prediction under MCAR setting with 30% missingness.

Q3. How does the runtime of M3-impute compare to the other methods? I know this is briefly discussed in appendix 5, but more details runtimes would be appreciated.

Thanks for the question. We now report the running time performance of each imputation method in Table 4 in the rebuttal PDF. We will incorporate it into the final manuscript.

Q4. Why are standard deviations not included in Table 1?

Thanks for the question. Due to space constraints, we were unable to include both MAE scores and the standard deviations of all methods in Table 1, but we provided the comprehensive results (including both MAE scores and standard deviations) in Table 8 in the appendix of the submitted manuscript.

We sincerely thank the reviewer for recognizing that our masking schemes are novel, the proposed correlation units are helpful, and our experiments demonstrate good performance. We also appreciate the constructive comments. Below we provide our response to the concerns raised.

Q1. The initialization of $\mathcal{P}$ in SCU remains unchanged once set. Why not select peers directly based on the degree of similarity rather than introducing randomness?

Sorry for the confusion caused. $\mathcal{P}$ is obtained from a sampling strategy where the probability of a peer being selected is proportional to its cosine similarity with the target node. The set $\mathcal{P}$ is updated every epoch and does not remain unchanged.

We would also like to point out that the set $\mathcal{P}$ balances both high-similarity peers and potential peers that serve for regularization and generalization purposes since directly selecting peers based on the degree of similarity, such as using a $k$-nearest neighbors method, may introduce extra computational complexity and hinder the method's generalization ability. We will incorporate the clarification into our final manuscript. Thanks.

Q2. Table 2 in Section 4.4 for sampling strategy comparison is confusing; the caption is unclear.

The cosine-similarity sampling strategy (introduced in Section 3.4) is integrated into M$^3$-impute for the main results across all experiments. Section 4.4 presents an ablation study where the original cosine-similarity sampling strategy is replaced with a uniform sampling strategy. The performance comparison is shown in Table 2, where M$^3$-impute indicates the cosine-similarity sampling strategy and M$^3$-uniform represents the uniform sampling strategy. As uniform sampling does not consider peer similarities, its performance is expected to be inferior to the original strategy, which is verified in Table 2. We will elaborated more in the table caption and incorporate the changes into the final manuscript.

Q3. Table 3 does not clearly present the trend of $\mathcal{P}$.

As explained in our response to Q1, the set $\mathcal{P}$ is updated every epoch. A proper peer size should balance high-similarity peers and potential peers that serve for regularization and generalization purposes. In general, the trend across different datasets shows that a too-small peer size may only include high-similarity peers, while a too-large peer size may include too many noisy nodes and incur higher computational overhead. The small fluctuations indicate that our method is relatively robust to this parameter. From the extensive experiments on 25 datasets, we recommend a peer size of 5–10 for practical use. We will properly revise the manuscript, not only proving a more comprehensive table but also showing the running time. Thanks.

Q4. A more thorough discussion of the results under different missing types could enhance the motivation and clarity of the paper.

Thanks for the comments. As demonstrated in Tables 1, 5, and 6 in our manuscript, M$^3$-Impute consistently outperforms the baselines across all three missingness patterns. This superior performance is due to M$^3$-Impute's unique approach. Rather than assuming the data follows MCAR, MAR, or MNAR missingness patterns from the outset, we designed M$^3$-Impute to leverage missingness information directly, enabling it to learn feature-wise and sample-wise correlations. Specifically, in Feature Correlation Unit (FCU), M$^3$-Impute uses the missingness results (i.e., known masks) to learn the correlations between the imputation targets and observable features. In Sample Correlation Unit (SCU), M$^3$ employs the missingness results to better capture sample correlations. Since feature and sample correlations exist regardless of the cause of missingness, M$^3$-Impute is naturally adaptive and robust across all three missingness settings.

Another notable feature of M$^3$-Impute is its ability to learn the cause of missingness. In the FCU unit, M$^3$-Impute explicitly captures the correlations between observed and missing features. When data is missing under MAR and MNAR conditions, the missing values depend on the observed ones. Since FCU explicitly captures these relationships, M$^3$-Impute can potentially identify the cause of missingness and enhance imputation accuracy. This capability may explain why M$^3$-Impute significantly outperforms the baselines in the MAR and MNAR settings compared to the MCAR settings.

Q5. Two bars exceed the upper boundary of Figure 3.

We will update the figure to show the full range of the performance. Thanks.

Q6. Add independent labels to all subfigures.

The subfigure labels will be added to the final manuscript. Thanks.

Q7. The influence on performance by applying different GNN models.

We have conducted new experiments using GNN variants such as GraphSAGE, GAT, and GCN. The results are shown in Table 3 in the rebuttal PDF. The results indicate that different aggregation mechanisms may introduce varying errors, but our method consistently outperforms its GRAPE counterpart, demonstrating its effectiveness.

We sincerely thank the reviewer for recognizing that our work is novel with a unique mask-guided representation learning method and we demonstrate strong empirical performance with comprehensive experiments. Below we provide our response to the concerns raised.

Q1. Weakness: Computational complexity

Thanks for the comment. The nature of neural network architectures makes it difficult to analyze the computational complexity rigorously. Nonetheless, we provide numerical results of running time here to show the computation overhead of each method. Please refer to Table 4 in the rebuttal PDF.

Q2. Minor Points #1: ''l4: `Existing imputation methods, however, fall short of considering the ‘missingness’ information in the data during initialization and modeling the entangled feature and sample correlations explicitly during the learning process,' $\rightarrow$ This is not true. Many existing methods consider missingness patterns.''

We appreciate your feedback. Regarding the point about ''fall short of considering the missingness information in the data during initialization'' not being precise, our intention was to highlight that many existing methods ignore the missingness information during their initialization stage. For instance, in iterative imputation frameworks such as MICE and HyperImpute, the initialization stage often involves filling in missing values with some starting values (typically using the mean values of features) before learning. Similarly, graph-based imputation methods like GRAPE and IGRM do not explicitly consider missingness information when initializing their node and edge embeddings. Nonetheless, we agree that our statement can be misleading. We will carefully revise the statement in the final manuscript to avoid any confusions. Thanks.

Q3. Minor Points #2: ''The distinction between statistical and learning based methods seems off. Certainly most learning based methods are statistical, and vice versa.''

Thanks for the suggestion. We agree with this point and will restructure our related work section in the final manuscript.

Q4. Minor Points #3: ''l. 38 struggles $\rightarrow$ struggle''

Thank you very much for pointing out the typo. We will correct it in the final version.

Q5. Questions: (How) do you do HPO for competing methods?

For baselines, we followed the commonly used hyperparameter settings from the previous studies, such as GRAPE, including edge drop-out ratio during training, dropout rate, learning rate, and the number of GNN layers. For M$^3$-Impute, we used the same set of parameters for all the experiments. Additionally, we have now included an analysis of the impact of hyperparameters on the performance of M$^3$-Impute. Results can be found in Tables 5-7 in the rebuttal PDF. They will be incorporated into the final manuscript. Thanks.

Q6. Limitation1: ''Gives little insight into why it works better on some datasets than others.''

Thanks for the comment. We included a performance analysis on different datasets in `Section 4.2 Overall Performance' of the submitted manuscript. To summarize, we apply learnable masks to the data matrix with missing values via our novel units, Feature Correlation Unit (FCU) and Sample Correlation Unit (SCU) to better capture feature-wise and sample-wise correlations, respectively. These correlations are crucial for improving the accuracy of missing data imputation.

We admit that there are few datasets with extreme cases of correlations, e.g., having almost all independent features or all completely dependent features, where our method may not perform as good as on the other datasets. For instance, in the KIN8NM dataset where most features are independent of each other, M$^3$-Impute does not perform as effective as it can be. It is somewhat expected since knowing any of the non-missing features offers little help in imputing the missing ones due to their independence. Nonetheless, M$^3$-Impute achieves second-best results for such datasets, and it does remain competitive, with a difference of less than 0.02 in MAE compared to the top performers. Furthermore, we have considered 10 new open datasets in the experiments, together with the original 15 datasets, leading to a total of 25 open datasets. The results show that our method consistently outperforms the baselines for most cases, demonstrating its general applicability.

Q7. Limitation2: ''Would be interesting to understand better how robust results are under systematic changes in datasets, e.g., different types of missingness.''

Thanks for the comment. We indeed reported the performance of our method under three types of missingness in the submitted manuscript. The results for the MAR and MNAR settings are presented in Table 5 and Table 6 in the appendix of the submitted manuscript, respectively. In addition to the 8 UCI datasets discussed in the main text, we also included performance results on 7 additional datasets under all three types of missingness. These results can be found in Table 9 of the appendix of the submitted manuscript.

We thank the reviewer for recognizing our paper as well-written and the experiments as well-conducted. We also appreciate the constructive comments. Below we provide our response to the concerns raised.

Q1. Although well presented, the method is complicated to understand.

The main idea behind our method is to explicitly utilize the missingness information as input and apply learnable masks to the data matrix to better capture feature-wise and sample-wise correlations, thereby improving imputation accuracy. To this end, we propose two novel units, Feature Correlation Unit (FCU) and Sample Correlation Unit (SCU), to capture feature-wise and sample-wise correlations, respectively. We will improve the presentation of the final manuscript. Thanks.

Q2. Computation resource comparison.

We have added the comparison of running time with other methods in Tab.4 of the rebuttal PDF and will incorporate this into the final manuscript.

Q3. Difference from simple concatenation of the mask to the data matrix.

We would like to point out that the simple concatenation of the mask to the data matrix corresponds to our transformation of the masked data matrix into a bipartite graph, and the execution of an imputation method on the augmented matrix corresponds to running our M$^3$-Impute method on the bipartite graph.

In addition, while we are not quite sure in what sense the paper ''On the consistency...'' is referred to here, we point out that M$^3$-Impute is a task-general architecture, as it is not limited to any specific downstream task. In contrast, in their paper, a downstream task is involved in the design of their method. We further compare M$^3$-Impute with theirs (NeuMiss + MLP). As shown in Tab.1, M$^3$-Impute consistently outperforms. We will include the new results in the final manuscript. Thanks.

Tab.1: MAE of label prediction under MCAR setting with 30% missingness

Q4. Performance under MNAR setting and comparison with not-MIWAE.

We presented the results under the MNAR setting in Table 6 in the submitted manuscript and confirmed that M$^3$-Impute consistently outperforms the baselines. This superior performance is due to M$^3$-Impute's unique approach. Rather than assuming the data follows MCAR, MAR, or MNAR missingness patterns from the outset, we designed M$^3$-Impute to leverage missingness information directly, enabling it to learn feature-wise and sample-wise correlations. Specifically, in FCU, M$^3$-Impute uses the missingness results (i.e., known masks) to learn the correlations between the imputation targets and observable features. In SCU, M$^3$ employs the missingness results to better capture sample correlations. Since feature and sample correlations exist regardless of the cause of missingness, M$^3$-Impute is naturally adaptive and robust across all three missingness settings.

Another notable feature of M$^3$-Impute is its ability to learn the cause of missingness. In FCU, M$^3$-Impute explicitly captures the correlations between observed and missing features. When data is missing under MAR and MNAR conditions, the missing values depend on the observed ones. Since FCU explicitly captures these relationships, M$^3$-Impute can potentially identify the cause of missingness and enhance imputation accuracy. This capability may explain why M$^3$-Impute significantly outperforms the baselines in the MAR and MNAR settings compared to the MCAR settings.

In addition, we have done new experiments for the comparison with not-MIWAE and report the results in Tab.2, showing that our method outperforms not-MIWAE substantially. We will include the results in the final manuscript. Thanks.

Tab 2: MAE of imputation under MNAR setting with 30% missingness

Q5. Add numbers to Figure 1.

The numbers will be included in the final manuscript. Thanks.

Q6. Hyperparameter setting in Algorithm 1 and other baselines.

For all the 25 datasets, we use the same hyperparameters for our method and follow the same setups as in the original papers for the baselines.

Q7. Comparison with missForest.

The results are shown in Tab.3 and will be added in the final manuscript. Thanks.

Tab. 3: MAE of imputation under MCAR setting with 30% missingness

Q8. Feature selection in MAR setting.

For the 30% missingness setup, we randomly selected 50% of the features to be observed (ensuring these features do not contain any missing values) and masked out values from the remaining 50% of the features until the desired missingness ratio is reached. We did not cherry-pick the subset of observed features; rather, we randomly selected them so that they could be different in each repeated run, as different random seeds are applied.

Q9. Minor comments:

They will be incorporated in the final manuscript. Thanks.

Again, we would like to thank all the reviewers and the AC for their invaluable time and constructive comments. We would like to provide our response below to the concerns that were newly raised after the discussion period.

(1) For the hyperparameters of the baselines, we used the optimal settings from Grape [1] and HyperImpute codebase [2] for a fair comparison. Though we agree that "Methods that produce results worse than mean imputation should be investigated and discussed, since this means exceptionally bad performance", we observe a similar trend of some baseline methods performing worse than the mean imputation in other works [1,2], which may be due to the small number of features or samples.

As suggested by the AC, we conducted a hyperparameter search for the baseline models, especially for MIWAE and MIRACLE. Results show that MIRACLE's performance improves substantially with a small embedding dimension, while MIWAE's performance improves with more iterations of training, albeit marginally. Despite such a hyperparameter tuning, the proposed method M -impute still achieves lowest MAE scores on average compared to all the other imputation methods. We have included the updated experiment results below, which are obtained after the hyperparameter tuning.

[1] You, Jiaxuan, Xiaobai Ma, Yi Ding, Mykel J. Kochenderfer, and Jure Leskovec. Handling missing data with graph representation learning. Advances in Neural Information Processing Systems 33 (2020): 19075-19087.

[2] Jarrett, Daniel, Bogdan C. Cebere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. Hyperimpute: Generalized iterative imputation with automatic model selection. In International Conference on Machine Learning, pp. 9916-9937. PMLR, 2022.

(2) For the comment of using MSE instead of MAE in training, we would like to point out that we indeed used MSE for training, which was clearly mentioned in Section 3.5 of our manuscript.

We acknowledge the efforts from all the reviewers and AC; however, we regret to know that these minor issues, which should be resolvable, were not raised during the discussion period.

Sincerely,

-- Authors