WaveSleepNet: An Interpretable Network for Expert-like Sleep Staging

Yan Pei,Wei Luo
2024-04-11
Abstract:Although deep learning algorithms have proven their efficiency in automatic sleep staging, the widespread skepticism about their "black-box" nature has limited its clinical acceptance. In this study, we propose WaveSleepNet, an interpretable neural network for sleep staging that reasons in a similar way to sleep experts. In this network, we utilize the latent space representations generated during training to identify characteristic wave prototypes corresponding to different sleep stages. The feature representation of an input signal is segmented into patches within the latent space, each of which is compared against the learned wave prototypes. The proximity between these patches and the wave prototypes is quantified through scores, indicating the prototypes' presence and relative proportion within the signal. The scores are served as the decision-making criteria for final sleep staging. During training, an ensemble of loss functions is employed for the prototypes' diversity and robustness. Furthermore, the learned wave prototypes are visualized by analysing occlusion sensitivity. The efficacy of WaveSleepNet is validated across three public datasets, achieving sleep staging performance that are on par with the state-of-the-art models when several WaveSleepNets are combine into a larger network. A detailed case study examined the decision-making process of the WaveSleepNet which aligns closely with American Academy of Sleep Medicine (AASM) manual guidelines. Another case study systematically explained the misidentified reason behind each sleep stage. WaveSleepNet's transparent process provides specialists with direct access to the physiological significance of its criteria, allowing for future adaptation or enrichment by sleep experts.
Signal Processing,Machine Learning
What problem does this paper attempt to address?
### Problems the paper attempts to solve The paper aims to solve the problem of limited clinical acceptance caused by the "black - box" characteristics of deep - learning models in automatic sleep staging. Although deep - learning algorithms have proven their effectiveness in automatic sleep staging, their "black - box" nature makes it difficult for doctors to understand the decision - making process of the models, thus limiting their application in clinical practice. To this end, the authors propose an interpretable neural - network model - WaveSleepNet, which can perform sleep staging in a manner similar to that of sleep experts and improve the transparency and interpretability of the model by visualizing waveform prototypes. ### Main contributions 1. **Case - based model**: WaveSleepNet is a case - based model that can automatically perform sleep scoring in a manner similar to the human cognitive process. To the best of the authors' knowledge, this is the first time that case - based learning has been applied to the sleep - staging task. The effectiveness of this model has been verified on three public datasets. 2. **Transparent reasoning process**: WaveSleepNet provides a transparent reasoning process. By calculating the distance between the input signal and the learned waveform prototypes, it determines the existence of specific waveforms and their relative proportions, thereby classifying sleep stages. Experimental case studies further confirm that the sleep - staging logic of this model is consistent with the standards of the American Academy of Sleep Medicine (AASM). 3. **Systematic analysis of misclassification reasons**: The authors also systematically analyze the misclassification errors of the model, which has not been explored in previous studies. By examining in detail the reasons for the misclassification of signals in different sleep stages, they provide directions for further improvement of the model to improve the accuracy of sleep staging. ### Method overview The architecture of WaveSleepNet mainly consists of three parts: the Feature Extraction Network, the WaveSensing Network, and the Decision Network. 1. **Feature Extraction Network**: This network extracts deep features from the original input signal and effectively captures information at different time scales and channel dimensions. The Feature Extraction Network is composed of a Multi - Resolution Convolutional Neural Network (MRCNN), Adaptive Feature Recalibration (AFR), and Stacked Convolutional Networks (stacked CNNs). 2. **WaveSensing Network**: This network uses the trained waveform prototypes to convert the features learned by the Feature Extraction Network into distance values between the features and the waveform prototypes. These distance values reflect the existence of specific waveforms in the input signal and their relative proportions. 3. **Decision Network**: This network is composed of a Waveform Estimator, a Proportion Estimator, and a Decision Block. The Waveform Estimator and the Proportion Estimator calculate the existence and proportion of waveforms respectively, and the Decision Block integrates these two scores to predict the final sleep stage. ### Loss function design In order to optimize classification accuracy and enhance interpretability, the loss function design includes the following parts: - **Classification loss**: Use the standard cross - entropy loss to punish misclassifications of sleep stages. - **Diversity loss**: Promote the diversity of waveform prototypes by minimizing the average squared L2 distance between waveform prototypes. - **Feature matching loss**: Two loss terms calculate the distance between feature patches and waveform prototypes in the latent space respectively, ensuring that each waveform prototype can correspond to an instance in the training set. - **L1 regularization**: Effectively improve the interpretability of the model by selecting important waveform prototypes for each sleep stage. ### Experimental results The authors conducted experiments on three public datasets (SleepEDF - 20, SleepEDF - 78, and SHHS) to verify the classification performance and interpretability of WaveSleepNet. The experimental results show that WaveSleepNet achieves performance comparable to that of the existing state - of - the - art models on multiple metrics, especially when multiple WaveSleepNet models are combined, its performance is even better. ### Conclusion WaveSleepNet successfully solves the "black - box" problem of deep - learning models in automatic sleep staging by introducing waveform prototypes and a transparent reasoning process.