Paper: Spatial-Temporal Mutual Distillation for Lightweight Sleep Stage Classification (ICLR-2024)

Reviewer comment (Experiments):
The experiments are not sufficient enough. The authors should compare their KD method with more popular sleep models.

Author response:
We agree that TinySleepNet is more lightweight and a better choice than DeepSleepNet. DeepSleepNet and TinySleepNet are both CNN-RNN architecture. The design of our CNN-RNN multi-channel teacher-student model is inspired by both DeepSleepNet and TinySleepNet.

In fact, as shown in Appendix A.5, our CNN-RNN model is more similar with TinySleepNet but using multiple channels.

The CNN part of the teacher model:

> | Layer | Layer Type | #Filters | Size | Stride | Activation | Mode |
> | ----- | ---------- | -------- | ---- | ------ | ---------- | ---- |
> | 1     | Input      | /        | /    | /      | /          | /    |
> | 2     | Conv1D     | 128      | fs/2 | fs/4   | Relu       | Same |
> | 3     | Dropout    | /        | 0.5  | /      | /          | /    |
> | 4     | Maxpool1D  | /        | /    | 8      | /          | /    |
> | 5     | Conv1D     | 128      | 8    | 1      | Relu       | Same |
> | 6     | Conv1D     | 128      | 8    | 1      | Relu       | Same |
> | 7     | Conv1D     | 128      | 8    | 1      | Relu       | Same |
> | 8     | Maxpool1D  | /        | 4    | 4      | /          | /    |
> | 9     | Dropout    | /        | 0.5  | /      | /          | /    |

The RNN part of the teacher model:

> | Layer | Layer Type  | Size | Activation |
> | ----- | ----------- | ---- | ---------- |
> | 1     | Encoder     | /    | /          |
> | 2     | Concatenate | /    | /          |
> | 3     | BiLSTM      | 128  | /          |
> | 4     | Dropout     | 0.5  | /          |
> | 5     | Dense       | 5    | Softmax    |

The CNN part of the student model:

> | Layer | Layer Type | #Filters | Size | Stride | Activation | Mode |
> | ----- | ---------- | -------- | ---- | ------ | ---------- | ---- |
> | 1     | Input      | /        | /    | /      | /          | /    |
> | 2     | Conv1D     | 32       | fs/2 | fs/4   | Relu       | Same |
> | 3     | Dropout    | /        | 0.5  | /      | /          | /    |
> | 4     | Maxpool1D  | /        | /    | 8      | /          | /    |
> | 5     | Conv1D     | 32       | 8    | 1      | Relu       | Same |
> | 6     | Conv1D     | 32       | 8    | 1      | Relu       | Same |
> | 7     | Conv1D     | 32       | 8    | 1      | Relu       | Same |
> | 8     | Maxpool1D  | /        | 4    | 4      | /          | /    |
> | 9     | Dropout    | /        | 0.5  | /      | /          | /    |

The RNN part of the student model:

> | Layer | Layer Type  | Size | Activation |
> | ----- | ----------- | ---- | ---------- |
> | 1     | Encoder     | /    | /          |
> | 2     | Concatenate | /    | /          |
> | 3     | BiLSTM      | 32   | /          |
> | 4     | Dropout     | 0.5  | /          |
> | 5     | Dense       | 5    | Softmax    |

Furthermore, to further show the generality of our framework, we add new experimentson a CNN+GCN architecture model and add the results to section 5.3 Table 4 :

> | Method                            | Accuracy | F1-score |
> | --------------------------------- | -------- | -------- |
> | Knowledge Distillation            | 75.07%   | 72.35%   |
> | Decoupled Knowledge Distillation  | 82.44%   | 80.26%   |
> | Fitnets                           | 81.88%   | 80.76%   |
> | Neuron Selectivity Transfer       | 83.31%   | 80.94%   |
> | Relational Knowledge Distillation | 76.68%   | 73.19%   |
> | Distilling Knowledge from GCN     | 82.65%   | 79.69%   |
> | Deep Mutual Learning              | 81.27%   | 77.84%   |
> | Ours                              | 84.26%   | 81.22%   |

Reviewer comment (Experiments):
There are some sleep staging models consisting of a less number of model parameters, e.g. TinySleepNet. When compared to DeepSleepNet which is a CNN-LSTM based model, TinySleepNet is both more lightweight and shares the same structure, so why not design the teacher-student model based on TinySleepNet.

Author response:
We agree that TinySleepNet is more lightweight and a better choice than DeepSleepNet. DeepSleepNet and TinySleepNet are both CNN-RNN architecture. The design of our CNN-RNN multi-channel teacher-student model is inspired by both DeepSleepNet and TinySleepNet.

In fact, as shown in Appendix A.5, our CNN-RNN model is more similar with TinySleepNet but using multiple channels.

The CNN part of the teacher model:

> | Layer | Layer Type | #Filters | Size | Stride | Activation | Mode |
> | ----- | ---------- | -------- | ---- | ------ | ---------- | ---- |
> | 1     | Input      | /        | /    | /      | /          | /    |
> | 2     | Conv1D     | 128      | fs/2 | fs/4   | Relu       | Same |
> | 3     | Dropout    | /        | 0.5  | /      | /          | /    |
> | 4     | Maxpool1D  | /        | /    | 8      | /          | /    |
> | 5     | Conv1D     | 128      | 8    | 1      | Relu       | Same |
> | 6     | Conv1D     | 128      | 8    | 1      | Relu       | Same |
> | 7     | Conv1D     | 128      | 8    | 1      | Relu       | Same |
> | 8     | Maxpool1D  | /        | 4    | 4      | /          | /    |
> | 9     | Dropout    | /        | 0.5  | /      | /          | /    |

The RNN part of the teacher model:

> | Layer | Layer Type  | Size | Activation |
> | ----- | ----------- | ---- | ---------- |
> | 1     | Encoder     | /    | /          |
> | 2     | Concatenate | /    | /          |
> | 3     | BiLSTM      | 128  | /          |
> | 4     | Dropout     | 0.5  | /          |
> | 5     | Dense       | 5    | Softmax    |

The CNN part of the student model:

> | Layer | Layer Type | #Filters | Size | Stride | Activation | Mode |
> | ----- | ---------- | -------- | ---- | ------ | ---------- | ---- |
> | 1     | Input      | /        | /    | /      | /          | /    |
> | 2     | Conv1D     | 32       | fs/2 | fs/4   | Relu       | Same |
> | 3     | Dropout    | /        | 0.5  | /      | /          | /    |
> | 4     | Maxpool1D  | /        | /    | 8      | /          | /    |
> | 5     | Conv1D     | 32       | 8    | 1      | Relu       | Same |
> | 6     | Conv1D     | 32       | 8    | 1      | Relu       | Same |
> | 7     | Conv1D     | 32       | 8    | 1      | Relu       | Same |
> | 8     | Maxpool1D  | /        | 4    | 4      | /          | /    |
> | 9     | Dropout    | /        | 0.5  | /      | /          | /    |

The RNN part of the student model:

> | Layer | Layer Type  | Size | Activation |
> | ----- | ----------- | ---- | ---------- |
> | 1     | Encoder     | /    | /          |
> | 2     | Concatenate | /    | /          |
> | 3     | BiLSTM      | 32   | /          |
> | 4     | Dropout     | 0.5  | /          |
> | 5     | Dense       | 5    | Softmax    |

Furthermore, to further show the generality of our framework, we add new experimentson a CNN+GCN architecture model and add the results to section 5.3 Table 4 :

> | Method                            | Accuracy | F1-score |
> | --------------------------------- | -------- | -------- |
> | Knowledge Distillation            | 75.07%   | 72.35%   |
> | Decoupled Knowledge Distillation  | 82.44%   | 80.26%   |
> | Fitnets                           | 81.88%   | 80.76%   |
> | Neuron Selectivity Transfer       | 83.31%   | 80.94%   |
> | Relational Knowledge Distillation | 76.68%   | 73.19%   |
> | Distilling Knowledge from GCN     | 82.65%   | 79.69%   |
> | Deep Mutual Learning              | 81.27%   | 77.84%   |
> | Ours                              | 84.26%   | 81.22%   |

