GraphSleepNet: Adaptive Spatial-Temporal Graph Convolutional Networks for Sleep Stage Classification

1 Introduction

Sleep stage classification is important for the assessment of sleep quality and the diagnosis of sleep disorders. Sleep experts identify sleep states based on sleep standard and observations recorded in polysomnography (PSG), which contains electroencephalography (EEG) at different positions on the head, electromyography (EMG), and electrooculography (EOG), etc. Currently, the two commonly used sleep standards are the Rechtschaffen and Kales (R&K) standard [Wolpert, 1969] and American Academy of Sleep Medicine (AASM) standard [Berry et al., 2012]. The transition rules between different sleep stages, which can assist sleep experts in identifying the sleep stages, are also recorded in the AASM standard. Although these rules provide very valuable information, classifying the sleep stages by human sleep experts is still a tedious and time-consuming task. Moreover, the classification results are affected by the variability and subjectivity of sleep experts.

∗ corresponding author: wj@bjtu.edu.cn

Automatic sleep stage classification can greatly improve the efficiency of traditional sleep stage classification and has important clinical value. A large number of researchers have made great contributions to automate this classification task. At first, traditional machine learning methods based on time domain, frequency domain, and time-frequency domain features are adopted. However, the classification accuracy of these methods depends heavily on feature engineering and feature selection, which require a lot of expert knowledge. Recently, deep learning methods have been widely applied to automatically classify sleep stage thanks to its powerful ability of representation learning. For example, Convolutional Neural Network (CNN) and Recurrent Neural Network (RNN) are often utilized to learn appropriate feature representations from transformed data or directly from raw data.

Although CNNs and RNNs can achieve high accuracy for sleep stage classification, their limitation is that the model's input must be grid data (themselves or after conversion to image-like representations) ignoring the connection among brain regions. Due to brain regions are in non-Euclidean space, graph is the most appropriate data structure to indicate brain connection. Motivated by the success of the graph convolution network (GCN) model in graph data, the sleep stage classification is studied by the graph representation method, where each EEG channel corresponds to a node of the graph, and the connection between the channels correspond to the edge of the graph. But an important input of GCN is the graph connection representation, which is generally a fixed graph structure which is set in advance. Although existing methods, such as the phase-locking value (PLV) method, can obtain a fixed graph structure, due to the limited understanding of the brain, it is still challenging to determine a suitable graph structure in advance for sleep stage classification.

Besides, how to exploit the sleep transition rules between neighboring stages is another challenge. Actually, sleep experts usually identify the class label of one sleep state according to both the characteristic EEG waves of the current state and the class labels of its neighbors. Hence, it is reasonable to take transition rules into account to improve classification accuracy.

In order to address the above challenges, we propose GraphSleepNet, a novel deep graph neural network, for automatic sleep stage classification. The GraphSleepNet is used to learn an adaptive graph structure representation that best serves spatial-temporal-GCN (ST-GCN) for sleep stage classification. The code has been released at: https://github.com/jingwang2020/GraphSleepNet.

Overall, the main contributions of the proposed GraphSleepNet for automatic sleep stage classification are summarized as follows:

• To the best of our knowledge, it is the first attempt to apply ST-GCN for automatic sleep stage classification. Moreover, we propose a novel adaptive sleep graph learning mechanism, which is integrated with ST-GCN simultaneously in a unified network architecture.
• We design a spatial-temporal convolution, which consists of graph convolutions for capturing spatial features and temporal convolutions for capturing the transition among different sleep stages.
• Experimental results demonstrate that the GraphSleepNet achieves state-of-the-art performance in sleep stage classification.

2 RelatedWork

2.1 Sleep Stage Classification

Identifying sleep stages plays an important role in diagnosing and treating sleep disorders. In earlier research, machine learning methods such as Support Vector Machine (SVM) and Random Forests (RF) are utilized to classify sleep stages. However, these methods need to extract hand-crafted features, which requires a lot of prior knowledge. Therefore, some researchers have turned to use deep learning methods for sleep stage classification.

