Toward Generalizability of Graph-based Imputation on Bio-Medical Missing Data

We sincerely appreciate your insightful feedback and recognition of the contributions made in our paper. We are committed to addressing your concerns comprehensively.

[Cons 1: Distinction from Previous Work]
The major difference of GRASS from other work is the generation of the graph information when the graph structure is not readily available and there are missing input features such as in the biomedical missing data scenario. We thus generate the graph information by utilizing the feature gradient obtained from the given tabular data, which does not require any external information. Compared to existing graph-based imputation methods, GRASS newly leverages the feature gradient to construct the graph structure by supplementing the noisy missing features. As Figure 3 demonstrates, the existing graph-based methods suffer from the low quality of the graph structure, and thus they occasionally underperform even basic methods like Mean Imputation where their graph construction depends solely on the noisy missing features. Therefore, GRASS has a clear contribution to better handle the challenging and realistic scenario (i.e., no given graph structure and missing input features) than the existing work that does not consider the characteristics of the scenario.

[Cons 2: Inheriting the Merit of Graph-based Methods in GRASS]
As pointed out, GRASS may not fully benefit from graph-based methods when the initial graph structure is absent. However, as Figure 1 illustrates, when the graph structure is well-defined, graph-based imputation methods can fully utilize the message-passing scheme (this would belong to the upper bound case of GRASS), an advantage not present in traditional tabular-based methods. In cases where the graph structure is not predefined, such as in the biomedical domain, the effectiveness of this message-passing scheme is uncertain. Our work investigates whether graph-based imputation can reach its full potential with a better-provided feature matrix and relevant graph structure. Figure 3 shows that when the graph structure is derived solely from the original missing feature matrix, it underperforms compared to classic tabular-based methods. However, utilizing GRASS as an initializer allows graph-based methods to surpass tabular-based methods by effectively leveraging the warmed-up matrix and adjacency matrix. This supports our research question: "Can graph-based imputation methods inherit the merit of the message-passing scheme with better-provided feature matrices and adjacency matrices?"

[Cons 3: Addressing Incomplete Features with GRASS]
Our approach specifically addresses scenarios with initially missing features. The term 'using incomplete features' might not fully align with our methodology. Rather than directly using incomplete features, we opt for task-relevant supplements to build the graph structure. As highlighted in the introduction, generating a graph from incomplete features is challenging and often leads to suboptimal results. This challenge motivates our approach to obtain a feature gradient by training an MLP layer, requiring only the input feature matrix. This gradient plays a crucial role as a supplement in constructing a column-wise graph. Thus, GRASS avoids using incomplete features directly and instead equips the model with task-relevant feature gradients for imputation via the generated column-wise graph.

This is a gentle last-minute reminder regarding the rebuttal period. We have diligently addressed the concerns raised by the reviewers, as outlined above, and are hopeful for a positive response in the final stage of evaluating our work. We would like to highlight that the aforementioned better-graph structure in [Cons 2: Inheriting the Merit of Graph-based Methods in GRASS] can be evaluated by the below experiments.

We have also included this experiment in Appendix A.10, on Page 26.

As shown in the above table, employing GRASS to obtain refined adjacency matrix ($\hat{\mathbf{A}}$) leads to an improved edge homophily ratio, indicating an increase in edges connecting nodes with the same labels. This improvement is attributed to the incorporation of feature gradients, which introduce task-relevant information through supervision signals from training the MLP. Consequently, a more refined graph structure, enriched with task-relevant information, allows current graph-based imputation methods to further enhance their performance by smoothing the representation of their nearest neighbors. In summary, this improvement in the graph structure is a key factor contributing to the performance boost observed with GRASS.

Overall, we hope that these experimental results will provide the reviewer with additional clarity and be taken into consideration while making the final decision on our work. Thank you for your invaluable time.

Best,
Authors

We are immensely grateful for your recognition of our rebuttal, particularly the discussion on convergence with Reviewer T1VM, which we found truly inspiring and significant from our perspective as authors. To address your remaining concerns, we have broadened our experiments to illustrate the impact of hyperparameters, encompassing not only the hidden dimensions in range {16, 64, 256} but also the dropout rates in range {0.0, 0.25, 0.5}, as initially pointed out in [Cons 2.] by the reviewer. We will begin by presenting the experimental results, followed by a summary of the key findings. For dataset, we used Pancreas dataset in bio domain with model variants in GCNMF (graph-perspective), Mean (tabular-perspective), scFP (domain-perspective).

