A Flexible Generative Model for Heterogeneous Tabular EHR with Missing Modality

We are very grateful for your constructive feedback about our work. We appreciate that you recognized the core contributions of our work, which enrouaged us a lot. We respond to your concerns about our work in the following comment.

Ablation Studies

Thank you for bringing up the ablation study part. Originally, we intended to include comprehensive ablation analyses in our submission; however, due to time constraints, this aspect was not fully developed. In response to your suggestion, we have now conducted an extensive ablation study focusing on four key configurations of FlexGen-EHR: the inclusion and exclusion of Optimal Transport (OT), and its performance on both complete and incomplete datasets, with the latter involving KNN imputation. This study aims to provide a detailed examination of the model's robustness and effectiveness under varying conditions. Results are summarized in the table below. It demonstrates that our FlexGen-EHR offers a more effective solution for handling missing modalities compared to traditional imputation methods, like kNN.

Questions

1. Typos: We thank the reviewer for pointing out these typos and errors in our submission. We have made the corresponding changes and additions to the text to fix these error.
2. Clarification on R^2, MMD, and KS-Statistics: R^2 is a statistical measure reflecting the proportion of variance in the dependent variable that is predictable from the independent variable. In our model, R^2 serves as an indicator of how well the synthetic data generated by FLEXGEN-EHR corresponds to the real-world data it aims to replicate. This metric is particularly relevant for evaluating the model's performance in terms of data fidelity. In contrast to R^2, MMD/KS-statics are both non-parametric tests used for determining if two samples are drawn from different distributions. By quantifying the maximum distance between these functions, MMD/KS-Statistics allow us to rigorously test the hypothesis that the synthetic and original data come from the same distribution. In our revision, we provided a more detailed discussion on the implementation and interpretation of both MMD and KS-Statistics in the context of our study.
3. Decoder Architectures: MLP only.
4. Sample Size Alignment in Equation 5 (Section 4.3): Yes, it means the number of samples with missing modalities does not necessarily to be equal to that of non-missing data becasue the learning problem is not a supervised correspondence problem.
5. Details on Membership Attack Implementation We label the real records with the lowest hamming distance to the closest record in the synthetic dataset as positive. We pick a hamming distance as our distance metric between patient records throughout our privacy evaluations in accordance with [1].This attack allow us to test the ability of the synthetic dataset to prevent an attacker from inferring whether a real record was used in the training dataset. Ideally, the accuracy is around 50%, which is similar to a random guess. This shows that neither the model nor the synthetic dataset reveals any meaningful or compromising information about the patient identity in the training dataset.

[1] Yan, C. et al. A multifaceted benchmarking of synthetic electronic health record generation models. Nat. Commun. 13, 7609 (2022).

Thanks for your constructive feedback about our work. We are encouraged by your recognition of the practical problem definition and the innovative approach in addressing missing modality in EHR data through FlexGen-EHR. We respond to your concerns about our work in the following comment.

Comparative Analysis with kNN and Optimal Transport

To address your valuable concern, we have incorporated an ablation study into our revised work. In this study, we systematically compared the performance of our approach with and without Optimal Transport (OT), specifically focusing on its effectiveness in handling incomplete data. This analysis serves as a direct comparison between the traditional k-nearest neighbors (kNN) imputation method and the novel optimal transport approach integrated into the FLEXGEN-EHR framework. For your convenience and clarity, we have summarized the results of this ablation study in the table below:

It demonstrates that our FlexGen-EHR offers a more effective solution for handling missing modalities compared to traditional imputation methods, like kNN. In addition, our study already encompassed substantial experimental validation. We have employed R², Maximum Mean Discrepancy (MMD), and Kolmogorov-Smirnov (KS) Statistics, which collectively offer a comprehensive evaluation of our method in addressing missing modalities. Combined with new experiments, we hope this response addresses your concern.

Questions

