Paper: A Flexible Generative Model for Heterogeneous Tabular EHR with Missing Modality (ICLR-2024)

Reviewer comment (Experiments):
The paper needs ablation studies. For example, it would be more convincing if the authors were to include a comparison of the performance gap: (1) between scenarios with and without the use of OT on complete data; (2) between FLEXGEN-EHR and FLEXGEN-EHR - OT + KNN imputation on partially missing data.

Author response:
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.

| Model        | Imputation | AUC Score| MMD |
|--------------|-----------|----------|----------|
| FlexGen-EHR  | OT| 74.86 ± 0.30 | 0.65 ±0.018 |
| FlexGen-EHR |kNN| 71.86 ± 0.29 | 0.69 ±0.016 |
| EHR-M-GAN    | OT| 72.86 ± 0.32 | 0.67 ±0.014 |
| EHR-M-GAN    |kNN| 70.86 ± 0.28 | 0.71 ±0.017 |

Reviewer comment (Writing):
There seem to be a few typos in the manuscript (e.g., column names in Table 1, MMCAR, the notation for time-invariant features, adversarys).

Author response:
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.

Reviewer comment (Writing):
Could you explain what R^2 is? The papers referred to use MMD and KS-statistics.

Author response:
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.

Reviewer comment (Reproducibility):
What are the architectures of the decoder? Are they MLP and LSTM?

Author response:
3. Decoder Architectures: MLP only.

Reviewer comment (Writing):
Could you provide more details about the statement "[Equation 7] can achieve alignment when the sample sizes between static and temporal aren’t the same." Does this refer to the sample sizes of the non-missing data?

Author response:
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.

Reviewer comment (Reproducibility):
It would be great if the authors could provide details about the membership attack implementation (possibly in the Appendix).

Author response:
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).

Reviewer comment (Experiments):
While the model showcases empirical strength, discussions or case studies on how it might be applied in real-world healthcare settings would enhance its practical value. The paper lacks an important motivation beyond patient privacy on why generating synthetic EHR is useful in practice. Providing this motivation is essential for the paper to be useful in the healthcare setting and should not be high-level only and be supported by experiments. For example, is this useful for data augmentation? If so does having data augmentation provided by this paper’s model improve downstream prediction tasks? And other questions in a similar vein.

Author response:
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,

| Without Augmentation |85.63 ± 0.32 | - | - |
|-----------------|-----------------|-----------------|-----------------|
| # of synthetic samples  | 1k  | 2k  | 3k  | 4k  |
| LDM             | 85.92 ± 0.22 | 86.14 ± 0.28 | 86.63 ± 0.31 |
| EHR-M-GAN       | 86.05 ± 0.27 | 86.33 ± 0.29 | 86.85 ± 0.28 |
| FlexGen-EHR     | 86.56 ± 0.27 | 86.88 ± 0.31 | 87.02 ± 0.27 | 

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.

Reviewer comment (Experiments):
It would be beneficial to understand the model's performance across diverse EHR datasets to gauge its broad applicability.

Author response:
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.

Reviewer comment (Evaluation):
One big challenge in generating synthetic data for EHR is the heterogeneity of the data with respect to specific groups (based on e.g. gender, age, race, etc). The paper is missing this analyses completely as it would benefit the validity of the generated data. Without identifying these spurious correlations, it is not at all clear where such fake patient data would be useful.

Author response:
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.

Reviewer comment (Reproducibility):
The paper combines various components, like optimal transport and diffusion models, which might make the model computationally intensive or challenging to implement.

Author response:
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.

Reviewer comment (Evaluation):
The final results although beating the baselines, the uncertainty is large enough that the model’s performance improvement relative to LDM, EHR-M-GAN is not significant.

Author response:
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. To establish the robustness of our improvement, we have taken the following steps:

1. Multiple Runs: We conducted our experiments over 5 independent runs, ensuring that our reported results are not driven by chance or random variation.

2. Averaging Metrics: To provide a more comprehensive perspective, we report utility and fidelity metrics averaged across these multiple runs. This approach helps mitigate the impact of potential outliers or extreme values, giving a more reliable assessment of the model's true performance.

By incorporating these measures, we aim to provide a clearer and more confident evaluation of the improvement achieved by our model, considering the inherent uncertainty that can be associated with performance metrics in healthcare datasets.

Reviewer comment (Evaluation):
What are the computational costs associated with FLEXGEN-EHR, especially when scaling to larger EHR datasets?