Reviewer comment (Experiments):
The main purpose of this paper is to achieve student model lightweighting through knowledge distillation. However, as shown in Tab.2, only the computational costs of teacher and student models are listed. It lacks the comparison with other classical sleep staging models.

Author response:
The method we proposed is a generic knowledge distillation framework that can be applied on most multi-channel sleep models. It is not model-specific. The teacher and student model we use are only a representitive of multi-channel sleep models and their corresponding student models. We compare computational costs of teacher and student models to show the effectiveness of the knowledge distillation framework. Since other classical sleep staging models can be compressed by our knowledge distillation framework, it is meaningless to compare the student model with other models.

Reviewer comment (Novelty):
Despite its novelty in the sleep classification domain, the technical contributions of this paper might be relatively incremental in the ML/AI community. For example, the second contribution ("how to transfer") is just a combination of [1] and [2].

Author response:
A1: We agree that our novelty is limited to sleep stage classification domain, and mutual distillation is widely studied in other area. Nevertheless, it cannot be ignored that the application of sleep models still lacks of exploration and  sleep model compression framework is needed. We root on the spatial and temporal knowledge in sleep stage classification area and adopt knowledge distillation into this area to meet the demand for compression and application in sleep stage classification. As we can see from the experiment results, knowledge distillation frameworks from other area cannot adapt well in the sleep stage classification area. We specifically design the knowledge distillation framework based on the spatial and temporal semantic prior knowledge in the sleep stage classification.

Reviewer comment (Theory):
The paper has missing discussions on why the proposed knowledge types are reasonable. This paper encodes spatial & temporal knowledge in a separate manner. However, is it not true that spatiotemporal knowledge should be jointly modeled as they are highly entangled? Given this awareness, the proposed knowledge types ("Channel-to-channel distances" and "Epoch-to-epoch similarities") might be outdated, unless the authors can provide the evidence why the separate knowledge types are better than the jointly encoded knowledge.

Author response:
A3: The proposed knowledge types are based on the knowledge of sleep stage classification. The spatial knowledge extracted from multi-channel sleep signals represents the functional relationship of human body. The temporal knowledge denotes the contextual relationship over time and the transition rules of multiple sleep epochs. These two kinds of knowledge are widely used and tested in the sleep stage classification area. 

To help better understand the spatial-temporal knowledge, we visualize the spatial and temporal connections between multi-channel sleep signals to better demonstrate the transfer of spatial and temporal knowledge. It is shown in the figure.3 and 4 in the new manuscript we upload.

Reviewer comment (Theory):
In Eq.(11-12), shouldn't the Loss terms be gradient?

Author response:
A4: It is true that spatial and temporal knowledge is highly entangled. We focus on the explicit extraction and transfer of spatial-temporal knowledge and visualize them to present a more pratical implications, making it easier to understand and explain in the clinical scenerios. As we can see in the figure 3 and 4 in the new manuscript, the spatial and temporal knowledge are visualized as the connections, which is more easier to understand. Despite that, it is free to classify the sleep stage after the fusion of spatial-temporal knowledge in the model we choose to compress.

Reviewer comment (Presentation):
Eq.(13-14) never appear in objectives. Where should they belong to?

Author response:
A5: Thank you for your correction. We have corrected them in the new version of manuscript:

> On the training epoch i, the update of both teacher and student model can be expressed as follows:
> $$
> L^t_{c} = CE(f^t_i(x), y) = -\sum_{i} y \cdot \log(f^t_i(x))
> $$
>
> $$
> L^s_{c} = CE(f^s_i(x), y) = -\sum_{i} y \cdot \log(f^s_i(x))
> $$
>
> $$
> Loss^t_i = \alpha L^t_{c} + \beta L_{spatial} + \gamma L_{temporal}
> $$
>
> $$
> Loss^s_i = \alpha L^s_{c} + \beta L_{spatial} + \gamma L_{temporal}
> $$