Teacher Assistant-Based Knowledge Distillation Extracting Multi-level Features on Single Channel Sleep EEG

1 Introduction

In recent years, sleep disorders are becoming a worrying problem that affects human health. Sleep stage classification is helpful for the diagnosis of sleep disorders. The experts complete the analysis of sleep quality by inferring each sleep stage with the signals from sensors on different parts of the body. Specifically, the signals include Electroencephalogram(EEG), Electromyography(EMG), Electrooculography(EOG), etc. Then, these signals are segmented into 30-second data samples which are called sleep epoch for the sleep stage classification. Finally, experts classify each sleep epoch into a specific stage according to the criteria of the American Academy of Sleep Medicine (AASM) [Berry et al., 2012] or other sleep manuals such as the Rechtschaffen and Kales(R&K) [Rechtschaffen, 1968]. Therefore, manual stage classification is a very time-consuming task.

∗Ziyu Jia is the corresponding author.

To automate the sleep stage classification, some deep learning methods are applied [Sekkal et al., 2022; Jia et al., 2022a; Liu and Jia, 2023]. For example, DeepSleepNet [Supratak et al., 2017] and SailentSleepNet [Jia et al., 2021b] are used to automatically extract multi-level features from sleep signals. Specifically, there are two kinds of important features in the sleep signals, which are epoch-level features and sequencelevel features. The epoch-level features represent the local characteristics of a single sleep epoch. For example, the N2 stage includes mainly sleep spindles and K complexes. The sequence level features are the transition rules between multiple sleep epochs. For instance, the N1 stage often serves as a transition stage between the W stage and other stages. To capture these features, the intermediate layers of existing sleep models are usually relatively large in size (wider and deeper). To the best of our knowledge, some sleep stage classification models usually have parameters up to the order of 100k or even 1M. The deployment of models may cost a large amount of computing resources.

In order to lighten the large-scale model, some knowledge distillation methods are applied [Wang and Yoon, 2021]. The teacher model (complex and large-scale model) can transfer knowledge to the student model (lightweight model) with the knowledge distillation framework. However, the performance of most knowledge distillation methods directly applied to sleep stage classification is unsatisfactory. These methods ignore the valuable information in the multiple levels of sleep epoch features and sleep sequence features shown in Figure 1. Therefore, it is a challenge to design a distillation framework that transfers multi-level sleep knowledge.

Another challenge is how to bridge the gap between teacher and student network with the smallest loss of knowledge. Specifically, in most cases, the teacher network is deep while the student network is shallow as shown in Figure 2. In some circumstances where teacher and student network have too much difference, knowledge may be transferred inefficiently [Mirzadeh et al., 2020]. Moreover, intermediate features extracted by the sleep stage classification model is relatively more complex and can not be fully conveyed to student because of the dimension alignment [Aguilar et al., 2020].

To address the above challenges, we propose a general knowledge distillation framework called SleepKD to lighten the EEG-based sleep stage classification model as shown in Figure 3. This framework consists of the multi-level distillation module, the teacher assistant module, and the other knowledge distillation module. The main contributions are summarized as follows:

• To the best of our knowledge, it is the first attempt to use multi-level knowledge distillation in the sleep stage classification model using EEG. Our proposed multilevel knowledge distillation module can efficiently boost epoch-level and sequence-level knowledge transfer from the teacher network to the student.
• We design the teacher assistant module for different kinds of sleep stage classification models. It bridges the excessive gap between teacher and student network for the gap-sensitive multi-level knowledge transfer.
• Experimental results show that our knowledge distillation framework achieves SOTA performance compared to existing knowledge distillation methods. In addition, the proposed knowledge distillation framework can significantly lighten the sleep model while maintaining its classification performance.

2 RelatedWorks

2.1 Sleep Stage Classification

Sleep stage classification is widely used to diagnose diseases such as sleep disorders. In early studies, machine learning methods are utilized to classify the sleep stages [Tzimourta et al., 2018; Basha et al., 2021; Sundararajan et al., 2021; Jia et al., 2022b]. However, these methods need a large amount of prior knowledge, which means they require a lot of manual costs to extract features. Therefore, many researchers start to implement automatic sleep stage classification by using deep learning methods.

