Data Augmentation for Seizure Prediction
with Generative Diffusion Model

Epilepsy is caused by abnormal discharge of brain neurons and is one of the world’s most common neurological diseases [1]. It is reported by the World Health Organization (WHO) that about 50 million people are suffered from epilepsy. With the development of medical science, 70% of the patients can be seizure-free after proper treatment. However, the remaining 30% of them still cannot be controlled by drugs [2]. Therefore, seizure prediction has great value for them. A precise prediction can let patients get timely protection, improving the quality of their lives by reducing their mental stress.

Electroencephalogram (EEG) is an efficient method for capturing and monitoring electrical activity of the brain [3]. EEG recordings can be divided into three kinds depending on the position where the signals are collected. There is non-invasive scalp electroencephalogram (sEEG) with electrodes attached to the subject’s scalp, semi-invasive electrocorticogram (ECoG) with electrodes placed between the cerebral cortex and the dura mater, and invasive electroencephalogram (iEEG) with implanted electrodes [4]. EEG is widely used to diagnose epilepsy in clinical [5]. Over the last few years, studies have demonstrated that EEG can be used to predict upcoming seizures [6][7]. In practice, physicians have divided the epilepsy EEG signals into four states: preictal (the period before seizures), ictal (the duration of seizures), postictal (the period following seizures) and interictal state (the interval between seizures) [8]. According to the definitions above, seizure prediction is transferred to a binary classification task which distinguishes preictal states from interictal ones. In this way, a variety of methods have been proposed to achieve the prediction.

In the traditional machine learning methods, many kinds of manually extracted features have been used to predict seizures, including spatial domain, time domain, frequency domain, and time-frequency domain features. After feature extraction, a classifier is used to do the classification by these features [9]. For instance, Chisci et al. utilized the autoregressive coefficient as features, and a support vector machine (SVM) as the classifier [10]. However, extracting features manually requires extensive expertise, and the efficacy of seizure prediction is heavily reliant on the quality of these extracted features. Consequently, deep learning (DL) methods have gained growing attention within the last few years, where representative features can be extracted automatically instead of manually [11][12]. Researchers have designed various DL architectures to achieve high-performance seizure prediction. Li et al. applied the multi-layer perceptrons (MLP) blocks to extract spatial and temporal features [13]. Gao et al. employed a multi-scale convolutional neural network (CNN) by using the dilated convolutions to capture information at different levels [14]. In another work [15], Gao et al. proposed a Transformer-based network. The model utilized the attention mechanism to establish associations between information at different locations and strengthen the role of key information. The multiple features extracted by these networks had great effects on seizure prediction and these DL methods all obtained satisfactory performance.

It is worthy noting that DL algorithms typically need a substantial amount of labeled data to achieve good performance. However, the low-frequency property of seizure onsets causes serious insufficiency of the preictal data for a specific patient [16][17]. Thus, the composition of dataset is seriously imbalanced, leading to the overfitting during model training [18].

To solve the issue of imbalanced data, various methods have been employed. Khan et al. selected interictal samples randomly to down-sample the interictal data [19]. However, down-sampling discards too many interictal samples, which could lead to the loss of useful information. Besides, the small quantity of data after down-sampling makes it hard for models to reach the optimal result. Therefore, some researchers proposed data augmentation (DA) methods. DA is the process that generates new samples to augment a small or imbalanced dataset. Troung et al. generated more preictal segments by using an overlapping technique which slid a 30-s window along the EEG signal [20]. Zhang et al. came up with an idea of recombination. For each seizure, they split every training EEG sample into three segments, and generated new samples as a recombination of the randomly selected segments [21]. Although those traditional DA methods can increase the number of preictal samples, they are just simple transformations of the existing data, with distribution constrained by original data due to the presence of numerous repetitive parts [22]. Because of the nonstationary signals in the brain, the epileptic representation of EEG varies over time [23] and an upcoming seizure may have representation biased from the previous ones. The simple transformations of existing samples can hardly expand the data distribution to cover the variable representation.

With the advancement of generative model, an increasing number of researchers have utilized these models to address the imbalanced data [24]. Rasheed et al. used a generative adversarial network (GAN) to produce synthetic EEG feature maps and evaluated the performance of generated data by a validation method [25]. However, the training process of GAN is difficult. Once the design is improper, the gradient may disappear or explode, resulting in unstable generation quality. In addition, the model may fall into local minimum during training and can only generate samples similar to a limited number of initial data. This is known as the mode collapse problem and will lead to the lack of generation diversity [26].

Recently, diffusion models have been widely used in various fields such as computer vision (CV) and natural language processing (NLP) [27][28]. Diffusion models are a class of promising generative models, which regard data as a step-by-step diffusion of a series of random noises. They use a Markov chain to gradually add Gaussian noise to data, getting the posterior distribution probability and learning the reverse denoised process. Thus they could synthesise new samples from random noise with the learned distribution of original data [29]. One important use of the diffusion model is DA due to its strong generation ability [30][31][32]. For example, Chambon et al. utilized a pre-trained latent diffusion model to generate high-fidelity, diverse synthetic chest x-rays (CXR) and measured a 5% improvement by jointly training a classifier on synthetic and real images [33]. In this paper, we proposed a diffusion-based model named DiffEEG. In specific, DiffEEG consists of two opposite processes, namely the forward/diffusion process and the reverse/denoised process. In the diffusion process, the model adds noise with different scales to the EEG sample step by step and converts it into random noise. This procedure is conditioned on the short-time Fourier transform (STFT) spectrogram of the input EEG to provide the guiding time-frequency features. By minimizing the loss between the output of the network and the added noise, the feature distribution of preictal data could be fully explored. In the denoised process, based on the learned conditional probability distribution, the synthetic data could be sampled from random noise input by gradually removing the noise with different scales. With the randomness of initial noise and the precise representation, the generated samples could have new information and diversity while ensuring the consistency of feature space, narrowing the distance between data clusters composed by initial samples from each seizure [34]. The distribution illustration of down-sampling, sliding windows, recombination and DiffEEG is presented in Fig. 1. Different from the existing DA methods, DiffEEG is capable of utilizing the learned distribution information to generate samples of diversity and provide additional information for model training.

We train DiffEEG with the patient’s preictal EEG signals so that our model can generate synthetic preictal data until we have an equal number of samples per class. The generation quality is improved with the reduction of the training loss. The samples are then given to a classifier to conduct seizure prediction task. We used five representative network frameworks (SVM, Spatio-temporal MLP, EEGNet, Multi-scale CNN, Transformer) as classifiers on the CHB-MIT and Kaggle dataset to verify the universality and generality of our model. In addition, we compared our model with that using just the original data and those using three existing methods for solving the data imbalance, including down-sampling, sliding windows and recombination, to evaluate the effectiveness and superiority of DiffEEG.

The main contributions of our work are as follows:

• •

  We propose a method of using diffusion model to solve the imbalanced data problem, which is the first work to introduce diffusion model to the seizure prediction field.

• •

  We verify that our diffusion model is superior to the existing DA methods and can be utilized in almost any other classification network to substantially improve the seizure prediction performance.

• •

  With the improvement of DiffEEG, the best prediction performance in our experiments outperforms that of most state-of-the-art (SOTA) methods.

The rest of the paper is structured as follows. Section II details our proposed methods, including the theory of diffusion model, the structure of DiffEEG and five classifiers. Section III shows the datasets, our experimental results and comparison with SOTA works. Section IV discusses the results of our methods. Finally, we conclude our work in Section V.

In this section, comparative experiments are conducted on the CHB-MIT and Kaggle database. The details of the two datasets and the evaluation metrics of our experiments are presented, followed by the experimental results and comparisons with several SOTA methods.

