EHRDiff : Exploring Realistic EHR Synthesis with Diffusion Models

In the literature on EHR synthesis, researchers are usually concerned with the generation of discrete code features such as ICD codes rather than clinical narratives. Researchers have developed various methods to generate synthetic EHR data. Early work was usually disease-specific or covered a limited number of diseases. Buczak et al. (2010) developed a method that generates EHR including visit records, clinical activity, and laboratory results for 203 synthetic tularemia outbreak patients of tularemia. The features in synthetic EHR data are generated based on retrieving similar real-world EHR which is inflexible and prone to privacy leakage. Walonoski et al. (2017) developed a software named Synthea which generates synthetic EHRs with various patient information based on publicly available data. They build generation workflows based on biomedical knowledge and real-world feature statistics. Various models are aggregated for different feature synthesis. However, Synthea only covered the 20 most common conditions.

Recently, researchers mainly applied generative modeling methods for EHR synthesis (Ghosheh et al., 2022). Medical GAN (medGAN) (Choi et al., 2017) introduced GAN to EHR synthesis. medGAN can generate synthetic EHR data with good quality and is free of tedious feature engineering. Following medGAN, various GAN-based methods are proposed, such as medBGAN (Baowaly et al., 2018), EHRWGAN (Zhang et al., 2019), CorGAN (Torfi & Fox, 2020a), etc. These GAN-based methods advance synthetic EHR to higher quality. However, a common drawback of GAN-based methods is that these methods suffer from the mode collapse phenomenon which results in a circumstance where a GAN-based model is capable of generating only a few modes of real data distribution (Thanh-Tung et al., 2018). To mitigate the problem, GAN methods for EHR generation rely on pre-trained auto-encoders to reduce the feature dimensions for training stability. However, inappropriate hyper-parameter choices and autoencoder pre-training will lead to sharp degradation of synthetic EHR quality or even failure to generate realistic data. There is also research that uses GAN-based models for conditional synthetic EHR generation to model the temporal structure of real EHR data (Zhang et al., 2020a). Since diffusion models are less studied in EHR synthesis, we focus on the unconditional generation of EHR and leave modeling conditional temporal structure with diffusion models to future work.

Besides GAN-based methods, there also exists research that explores generating synthetic EHR data through variational auto-encoders (Biswal et al., 2020) or language models (Wang & Sun, 2022). Concurrently, MedDiff (He et al., 2023) is proposed and explores diffusion models for synthetic EHR generation, and they propose a new sampling technique without which the diffusion model fails to generate high-quality EHRs. Ceritli et al. (2023) directly apply TabDDPM (Kotelnikov et al., 2022) to synthesizing EHR data.

Diffusion models are formulated with forward and reverse processes. The forward process corrupts real-world data by gradually injecting noise, and harvesting training data with different noise levels for a denoising distribution, while the reverse process generates realistic data by removing noise using the denoising distribution. Sohl-Dickstein et al. (2015) first proposed and provided theoretical support for diffusion models. Denoising diffusion probabilistic models (DDPM) (Ho et al., 2020) and noise conditional score networks (NCSN) (Song et al., 2021) have discovered the superior capability in image generation, and diffusion models have become a focused research direction since then. Recent research generalizes diffusion models to the synthesis of other data modalities and achieves excellent performance (Li et al., 2022; Kong et al., 2021). Our work is among the first research that introduces diffusion models to realistic EHR synthesis (He et al., 2023; Ceritli et al., 2023).

Real-world EHR contains diverse patient information including demographics, physiological conditions, laboratory results, ICD codes, etc. Such information may be presented in various data formats (e.g., categorical, binary, continuous, etc). Features of different information may also have different scales, such as ages and heights. In this work, we use a vector $\bm{x}_{0}$ to represent the EHR, where $\bm{x}_{0}\in\mathcal{R}^{|\mathcal{C}|}$ and $\mathcal{C}$ represents the dimension of EHR features of interest. We treat all features of different formats as real numbers ranging from $0$ to $1$. Following previous research (Choi et al., 2017; Baowaly et al., 2018) which only focused on binary code data, we use $1$ standing for code occurrence and $0$ otherwise. For continuous and categorical features, we normalize each feature within the range.