Currently, there are two typical deep learning architectures that are widely used for sleep stage classification, CNNbased [Zhang and Wu, 2017; Cui et al., 2018; Phan et al., 2019] and a hybrid of CNN and RNN [Yang et al., 2018; Back et al., 2019; Fan et al., 2021]. CNN-based architectures are widely applied to sleep stage classification models. CNN is used to extract sequence level information in sleep signals [Chambon et al., 2018]; SleepUtime [Perslev et al., 2019] is proposed with a fully feed-forward deep learning method based on physiological time series segmentation by using the U-Time module; SalientSleepNet [Jia et al., 2021b] is devised with a U 2 -structure stream by nesting multiple U-units and capturing more useful information for sleep stage classification. It extracts multi-level features with Multi-Scale Extractor (MSE) for sleep epoch and sequence features.

Also, researchers propose a series of sleep stage classification models based on the hybrid architecture of CNN and RNN. DeepSleepNet [Supratak et al., 2017] is applied to extract sleep epochs features and sequence features by using CNN and Bi-directional Long Short-Term Memory (BiLSTM); A hierarchical neural network is designed to learn both comprehensive features and sequential features for sleep stage classification [Sun et al., 2019]; SleepEEGNet [Mousavi et al., 2019] is devised to extract time-invariant features from the original signal, and to capture both long-term and short-term context dependencies by using Bidirectional Recurrent Neural Network (BRNN). In addition, there are other architectures that can be used for sleep stage classification. For example, GraphSleepNet [Jia et al., 2020] is composed of a deep graph neural network for sleep stage classification. MSTGCN [Jia et al., 2021a] is constructed to extract the sleep features by using ST-GCN. An improved real-time sleep stage estimation is devised to have a better sleep stage classification [Harada et al., 2017]. An evolutionary algorithm is used to complete the age-based sleep stage estimation [Matsushima et al., 2012]. Relative evaluation is employed to score the sleep stage by heart rate [Tobaru et al., 2019].

Although these approaches yield good results in the field of sleep stage classification, the parameters of the network model grow. This leads to high computational and storage costs for the models at the industrial level, making deployment difficult to be achieved. Therefore, lightweight sleep stage classification models are particularly important. There have been a few related studies such as [Joshi et al., 2021], which ignore the intermediate features of the sleep signals or the large gap between teacher and student.

2.2 Knowledge Distillation

Knowledge distillation can lighten large-scale models. Most of the knowledge distillation approaches can be classified into two types: distillation from logits [Furlanello et al., 2018; Cho and Hariharan, 2019; Tian et al., 2019; Zhao et al., 2022]

and distillation from intermediate layers [Yim et al., 2017; Kim et al., 2018; Heo et al., 2019; Chen et al., 2021]. Specifically, the distillation from logits focuses on the logits distribution. The concept of Knowledge Distillation is proposed [Hinton et al., 2015], which aims to transfer knowledge through the difference of the logits distribution of the classification from teacher and student network.

In addition, there are some studies focusing on distilling intermediate features of teacher network. The features of intermediate layers are used to guide the student model during distillation [Romero et al., 2014]. It uses a wide and shallow teacher model to train a narrow and deep student model from intermediate layers with the mean squared error.

Most of the studies are based on the above types of methods. However, there are a small number of special distillation approaches [Zhang et al., 2018; Mirzadeh et al., 2020; Xu et al., 2020; Son et al., 2021]. Actually, different types of distillation methods are complementary. Therefore, the proposed distillation framework combines two major types of distillation methods. We train the student model with knowledge not only from single and multiple epochs of sleep signals but also from the true label and the logits distribution of the teacher network. To improve the performance of distillation even further, we also introduce the teacher assistant to the distillation to help bridge the gap between the teacher and student network.

3 Preliminary

The input of the model is a sequence of sleep epochs, and the output is a sequence of corresponding labels. Each sleep epoch is defined as x ∈ Rn , where n denotes the number of samples in a sleep epoch. The input sequence of sleep epochs is defined as S = {x1, x2, · · ·, xL}, where xi (i ∈ [1, 2, · · ·, L]) denotes a sleep epoch and L is the number of sleep epochs from the input sequence.

We distill different sleep stage classification models based on the proposed distillation framework. Then, we evaluate them based on the final classification performance and model compression ratio. We define the predicted output of the model as ˆ Y = { ˆ y1, ˆ y2, · · ·, ˆ yL}, where ˆ yi ∈ {0, 1, 2, 3, 4} denotes the classification result of xi, corresponding to the five sleep stages W, N1, N2, N3, and REM in the AASM manual, respectively.