In the GCNMF Table, it reveals that a hidden dimension size of 256 with a dropout rate of 0.25 is more advantageous than the conventional setting of a 64-dimensional hidden layer with a 0.5 dropout rate. Moreover, in the Table of Mean and scFP, the case of an imputation model equipped solely with a logistic classifier (i.e., a 1-layer MLP) and thus invariant to hyperparameter flexibility, we found performance gains when using our proposed method, GRASS, as an initializer. This improvement is attributed to the trainable parameters in a 2-layer MLP and the incorporation of feature gradient. In summary, while setting constant hyperparameters ensures reproducibility and consistency across different datasets in the biomedical domain, allowing flexibility in these parameters can potentially enhance overall performance, as the reviewer suggested. For our finalized version, we would incorporate more flexibility, such as the number of layers.

We have also created a new section in the updated paper to facilitate this discussion, which you can find in Appendix A.9, on Page 26.

Once again, the enhancements made to this paper were possible due to the reviewer's invaluable feedback.
We hope our latest efforts will further alleviate any remaining concerns.

[Cons 3: Convergence Issue]
We greatly value the reviewer's expertise in feature propagation and related works concerning missing features, and we are thankful for bringing this crucial discussion to the forefront. We concur that the convergence property is of paramount importance in the missing feature community, and we have delved deeper to articulate our perspective on this issue. As stated in our initial submission, we acknowledge that strong connectivity in kNN graph generation is not always guaranteed, and achieving convergence may not always be advantageous.

To examine whether the convergence property contributes to performance improvement, we conducted an empirical analysis. In the table below, we extended the range of $k_\text{col}$ and $k_\text{row}$ values from {1,3,5,10} to as high as 50, to test the connectivity of the resulting graphs. We observed that when $k_\text{col}$ and $k_\text{row}$ exceed 10, the graph becomes strongly connected, implying that every node is reachable from every other node. Interestingly, while strong connectivity facilitates the selection of neighbor numbers and meets the necessary condition for FP to converge, it does not automatically lead to performance improvements. Optimal performance was actually achieved within the initially proposed smaller range of $k$ values. Further analysis revealed that increasing $k_\text{col}$ resulted in an oversmoothing issue in the output, particularly noticeable in the warmed-up matrix. This effect was quantified using the MADGap metric [1], which assesses the representational differences between remote and neighboring nodes. Our findings suggest that in bio-medical tabular data, where an initial graph structure is not provided and a $k$NN graph must be manually constructed, opting for a larger $k$ value to exploit the convergence property of FP may not be the most effective strategy, especially in scenarios with significant missing data.

We have provided more visual illustrations in Figure 7, Appendix A.2, on Page 16.

