DistillSleepNet: Heterogeneous Multi-Level Knowledge Distillation via Teacher Assistant for Sleep Staging

Abstract

Accurate sleep staging is crucial for the diagnosis of diseases such as sleep disorders. Existing sleep staging models with excellent performance are usually large and require a lot of computational resources, limiting their application on wearable devices. Therefore, it is a key issue to distil the knowledge embedded in large models into small heterogeneous models for better deployment. In the process of knowledge distillation of heterogeneous models for sleep electroencephalography (EEG) signals, we mainly deal with three major challenges: 1) There are large structural differences between heterogeneous sleep staging models; 2) What kind of knowledge should be conveyed in sleep EEG signals in the knowledge distillation of heterogeneous models; 3) Significant scale differences exist between heterogeneous models. To address these challenges, we design a generic heterogeneous model knowledge distillation framework for sleep staging. Specifically, we first propose a knowledge distillation strategy for heterogeneous models that addresses the large structural differences between heterogeneous models. Then, a multi-level knowledge distillation module is designed to effectively transfer important multi-level feature knowledge. In addition, the teacher assistant module is introduced to ease the scale difference between the heterogeneous models which further enhances the knowledge distillation performance. Experimental results on both Sleep-EDF and ISRUC datasets show that our distillation framework achieves state-of-the-art performance.