4 The Proposed SleepKD

Figure 3 presents our SleepKD framework for sleep stage classification. It can be summarized into three key points. 1) We develop a novel multi-level knowledge distillation that simultaneously conveys epoch-level knowledge and sequencelevel knowledge in sleep EEG signals. 2) We design the teacher assistant module to bridge the excessive gap between the teacher and student network for the gap-sensitive multilevel knowledge transfer, which makes the distillation more effective. 3) We employ the Kullback-Leibler divergence between the output of teacher and student to transfer the knowledge from logits distribution to the student network.

4.1 Multi-Level Module

There are two kinds of important features in the EEG, which are epoch-level features and sequence-level features. They represent the local characteristics of a single sleep epoch and the transition rules between multiple sleep epochs, respectively. In order to capture these two types of features, the intermediate layers of existing models are usually designed to be large. In this paper, we use the knowledge distillation technique to transfer the knowledge extracted from the teacher network to the student network.

Because epoch-level and sequence-level knowledge are extracted from the intermediate layers, we distill these two kinds of knowledge in the intermediate layers of the network. Figure 4 shows the epoch-level knowledge distillation. Specifically, we minimize the difference between teacher's and student's epoch features at the epoch level. This enables the student to learn epoch features from the teacher network. Thus, the student can better capture the features of each single sleep epoch and improve classification with them. The loss at the epoch level is defined as follows:

where F Te denotes the epoch features of the teacher and F Se denotes the epoch features of the student. Because of the dimension difference between the teacher and student features, we use an alignment function Φ. It can be max-pooling or 1 × 1 convolution. LMSE denotes the loss function calculated by mean square error.

In addition, as shown in Figure 4, we calculate the difference between the teacher's and student's sequence features at the sequence level. By minimizing the difference, the student is allowed to learn the sequence-level features from the teacher. The student is able to learn sleep transition rules, further improving classification performance. The loss at the sequence level is defined as follows:

where F Ts denotes the sequence features of the teacher. F Ss denotes the sequence features of the student. Φ is the alignment function mentioned above. To measure the difference, we choose mean square error as the loss function, which is denoted as LMSE.

4.2 Teacher Assistant Module

Existing studies have shown that knowledge transfer is hindered when teacher and student network are too much different. In the multi-level module, we introduce multi-level knowledge from intermediate layers in the network. However, the dimensions of intermediate features are mismatched because teacher and student have different architectures. To transfer the knowledge, dimension alignment between multilevel features is necessary. During the dimension alignment, the complex knowledge in these features could be lost, which makes it more sensitive to the gap. To smooth the transfer of multi-level knowledge, we design a teacher assistant module between teacher and student network to bridge the gap. It enhances the transfer of multi-level knowledge and makes distillation more effective for sleep stage classification.

We design the teacher assistant module for two kinds of typical sleep stage classification architectures. One is CNNbased architectures and the other is hybrid architectures based on CNN and RNN. For the model using CNNs to extract epoch and sequence features of EEG signals, we design a medium-sized teacher assistant model by reducing the convolution layers of the teacher assistant model between teacher and student network. For instance, the number of CNN layers of the teacher model is 6, and the number of CNN layers of the student model is 2. We set the number of CNN layers of the teacher assistant model to 4. The details of the design are shown in Figure 5.

For the hybrid architecture with CNN and RNN, we design a teacher assistant module with a scale between the teacher model and the student model. For the CNN layers of the teacher assistant module, we reduce it as CNN-based architectures above, which means the CNN depth of the teacher assistant model is between the teacher and student network. Meanwhile, we design a medium-sized RNN component to extract the sequence-level features. For example, the number of units of the RNN layer in the teacher network is 512, and the student network is 128. We set the number of units of the RNN layer in the teacher assistant network to 256. We design the teacher assistant module through such a strategy, and the details are shown in Figure 6.

4.3 Other Knowledge DistillationModule

The soft labels, which are the probability distribution for each stage from the teacher model output, also contain useful knowledge. Therefore, we introduce Kullback-Leibler divergence to compute Lsoft between the teacher and student network. This allows the teacher network to transfer knowledge from its logits distribution to the student network. Lsoft is