Despite decades of development in medical technology, the collection of high-quality EEG signals still demands expensive medical devices and plenty of expert time [25]. It is impractical to increase the amount of preictal EEG data through continuous collection. Data augmentation is the process that generates new samples to augment a small or imbalanced database by converting existing samples. Exposing the model to various augmented data could raise the precision and stability of classifier, making it more robust to changes [58]. According to the results in Table III $\sim$ Table XII, it can be observed that all the classification models using the augmented data generated by DiffEEG achieve higher Sens and AUC than the models using just the original data, indicating that our method can improve the predictive precision of epilepsy. However, the models using the original data exhibited lower FPR. This is because the interictal samples in the original data are far more than the preictal data, leading to the overfitting of the model. As a result, the model can accurately identify the interictal samples but struggles to recognize the preictal samples. Besides, compared with the existing methods for addressing the imbalanced data problem, we can conclude that the DiffEEG has obvious advantage over them. There has been a notable enhancement in the performance of seizure prediction, proving that the data generated by DiffEEG has high quality and diversity. From a statistical perspective, the augmented data can increase the number of subjects whose improvement over chance is statistically significant, suggesting that our DA technique is capable of enhancing the reliability of the model’s predictions. In practice, the distribution of EEG signals varies over time due to the changes in external environment and mental states of patients. The epileptic representation in different preictal time varies as well if the location or mode of abnormal discharge is different [23]. When the data cluster of an upcoming seizure has few nearby distribution with the existing clusters, the seizure will be difficult to predict. Based on the distribution learned during the diffusion process, the data produced by DiffEEG can narrow the distance between different clusters and provide additional information to classifiers. Instead, the distribution of data generated by sliding windows or recombination is restricted by the original data and cannot be fully explored, so their improvement to prediction performance is limited.

By comparing the results of five classification networks, it is clear that no matter what DA methods we used, the performance of Multi-scale CNN was the best. SVM is a powerful machine learning classifier, but its performance typically relies on the handcrafted feature engineering. The manually extracted features are hard to capture all the complex representations from raw data, making the performance of SVM inferior to that of deep learning networks. Although the blocks in Spatio-temporal MLP can extract information within and between channels, the fully connected structure of MLPs determines that the efficiency of data processing is not high. The lack of weight sharing mechanism also causes a huge amount of parameters, making it difficult to converge to optimal result. The structure of EEGNet is lightweight and computationally efficient. The multi-layer 1D convolution enables the model to effectively extract the time-frequency information. However, EEGNet lacks consideration of the inter-channel correlation, leading to the insufficient extraction of spatial features. For Transformer, the input embedding and position encoding can capture the global information of data. The attention module can learn the relation between features and strengthen the role of key information. For Multi-scale CNN, the local perception ability of convolution makes it easier to extract local information. By the use of dilated convolution, spatial and temporal features of different scales can be extracted, providing multi-scale information for seizure prediction. In the meantime, the model also utilizes attention mechanism as Transformer to assign weights to features and focus on the critical representation. Therefore, the Multi-scale CNN network obtains the best effects. SVM is a widely used traditional classification method. MLP, CNN and Transformer are the most representative frameworks in deep learning and almost all classifiers are based on them. As our proposed DiffEEG shows superiority on these networks, it is proved that DiffEEG can be utilized in almost any other classification network to substantially improve the seizure prediction performance.

The generation of multichannel EEG signals has been a challenge to researchers for a long time. Therefore, the existing generative models either generate the feature maps of EEG signals [25] or generate each channel respectively before combining them [59]. This is done in case the model cannot fully learn the representations of multichannel EEG signals. Our proposed model, DiffEEG, can overcome the challenge and directly generate signals of the whole channels. The real and synthetic samples of 18 channels and 30-s duration generated by DiffEEG are shown in Fig. 8. To further inspect the details of the samples, we only select the first three channels of 10-s duration as shown in Fig. 9. Through visual inspection, we can see that the augmented data is not the simple superposition of existing data, but diversified data generated based on the learned distribution, which could provide additional information to improve the performance.

