Mutual Distillation Extracting Spatial-temporal Knowledge for Lightweight Multi-channel Sleep Stage Classification

Sleep stage classification has important clinical significance for the diagnosis of sleep-related diseases. To pursue more accurate sleep stage classification, multi-channel sleep signals are widely used due to the rich spatial-temporal information contained. However, it leads to a great increment in the size and computational costs, which constrain the application of multi-channel sleep models on hardware devices. Knowledge distillation is an effective way to compress models, yet existing knowledge distillation methods cannot fully extract and transfer the spatial-temporal knowledge in the multi-channel sleep signals. To solve the problem, we propose a general knowledge distillation framework for multi-channel sleep stage classification called spatial-temporal mutual distillation. Based on the spatial relationship of human body and the temporal transition rules of sleep signals, the spatial and temporal modules are designed to extract the spatial-temporal knowledge, thus help the lightweight student model learn the rich spatial-temporal knowledge from large-scale teacher model. The mutual distillation framework transfers the spatial-temporal knowledge mutually. Teacher model and student model can learn from each other, further improving the student model. The results on the ISRUC-III and MASS-SS3 datasets show that our proposed framework compresses the sleep models effectively with minimal performance loss and achieves the state-of-the-art performance compared to the baseline methods.