Author response:
Like other generative models, the computational requirements of FlexGen-EHR are subject to variation based on the scale and intricacy of the employed electronic health record (EHR) dataset. When expanding the model to accommodate larger EHR datasets, the key computational aspect to account for primarily centers around the complexity of solving an optimal transport problem. However, it's worth noting that this complexity is not a significant concern in our work, as we address the optimal transport within the embedding space, which typically exhibits a compact size (e.g., 128 dimensions).

Reviewer comment (Evaluation):
Given the significance of missing data in EHRs, how does the model ensure that the synthetic data generated doesn't inadvertently introduce biases or inaccuracies, especially in predictive tasks?

Author response:
Like other generative models, the primary objective of FLEXGEN-EHR is to learn the data distribution to the best of its ability, with the evaluation hinging on downstream tasks. While bias mitigation remains an active area of research and development, it is essential to underscore that FLEXGEN-EHR, characterized by its flexibility, is well-positioned to integrate fairness-aware techniques. By incorporating such fairness-aware approaches into its framework, FLEXGEN-EHR strives to minimize any unintended biases that may emerge during the synthesis of data. We intend to quantify these impacts and compare the performance of different generative methods in the future.

Reviewer comment (Experiments):
How can the model handle time-varying labels? (e.g. health status of patient at any point in time instead of mortality just at the end of the trajectory.)

Author response:
We appreciate the important consideration of handling time-varying labels, such as the health status of a patient at any given time rather than just assessing mortality at the end of a patient's trajectory. While our current work does not focus on addressing this aspect, we acknowledge its practical relevance. One potential avenue for handling time-varying labels could involve the utilization of advanced conditional encoder network architectures. These sophisticated architectures could allow us to model and predict health status changes throughout a patient's journey. By feeding (records, time $t$, label at time $t$) into the network, the network could implicitly learn more comprehensive understanding of their evolving medical condition.

Reviewer comment (Writing):
Please provide more information for Equation 6 and what is trying to be achieved in the paragraph above it.

Author response:
The Gromov-Wasserstein problem aims to find an optimal way to match or align two sets of data, in our work, referred to as "temporal" and "static" latent embeddings. This matching is not a strict one-to-one correspondence but rather a probabilistic or likelihood-based translation. In simpler terms, it's about determining how likely it is that specific pairs of data points in these two sets correspond to each other.
Imagine you have two sets of data, one representing temporal information (like time-series data) and the other representing static information (like attributes of objects). The Gromov-Wasserstein problem helps us find the best way to associate or match the elements from these two sets based on how likely they are to be related.
In essence, it's a mathematical approach to quantify the likelihood of correspondence between elements from two different datasets, helping us make sense of complex relationships in data.

Reviewer comment (Reproducibility):
Have an algorithm section either in main text or appendix that shows a step by step overview of the model both for training and for generation.

Author response:
We greatly appreciate your suggestion, and in response, we have incorporated a comprehensive algorithm section into our revision. This section provides a step-by-step overview of our model, encompassing both the training and generation processes.

Reviewer comment (Experiments):
How can the model handle more realistic missingness? (see W5).

Author response:
We understand the reviewer’s concern about the challenge of missing modalities in EHR data. However, we wish to clarify that the type of missingness addressed in our study is akin to structured missingness, where specific patterns or conditions lead to data missingness. This scenario is frequently encountered in real-world healthcare datasets [1,2]. Additionally, in response to Reviewer ULHZ's valuable suggestion, we have expanded our experimental framework to include scenarios with missing features. This extension allows us to more comprehensively demonstrate the robustness of FlexGen-EHR in handling realistic and complex missing data situations. The enhanced experiments showcase FlexGen-EHR's superior performance in comparison to other models like LDM and MedGAN. The results, clearly depicted in the table below, highlight our model’s effectiveness in managing various missing data patterns, thereby reinforcing its practical utility in real-world healthcare applications."
| Model        | AUC Score | 
|--------------|-----------|
| Real         | 85.63 ± 0.32 | 
| TabDDPM      | 68.64 ± 0.29 | 
| EHR-M-GAN    | 70.48 ± 0.33 |  
| FLEXGEN-EHR  | 74.86 ± 0.30 |  

[1] Jones et.al., Characterizing and Managing Missing Structured Data in Electronic Health Records: Data Analysis, 2018

[2] Ferri et.al., Extremely missing numerical data in Electronic Health Records for machine learning can be managed through simple imputation methods considering informative missingness: A comparative of solutions in a COVID-19 mortality case study, 2023