Furthermore, statistical tests are performed for the synthetic data to validate that the augmented data increase diversity and new information while preserving the characteristics of the genuine EEG data. As the augmented data are generated conditioned on the STFT spectrum of real data, a paired t-test can be appropriately performed for the augmented data and the corresponding real data. T-tests are usually carried out on the feature space instead of the high-dimensional time-domain data [60]. Hence, we select variance to reflect the amplitude distribution, and extract the spectral power of $\delta$, $\theta$, $\alpha$, $\beta$, $\gamma$ bands to give information about the energy variations. We utilized these 6 features to conduct the paired t-test at a significance level of 0.05. The p-values on each subject are shown in Table XIV. Based on the theory of t-test, if the p-value is less than 0.05, there is considered to be a significant difference between the two sample groups; otherwise, there is no significant difference. From the results, it can be observed that the variance of both groups is not significantly different across all subjects except for Dog3, indicating that the amplitude distribution of the augmented data is similar to that of the real data for most patients. For each spectral feature, some of the patients exhibit significant difference between the real and augmented groups, some of the patients not. This reflects the diversity introduced by our DiffEEG model in the generated samples. The synthetic data retain the temporal rhythms of the original data while bringing in diverse variations in the frequency domain.

The DiffEEG model is trained with the guidance of STFT condition. Compared to directly using the raw EEG signal, the conditional model can generate samples with higher quality. Numerous articles have demonstrated the advantages of conditional diffusion models over unconditional ones. Kong et al. mentioned that generating audios in the time domain without conditional information is a challenging task [28]. Their conditional model achieved higher Mean Opinion Score (MOS) values than the unconditional model, indicating that the generated samples of the conditional model have higher quality. Ho et al. also showed that the sample quality and fidelity can be greatly improved by the conditional guidance [34]. Certainly, the conditional models also have their limitations. Ho et al. explained that the presence of condition could reduce the diversity of generation, resulting in a trade-off between quality and diversity. In our EEG generation task, we prioritize the quality because the complicated distribution of EEG data is difficult to learn. The loss of authenticity in the augmented data will lead to a decline in the seizure prediction performance. Besides, we have made use of the recombined spectrograms to increase the diversity as mentioned in Section II-D.

Finally, seizure prediction experiments are carried out to test the contributions of the synthetic data. The results verify that the data generated by DiffEEG can greatly improve the prediction performance and beat the existing DA methods in various situations.

The outstanding effects achieved by DiffEEG on public datasets have revealed multiple potential impacts on the clinical practice of seizure prediction.

• •

  Rapidly establishing high-performance and generalizable epilepsy prediction models for patients. In clinical practice, the insufficient preictal data often restricts the performance of seizure prediction models [59]. Relying on collecting more preictal data to build high-performance models is extremely time-consuming. With DiffEEG, preictal data with additional information can be generated in a short time, enhancing the predictive effect of the model. Moreover, the strong generalization capability of the synthetic data allows them to be widely applicable to various prediction models.

• •

  Assisting doctors in studying biomarkers for epilepsy prediction and enhancing the interpretability of predictions. It is challenging to capture the EEG biomarkers before seizures because of the complex spatial-temporal dynamics present in the epileptic brain [61]. The DiffEEG model can fully explore the feature distribution of EEG signals and generates data with distinct epileptic representation. These data provide a possible reference for doctors in identifying biomarkers for seizure prediction, improving the interpretability of deep learning-based prediction models.

• •

  Protecting the data privacy of patients. In the medical field, EEG data is typically sensitive information. Utilizing synthetic samples for data augmentation can reduce the reliance on original data, thereby lowering the risks of data leakage and privacy concerns.