Diffusion models are characterized by forward and reverse Markov processes with latent variables where the forward process transforms the real samples to random Gaussian noise while the reverse process generates synthetic samples by gradual denoising the noise. As demonstrated by Song et al. (2021), the forward and reverse processes can be described by stochastic differential equations (SDE). The general SDE form for modeling the forward process follows:

where $\bm{x}$ represents data points, $\bm{w}$ represents the standard Wiener process, $t$ is diffusion time, and ranges from $0$ to $T$. At $t=0$, $\bm{x}_{0}$ follows real data distribution while at $t=T$, $\bm{x}_{T}$ asymptotically follows a random Gaussian distribution. Functions $f$ and $g$ respectively define the sample corruption pattern and the level of injected noises. $f$ and $g$ together corrupt real-world samples to random noise. Based on the forward SDE, the SDE for the reverse process can be derived as follows:

where $p_{t}(\bm{x})$ is the marginal density of $\bm{x}$ at time $t$, and $\nabla_{\bm{x}}\log p_{t}(\bm{x})$ is the score function, which indicates a vector field of which the direction is pointed to the high-density data area. With reparameterization proposed in Karras et al. (2022), the reverse generation process can also be described with the probability flow ordinary differential equations (ODE) instead of SDEs (Song et al., 2021):

where $\dot{\sigma}(t)$ represents the derivative of $\sigma(t)$. SDE and probability flow ODE indicate stochastic and deterministic generation processes respectively. To generate synthetic samples following reverse processes, it is required to learn a score function $s_{\theta}(\bm{x})$ parameterized by $\theta$ by score matching $\min_{\theta}\mathbb{E}_{p(\bm{x}_{0})p_{\sigma_{t}}(\bm{x}|\bm{x}_{0})}\left[\|s_{\theta}(\bm{x})-\nabla_{\bm{x}}\log p_{\sigma_{t}}(\bm{x}|\bm{x}_{0})\|^{2}_{2}\right]$.

For the diffusion process of EHRDiff , we use $h(t)=1$ and $\sigma_{t}=t$ from previous research (Karras et al., 2022). Therefore $\nabla_{\bm{x}}\log p_{\sigma_{t}}(\bm{x}|\bm{x_{0}})=-\frac{\bm{x}-\bm{x}_{0}}{\sigma_{t}^{2}}$ and we reparameterize $s_{\theta}(\bm{x})=-\frac{\bm{x}-D_{\theta}(\bm{x},\sigma_{t})}{\sigma_{t}^{2}}$, then the objective can be derived as:

and with further simplification of ignoring ${\sigma_{t}^{2}}$, the final objective becomes:

Generally, $D_{\theta}(\bm{x},\sigma_{t})$ can be modeled by neural networks, while such direct modeling may cause obstacles for optimization because the variance of $\bm{x}_{t}$ and the scale of $\sigma_{t}$ are diverse at different time step $t$. Therefore, we chose the pre-conditioning design of $D_{\theta}(\bm{x},\sigma_{t})$ (Karras et al., 2022) where $D_{\theta}(\bm{x},\sigma_{t})$ is decomposed as:

$F_{\theta}$ is modeled with neural networks and such designs of $c_{\text{in}}$ and $c_{\text{noise}}$ regulate the input to the network to be unit variance across different time step $t$, and $c_{\text{out}}$ and $c_{\text{skip}}$ together set the neural model prediction to be unit variance with minimized scale. Therefore in EHRDiff , we chose $c_{\text{out}}=\sigma\sigma_{\text{data}}/\sqrt{\sigma^{2}+\sigma_{\text{data}}^{2}}$ and $c_{\text{skip}}(\sigma)=\sigma_{\text{data}}^{2}/(\sigma^{2}+\sigma_{\text{data}}^{2})$. $c_{\text{noise}}(\sigma)=0.25\ln{\sigma}$ which is designed empirically with the principle of constraining the input noise scale from varying immensely and $\ln(\sigma)\sim\mathcal{N}(P_{\text{mean}},P_{\text{std}}^{2})$ following Karras et al. (2022), where $P_{\text{mean}}$ and $P_{\text{std}}$ are hyper-parameters to be set.