Reviewer comment (Presentation):
The author should bold the best results in Table 4 as Table 2 and Table 3. There exists a trade-off between generation quality and privacy guarantee. The method achieves worse performance on some datasets under some criteria is not a serious problem. However, the author should not neglect it.

Author response:
## 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.

Reviewer comment (Theory):
In the second line from the bottom, the author states "their tendency to treat numerical and categorical features independently". Although TabDDPM deals with numerical and categorical features separately, TabDDPM learns the dependency between them. This is because although the diffusion process is different, they concatenate the numerical and categorical features to the network, which is the same as the latent diffusion model used in this paper.

Author response:
## 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.”

Reviewer comment (Writing):
Typo in the second line of paragraph "Heterogeneous Tabular EHR Generation": x^{S}_i instead of x^{T}_i.

Author response:
## Minor Comments. 
Thank you for pointing out several typos and errors in our submission. We have made the corresponding edits within the text.

Reviewer comment (Experiments):
While the manuscript underscores the significance of proficiently addressing missing modality, it lacks experimental validations. To effectively highlight the merits of optimal transport, I would suggest incorporating a comparative analysis between doing an imputation for missing modalities via kNN in the FLEXGEN-EHR framework, and leveraging Optimal Transport for missing modality imputation.

Author response:
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:
| Model        | Imputation | AUC Score (the higher, the better) | MMD (the lower, the better) |
|--------------|-----------|----------|----------|
| FlexGen-EHR  | OT| 74.86 ± 0.30 | 0.65 ±0.018 |
| FlexGen-EHR  |kNN| 71.86 ± 0.29 | 0.69 ±0.016 |
| EHR-M-GAN    | OT| 72.86 ± 0.32 | 0.67 ±0.014 |
| EHR-M-GAN    |kNN| 70.86 ± 0.28 | 0.71 ±0.017 |

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.

Reviewer comment (Experiments):
In Table 2, could you elucidate how your method contrasts with baseline models that are solely focused on synthesizing Discrete codes? Specifically, how did you employ MedGAN in generating the Labevents data for MIMIC-III?

Author response:
*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.

Reviewer comment (Evaluation):
I noted the statement, "we observed that latent space embedding models, trained on disparate features, manifested analogous geometric patterns and behaviors." Could you provide further clarity on the specific geometric patterns and behaviors that were identified during this observation?

Author response:
*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.

Reviewer comment (Experiments):
In sections 4.1 and 4.2, the authors assume an individual patient EHR to take one of three forms: all static features exclusively present, all temporal features exclusively present, or all features present. In real EHR data, there are often patients with some (not all) static features present and some (not all) temporal features present.

Author response:
## 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.

| Model        | AUC Score | 
|--------------|-----------|
| LDM          | 80.28 ± 0.65 | 
| MedGAN         | 79.28 ± 0.13 | 
| FLEXGEN-EHR  | 82.87 ± 0.15 | 

[1] Interpolation-Prediction Networks, ICLR 2019

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

Reviewer comment (Experiments):
Since addressing the problem of missing modality seems to be the primary goal of the paper, more experiments in general would be helpful. This could include showing more baselines beyond EHR-M-GAN, additional evaluation metrics or scores for fidelity and utility, and more results in the Appendix to support the claims. Currently, only R-squared values are depicted for generation fidelity and only eICU for utility.

Author response:
## 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). 

| Model        | AUC Score | 
|--------------|-----------|
| Real         | 85.63 ± 0.32 | 
| TabDDPM      | 68.64 ± 0.29 | 
| EHR-M-GAN    | 70.48 ± 0.33 |  
| FLEXGEN-EHR  | 74.86 ± 0.30 |

Reviewer comment (Experiments):
It would be very helpful to include results from training on real data for Figures 2-4 to provide more context. Currently, EHR-M-GAN is the only comparison.

Author response:
## 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). 

| Model        | AUC Score | 
|--------------|-----------|
| Real         | 85.63 ± 0.32 | 
| TabDDPM      | 68.64 ± 0.29 | 
| EHR-M-GAN    | 70.48 ± 0.33 |  
| FLEXGEN-EHR  | 74.86 ± 0.30 |

Reviewer comment (Presentation):
A couple of brief mentions of future work or limitations in the conclusion would help clarify and reaffirm the authors’ current goals and progress. For example, addressing lower privacy scores or the lack of ablation studies and their implications.

Author response:
## 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.