Failure prediction of a deep learning model for PSG-based sleep stage classification problem through dropout confidence

Abstract

Deep learning models have made predicting sleep stages from polysomnography records as accurate as sleep technicians. However, these black-box classication models are seldom used in practice because the risk of having a wrong prediction is high. We aim to help clinicians understand when to trust the model's predictions and when to manually correct them by providing reliable condence values to the predictions of the sleep staging model. We use the state-of-the-art TinySleepNet architecture to train our classication model and adapt the CondNet architecture to train an auxiliary condence model that will provide condence values to the predictions of the classication model. We train the condence model with our proposed Dropout Correct Rate (DCR) loss using a raw, single EEG channel in the PSG records from the public SHHS dataset and the local clinical records from SNUBH. Removing 20% of the input data with the least condence values, rovided by our DCR-trained condence model, improves the accuracy from 77% to 84%, F1-score from 0.79 to 0.86, and Cohen's kappa from 0.72 to 0.81 of the classication model. As the rst study to introduce condence estimation in PSG-based automated sleep staging models, we demonstrate that providing condence values to the predictions of the classication model improves both the accuracy and the interpretability of the model. With our proposed method, we make automated sleep staging systems not only fast and ecient but also reliable and usable in practice.