With the aforementioned formalization of $h(t)$, $\sigma(t)$ and the learned score function $s_{\theta}$, the generation process of EHRDiff can be expressed as the ODE:

Solving the ODE numerically requires discretization of the time step $t$ and a proper design of noise level $\sigma_{t}$ along the solution trajectory. Therefore, following previous research (Karras et al., 2022), we set the maximum and minimum noise levels as $\sigma_{\text{max}}$ and $\sigma_{\text{min}}$, and use the following form of discretization:

where $i$’s are integers and range from $0$ to $N$, $\sigma_{t_{N}}=0,\sigma_{t_{N-1}}=\sigma_{\text{min}}$, and $\rho$ controls the schedules of discretized time step $t_{i}$ and trades off the discretized strides $t_{i}-t_{i-1}$ the larger value of which indicates a larger stride near $t_{0}$. In order to solve the ODE more precisely and generate synthetic EHR with higher quality, we use Heun’s $2$nd order method, which adds a correction updating step for each $t_{i}$ and alleviates the truncation errors compared to the $1$st order Euler method. We leave the detailed sampling procedure to Appendix A.

Many previous research uses in-house EHR data which is not publicly available for method evaluation (Zhang et al., 2019; Yan et al., 2020). Such experiment designs set obstacles for later research to reproduce experiments. In this work, we use a publicly available EHR database, MIMIC-III, to evaluate EHRDiff .

Deidentified and comprehensive clinical EHR data is integrated into MIMIC-III (Johnson et al., 2016). The patients are admitted to the intensive care units of the Beth Israel Deaconess Medical Center in Boston. For each patient’s EHR for one admission, we extract the diagnosis and procedure ICD-9 code and truncate the ICD-9 code to the first three digits. This preprocessing can reduce the long-tailed distribution of the ICD-9 code distribution and results in a 1,782 code set. Therefore, the EHR for each patient is formulated as a binary vector of 1,782 dimensions. The final extracted number of EHRs is 46,520 and we randomly select 41,868 for model training while the rest are held out for evaluation.

Although most of the existing research focused on synthesizing discrete code features, real-world EHR data contains various data formats such as continuous test results values or time series of electrocardiograms (ECG). In this work, we extend the previous research and explore applying EHRDiff to the synthesis of EHR data other than binary codes. We use the following two datasets: CinC2012 Data and PTB-ECG Data. CinC2012 Data is a dataset for predicting the mortality of ICU patients and contains various feature formats such as categorical age, or continuous serum glucose values. PTB-ECG Data contains ECG signal data for heart disease diagnosis. Detailed introductions of both datasets are left in C. All the categorical features are converted into binary columns by one-hot encoding, and continuous features are normalized to values in the range of $[0.0,1.0]$. We use sets A and B in CinC2023 Data as training and held-out testing sets respectively. The PTB-ECG Data is split with a ratio of 8:2 for training and held-out testing.

To better demonstrate EHR synthesis performance, we compare EHRDiff to several strong baseline models as follows.

In our experiments, we evaluate the generative models’ performance from two perspectives: utility and privacy (Yan et al., 2022). Utility metrics evaluate the quality of synthetic EHRs and privacy metrics assess the risk of privacy breaches. In the following metrics, we generate and use the same number of synthetic EHR samples as the number of real training EHR samples.