1. Clarification in Table 2 on Baselines for Discrete Codes Generation: MedGAN, as well as CorGAN, their original implementations are capable of generating discrete samples. To generate temporal features like labevents, we modify the last layer of MLP so that their output is not limited to discrete integers.
2. Clarifying Observed Geometric Patterns in Latent Space Embeddings: By saying “geometric patterns”, we mean “proximity relationships”, where specific distances or relationships remain consistent across two latent space embedding models. By exploiting this observation, we are able to align two latent spaces learned by two different encoders.

We thank the reviewer for raising this interesting question and allowing us to discuss it. We hope this response helps validating FlexGen-EHR as a novel and effective generative model for EHRs with missing modality.

Thank you for your extensive comments, critiques, and praises for our work. Following your suggestions, we added new experiments, clarifications of utility, and analyses. In addition, we have provided detailed responses to all the points you raised. We believe there is some confusion regarding our work. We have worked hard to improve the communication of our method, and we would greatly appreciate your considering increasing the score. Many thanks!

RE to Weakness 1: Discussions or case studies on how it might be applied in real-world healthcare settings.

We appreciate the emphasis on practical application. First, we demonstrated it in our submission that FlexGen-EHR is a useful and practical tool for preserving patient privacy by training ML models using synthetic EHRs only. Results showed our method consistently outperformed baselines and is comparable with training using real data. Also, we agree that data augmentation would enrich the paper's practical value. In response, we added a section detailing data augmentation, where our model can contribute to the quality of ML models. A brief table demonstrating the benefit of FlexGen-EHR is summarized as follows,

Although not notable, it is worth noting that existing literature has consistently shown that performance improvements on datasets like MIMIC-III and eICU tend to be relatively modest, unlike those in other datasets.

RE to Weakness 2: It would be beneficial to understand the model's performance across diverse EHR datasets to gauge its broad applicability.

We appreciate the suggestion to explore this aspect further. Due to time constraints, we leave it as future work for a more in-depth investigation.

RE to Weakness 3: handle heterogeneity of the EHRs

FlexGen-EHR is able to offer a solution to the challenge of EHR heterogeneity through the utilization of conditional sampling. This approach empowers the model to generate group-specific samples by incorporating crucial group information and prediction labels directly into the network. It's worth emphasizing that evaluating the effectiveness of this technique can indeed be a substantial research endeavor in its own right. Given the significance of fairness in synthetic EHR generation, we recognize the need for a deeper exploration of this aspect. In our future research efforts, we are committed to placing a stronger emphasis on addressing the issue of fairness to ensure that the outcomes are not only robust but also unbiased and equitable across diverse patient groups.

RE to Weakness 4: the model computationally intensive or challenging to implement.

We understand the concerns regarding the computational demands and implementation complexity of FlexGen-EHR. Our model, however, is straightforward, comprising only two main components: optimal transport and diffusion. Both these elements have manageable computation costs and are not overly challenging to implement. We emphasize the stability of training our diffusion model compared to Generative Adversarial Networks (GANs). While diffusion models may generate samples more slowly than GANs, this is offset by their increased stability. Regarding implementation, the optimal transport component can be efficiently solved and implemented using the well-established Python Optimal Transport (POT) package.

I thank the authors for their detailed responses. I really appreciate the effort that has been put into the rebuttal and have increased my score from 4 to 5. On a separate note, I strongly believe the response to RE5 should be incorporated in the paper and further analyzed in the appendix as it is a practical missing piece of the work.

Thank you for your constructive feedback and for acknowledging the novelty of our approach. Your insights are crucial in guiding the improvement of our manuscript. We meticulously address your suggestions in the following discussion, in order to clarify your concern.

Enhancing Table 4 Presentation

We have addressed this in our revised manuscript, ensuring that the best results are now boldly highlighted, as done in Tables 2 and 3.

Trade-off Between Generation Quality and Privacy Guarantee

We sincerely appreciate your observation regarding the delicate trade-off between generation quality and privacy guarantees within diffusion-based models. In response, we have explicitly incorporated this observation into the text to ensure clarity. However, we wish to underscore our position on this matter and emphasize the meticulous consideration we have devoted to this facet in our research.