where DKL denotes the Kullback-Leibler divergence, which is used to calculate the relative entropy of the output distribution between the teacher model and the student model. pT denotes the output of the teacher model and pS denotes the output of the student model. Moreover, we calculate Lhard using the cross-entropy loss function to obtain knowledge of hard labels. Lhard is defined as follows:

where LCE denotes the cross-entropy loss function and y denotes the true label. Finally, the total loss LTotal is defined as follows:

where α, β, γ, δ denote the weights of Lepoch, Lseq, Lsoft, Lhard, respectively.

5 Experiments

5.1 Datasets and Data Processing

We evaluate our method on two public datasets: ISRUC-III [Khalighi et al., 2016] and Sleep-EDF [Kemp et al., 2018].

ISRUC-III collects the PSG data samples from 10 subjects (1 for males and 9 for females) for a whole night in 8 hours. The annotations of this dataset are scored by two professional experts.

Sleep-EDF is a very famous public dataset that contains the PSG data samples from 20 subjects (10 for males and 10 for females) in 2 days. The ages of the subjects range from 25 to 34 years old. These recordings were manually classified into one of the eight classes (W, N1, N2, N3, N4, REM, Movement, Unknown) by sleep experts according to the R&K

standard. For a fair comparison, we remove the Movement and Unknown stage, and merge the N3 and N4 stage into a single N3 stage according to the AASM manual.

In the experimental data, the EEG is typically segmented into sleep epochs of 30 seconds. We finally extract the sleep epochs of the EEG signal from the Fpz-Cz channel in the Sleep-EDF and the sleep epochs of the EEG signal from F3A2 in ISRUC-III. The EEG data from each dataset is downsampled to 100Hz.

5.2 Baseline Methods

We select some classical and well-performing knowledge distillation methods as baseline methods from three aspects: distillation from logits, distillation from intermediate features, and distillation with teacher assistants.

• KD [Hinton et al., 2015]: Propose a simple way to improve the performance by distilling the knowledge of the complex model into a compact model with the output of the former.
• Fitnets [Romero et al., 2014]: Extend the idea of the traditional knowledge distillation by using both the output of the teacher network and the intermediate representation as a hint to the student.
• NST [Huang and Wang, 2017]: Implement a knowledge transfer loss function by minimizing the Maximum Mean Discrepancy between the feature map of the sophisticated model and the slimming model.
• TAKD [Mirzadeh et al., 2020]: Introduce a multi-step knowledge distillation by using teacher assistant (TA) whose size is between the teacher and student model.
• DGKD [Son et al., 2021]: Devise the densely-guided knowledge method using multiple teacher assistant to fill the large gap between teacher and student model gradually.
• DKD [Zhao et al., 2022]: Reformulate the classical KD method with non-target class knowledge distillation (NCKD) and target class knowledge distillation (TCKD).

5.3 Experiment Settings and Implementation

We split the datasets into the train, validation, and test sets by a ratio of 8:1:1 on Sleep-EDF and ISRUC-III separately. The input sleep sequence is 20-epoch long. Each epoch lasts 30 seconds. We implement the teacher, teacher assistant, and student models with TensorFlow. We use Adam as the optimizer in each experiment. In experiments of the CNN framework, we choose SalientSleepNet as a representative. The learning rate of SalientSleepNet is 0.001. The number of training epochs is 60 and the batch size is 8. The weights are α = 0.3, β = 0.2, γ = 0.4 and δ = 0.1. In experiments of the CNN and RNN framework, we choose DeepSleepNet as a representative. DeepSleepNet has a learning rate of 0.00001. The number of training epochs is 200 for SleepEDF, 300 for ISRUC-III and the batch size is 20. The weights are α = 1.0, β = 0.1 and γ = δ = 1.0. We set these hyperparameters according to the performance on validation sets.

We design corresponding TA and student network for SalientSleepNet and DeepSleepNet. As for SalientSleepNet, we take the number of U-units and Multi-Scale Extractors (MSE) as the proxy of the model capacity since they are designed to extract epoch-level and sequence-level knowledge, respectively. In the implementation of DeepSleepNet, we consider the number of convolution layers in each stream's epoch extractor and units in BiLSTM, which is employed to capture the sequence features from a series of epoch features, as the complexity of our models. As a result, we decide to reduce the number of these components above to design corresponding teacher assistants and students. The specific numbers of layers are presented in Table 1.

In order to evaluate the classification performance of the model and compare it with the other baseline models, we employ Accuracy and F1-Score as the evaluation metrics, which are defined as follows:

where TP is a true positive, TN is a true negative, FP is a false positive, and FN is a false negative.

Besides, we also use compression ratio to evaluate our method, which is calculated as follows:

where PTeacher and PStudent represent the number of parameters in teacher and student network, respectively.

5.4 Experiment Results

As shown in Table 2 and Table 3, we perform several experiments with SleepKD and baseline methods on SalientSleepNet and DeepSleepNet, which are classic models with a CNN framework and a hybrid framework based on CNN and RNN, respectively. Our SleepKD achieves the SOTA knowledge distillation results.

As for KD and DKD, they only focus on the knowledge from logits distribution. These kinds of approaches introduce the output of the teacher network as the extra label to help the student to reach better performance. By contrast, there is a limitation of information transfer in these kinds of methods. Hence, their classification performance is relatively lower than the other approaches. Besides, Fitnets and NST concentrate on the knowledge in the feature map of intermediate layers. However, the knowledge from intermediate layers is gap-sensitive. When facing a student network with a high compression ratio, these methods are restricted by the significant gap between the teacher and student network. As a result, information in the teacher network may not be conveyed efficiently. In addition, the TAKD and DGKD present a great performance because they realize that the gap between teacher and student can be bridged by teacher assistant. However, these types of knowledge distillation ignore the multi-level features from intermediate layers. Student can not learn to extract information in epochs and sequences from the teacher which limits their performance. Because of the consideration of the knowledge at multiple levels (epoch knowledge and sequence knowledge) and the huge gap between the teacher and student network, our SleepKD achieves the best performance. Take the performance on Sleep-EDF as an example: the accuracy of SleepKD reaches up to 87.05% and 85.66% on SalientSleepNet and DeepSleepNet.

Furthermore, we evaluate SleepKD in different aspects (which include parameters, compression ratio, acceleration, etc.). Table 4 and Table 5 present that the student models achieve 74.68% and 71.78% on the compression ratio while the reduction of the accuracy are less than 1%. These data reveal that our framework is able to compress the model the most with the least cost of accuracy. Therefore, the performance in different aspects gets a significant improvement.

5.5 Ablation Experiments

To evaluate the effectiveness of each module, the ablation experiments are designed. With the same experiment settings, we select different combinations of loss terms to verify the effectiveness of each loss term in SleepKD loss function LT otal .

• L1 =LTotal -Lseq
• L2 =LTotal -Lepoch
• L3 =LTotal -Lsoft
• L4 =LTotal -Lhard

Figure 7 demonstrates that each loss term in SleepKD is useful and effective. These modules transfer valuable information. Multi-level information of EEG signals transferred by Lepoch and Lseq significantly improve student performance. Students learn the teacher's probability distribution by Lsoft and knowledge from labels by Lhard.

To demonstrate the effectiveness of the teacher assistant module, the following distillation methods are designed:

• varient a): Multi-Level Module
• varient b): Multi-Level Module + TAModule

Figure 8 illustrates that the teacher assistant module improves the efficiency of distillation knowledge transfer and enhances the performance of the student model by about 1%. The teacher assistant module protects the gap-sensitive multi-level knowledge and helps transfer logits knowledge by smoothing the gap between the teacher and student network.

6 Conclusion

In this paper, we propose a knowledge distillation framework for the sleep stage classification model. We employ knowledge distillation on the multi-level sleep stage classification model for the first time and introduce the teacher assistant module to improve the distillation. The proposed SleepKD framework can adapt well to the current mainstream multilevel sleep stage classification model. It is able to transfer the features of single sleep stages and transition rules between multiple sleep stages in sleep signals. Meanwhile, we design corresponding teacher assistant modules for different architectures. This can bridge the excessive gap between teacher and student network and further enhance knowledge distillation. Experiments show that our distillation framework achieves excellent results on two popular architectures (CNNbased and hybrid of CNN and RNN). Moreover, SleepKD achieves state-of-the-art distillation performance compared to other distillation methods. The proposed method is a general distillation framework for time series classification. In the future, we can apply the proposed method to other large-scale time series models.

Acknowledgments

This work was supported by STI2030-Major Projects 2021ZD0200200 and China Postdoctoral Science Foundation 2023M733738.

Contribution Statement

Heng Liang, Yucheng Liu, and Haichao Wang have equal contributions to this paper.