Our work also presents vast application prospects in other EEG-based tasks. Issues such as data imbalance and scarcity are prevalent in EEG-based tasks. For instance, in sleep staging tasks, Stage N2 generally occupies a majority of the sleep time (45%-55%), while stage N1 only accounts for very little time (2%-5%) [62]. In emotion recognition tasks, the AMIGOS dataset exhibits a higher arousal ratings, whereas in the DEAP dataset, the proportion of positive emotions is larger than that of negative ones [63]. In the diagnosis of brain diseases such as depression, Alzheimer’s disease and mild cognitive impairment, it is typical to face a lack of data due to the equipment requirements, time constraints and privacy concerns. The EEG signals for all these tasks exhibit significantly different feature distribution. However, DiffEEG can leverage the STFT to provide time-frequency information, guiding the model to learn task-specific representations for each distribution. Therefore, our approach enables the generation of diverse and high-quality EEG samples for data augmentation under various tasks, raising the generalization, accuracy, and robustness of these tasks.

Despite the promising influence and applications demonstrated by DiffEEG, it is important to acknowledge its potential limitations and pursue the directions for future research. One of the primary limitations is the slow training process. Although we speed up the training process by randomly selecting the diffusion steps, it still takes plenty of time to train the model. In future work, the input block could be modified to transform the multichannel EEG data into latent space. By performing the diffusion process in a low-dimensional latent space, the computation cost and training time can be greatly reduced. Another limitation is the lack of cross-subject capabilities. To ensure the generation quality, DiffEEG is trained under a subject-specific strategy. When a new patient arrives, we need to retrain a model rather than utilizing the existing data and models, which is both time-consuming and resource-wasting. In the future, the structure of DiffEEG could be improved for pretraining on a large dataset. In this way, we only need to fine-tune the pretrained model to fit the data distribution of each patient, enabling the model with few-shot learning capability.

In addition to the inherent limitations of the DiffEEG itself, there are several challenges and considerations when integrating it into real-world applications. The first obstacle comes from the acquisition quality of the EEG signals. In real-world scenarios, factors such as the loose or offset of the collecting electrodes, the patient’s chewing and vigorous movement can introduce severe noise in the EEG signals [64]. The noise interference may mask or distort the epileptic representation, hindering the model from learning the distribution characteristics of epileptic EEG data. We cannot guarantee that a model trained on low-quality EEG signals is capable of producing high-quality EEG data. The second challenge is the difficulty in continuously optimizing the model. Patients employing the epilepsy prediction system still experience seizures, and the consequent preictal data could be utilized to update and enhance the DiffEEG model. However, retraining the model with all the data each time new data arrives is too costly, and training the model solely on the new data could lead to the catastrophic forgetting problem. Hence, integrating incremental learning into DiffEEG is worth considering. At last, the highly realistic content generated by DiffEEG can raise ethical concerns. The model has the potential to create deceptive EEG signals including deepfakes and misinformation [65]. Therefore, ensuring responsible use and preventing malicious applications are ongoing challenges.

In this article, we propose a novel and effective data augmentation method by the use of diffusion model. We conduct contrast experiments with just using the original data and with three existing methods of solving the imbalanced data problem: down-sampling, sliding windows and recombination. We evaluate the prediction performance on five representative classifiers to verify the effectiveness and generality of DiffEEG. The results prove that the data generated by DiffEEG achieve better performance on all five classification networks. Among them, the Multi-scale CNN utilizing the data augmented by DiffEEG obtained the best performance, with an average Sens, FPR, AUC of 95.4%, 0.051/h, 0.932 on the CHB-MIT database and an average Sens, FPR, AUC of 93.6%, 0.121/h, 0.822 on the Kaggle database. Both results reached the SOTA level. To the best of our knowledge, our work is the first diffusion model used in epileptic EEG generation, creating a new and effective way to solve the imbalanced data. This model can be widely applied to all kinds of seizure prediction classifiers to improve the classification performance significantly.