Indeed, our findings substantiate that FlexGen-EHR maintains a level of privacy preservation that aligns with existing methods, all while achieving superior data fidelity. It is essential to recognize that this trade-off between generation quality and privacy guarantee is a shared characteristic among all diffusion-based generative models.

The intrinsic nature of diffusion-based generative models involves learning the distribution of the original data. In Section 5.4 of our manuscript, specifically in the Generation Privacy Table (Table 4), we provide empirical evidence that underscores this phenomenon. Our observations reveal that diffusion-based methods, including FLEXGEN-EHR, TabDDPM, and LDM, exhibit a distinct performance profile when compared to non-diffusion-based methods like GANs or VAEs.

This distinction highlights the unique challenges faced by diffusion-based models in striking the delicate balance between data utility and privacy preservation. This dynamic presents an intriguing area for further investigation and exploration within the realm of generative modeling.

Clarification on Handling static and temporal Features using TabDDPM

We have modified this claim by saying “Different from TabDDP that concatenates the numerical and categorical features to the network, FlexGen-EHR adeptly discerns and represents the underlying relationships between static and temporal features.”

Minor Comments.

Thank you for pointing out several typos and errors in our submission. We have made the corresponding edits within the text.

We appreciate your valuable feedback, which has helped us improve our work.

Future Work and Limitations:

We have added additional discussion avenues for future research as follows. The discussion included the implications of our current findings and any limitations encountered during our research. For example, we have included:

Limitations

1. Privacy Concerns in Diffusion Models While FLEXGEN-EHR demonstrates remarkable capability in generating high-fidelity synthetic EHR data, it's important to address the subtle yet significant privacy trade-offs associated with diffusion models. Our model, like many in its category, slightly compromises on privacy to achieve the desired level of data utility and quality. This is a consequence of the inherent characteristics of diffusion models that aim to capture complex data patterns in great detail. This detailed representation, while beneficial for data utility, can inadvertently encode sensitive information, potentially making the model susceptible to privacy risks such as re-identification or membership inference attacks. We have undertaken measures to mitigate these risks, such as evaluating against membership inference attacks, yet the model's intrinsic properties suggest an inevitable balance between data utility and privacy. Future iterations of FLEXGEN-EHR could explore advanced privacy-preserving techniques, possibly integrating differential privacy or advanced anonymization methods to strengthen data confidentiality while maintaining the quality of synthetic data generation.

Future Work

1. Exploration of Federated Diffusion Models Another promising direction for future work is the exploration of federated diffusion models. Federated learning, a method of machine learning where the model is trained across multiple decentralized devices or servers holding local data samples, offers a unique opportunity to address privacy concerns inherent in traditional diffusion models. This approach aligns with the increasing emphasis on privacy-preserving techniques in healthcare data analysis. Adapting FLEXGEN-EHR to a federated learning framework, where a federated diffusion model is employed, could potentially enhance privacy safeguards. In this setup, the model would learn from a diverse and distributed dataset without needing to access or transfer sensitive patient data centrally. This would not only mitigate privacy risks but also allow for a more scalable and ethically compliant model, especially pertinent given the sensitive nature of EHR data. Additionally, a federated approach could facilitate the incorporation of a more diverse dataset, encompassing a wider range of patient demographics and conditions, thereby improving the generalizability and robustness of the model.

Questions

1. Table 1 is a bit confusing to read. Table 1 has been revised for better clarity. The temporal feature dimensions (d) for MIMIC-III and eICU is explicitly labeled, and any similarities or differences are clearly stated.
2. Is the paper also claiming novelty in joint static and temporal diffusion? We agree one of the claimed novelties is joint static and temporal diffusion by using optimal transport.
3. What about considering generating samples that include missing modalities? Should generating missing modalities also be a point of consideration in terms of data fidelity? Indeed, we have already included results on data fidelity evaluation experiments (Figure 2 and 3) in our submission. Results showed that FlexGen-EHR maintains high generation fidelity. We would be very grateful if you considered increasing the score. We have already incorporated reviewers’ feedback and integrated all changes into the revised manuscript