In our experiments, for the diffusion noise schedule, we set $\sigma_{\text{min}}$ and $\sigma_{\text{max}}$ to be $0.02$ and $80$. $\rho$ is set to $7$ and the time step is discretized to $N=32$. $P_{mean}$ is set to $-1.2$ and $P_{std}$ is set to $1.2$ for noise distribution in the training process. For $F_{\theta}$ in Equation 8, it is parameterized by an MLP with ReLU (Nair & Hinton, 2010) activations and the hidden states are set to $[1024,384,384,384,1024]$. For the baseline methods, we follow the settings reported in their papers. The reported standard errors marked with $\pm$ are calculated under 5 different runs.

Since both datasets are designed for classification, we inspect the utilities of synthesized data by evaluating the Area Under receiver operating characteristic Curve (AUC) of classifiers trained with synthetic data. We use LightGBM (Ke et al., 2017) as classifiers and train on synthetic data of the same size as real training data.

The results shown in Table 3 that classifiers trained by synthetic data from EHRDiff achieve the highest AUC values and are consistently better than GAN-based methods, reaching 0.8405 and 0.9898 on average on CinC2012 Data and PTB-ECG Data respectively. They also have on-par performance with classifiers trained by real data. The results show the great utility of EHRDiff generated EHR data, and the efficacy is consistently good across different EHR data feature formats. This demonstrates EHRDiff is practical in real-world scenarios and can approach EHR synthesis of diverse formats. The downstream biomedical models can benefit from training on synthetic EHR data by EHRDiff , potentially overcoming the obstacles of limited publicly available real EHR data.

In this section, we discuss and conduct ablation experiments on the effective designs of EHRDiff . The designs of EHRDiff include three major aspects: diffusion process, sampling method for inference, and neural network design for modeling denoising functions. We demonstrate the performance for each through the utility and privacy metrics on MIMIC-III. The results are shown in Table 4. For the diffusion process, we compare the diffusion process of EHRDiff to the choice of variance exploding (VE) and variance preserving (VP) diffusion process. VE and VP are diffusion processes originally proposed in Song et al. (2021). We can see that EHRDiff outperforms both choices in terms of utility metrics while regarding the privacy concerns, EHRDiff performs a little inferior to VP and VE diffusion processes. As discussed in Section 4.5.2, the privacy performance can be compromised to better utility performance. For the sampling process, Heun’s 2nd order method outperforms Euler’s 1st method in utility. A 2nd order sampling method can lead to a better step-wise estimation of each denoising step, and hence has better generation performance than 1st methods when using the same trained denoising function. For the network design, we can see that EHRDiff without pre-conditioning fails to generate high-quality EHR data as shown by the lowest NZC results of only 432. The possible explanation for this is that the binary code features may have diverse prevalence which leads to much difference in variance of each feature. Pre-conditioning sets the inputs and outputs of the neural networks to be unit variance and thus eases the modeling difficulty for the network in the denoising function.

Overall, we have demonstrated the effectiveness of the design choices in EHRDiff regarding diffusion processes, sampling methods, and network designs. EHRDiff achieves the best performance regarding utility metrics, indicating a superior generation quality, while such designs in EHRDiff only compromise the privacy concerns marginally.

Our main motivation is to use the synthetic data from EHRDiff to aid the downstream methodology development on EHR, emphasizing the limitation of scarce real-world EHR data. Therefore, we further demonstrate the influence of synthetic EHR data scales on training downstream models. We train classifiers of the same type on PTB-ECG and CinC2023 data with different scales of synthetic data from EHRDiff and the previous state-of-the-art EMR-WGAN for comparison. The results are shown in Figure 3.

As can be seen from the results, the downstream performance of models trained on synthetic data from EHRDiff improves as the size of synthesized data scales up. Models trained on synthetic data generated from EHRDiff outperformed those trained on data from EMR-WGAN across different sizes of synthetic training data for both CinC2012 and PTB-ECG, except on minimal data size of 100 on CinC2012. Notably, our results demonstrate that smaller sample sizes of synthetic EHR data from EHRDiff can lead to superior downstream model performance compared to larger sample sizes from EMR-WGAN.