Currently, two popular deep learning models, CNN and RNN, are widely used in sleep stage classification. A fast discriminative complex-valued CNN (FDCCNN) [Zhang and Wu, 2017] is proposed to capture the sleep information hidden inside EEG signals. A CNN model based on multivariate and multimodal physiological signals [Chambon et al., 2018] takes into account the transitional rules of sleep stages to assist classification. The extracted functional connectivity metrics are employed as multi-dimensional images to input into the CNN model for classification [Chriskos et al., 2019]. A hierarchical RNN named SeqSleepNet [Phan et al., 2019] tackles the task as a sequence-to-sequence classification task. At the same time, hybrid models are also used by some researchers. DeepSleepNet [Supratak et al., 2017] utilizes CNN to extract time-invariant features, and Bi-directional Long Short-TermMemory (BiLSTM) to learn the transition rules among sleep stages. A hierarchical neural network [Sun et al., 2019] implements comprehensive feature learning stage and sequence learning stage, respectively. Additionally, with the development of attention mechanisms, a deep Bi-directional RNN with attention mechanism is used for single-channel sleep stage classification [Phan et al., 2018].

Although the models above can extract spatial or temporal features from EEG signals, there is one obvious limitation: their input must be grid data. As different brain regions are not in the Euclidean space, grid data may not be the optimal data representation. Hence, the graph is the most appropriate data structure.

2.2 Graph Convolution Network

Recently, GCNs have shown advanced performance in addressing graph structure data. Graph convolution methods are classified into spatial convolution and spectral convolution methods. Spatial convolution methods usually perform graph convolution operations directly on the nodes and their neighbors of the graph. For example, an end-to-end CNN that operates on a graph is presented in [Duvenaud et al., 2015]. The GraphSAGE model [Hamilton et al., 2017] generates embeddings by sampling and aggregating features from a node's local neighborhood. Graph Attention Networks (GAT) [Veliˇckovi´c et al., 2018] uses masked self-attentional layers to assign different weights to different nodes in a neighborhood. Spectral convolution methods usually define graph convolution operation on the spectral representation of the graph. For example, a general graph convolution framework is proposed based on graph Laplacian [Bruna et al., 2014]. Chebyshev expansion of graph Laplacian is utilized to reduce computational complexity [Defferrard et al., 2016]. A simpler GCN for semi-supervised learning is proposed in [Kipf and Welling, 2016]. Although, the existing GCNs usually utilize a fixed graph, which is not optimal for sleep stage classification, because the graph structure of different sleep states may be different. Moreover, the GCNs are applied directly to classify sleep stages without the capture of sleep transition rules in the temporal dimension.

3 Preliminaries

In our study, a sleep stage network is defined as an undirected graph G = (V, E, A), where V denotes the set of vertices and each vertex in the network represents an electrode on brain; |V | = N is the number of vertices in sleep stage network; E denotes the set of edges and indicates the connection between vertices; A denotes the adjacency matrix of sleep stage network G. As shown in Figure 1, Gt is constructed from a 30s EEG signal sequence St. It is worth noting that the adjacency matrix A used in the proposed model is obtained through learning, while traditional GCNs usually use a fixed adjacency matrix.

The sleep feature matrix is the input of the GraphSleepNet. We define the raw signals sequence as S = (s1, s2, . . . , sL) ∈ RL×N×Ts , where L denotes the number of samples, Ts denotes the time series length of each sample si ∈ S(i ∈ {1, 2, · · · ,L}). For each sample si, we extract the differential entropy (DE) features [Hyv¨arinen, 1998] on various frequency bands and define each sample si's feature matrix Xi = (xi1 , xi2 , . . . , xiN )T ∈ RN×Fde , where xin ∈ RFde (n ∈ {1, 2, · · · , N}) denotes Fde features of node n at sample i. Figure 1 shows the sleep stage network and sleep feature matrix, which are constructed from the raw signals.

The research goal is to learn the mapping relationship between the encoded signals and sleep stage class based on the proposed model. The problem of sleep stage classification is defined as: given X = (Xi-d, . . . ,Xi, . . . ,Xi+d) ∈ RN×Fde×Tn identify the current sleep stage y, where X denotes the temporal context of Xi, y denotes the Xi's sleep stage class label, and Tn = 2d+1 denotes the length of sleep stage networks, where d ∈ N+ denotes temporal context coefficient.

4 Adaptive Spatial-Temporal GCN