(link to Figure 7: https://anonymous.4open.science/r/grass-iclr-41D5/figs/crop_discussion_k.pdf)

To summarize, in the context of manually generated $k$NN graphs, particularly in bio-medical datasets where an adjacency matrix is not inherently available, strong connectivity does not necessarily guarantee performance improvement. Such settings are prone to oversmoothing, and hence, convergence may not always be beneficial for classification tasks.

[1] Chen, Deli, et al. "Measuring and relieving the over-smoothing problem for graph neural networks from the topological view." Proceedings of the AAAI conference on artificial intelligence. Vol. 34. No. 04. 2020.

2-3. Downstream Neural Network: As the reviewer correctly noted, we have indeed appended a logistic classifier to methods primarily focused on imputation. For graph-based imputation methods like GCNMF, PaGNN, FP, and PFCI, which target downstream tasks rather than solely imputation, we utilized a standard 2-layer GCN. Conversely, for methods such as Mean, GAIN, kNN, GRAPE, and IGRM, which are more imputation-centric, we employed a logistic classifier (1-layer MLP) for downstream tasks. Notably, while GRAPE and IGRM incorporate a graph structure, our empirical findings indicated that using a logistic classifier post-imputation yielded better results than directly employing their GNN for downstream tasks. This was primarily due to the differing loss functions used in imputation (MSE loss) and classification (CE loss). Therefore, we decided to consistently use a logistic classifier for models where the primary focus is on imputation.

2-4. Error Bar: We appreciate the reviewer pointing out the need for clearer error bar notation. We conducted our experiments with five random splits, reporting the average scores. To provide a clearer picture of the variability in our results, we have now included the standard deviation alongside the average scores in the tables presented in the main text and Appendix A.5 on Page 22. Additionally, to highlight the best-performing and most-improved models, we have used bold formatting for the best model and underlining for the most improved model in these tables.

[Cons 2: Issues in Experimental Settings - Hyper-parameters, Baseline, Downstream Neural Network, Error Bar]

2-1. Hyper-parameters: We would like to respectfully assert that our hyper-parameter setup is fair and appropriate. Considering that our work is largely based on Feature Propagation, which in the original paper utilized a structure of 2 layers with a hidden dimension of 64 and a dropout rate of 0.5, our approach aligns with established practices. In terms of the downstream GNN model (GCN), which is widely used for graph-based imputation, adding complexity to this model could unnecessarily complicate the comparison of imputation methods. A fixed hyper-parameter setting for the downstream network acts as a control factor, allowing us to focus on the experimental factor, which is the method of imputation, such as GCN-Zero, GCN-NM, or GCN-FP. This approach enables a direct comparison of the impact of the imputation module.

Furthermore, we would like to note that the datasets mentioned by the reviewer, namely Cora, CiteSeer, and PubMed, are not within the scope of our study, as we focus on tabular data rather than graph-structured data. Therefore, the more relevant hyperparameters in our context pertain to the imputation step, particularly the number of connected neighbors. To this end, we have allowed a flexible search space for all graph-based models, with options among {1,3,5,10}.

Additionally, in the computational biology domain, which is our primary focus, there is a preference for a unified hyper-parameter configuration across different datasets to ensure reproducibility and generalizability across various types of biological data. Nonetheless, we acknowledge the importance placed on hyper-parameter tuning in the broader ML community. With this in mind, we will ensure a more thorough examination of hyper-parameters in our final version of the paper.

2-2. Baseline: As pointed out by the reviewer, we initially omitted scFP in our submission, focusing more on general approaches prevalent in graph-based imputation methods and conventional methodologies within the realm of tabular data. It is important to note, based on our understanding, that scFP is tailored specifically for the single-cell RNA-seq (scRNA-seq) domain and primarily evaluated on bio datasets, not covering medical datasets. However, recognizing the similarities in process with scFP, and in response to the reviewer's request, we have expanded our evaluation to include scFP alongside our current baselines across all datasets. This broader evaluation approach allows us to provide a more comprehensive comparison, particularly in bio datasets, where scFP's performance can be directly contrasted with other methods. The results of this extended evaluation, particularly focusing on bio datasets, are presented below. We can observe that GRASS-initialized graph-based imputation methods consistently outperform scFP. Moreover, interestingly, scFP with GRASS initialization enhances the performance up to 16.23%, which implies that GRASS can provide distinct knowledge from that of scFP, i.e., column-wise relationship.

We deeply appreciate the reviewer's dedicated time in reviewing our paper and providing valuable, thorough feedback, which has significantly contributed to the improvement of our work. We have done our utmost during the rebuttal period to address the concerns raised, and we now respectfully offer a point-by-point response to the reviewer’s valuable insights.

[Cons 1. Novelty regarding scFP]

The reviewer's concerns about novelty in comparison to scFP are acknowledged. We respectfully argue that our work, GRASS, is distinct from scFP in the following three key areas: Target Domain, Target Task and Imputation Methodology, and Usage of Feature Propagation. Here are the details for each:

(1) Target Domain: While scFP focuses on the single-cell RNA-seq (scRNA-seq) domain, particularly from a biological perspective, GRASS adopts a more general approach for the broader `bio-medical' domain, as indicated in the paper's title. This distinction is crucial as scRNA-seq datasets typically comprise numerical features where each element represents the count of a gene's RNA transcript sequenced by the sequencing machine. In contrast, medical datasets often include both numerical and categorical features, such as patient information. This versatility underscores the broader applicability of GRASS, capable of handling both numerical and categorical features, the latter through the clamping technique as discussed in Section 3.3. Therefore, we argue that the target domain of scFP, primarily focused on numerical matrix imputation in scRNA-seq, differs from that of GRASS, which extends to handling categorical data often encountered in patient data.

(2) Target Task and Imputation Methodology: Unlike scFP, which is unsupervised with its primary goal being effective imputation in sparse and noisy cell-gene count matrices, this work concentrates on supervised tasks, specifically on downstream applications like classification. Notably, the objective of imputation is often to enhance performance in relevant downstream tasks [1, 2]. In this context, while the unsupervised approach of scFP can align with supervised tasks through probing (i.e., attaching a classifier), it is important to note that since its imputation occurs prior to probing, scFP cannot incorporate any downstream task-related knowledge during the imputation process, potentially leading to shortcomings in classification tasks. Conversely, as GRASS is directly designed with downstream tasks in mind, it incorporates knowledge pertinent to these tasks during imputation. This is achieved by utilizing the feature gradient, which is obtained during training 2-layer MLP. This fundamental difference in the target task (classification vs. imputation) and the imputation process (incorporating relevant downstream knowledge or not) distinctly sets the two methodologies apart.

(3) Usage of Feature Propagation: Although both scFP and GRASS employ Feature Propagation (FP), their applications of this process differ significantly. Specifically, scFP utilizes FP from a row-wise perspective, i.e., focusing on cell-cell relationships while assuming gene-gene relationship independence. Although beneficial for smoothing similar and relevant samples, this approach does not capture interactions between columns (features), which are pivotal in the biomedical domain. For instance, in scRNA-seq, gene-gene relationships, such as co-expression networks, play a critical role in identifying key regulatory genes or pathways, offering insights into underlying biological or disease mechanisms [3,4]. Acknowledging this, GRASS initially employs column-wise FP to capture potential feature interactions, e.g., gene-gene relationships. It's also noteworthy that GRASS incorporates not only the feature matrix but also the feature gradient relevant to downstream tasks when generating the column-wise $k$NN graph. Consequently, before initiating row-wise (sample-wise) smoothing in the relevant GNN model, GRASS is able to consider feature relationships that scFP does not capture. This distinction is illustrated in Figure 8 in Page 23 and significantly differentiates the two methodologies.

(link to Figure 8: https://anonymous.4open.science/r/grass-iclr-41D5/figs/crop_scfp_grass_2.pdf)

[1] Rossi, Emanuele, et al. "On the unreasonable effectiveness of feature propagation in learning on graphs with missing node features." Learning on Graphs Conference. PMLR, 2022.
[2] Wang, Juexin, et al. "scGNN is a novel graph neural network framework for single-cell RNA-Seq analyses." Nature communications 12.1 (2021): 1882.
[3] Cochain, Clément, et al. "Single-cell RNA-seq reveals the transcriptional landscape and heterogeneity of aortic macrophages in murine atherosclerosis." Circulation research 122.12 (2018): 1661-1674.
[4] Galfre, Silvia Giulia, et al. "COTAN: scRNA-seq data analysis based on gene co-expression." NAR Genomics and Bioinformatics 3.3 (2021): lqab072.

For Medical datasets the overall performance is as below:

We have taken steps to enhance the visibility of the tables discussed in our responses. These tables have been presented in a more accessible format in Appendix A.5 on Page 22 of our manuscript. We kindly ask the reviewers to refer to this Appendix for a detailed view of the results.

We deeply appreciate your recognition of the core intuition behind our paper, particularly regarding the role of feature gradients in aiding the discriminative representation of data. We hope that the following responses adequately address your concerns.

[Cons 1: Different Scales in GCN Propagation]
Thank you for highlighting the importance of normalization in the context of feature matrix scaling. Indeed, as you pointed out, the scale of imputed values, interpreted as weighted sums of columns, is crucial. To address this, we preprocess the feature matrix before its multiplication with the adjacency matrix. For numerical features, we employed a MinMaxScaler, ensuring a range of [0,1], as shown in line 163 of our preprocess.py script on our anonymous GitHub repository (https://anonymous.4open.science/r/grass-iclr-41D5/preprocess.py). Categorical features are one-hot encoded to align with the same range as numerical columns, as demonstrated in line 187 of the same script. This preprocessing ensures that the feature matrix values fall within [0,1], preventing the explosion of imputed values, considering the symmetrically normalized adjacency matrix with off-diagonal elements also ranging from [0,1]. This approach effectively addresses scalability issues in our algorithm.

[Cons 2: Expression of Gradient and Pseudocode of Algorithm]
We are grateful for your constructive feedback on expressing the concept of the feature gradient and the need for a clear pseudocode of our algorithm. The feature gradient, as defined in Definition 3.1 and Proposition 1, is the partial derivative of the loss with respect to the feature matrix, with its value dependent on trainable weight parameters related to the downstream task. This gradient informs the construction of a column-wise graph embedded with task-relevant information, distinguishing our method from existing approaches. The complete pseudocode, encompassing the process of obtaining the feature gradient, column-wise feature propagation, clamping technique, and final output, is provided in Appendix A.4 on Page 20. Additionally, for a more implementation-focused understanding, we have included a PyTorch-style implementation for obtaining the feature gradient in Appendix A.4 on Page 21, achieved simply by enabling requires_grad=True. We respectfully request the reviewer to consult this PyTorch-style pseudocode for a more in-depth understanding.

[Cons 3: Choice of Threshold $\theta$]
The clamping technique is applied across all categorical features using a global threshold $\theta$. We acknowledge that different thresholds could be significant for distinct features, especially in reflecting the inherent uncertainty in imputing categorical values, as addressed in Equation (3). The softmax operation is applied sample-wise for each categorical feature set. For example, in a one-hot encoded patient race category with three options (European, American, Japanese), imputed values of 0.8, 0.9, and 0.4 would transform to probabilities of 0.377, 0.416, and 0.207, respectively, post-softmax. A threshold of 0.4 would then yield a clamped output of [0,1,0] or [?,?,?] if the threshold is over 0.5. This allows subsequent imputation methods to fill in the gaps based on their algorithms. However, in cases with a large number of possible values for a category, the probability distribution might be too spread, as you rightly pointed out. Introducing a temperature parameter in the softmax operation could be a potential solution in such scenarios. We will thoroughly explore this aspect in our final version and sincerely thank you for this valuable insight.

[Cons 3: Gradient Norm in Figure (4) and Pairwise Cosine Similarity]
We are grateful to the reviewer for their insightful feedback, which significantly contributes to the improvement of our paper. We agree with the reviewer's suggestion that normalizing gradients may render the gradient norm less indicative. The original intention behind Figure 4(b) was to showcase the gradient dynamics of key marker genes. However, upon reflection and in line with the reviewer's suggestion, we find that presenting the pairwise cosine similarity between the original input feature matrix and the feature gradients offers a more intuitive and insightful perspective. Our analysis of this similarity has yielded promising results, as detailed below.

Here, it shows the marker genes for `activated stellate' such as COL1A1, TIMP1, TGFBI, and PDGFRB exhibit higher cosine similarity in feature gradient compared to its original matrix. This is achieved by incorporating task-relevant information, brought from label supervision, into the feature gradient. Following the reviewer's valuable suggestion, we have updated Figure 4(b) with a clearer representation to better illustrate this point. We kindly invite the reviewer to review this new version in the updated PDF file.

(link to Figure 4: https://anonymous.4open.science/r/grass-iclr-41D5/figs/crop_confusion.pdf)

[Question 1: Reasoning Behind 10/10/80 Split]
Our choice of a 10/10/80 split was designed to mirror real-world scenarios, where training labels often form a small subset. This experimental setting aimed to evaluate the suitability of our proposed method in environments with limited training labels. However, we acknowledge the reviewer's suggestion that a larger training dataset could be beneficial. In such scenarios, we anticipate that GRASS would gain further from more substantial supervision, thereby enhancing the knowledge content in the feature gradient. We plan to explore and include a variety of training settings in our final version for a more comprehensive validation.

[Question 2. Descriptions regarding the baselines & Preprocessing of dataset]
We apologize for previously omitting detailed descriptions of baselines and preprocessing methods for the bio and medical datasets. In response to the reviewer's request, we provide the following additional information:

In our bio datasets, we utilize cell-gene matrices to predict relevant annotated cell types for each cell. This cell type information serves as a supervisory signal during training. For preprocessing, we filter out cells and genes that have not been transcribed in each row and column, respectively, and apply a log transformation to normalize the count values.

We extend our sincere thanks to Reviewer JSNu for their time and constructive comments. We are eager to engage in further discussion with Reviewer JSNu and hope that our response effectively addresses their concerns.

[Cons 1: Limited Dataset with Larger Samples]
We acknowledge and appreciate the reviewer’s concern regarding the scalability of our algorithm to larger datasets. In response, we have expanded our testing to include a substantially larger dataset from the scRNA-seq domain, known as the Macosko dataset. This dataset is notably extensive, encompassing 44,808 cells, 22,452 genes, and 14 distinct cell types, with an initial missing rate of 81.41%. For our analysis, we preprocessed 2,000 highly variable genes, a standard practice in scRNA-seq studies. The performance results obtained from this dataset are summarized below.

We acknowledge that GRASS integrates effectively with existing methods, except in situations where initial baselines, such as GCNMF and GRAPE, face Out-Of-Memory (OOM) issues. These issues arise due to the weights associated with the Gaussian Mixture Model in GCNMF and the construction of a heterogeneous node-feature graph in GRAPE. However, it is observed that graph-based methods generally enhance their performance when integrated with GRASS. In large graphs, where the number of samples typically exceeds the feature dimension, GRASS aligns well with current graph-based imputation methods, owing to its computational efficiency.

To summarize, our algorithm demonstrates good alignment with larger datasets due to its modest computational demands, which we will discuss further.

[Cons 2: Computational Demand]
The reviewer’s concern about computational demand alongside performance benefits is valid and important. We address this from two perspectives:

Memory Cost: The primary resource used by GRASS is the feature gradient ($\overline{\boldsymbol{\nabla}}_{\mathbf{X}} \in \mathbb{R}^{N \times F}$), which is instrumental in constructing a column-wise graph. This gradient has the same dimensions as the original feature matrix ($N$ nodes and $F$ features). In graph-based imputation models, the complexity of the row-wise adjacency matrix ($\mathbf{A} \in \mathbb{R}^{N \times N}$) often surpasses that of the feature matrix, resulting in a complexity of $\mathcal{O}(N^2)$, especially in biomedical datasets where $N >> F$. Therefore, the extra memory needed for the feature gradient is not excessively large. Moreover, the complexity of the generated column-wise graph ($\mathbf{A}^{feat} \in \mathbb{R}^{F \times F}$) is lower than that of the row-wise adjacency matrix, allowing GRASS to integrate with existing graph-based models without significant memory costs. Post computation of the warmed-up matrix ($\hat{\mathbf{X}}$) and adjacency matrix ($\hat{\mathbf{A}}$), the memory allocated for the feature gradient and column-wise graph can be freed, leaving only the training cost of the baseline model.

Time Cost: The time complexity of GRASS is comparable to training a conventional 2-layer MLP, efficient for tabular data. This involves training two weight matrices – one for the raw feature space to hidden space transformation, and another for mapping the hidden space to output space. The computation of the feature gradient, despite its apparent complexity as outlined in Proposition 1, is straightforward. The requires_grad switch enables automatic saving of gradient information, equating the time complexity for computing the feature gradient to that of training a 2-layer MLP. Moreover, column-wise Feature Propagation can be executed efficiently via sparse multiplication, ensuring that the time required for obtaining the warmed-up matrix and adjacency matrix is not significant.

Thank you for the time spent evaluating our rebuttal and for recognizing the promise in our work. We have made our utmost efforts to thoroughly address the concerns raised, as detailed above. Here, we would like to further clarify the origin of our performance gain, providing a clearer response to Reviewer r66i's [Cons 1.].

Below, we further explore how the performance enhancement is achieved by initializing current graph-based imputation methods with GRASS. Specifically, using GRASS as an initializer provides a warmed-up matrix along with a refined adjacency matrix. We assessed this adjacency matrix($\hat{\mathbf{A}}$) in terms of the homophily ratio, a crucial inductive bias in the graph domain.

As shown in the above table, employing GRASS to obtain refined adjacency matrix ($\hat{\mathbf{A}}$) leads to an improved edge homophily ratio, indicating an increase in edges connecting nodes with the same labels. This improvement is attributed to the incorporation of feature gradients, which introduce task-relevant information through supervision signals from training the MLP. Consequently, a more refined graph structure, enriched with task-relevant information, allows current graph-based imputation methods to further enhance their performance by smoothing the representation of their nearest neighbors. In summary, this improvement in the graph structure is a key factor contributing to the performance boost observed with GRASS.

We have also included this experiment in Appendix A.10, on Page 26.

Overall, we hope that these experimental results will provide the reviewer with additional clarity and be taken into consideration while making the final decision on our work. Thank you for your invaluable time.

Best,
Authors

We thank the authors for their response Their feedback is appreciated, nevertheless my score shall stay the same The authors are highly encouraged to extend their promising work

[Question 6: Meaning of Algorithmic Convergence]
We appreciate the mention of the convergence property discussed in our manuscript. The convergence concept primarily stems from the FP paper, which discusses the convergence of imputed values. A key difference between our work and FP is the nature of the graph structure—whether it is manually constructed via $k$NN (as in GRASS) or initially given (as in FP). This distinction is critical as the strong connectivity of the graph, where every node is accessible from others, is a prerequisite for convergence in imputed values. In our approach, utilizing $k$NN-based graph generation, convergence is not guaranteed, as noted by Reviewer T1VM. However, we argue that in scenarios with prevalent missing features, especially in tabular data formats requiring manual graph generation, convergence is not necessarily advantageous. Our rationale is further elaborated in Appendix A.2., on Page 16. Briefly, as demonstrated in the table below, the risk of oversmoothing (indicated by the MADGap, which measures the representational difference between remote and neighboring nodes) can detrimentally impact performance, even if convergence is achieved.

[Question 7: Employed Loss Function and Uncertainty of Imputed Values]
We appreciate the reviewer's inquiry regarding the loss function used in our study. Indeed, for the downstream classification task, we employed the CrossEntropy loss. While this work concentrates on classification, aligning with existing graph-based imputation methods, we acknowledge the reviewer's point about the potential of using different loss functions for various tasks, such as MSE loss, KL divergence loss, or self-supervised loss. We find this suggestion intriguing and consider it an excellent direction for future research. We plan to incorporate this aspect of versatility into the final version of our paper, exploring the generalizability of graph-based methods across diverse loss functions.

Regarding the implementation, obtaining task-relevant feature gradients is feasible using AutoGrad. By setting requires_grad=True, we can efficiently compute these gradients, underscoring the flexibility of our approach.

As for the uncertainty of imputed values, especially in categorical features, we recognize the significance of reflecting such uncertainty. For instance, in biological datasets, an imputed value of 0.5 for a tissue type, where '0' denotes epithelial and '1' denotes mesenchymal, would be biologically nonsensical. To address this, as detailed in Section 3.3, we implemented a clamping technique. This technique employs a softmax operation to sample categorical features based on their specific number. For example, in a one-hot encoded representation of a patient's race (European, American, Japanese), imputed values like 0.8, 0.9, and 0.4 would be transformed via softmax to 0.377, 0.416, and 0.207, respectively. Setting a threshold, such as 0.4, allows the clamping operation to output a definite category or leave room for ambiguity (e.g., [0,1,0] or [?,?,?]) when the threshold is higher than 0.5. This approach enables subsequent imputation methods to fill in the data based on their algorithms. The efficacy of this clamping process is demonstrated in Table 1 on Page 9, where incorporating the clamping technique significantly enhances overall performance.

[Cons 4: Inconsistencies in Notation]
We appreciate the reviewer's attention to detail regarding the notation. The symbol $Y$ is clearly defined in Proposition 1 as a label matrix used during the computation of the feature gradient, and is denoted in orange. This specific use of $Y$ is confined to the calculation of the feature gradient and does not recur elsewhere in the paper.

[Minor Comments]
Thank you for pointing out the redundancy in terms and the need for clarity in citations. Based on your valuable feedback, we have revised the redundant terms and clarified the citations in the updated PDF version.

[Question 1: “Clear of Presumptions Associated with Missing Completely at Random (MCAR)”]
We apologize for any ambiguity in our wording. By "Steering clear of presumptions associated with Missing Completely at Random (MCAR)," we mean that our method of zero imputation avoids assuming that the missing data occurred completely at random. This assumption may not always hold true in real-world datasets, making our approach more flexible and generalizable. This is especially pertinent in the biomedical domain, where zero often has a meaningful interpretation, such as the absence of a feature.

[Question 2: Meaning of Each Bar in Figure 3]
In Figure 3, each bar represents the performance (AUROC for medical datasets, Macro-F1 for bio datasets) of existing methods, as well as their enhanced performance when GRASS is used as an initializer.

[Question 3: Fewer Labels in Figure 5]
We are grateful to the reviewer for identifying the discrepancy in Figure 5 regarding the representation of cell-types. We have corrected the figure to accurately reflect all 14 cell-types. Recognizing the figure's small size, we have provided a link for a clearer view of the refined tSNE visualization. This figure demonstrates that, compared to the original feature matrix, the feature gradient imbued with cell-type classification knowledge significantly enhances the discernment of different cell-types in the warmed-up matrix through column-wise FP.

(link to Figure 5: https://anonymous.4open.science/r/grass-iclr-41D5/figs/crop_tsne_final.pdf)

[Question 4: Meaning of Numbers in Each Column in Table 1]
Table 1, presenting our ablation study, contains cells that denote the average AUROC score along with the standard deviation across five runs. To facilitate understanding, we have explicitly mentioned this metric in the manuscript.

[Question 5: Meaning of Each Cell in Figure 6 (a)]
We apologize for not providing detailed explanations for Figure 6 (a). This figure aims to illustrate the sensitivity of GRASS in terms of selecting the number of neighbors for column-wise and row-wise graphs, assessed through the AUROC metric. This metric is also described in the figure's caption.

For Medical datasets the overall performance is as below:

We have included a more visually detailed illustration of the table in the updated PDF, located in Appendix A.5 on page 22.
We respectfully request the reviewer to consider this illustration in the Appendix for a comprehensive understanding.

We sincerely appreciate the reviewer's dedicated time in reviewing our paper and acknowledging the use of feature gradients, specifically feature saliency, targeted toward the downstream task. We are greatly encouraged by the reviewer's insightful summary of our work. In response to Reviewer r66i's questions, we provide detailed answers below.

[Cons 1: Incremental and Consistency]
We apologize if our presentation inadvertently misled the reviewer's understanding of our work. We would like to clarify that our proposed method, GRASS, serves as an agnostic preprocessing step. It generates a warmed-up feature matrix and adjacency matrix that seamlessly integrate with existing graph-based imputation methods. This is illustrated in each performance baseline, where the additional colored bars in green and brown (symbolizing 'grass') denote improvements when GRASS is utilized as an initializer. The extent of performance gain varies with the initial missing rate of the dataset and the size of the training sample. This variation explains the relatively lower performance gain in medical datasets compared to bio datasets, given the smaller number of training samples in a 10/10/80 train/val/test split. Overall, Figure 3 highlights the improved generalizability of current graph-based imputation methods when equipped with GRASS as an initializer. In the Introduction, we observed that graph-based imputation can sometimes underperform classical tabular-based methods, as seen with the Mouse ES dataset in Figure 3. However, upon integrating GRASS, its performance significantly increases, demonstrating GRASS's effectiveness within current graph-based imputation methods.

[Cons 2: Computational Trade-off]
Regarding the reviewer's concern about computational trade-offs versus performance gains, the computational demands of GRASS can be broken down into two aspects:

Memory Cost: GRASS primarily uses the feature gradient ($\overline{\boldsymbol{\nabla}}_{\mathbf{X}} \in \mathbb{R}^{N \times F}$), which supplements the construction of a column-wise graph and shares the same dimensions as the original feature matrix ($N$ nodes and $F$ features). In the context of graph-based imputation models, which typically use a row-wise adjacency matrix ($\mathbf{A} \in \mathbb{R}^{N \times N}$), the complexity of the adjacency matrix often exceeds that of the feature matrix ($\mathcal{O}(NF) + \mathcal{O}(N^2) = \mathcal{O}(N^2)$). This is particularly pertinent in biomedical domains, where datasets are tabular and the number of samples ($N$) significantly surpasses the number of features ($F$). Therefore, the extra memory needed for the feature gradient is relatively modest. Moreover, the complexity of the generated column-wise graph ($\mathbf{A}^{feat} \in \mathbb{R}^{F \times F}$) is also lower than the row-wise adjacency matrix. Thus, GRASS aligns with existing graph-based models without excessive memory costs. After computing the warmed-up matrix ($\hat{\mathbf{X}}$) and adjacency matrix ($\hat{\mathbf{A}}$), memory allocated for the feature gradient and column-wise graph can be freed, leaving only the cost of training the original baseline model.

Time Cost: The process resembles training a conventional 2-layer MLP, efficient for tabular data and involving two weight matrices. The calculation of the feature gradient, as detailed in Proposition 1, is straightforward in practice. By enabling the requires_grad switch, gradient information is automatically saved, equating the time complexity for computing the feature gradient to training a 2-layer MLP. Furthermore, column-wise Feature Propagation is executed efficiently via sparse multiplication, as detailed in our methodology. Therefore, the time required to obtain the warmed-up matrix and adjacency matrix is not significant.

In conclusion, considering the relatively low memory and time complexities, we assert that the performance gains achieved with GRASS are substantial and meaningful.

[Cons 3: Presentation of Results]
We apologize to the reviewer for any confusion caused by our initial presentation of the results. To provide a clearer understanding of how GRASS improves existing baselines, we have reformatted the results into a table, as shown below: