Multi-view Spatial-Temporal Graph Convolutional Networks with Domain Generalization for Sleep Stage Classification

Sleep stage classification is crucial for sleep quality assessment and disease diagnosis. However, existing classification methods still face numerous challenges in handling the spatial and temporal features of time-varying multi-channel brain signals, coping with individual differences in biological signals, and model interpretability. Traditional machine learning methods rely on complex feature engineering, while deep learning methods, though performing well in feature representation learning, still need improvements in spatial-temporal feature utilization, cross-individual generalization ability, and model interpretability.

To address these challenges, Ziyu Jia and others from Beijing Jiaotong University and Li-Wei H. Lehman from the Massachusetts Institute of Technology proposed a Multi-View Spatial-Temporal Graph Convolutional Network (MSTGCN) combined with domain generalization for sleep stage classification.

Paper Source

This paper was co-authored by Ziyu Jia, Youfang Lin, Jing Wang (corresponding author), Xiaojun Ning, Yuanlai He, Ronghao Zhou, and Yuhan Zhou from the School of Computer and Information Technology, Beijing Jiaotong University, and Li-Wei H. Lehman from the Institute for Medical Engineering and Science, Massachusetts Institute of Technology. It was published in the 2021 issue of IEEE Transactions on Neural Systems and Rehabilitation Engineering.

Research Details

Research Process

1. Construct Brain Perspective Graph: Construct two types of brain perspective graphs based on brain region functional connectivity and physical distance proximity. Each EEG channel corresponds to a node in the graph, and specific connections between channels correspond to the graph’s edges.

2. Spatial Graph Convolution: Extract rich spatial features using graph convolution, which aggregates information from neighboring nodes to capture spatial dimension features.

3. Temporal Convolution: Use temporal convolution to capture transition rules across different sleep stages, helping to identify the current sleep stage.

4. Spatial-Temporal Attention Mechanism: Design a spatial-temporal attention mechanism to automatically capture the spatial-temporal information most relevant to sleep stage classification.

5. Domain Generalization: During model training, treat each individual as a specific source domain for sleep feature extraction. Use adversarial domain generalization methods to extract sleep features that do not vary with individuals, thereby enhancing the model’s cross-individual generalization ability.

Experiments and Results

• Datasets: The experiments used two public sleep datasets, ISRUC-S3 and MASS-SS3.

• Experimental Setup: Use 10-fold cross-validation and 31-fold cross-validation, adopting an individual-independent strategy for cross-validation. The model was implemented using TensorFlow.

The experimental results show that MSTGCN outperforms existing baseline models on multiple metrics (overall accuracy, F1 score, and Kappa value). Moreover, the experiments also verified the classification effect of the model in different sleep stages. In the ISRUC-S3 dataset, the classification accuracy was highest in the Wake and N3 stages, while in the MASS-SS3 dataset, the accuracy was highest in the REM and N2 stages. Although the N1 stage classification performance did not meet expectations, it still performed excellently among many baseline models.

Conclusion

This study proposes a Multi-View Spatial-Temporal Graph Convolutional Network (MSTGCN) combined with domain generalization for sleep stage classification. The specific contributions are as follows:

1. Constructed different brain perspectives based on functional connectivity and physical distance proximity, providing rich spatial topological information for the classification task.
2. Designed a spatial-temporal graph convolution based on an attention mechanism to capture both spatial and temporal features, improving classification performance.
3. Combined domain generalization methods with spatial-temporal graph convolutional networks to extract sleep features that do not vary with individuals, enhancing the model’s cross-individual generalization ability.
4. Conducted experiments on two public datasets, ISRUC-S3 and MASS-SS3, showing that the model achieves state-of-the-art performance.
5. Explored the interpretability of key modules in the model, especially the functional connectivity results obtained through adaptive graph learning, indicating that functional connectivity during light sleep is more complex than during deep sleep.

Compared to previous studies like graphsleepnet, MSTGCN has the following significant improvements:

1. Brain Network Construction: The brain network constructed based on physical distance proximity, combined with the previous adaptive functional connectivity brain network, forms a multi-view brain network, providing richer spatial topological information.
2. Domain Generalization: Combined domain generalization with spatial-temporal graph convolutional networks, enhancing the model’s generalization ability.
3. Comprehensive Experimental Evaluation: Evaluated the effectiveness of MSTGCN on two sleep datasets and conducted ablation experiments to verify the impact of each component on performance.
4. Module Interpretability: Explored and discussed the interpretability of key modules in MSTGCN.

Research Significance

The MSTGCN model proposed in this study not only achieves state-of-the-art levels in the sleep stage classification task but also demonstrates a general framework for spatial-temporal modeling of multivariate physiological time series data, taking into account cross-individual generalization ability and model interpretability. Compared to traditional methods, this model provides a more accurate and efficient sleep stage classification method, with significant clinical application value.