The overall architecture of the GraphSleepNet is shown in Figure 2. We summarize three key ideas of the GraphSleepNet: 1) Represent the pairwise relationship between nodes to dynamically construct adjacency matrix A. 2) Combine spatial graph convolution and temporal convolution to extract both spatial and temporal features. 3) Employ a spatial-temporal attention mechanism to automatically pay relatively more attention to valuable spatial-temporal information. The overall architecture is elaborately designed to accurately identify sleep stages.

4.1 Adaptive Sleep Graph Learning

Adaptive sleep graph learning is the first key part of the GraphSleepNet as shown in Figure 3. It dynamically learns the graph structure, rather than the graph constructed by prior knowledge or artificially (e.g., k-nearest neighbor graph [Jiang et al., 2013]). Therefore, we define a non-negative function Amn = g (xm, xn) (m, n ∈ {1, 2, · · · ,N}) to represent the connection relationship between nodes xm and xn based on the input feature matrix Xi = (xi1 , xi2 , . . . , xiN )T ∈ RN×Fde . g (xm, xn) is implemented through a layer neural network, which has the learnable weight vector w = (w1,w2, . . . ,wFde )T ∈ RFde×1 . The learned graph structure (adjacency matrix) A is defined as:

where rectified linear unit (ReLU) is an activation function to guarantee that Amn is non-negative. The softmax operation normalizes each row of A. The weight vector w is updated by minimizing the following loss function,

That is, the larger distance ‖xm - xn‖2 between xm and xn, the smaller Amn is. Due to the brain connection structure is not a fully connected graph, we utilize the second term in the loss function to control the sparsity of graph A, where λ ≥ 0 is a regularization parameter.

The proposed graph learning mechanism automatically constructs the neighborhood connection of the nodes. To avoid the trivial solution (i.e., w = (0, 0, · · · , 0)), which is due to minimizing the above loss function Lgraph learning independently, we utilize it as a regularized term to form the final

where Lcross entropy denotes the original loss function of the multi-classification task, L denotes the number of samples, R denotes the number of classes. y is the true label and ˆ y is the value predicted by the model.

4.2 Spatial-Temporal Graph Convolution

Spatial-temporal graph convolution is a combination of spatial graph convolution and temporal standard convolution, which is used to extract both spatial and temporal features. As shown in Figure 2, the spatial features are extracted by aggregating information from neighbor nodes for each sleep stage network and the temporal features are captured by exploiting temporal dependencies from neighbor sleep stages.

Spatial Graph Convolution

We use graph convolution based on spectral graph theory to extract spatial features in the spatial dimension. For each sleep stage to be identified, the adaptive sleep graph learning module provides an adjacency matrix A for graph convolution. In addition, we employ the Chebyshev expansion of graph Laplacian to reduce computational complexity. Chebyshev graph convolution [Defferrard et al., 2016] using the K - 1 order polynomials is defined as:

where gθ denotes the convolution kernel, ∗G denotes the graph convolution operation, θ ∈ RK denotes a vector of Chebyshev coefficients and x is the input data. L = D-A denotes Laplacian matrix, where D ∈ RN×N is degree matrix. ˜L = 2 λmax L - IN, where λmax is Laplacian matrix's maximum eigenvalue and IN is a identity matrix. Tk(x) = 2xTk-1(x) - Tk-2(x) denotes the Chebyshev polynomials recursively, where T0(x) = 1 and T1(x) = x.

By using the approximate expansion of Chebyshev polynomial, the information of neighboring 0 to K - 1 order neighbors centered at each node is extracted.

In this work, we generalize the above definition to the nodes with multiple channels. The l-th layer's input is ˆ X (l-1) = ( ˆX1, ˆX2, . . . , ˆXTl-1 ) ∈ RN×Cl-1×Tl-1 , where Cl-1 denotes channel's number of each node, i.e., l = 1, C0 = Fde. Tl-1 denotes the l-th layer's temporal dimension. For each ˆXi, we obtain gθ ∗G ˆXi by using Cl filters on ˆXi, where Θ = (Θ1,Θ2, . . . ,ΘCl ) ∈ RK×Cl-1×Cl is the convolution kernel parameter [Defferrard et al., 2016]. Hence, the information of the 0 ∼ K - 1 order neighbors is aggregated to each node.

Temporal Convolution

To capture the sleep transition rules, which are used by sleep experts to classify the current sleep stage in combination with neighboring sleep stages, we employ CNN to perform convolution operation in the temporal dimension. Specifically, after graph convolution operation has sufficiently extracted the spatial features from each sleep stage network, we implement a standard 2D convolution layer to extract the temporal context information of the current sleep stage. The temporal convolution operation on the l-th layer is defined as:

where ReLU is the activation function, Φ denotes the convolution kernel's parameters, * denotes the standard convolution operation.

4.3 Spatial-Temporal Attention

The attention mechanism is often used to automatically extract valuable information. In this study, we employ a spatialtemporal attention mechanism [Guo et al., 2019] to capture valuable spatial-temporal information on the sleep stage network. The spatial-temporal attention mechanism contains spatial attention and temporal attention.

Spatial Attention

In the spatial dimension, different regions have different effects on the sleep stage which is dynamically changing during sleep. To automatically extract the attentive spatial dynamics, we utilize a spatial attention mechanism, which is defined as

where V p, bp ∈ RN×N , Z1 ∈ RTl-1 , Z2 ∈ RCl-1×Tl-1 , Z3 ∈ RCl-1 denote learnable parameters, σ denotes the sigmoid activation function. X (l-1) = (X1,X2, . . . ,XTl-1 ) ∈ RN×Cl-1×Tl-1 denotes the l-th layer's input. P denotes spatial attention matrix, which is dynamically computed by current layer's input. Pm,n denotes the correlation between node m and n. The softmax operation is utilized to normalize the attention matrix P. In the proposed model, when the graph convolution is performed, the learned adjacency matrix A and spatial attention matrix P can dynamically adjust the update of nodes.

Temporal Attention

In the temporal dimension, there are correlations among neighboring sleep stages, and the correlations change in different situations. Therefore, a temporal attention mechanism is utilized to capture dynamic temporal information among sleep stage networks.

The temporal attention mechanism is defined as follows:

where V q, bq ∈ RTl-1×Tl-1 , M1 ∈ RN , M2 ∈ RCl-1×N , M3 ∈ RCl-1 denotes learnable parameters. Qm,n denotes the strength of correlation between sleep stage network m and n. Finally, the softmax operation is utilized to normalize the attention matrix Q. The input of the ST-GCN is tuned by the temporal attention: ˆ X (l-1) = ( ˆX1, ˆX2, . . . , ˆXTl-1 ) = (X1,X2, . . . ,XTl-1 )Q′ ∈ RN×Cl-1×Tl-1 to pay more attention to valuable temporal information.

5 Experiments

