EEG-based sleep staging via self-attention based capsule network with Bi-LSTM model

Sleep staging via electroencephalogram is essential for determining the quality of sleep. Manual sleep staging is expensive and time-consuming. Recently, many deep learning-based sleep staging methods are demonstrated to outperform traditional methods. However, most methods do not fully exploit the temporal correlation between features of electroencephalogram signals. In this paper, we propose a self-attention routing-based capsule network with bi-directional long short-term memory model to extract more discriminative features from electroencephalogram signals and improve the accuracy of sleep staging. First, a convolutional neural network is used to extract salient features from the electroencephalogram signal. Second, to learn the transition rules between different sleep epochs, a bi-directional long short-term memory is used to capture the temporal dependence between the encoded electroencephalogram signals. Finally, to fully explore the temporal correlation between the features from the electroencephalogram signals, a self-attention routing-based capsule network is utilized to recode the importance based on the intrinsic temporal similarity of electroencephalogram signals. We evaluated our model by two different single-channel electroencephalogram signals (i.e., Fpz-Cz and Pz-Oz electroencephalogram channels) from two public sleep datasets, named Sleep-EDF-39 and Sleep-EDF-153. Our overall accuracies on the Sleep-EDF-39 and Sleep-EDF-153 datasets are 85.8% and 83.4%, with a kappa of 0.8 and 0.77, respectively. The results show that our proposed method achieves the state-of-the-art level of sleep staging using a single-channel electroencephalogram and offers the possibility of widespread application of capsule networks for sleep staging.