Source-linked AI summary
Multi-View Spatial-Temporal Graph Convolutional Networks with Domain Generalization for Sleep Stage Classification
Ziyu Jia, Youfang Lin, Jing Wang, Xiaojun Ning, Yuanlai He, Ronghao Zhou, Yuhan Zhou, Li-wei H. Lehman
TL;DR
Sleep-stage classification still needs better use of spatial-temporal brain structure, cross-subject generalization, and interpretability. MSTGCN combines complementary functional-connectivity and physical-distance graphs with attention-based spatial-temporal convolutions and domain generalization. Experiments on two public datasets report performance exceeding state-of-the-art baselines.
Problem
Existing methods incompletely model spatial-temporal brain topology, subject variability, and brain-level interpretability in sleep-stage classification.
Method
MSTGCN combines functional-connectivity and physical-distance brain graphs, spatial-temporal convolutions with attention, and adversarial domain generalization.
Results
Experiments on ISRUC-S3 and MASS-SS3 report that MSTGCN outperforms baseline methods on multiple overall metrics and achieves state-of-the-art performance.
Takeaways & Limitations
The framework extracts subject-invariant sleep features while providing complementary spatial topology information and interpretable functional-connectivity analysis.
Abstract
from arXiv · showhide
Sleep stage classification is essential for sleep assessment and disease diagnosis. Although previous attempts to classify sleep stages have achieved high classification performance, several challenges remain open: 1) How to effectively utilize time-varying spatial and temporal features from multi-channel brain signals remains challenging. Prior works have not been able to fully utilize the spatial topological information among brain regions. 2) Due to the many differences found in individual biological signals, how to overcome the differences of subjects and improve the generalization of deep neural networks is important. 3) Most deep learning methods ignore the interpretability of the model to the brain. To address the above challenges, we propose a multi-view spatial-temporal graph convolutional networks (MSTGCN) with domain generalization for sleep stage classification. Specifically, we construct two brain view graphs for MSTGCN based on the functional connectivity and physical distance proximity of the brain regions. The MSTGCN consists of graph convolutions for extracting spatial features and temporal convolutions for capturing the transition rules among sleep stages. In addition, attention mechanism is employed for capturing the most relevant spatial-temporal information for sleep stage classification. Finally, domain generalization and MSTGCN are integrated into a unified framework to extract subject-invariant sleep features. Experiments on two public datasets demonstrate that the proposed model outperforms the state-of-the-art baselines.
I. INTRODUCTION
Sleep stage classification remains difficult because existing methods incompletely model brain spatial-temporal structure, subject variability, and brain-level interpretability. MSTGCN addresses these challenges with multi-view graph modeling, attention, domain generalization, and interpretability analysis.
- Challenges: Existing sleep-stage methods do not fully capture spatial-temporal features, brain-region topology, subject variability, or model interpretability.These gaps persist despite high reported classification accuracy.
- Proposed approach: MSTGCN constructs complementary brain views from functional connectivity and physical distance proximity to represent spatial topology.Each EEG channel is a graph node, with channel relationships represented as edges.
- Proposed approach: Spatial graph convolution extracts spatial features, temporal convolution captures sleep-stage transitions, and attention selects relevant spatial-temporal information.The model uses neighboring sleep-stage labels as part of the temporal context for classifying the current state.
- Proposed approach: Domain generalization is integrated with MSTGCN to extract subject-invariant sleep features without requiring information from the target domain.Each training subject is treated as a source domain.
- Results and analysis: Experiments on ISRUC-S3 and MASS-SS3 report state-of-the-art performance, while interpretability analysis examines learned functional connectivity.The reported analysis indicates more complex functional connectivity during light sleep than deep sleep.
II. RELATED WORK
Related work covers CNN/RNN, graph-based, and transfer-learning approaches for sleep staging, while the paper defines graph-based sleep representations and the classification input. The proposed setup combines functional and spatial-distance brain networks with epoch-level node features and temporal context.
- II. RELATED WORK: CNN and RNN methods extract spatial or temporal features but require grid-like inputs that ignore non-Euclidean brain-region connections.GraphSleepNet addresses this limitation using functional-connectivity brain networks, but does not model physical proximity.
- II. RELATED WORK: Transfer-learning methods address subject differences but commonly require collecting new-subject data and fine-tuning before evaluation.This can be inefficient for clinical systems involving unknown subjects.
- III. PRELIMINARIES: Polysomnography records physiological signals during sleep, which are segmented into 30-second epochs for sleep-stage classification.Sleep experts assign stages according to the AASM sleep-staging standard.
- III. PRELIMINARIES: A sleep brain network is represented as G = (V, E, A), where vertices are electrodes, edges indicate connections, and A is the adjacency matrix.The paper uses functional-connectivity and spatial-distance graphs as distinct network views.
- III. PRELIMINARIES: The model extracts node features from each sleep epoch and uses a temporal context of neighboring epochs to identify the current sleep stage.Functional-connectivity and spatial-distance graphs are constructed for each epoch and supplied to the model.
IV. MULTI-VIEW SPATIAL-TEMPORAL GCN
MSTGCN models sleep-stage brain signals using complementary functional-connectivity and spatial relationships, with adaptive graph learning and graph-based spatial feature extraction.
- IV. MULTI-VIEW SPATIAL-TEMPORAL GCN: MSTGCN constructs multiple brain views to represent complementary spatial relationships for sleep-stage classification.The views include functional connectivity and spatial distance-based brain graphs.
- 1) Functional Connectivity-based Brain Graph:: Adaptive functional-connectivity learning builds graph edges from feature relationships between pairs of brain nodes.The mechanism uses node features and learnable weights to generate the adjacency structure.
- 1) Functional Connectivity-based Brain Graph:: More distant node features receive smaller learned functional-connectivity weights, encouraging connections between more similar nodes.The graph-learning loss also includes a sparsity term with λ = 0.001.
- 1) Functional Connectivity-based Brain Graph:: The learned functional-connectivity graph automatically constructs node neighborhoods and is used as a regularized component to avoid a trivial zero-weight solution.The graph-learning objective is incorporated into the broader model loss.
2) Spatial Distance-based Brain Graph:
The model adds a physical-distance brain graph and uses spatial-temporal attention to emphasize dynamically informative brain regions and network information.
- 2) Spatial Distance-based Brain Graph:: The spatial distance-based brain graph represents the influence of physically adjacent brain regions, with closer regions assigned greater influence.This graph complements functional connectivity because physical proximity can capture interactions not reflected by functional adjacency.
- 1) Spatial Attention:: The spatial-temporal attention mechanism captures valuable information through separate spatial and temporal attention components.Different brain regions can have changing effects on sleep-stage classification during sleep.
- 1) Spatial Attention:: Spatial attention dynamically computes correlations between brain nodes from the current layer’s input.The attention matrix is normalized with softmax and can adjust node updates during graph convolution.
2) Temporal Attention:
Temporal attention models changing correlations among neighboring sleep brain networks so the MST-GCN can emphasize informative temporal information.
- 2) Temporal Attention:: Temporal attention captures dynamic correlations among neighboring sleep stages and sleep brain networks.The method accounts for temporal correlations that vary across situations.
- 2) Temporal Attention:: The temporal attention matrix Q represents correlation strengths between sleep brain networks and is normalized with softmax.Q is computed from the preceding layer’s input using learnable parameters.
- 2) Temporal Attention:: The input features are multiplied by the normalized attention matrix to emphasize informative temporal information.This tuning produces the temporal-attended input for the MST-GCN.
- Spatial-Temporal Graph Convolution: Spatial-temporal graph convolution combines graph convolution for neighboring-node information with temporal convolution for dependencies among neighboring sleep stages.Chebyshev graph convolution aggregates information from 0 to K−1 order neighbors and reduces computational complexity through polynomial expansion.
2) Temporal Convolution:
The model applies temporal convolution after spatial feature extraction, fuses the two brain-view representations, and uses adversarial domain generalization to learn subject-invariant features.
- 2) Temporal Convolution:: Temporal convolution extracts context for the current sleep stage after graph convolution has extracted spatial features.A standard 2D convolution layer is used for temporal-context extraction.
- 2) Temporal Convolution:: Features from the functional-connectivity and spatial-distance views are concatenated for feature fusion.The two inputs represent features extracted from their respective brain graphs.
- D. Domain Generalization: Adversarial domain generalization trains the model to make source subjects indistinguishable while retaining sleep-stage classification performance.The feature extractor seeks common subject-invariant features rather than subject-specific features.
- D. Domain Generalization: The framework uses a feature extractor, label predictor, domain classifier, and cross-entropy losses for sleep-stage and domain classification.The transferred features are sent to both the label predictor and domain classifier.
- D. Domain Generalization: A Gradient Reversal Layer integrates feature learning and domain generalization in one backpropagation framework instead of separate classifier and discriminator training.The feature extractor minimizes label loss while maximizing domain-classification loss, whereas the domain classifier minimizes its own loss.
A. Dataset and Experiment Settings
Experiments use subject-independent cross-validation on ISRUC-S3 and MASS-SS3 to compare MSTGCN with baselines and assess its modules. MSTGCN achieves strong overall and stage-specific performance, while ablations support attention, multi-view fusion, and domain generalization.
- Experiment settings: ISRUC-S3 and MASS-SS3 are evaluated with subject-independent 10-fold and 31-fold cross-validation, respectively.The experiments compare MSTGCN with seven baselines under the same settings.
- Comparison with baselines: MSTGCN outperforms baseline methods on overall Accuracy, F1-score, and Kappa across both datasets.The comparison includes traditional machine learning and deep learning models.
- Stage-specific performance: MSTGCN accurately identifies most sleep stages, with highest accuracy for Wake and N3 on ISRUC-S3 and REM and N2 on MASS-SS3.N1 remains difficult on both datasets, although its MASS-SS3 F1-score is 4% higher than the sub-optimal result.
- Ablation experiment: The ablation study progressively adds temporal convolution, attention, a functional-connectivity view, and domain generalization to a spatial-distance graph-convolution base model.The variants isolate the contributions of the model’s main modules.
- Ablation experiment: Attention, multi-view fusion, and domain generalization are reported to improve valuable feature capture, complementary information, and subject-invariant feature learning.Figure 6 identifies the key modules as effective, especially the attention variant.
2) Adaptive Functional Connectivity Graph:
The study evaluates adaptive functional-connectivity graph learning against fixed adjacency constructions and visualizes the learned graphs across sleep stages. Adaptive learning performs best, while connectivity patterns vary by stage, with more connections in Wake and N1 than in N3.
- Graph comparison: Adaptive functional-connectivity graph learning is compared with fully connected, KNN, PCC, PLV, and MI adjacency matrices.The fixed alternatives include common graph constructions and neuroscience-based connectivity measures.
- Graph comparison: The adaptive learned adjacency matrix achieves the highest sleep-stage classification accuracy among the compared adjacency matrices.PCC, PLV, and MI achieve suboptimal effects, while the fully connected matrix performs poorly.
- Interpretability: The learned adjacency matrices visualize stage-specific functional connectivity for Wake, N1, N2, N3, and REM sleep.The visualization is intended to expose how functional connectivity changes across stages.
- Interpretability: Wake and N1 show more functional connectivity, whereas N3 shows the least connectivity.The reported pattern is consistent with existing neuroscience research.
- Temporal attention: Temporal attention assigns the greatest weight to the current sleep stage and some weight to adjacent stages.This visualization links the attention pattern to sleep-stage classification.
3) Attention Mechanism:
The attention mechanism identifies which sleep epochs and EEG channels contribute most to classification. It emphasizes the current stage and consistently highlights C3 and C4 across sleep stages.
- Temporal attention: Temporal attention assigns the highest weight to the current sleep stage T, with adjacent epochs receiving similar but lower attention.This pattern is consistent with the AASM standard and how sleep experts judge the current stage.
- Spatial attention: Spatial attention varies across sleep stages because their EEG patterns differ.F3 and F4 receive the lowest attention weights, whereas C3 and C4 receive the highest.
- Spatial attention: C3 and C4 may be the most informative EEG channels for sleep stage classification.They are located near the middle of the scalp and may contain richer EEG information while being less affected by external factors.
- Model role: MSTGCN combines attention with spatial-temporal graph convolution to capture relevant information for sleep stage classification.The framework models spatial and temporal dimensions while incorporating subject differences in sleep data.
S.1 The detailed construction process of two brain view graphs
The model represents EEG channels as graph nodes and brain-region connections as edges, using functional connectivity and spatial distance to construct complementary brain views.
- Graph construction: Each EEG channel is treated as a node, while connections among brain regions are represented as graph edges.Each node has its own attributes or features.
- Graph construction: The two brain views are a functional connectivity-based graph GF C and a spatial distance-based graph GDC.Figure S.1 provides the detailed construction process for both graphs.
- Feature extraction: The feature-extraction network is a two-branch CNN in which small kernels capture temporal information and large kernels capture frequency information.This design separates temporal and frequency-oriented feature extraction.
- Baselines: The proposed model is compared with machine-learning, recurrent, convolutional, sequence-to-sequence, and spatial-temporal graph baselines.The listed baselines include SVM, random forests, MLP-LSTM, CNN-BiLSTM, multivariate time-series models, SeqSleepNet, and adaptive graph learning.
S.4 Confusion matrices
The supplementary figures provide confusion matrices and hypnogram comparisons for MSTGCN, alongside an analysis of network configuration sensitivity.
- Hypnogram: The proposed MSTGCN hypnogram is compared with the ground truth to assess automatically scored sleep stages.The hypnogram indicates that most sleep stages are classified correctly.
- Network configuration: MSTGCN is not sensitive to the tested hyper-parameter settings.Fewer ST-GCN layers and more convolution kernels slightly improve performance, while Chebyshev polynomial K and regularization parameter λ are also varied.