We hope our response and experiments further convince you of our work’s novelty and effectiveness.

We thank the reviewer for these comments and questions. We appreciate that the reviewer recognized the strengths of FlexGen-EHR. To address the reviewer’s questions and concerns, we offer several clarifications as well as additions to our paper:

1. We have added more detail to the text in our revision, including the novelty of our method, clarification on Table 1, and discussion of future work and limitations.
2. We conducted ablation studies and show that the optimal transport component improves the performance of other methods, while FlexGen-EHR remains a superior approach for missing data generation.
3. We conducted additional experiments regarding generation when partial features are missing in EHR data. The results indicate a superior performance of FlexGen-EHR compared to LDM and MedGAN

Please let us know if there is anything else we can do to address your comments. We would be very grateful if you considered increasing the score.

Clarification on Novelty and TabDDPM:

FlexGen-EHR's novelty primarily lies in its unique approach of learning a unified latent space to capture the intricate dependencies between static and temporal EHR data, rather than being the first diffusion model to capture joint relationships. This distinction, however, was not adequately emphasized, and we thank the reviewer for pointing out this ambiguity. We’ve revised the introduction section in muniscript by clarifying that our contribution is focused on the novel application of unified latent space modeling in EHR data, rather than being the first to capture joint relationships among diffusion models.

Detailing Feature Dimensions in MIMIC-III and eICU:

We have incorporated your suggestion to clarify the feature dimensions corresponding to static and temporal data. We have updated the Appendix with a more detailed table that explicitly maps each dimension to its respective feature in both datasets, providing a clearer understanding of the data structure we are working with.

Addressing Partial Feature Presence in EHR Data:

Thank you for your insightful suggestion to include results from a randomly sampled patient cohort, ensuring a more realistic representation of missing data and modalities in Electronic Health Record (EHR) datasets. We acknowledge the importance of reflecting real-world scenarios where patient records often contain a combination of static and temporal features. In response, we address your concern by conducting additional experiments (preliminary due to time limitations). This experiment on MIMIC-III demonstrated that our model, FlexGen-EHR, outperforms baseline methods in scenarios with missing temporal features. For this comparison, we followed [1] and [2] which utilized 12 temporal variables, each recorded within the first 48 hours of a patient's admission. The average missing data rate was approximately 83%. The objective of our study was to predict in-hospital mortality as a binary classification problem. We divided the dataset into 64% for training, 16% for validation, and 20% for testing. The performance metric used was the area under the Receiver Operating Characteristic (ROC) curve (AUC). The results, indicating a superior performance of FlexGen-EHR compared to LDM and MedGAN are summarized in the table below.

[1] Interpolation-Prediction Networks, ICLR 2019

[2] Learning from Irregularly-Sampled Time Series: A Missing Data Perspective, ICML 2020

Expansion of Experiments and Metrics:

We also believe more baselines and evaluation metrics is more convincing. Indeed, we have already included results on data fidelity evaluation experiments (Figure 2 and 3) in our submission. Results showed that FLEXGEN-EHR maintains high generation fidelity. In addition, we have extended our utility evaluations on missing modality generation (10% missing in each modality ) with more baselines: 1: training with real samples and 2: Training with synthetic samples generated by TabDDPM. Results demonstrated that FlexGEN-EHR consistently outperforms baselines and is comparable to the gold standard (real samples).

I thank the authors for their comprehensive edits, revisions, and clarifications. I have thoroughly looked through the authors’ revisions and their replies to my questions. I will adjust my rating to a 6, above the acceptance threshold.

Thank you for your great question. In our work, we take a different preprocessing way. We flatten temporal features into a 1-d vector ($T \times d$), where each value represents a quantity of interest at a time stamp. We agree the decoder step might not be generalizable to video generation.