To verify the effectiveness of the GraphSleepNet, we evaluate it on the Montreal Archive of Sleep Studies (MASS)-SS3 dataset [O'Reilly et al., 2014].

5.1 Dataset

We use the MASS-SS3 dataset, which contains the PSG recordings from 62 healthy subjects (28 male and 34 female). Each recording contains 20 EEG channels, 2 EOG channels, 3 EMG channels, and 1 ECG channel. In addition, experts classify these PSG recordings into five sleep stages (W, N1, N2, N3, and REM) according to AASM standard [Berry et al., 2012].

5.2 Preprocessing

The PSG recordings are pre-processed with the bandpass filters of 0.30-100 Hz (EEG), 0.10-100 Hz (EOG, ECG) and 10-100 Hz (EMG). We extract DE features for each channel from 9 crossed frequency bands: 0.5-4 Hz, 2-6 Hz, 4-8 Hz, 611 Hz, 8-14 Hz, 11-22 Hz, 14-31 Hz, 22-40 Hz, 31-50 Hz. To make a fair comparison of existing sleep classification methods, we employ all signals in our experiments.

5.3 Baseline Methods

We compare the GraphSleepNet with the following baselines:

• [Dong et al., 2017]: A mixed neural network, which combines multilayer perceptron (MLP) and LSTM, and also compare its performance with RF and SVM.
• [Supratak et al., 2017]: A model combines CNN and BiLSTM to capture both time-invariant features and transition rules among sleep stages.
• [Chambon et al., 2018]: A temporal sleep stage classification use multivariate and multimodal time series.
• [Phan et al., 2019]: SeqSleepNet changes the single sleep stage classification problem into a sequenceto-sequence classification problem by using attentionbased bidirectional RNN (ARNN) and RNN.
• [Sun et al., 2019]: A hierarchical neural network, which learns comprehensive features and sequence respectively.
• [Jiang et al., 2019]: Robust sleep stage classification which uses multimodal decomposition and Hidden Markov Model (HMM) -based refinement.

5.4 Experiment Settings

We employ 31-fold cross-validation to evaluate the performance of the GraphSleepNet. In addition, we use the subjectindependent schemes to randomly divide the training set and test set. The hyperparameters of the experiment are shown in Table 1. We implement the GraphSleepNet using TensorFlow. Please refer to the code for more detailed.

5.5 Comparison with State-of-the-Art Methods

We compare our GraphSleepNet with the other eight baseline methods for sleep stage classification on the MASS-SS3 dataset, as shown in Table 2. The results show that our GraphSleepNet consistently achieves the best performance on all evaluation indicators compared to other baseline methods.

The traditional machine learning methods (SVM and RF) cannot learn the complex spatial or temporal features well. However, existing deep learning models such as CNN and RNN [Dong et al., 2017; Supratak et al., 2017; Chambon et al., 2018; Phan et al., 2019; Sun et al., 2019] can directly extract the spatial or temporal features. Therefore, their performance is better than the methods based on traditional machine learning.

Although CNN and RNN achieve high accuracy, their limitation is that the model's input must be grid data ignoring the connection among regions. Due to brain regions are in nonEuclidean space, graph is the most appropriate data structure to indicate the connection. The proposed GraphSleepNet extracts both spatial and temporal features based on an adaptive sleep graph structure. Therefore, the proposed GraphSleepNet is superior to other baseline methods.

5.6 Experimental Analysis

To further investigate the effectiveness of the adaptive sleep graph learning, we design five fixed adjacency matrices to compare with it. We compare these different adjacency matrices with an adaptive adjacency matrix. The different adjacency matrices are defined as below:

• Fully Connected Adjacency Matrix: A matrix whose elements are all 1. It represents that there are all connections among all nodes and each node also has selfconnection in the graph.
• Random Adjacency Matrix: A randomly generated matrix with each element's value between 0 and 1 indicating that the connections between nodes are random.
• Space Distance Adjacency Matrix: A matrix constructed based on the actual spatial distance between electrodes.
• K-Nearest Neighbor Adjacency Matrix [Jiang et al., 2013]: A matrix, which represents a k-nearest neighbor graph. That is, each node has k neighbor nodes.
• PLV Adjacency Matrix [Aydore et al., 2013]: A matrix generated by PLV method between each pair of nodes.

Figure 4(a) illustrates that the adaptive (learned) adjacency matrix achieves the highest accuracy for sleep stage classification. In addition, the adjacency matrix combined with prior knowledge also achieves a suboptimal effect, such as the PLV adjacency matrix. The fully connected adjacency matrix does not work well because the brain network is not a fully connected graph. In general, the adjacency matrix can significantly affect the classification effect. The proposed adaptive sleep graph learning for classification tasks is superior to all fixed graphs.

In addition, we study the impact of the number of input sleep stage networks (temporal context) on the classification performance. Figure 4(b) presents the performance of the classification improves as the number of input sleep stage networks Tn increases, and the best accuracy is achieved when Tn = 5. It is worth noting that according to the AASM standard, sleep experts also judge the current sleep state based on 5 sleep stages, which shows the practicability of the proposed model. However, too long time context is input, resulting in degradation of classification performance, which may be caused by information redundancy.

6 Conclusion

In this paper, we propose a novel deep graph neural network for automatic sleep stage classification. The main advantage of the proposed model is to adaptively learn the sleep connection structure that best serves ST-GCN for the classification task. Furthermore, this model combines spatial-temporal convolution and spatial-temporal attention mechanisms to simultaneously capture attentive spatial-temporal features of the sleep data. Experimental results demonstrate our model achieves state-of-the-art performance in the existing models. In addition, our proposed model is a general-framework for multivariate physiological time series.

Acknowledgments

This work was supported by the Natural Science Foundation of China (No. 61603029).