ST-USleepNet: A Spatial-Temporal Coupling Prominence Network for Multi-Channel Sleep Staging

1 Introduction

Sleep disorders are associated with a wide range of health issues [Mahowald and Schenck, 2005], including cardiovascular diseases [Redline et al., 2023], metabolic disorders [Gileles-Hillel et al., 2016], and neurodegenerative conditions [Abbott and Videnovic, 2016]. Accurate diagnosis and management of these disorders heavily rely on sleep staging

∗Corresponding Authors

[Malhotra and White, 2002], a process that involves categorizing sleep into distinct stages based on physiological signals. This classification provides essential insights into sleep patterns and plays a critical role in identifying abnormalities.

Sleep staging is traditionally performed using polysomnography (PSG), which primarily analyzes electroencephalogram (EEG) signals, but also incorporates other physiological measurements such as electrooculogram (EOG) and electromyogram (EMG). These signals are segmented into 30-second epochs and manually classified by experts into standard stages, including wake, non-REM

(N1, N2, N3) sleep stages and rapid eye movement (REM) [Berry et al., 2012]. However, this manual process is labor-intensive and challenging to scale for large datasets [Danker-Hopfe et al., 2009]. To improve the efficiency and accuracy of sleep staging, researchers have turned to artificial intelligence to develop automated sleep staging models [Supratak et al., 2017; Perslev et al., 2019; Supratak and Guo, 2020]. Despite advances, there are two unresolved challenges in this field.

How to simultaneously capture the characteristic sleep waveforms and salient spatial brain networks from multichannel raw signals? Different sleep stages exhibit prominent temporal and spatial features, primarily the characteristic sleep waveforms and salient spatial brain networks as outlined by the AASM sleep standards [Berry et al., 2012]. Figures 1 (a) and (b) show characteristic sleep waveforms over time. Panel (a) features spindle waves from the N2 stage, while panel (b) displays delta waves from the N3 stage. Figure 1 (c) presents the brain networks and salient spatial subnetworks during both N2 and N3 stages, with the salient subnetworks highlighted by retaining only the green nodes. Given the crucial information contained in these prominent sleep features, many studies [Jia et al., 2021b; Pei et al., 2024b; Wang et al., 2024a] have effectively extracted salient waveforms from sleep signals. However, these approaches are limited to single-channel signals, whereas PSG signals are predominantly multi-channel. The prominent sleep features in multi-channel signals are significantly more complex than those in single-channel signals. This complexity has two main sources. One arises from the characteristic sleep waveforms observed across different channels, as illustrated in Figures 1 (a) and (b), which show simultaneous waveforms in two channels. The other source comes from the formation of distinct functional connections at various sleep stages due to interactions between signals [Stevner et al., 2019; Parente and Colosimo, 2020]. These interactions result in intricate spatial patterns that further enrich the complexity of multi-channel signals. Consequently, simultaneously capturing both multi-channel characteristic sleep waveforms and salient spatial brain networks is crucial for accurate sleep staging and presents a significant challenge.

How to capture the spatial-temporal coupling patterns to better understand brain activities? Signals from different brain regions exhibit complex yet crucial spatial-temporal coupling patterns as they evolve over time [H¨ovel et al., 2020]. Figure 1 (d) illustrates the spatial-temporal relationships within the brain network over time. In the figure, each circle represents a signal channel, corresponding to an electrode placement on the brain; each arrow signifies a connection within the brain network over time. Using the green node as an example, different arrows indicate distinct types of relationships: red arrows denote temporal correlations within the same channel across different time steps, black arrows symbolize spatial dependencies at the same time step across various channels, and blue arrows represent spatialtemporal coupling patterns between different channels at different times. These spatial-temporal connections within the brain are crucial for revealing brain activities [Pang et al., 2023]. Most existing models can only capture temporal correlations [Phan et al., 2018; Eldele et al., 2021; Ji et al., 2023; Ji et al., 2024] or spatial dependencies [Jia et al., 2020; Ji et al., 2022]. Therefore, capturing these spatial-temporal coupling patterns for sleep staging presents a challenge.

To address the above challenges, we propose STUSleepNet, a raw-signal-based novel spatial-temporal coupling prominence framework for sleep staging. It consists of two main components: a spatial-temporal graph construction module (ST) and a U-shaped sleep network (USleepNet). Inspired by studies from other domains [Song et al., 2020; Gong et al., 2023], the ST module converts multi-channel raw signals into a spatial-temporal graph, allowing for the modeling and capturing of spatial-temporal coupling patterns. On the other hand, the USleepNet, which leverages a U-shaped structure originally developed for image segmentation [Ronneberger et al., 2015], segments the signals and the graph constructed by the ST to obtain prominent temporal and spatial sleep features like how image segmentation isolates the main target of an image. Therefore, the USleepNet has two streams: for the temporal stream, we utilize a fully convolutional U-Net to segment characteristic sleep waveforms; for the spatial stream, we employ a graph convolutional U-Net to isolate salient spatial brain networks. The key contributions of this work include the following:

• We develop the USleepNet module, comprising a temporal prominence network and a spatial prominence network to simultaneously extract multi-channel characteristic sleep waveforms and salient spatial brain networks.
• We design a spatial-temporal graph construction module to model the spatial-temporal coupling patterns.
• Experimental results indicate that our ST-USleepNet model reaches state-of-the-art performance across three sleep staging datasets.
• Visualization Studies demonstrate that our method effectively extracts prominent sleep features and spatialtemporal coupling patterns across different sleep stages, providing a certain degree of interpretability.

2 RelatedWork

Traditional sleep staging relied on machine learning methods like Random Forest (RF) [Memar and Faradji, 2018] and Support Vector Machine (SVM) [Lv et al., 2016; Alickovic and Subasi, 2018], which required manually crafted features based on expert knowledge, making the process costly and inefficient. To overcome this, deep learning approaches such as convolutional neural networks (CNNs) [Eldele et al., 2021; Jia et al., 2021a; Ji et al., 2023; Ji et al., 2024] and Recurrent Neural Networks (RNNs) [Pei et al., 2024a] were introduced to automatically extract high-dimensional features. However, these methods still follow a two-stage pipeline and lack endto-end training capabilities.

Recent efforts have focused on end-to-end architectures for processing temporal correlations in sleep staging without preliminary steps. Temporal modeling approaches include CNNs [Perslev et al., 2019; Jia et al., 2021b; Perslev et al., 2021], RNNs [Phan et al., 2018], their combinations

[Supratak et al., 2017; Supratak and Guo, 2020], and Transformers [Phan et al., 2022]. For instance, U-time [Perslev et al., 2019] employed a CNN-based Unet for single-channel signals, while SalientSleepNet [Jia et al., 2021b] integrated a U 2 -structured CNN Unet with attention to identify typical waveforms. Some studies have also focused on extracting typical waveforms to perform expert-like classification [Pei et al., 2024b; Wang et al., 2024a]. Despite their advancements, these methods primarily focus on single-channel or abstract temporal features, overlooking spatial dependencies.

Spatial dependencies among brain regions are critically important [Sakkalis, 2011]. GraphSleepNet [Jia et al., 2020] was the first to apply a spatial-temporal graph convolution network (ST-GCN) for sleep staging. Subsequently, researchers introduced more graph-based methods [Jia et al., 2021a; Ji et al., 2022]. Besides, Liu et al. proposed the BSTT [Liu and Jia, 2023] model, which uses a spatial-temporal transformer to capture spatial brain features. However, these spatial methods mainly focus on capturing global spatial features instead of the salient spatial brain networks and ignore the spatial-temporal coupling patterns.

To address the limitations of existing methods and effectively capture characteristic sleep waveforms, salient spatial brain networks, and spatial-temporal coupling patterns across different sleep stages frommulti-channel raw signals, we introduce ST-USleepNet for sleep staging.

3 Methodology

The framework of ST-USleepNet is illustrated in Figure 2, consisting of two modules. First, a spatial-temporal graph construction module is developed to capture the spatialtemporal coupling patterns, converting raw signals into a spatial-temporal graph. Next, we introduce USleepNet to extract characteristic sleep waveforms and identify spatial salient brain networks from these signals. USleepNet is composed of two interwoven streams: the temporal stream, also called the temporal prominence network, represented by yellow blocks in Figure 2, and the spatial stream, also called the spatial prominence network, depicted by orange blocks. Each stream comprises l encoder and l decoder blocks, with l = 2 exemplified in the figure. Additionally, the temporal stream integrates a segment classifier. To effectively combine the two streams, a fusion strategy is implemented.

3.1 Spatial-Temporal Graph Construction

The spatial-temporal graph construction module aims to transform the raw signals into a spatial-temporal graph. Given a multi-channel physiological signal data sC×fsT , where C represents the number of channels, fs denotes the sampling frequency, and T indicates the duration of each segment, the data sC×fsT is divided into n patches, resulting in XnC×fst . We use the raw signals X directly as the initial features of the graph, then we have Gs = (V, AnC×nC , XnC×fst ), where V represents the set of nodes, A is the adjacency matrix. The nodes of the graph Gs are constituted by patches: node Ncitj represents a channel ci, where i ∈ {k ∈ Z | 0 ≤ k < C}, at a certain time step tj, where j ∈ {k ∈ Z | 0 ≤ k < n}. Next, we will discuss how to construct the adjacency matrix AnC×nC .

The adjacency matrix A is derived from the multiplication of a similarity matrix with several weighted masks, as shown in Figure 3. First, an adjacency matrix is established based on the cosine similarity between different features, as highly similar signals often indicate similar functional states. To avoid assigning weights to unrelated nodes, a threshold δ is applied, establishing edges only between nodes with similarity exceeding δ. The similarity matrix S is then defined as follows:

Here, Xi and Xj represent features of two different nodes; sim cos represents the cosine similarity function. Next, we add some prior knowledge using other weighted masks. One is to add the temporal information: patches closer in time should have higher similarity. A temporal weighted mask Mt is constructed based on the time interval between patches. With a hyperparameter δt ∈ (0, 1), the temporal weight mtij between the feature Xciti and Xcjtj is defined as:

where ti and tj are the time steps of the patches, and ti, tj ∈ {k ∈ Z | 0 ≤ k < n}.

Similarly, positional features between channels should be considered because patches from different channels should have lower similarity than those from the same channel. A hyperparameter δp ∈ (0, 1) is defined to penalize edge weights between different channels, forming a positional weighted mask Mp. Thus, the positional weight mpij between the feature Xciti and Xcjtj can be expressed as:

Therefore, the adjacency matrix A for Gs is:

According to the above rules, the multi-channel signal segment sC×fsT is transformed into a spatial-temporal graph Gs = (V, AnC×nC , XnC×fst ).

3.2 U-shaped Sleep Network

The spatial-temporal graph Gs constructed from the raw signals is then input into the U-shaped sleep network (USleepNet). There are two interwoven networks, the temporal prominence network for the temporal stream and the spatial prominence network for the spatial stream.

Temporal Prominence Network

This network consists of three logical submodules: the temporal encoder block Ten, the temporal decoder block Tde, and a segment classifier Tc. Ten extracts multi-scale temporal features from the feature matrix X of the graph Gs, outputting a compressed representation, which Tde further processes. Lastly, Tc generates a segment-wise label.

Temporal Encoder Block Ten consists of two 1D convolution layers. The number of channels in the second convolution layer is greater than that in the first layer, since deeper networks tend to capture more complex features. Each convolution uses the same padding, followed by batch normalization. The kernel size remains consistent within the two convolution operations of each Ten, while it may vary across different Ten. After these two convolution operations, max pooling is performed to reduce the feature dimensions.

Temporal Decoder Block Tde consists of two transposed1D convolution layers. Each Tde corresponds to a Ten, performing inverse operations to recover features. It unpools the representation, concatenates it with the corresponding Ten output, and applies two sequential convolutions. Kernel sizes are consistent within each Tde but may vary across blocks, with batch normalization applied after each convolution.

Segment Classifier Tc generates segment-wise labels. It takes embedding Hch×nC×fst as input, reshapes it to ˆHch×C×nfst , and applies a softmax function to produce point-wise labels for each time step and channel. Mean pooling combines the channels into one, and a fully connected layer outputs the segment label.

Spatial Prominence Network

For the spatial stream, inspired by Graph U-Net [Gao and Ji, 2019], we propose the spatial prominence network. This network follows an encoder-decoder structure and consists of two logical submodules, the spatial encoder block Sen and the spatial decoder block Sde. Sen takes the spatial-temporal graph Gs as input and outputs a compressed feature subgraph, while Sde takes the subgraph and recovers a supergraph containing the learned spatial relationships.

Spatial Encoder Block Sen is to extract spatial features between multiple channels and sample important channels to form a new subgraph for high-level feature encoding. It consists of two GCN layers and a graph pooling layer.

The input to Sen consists of the adjacency matrix Av×v and the feature matrix Xv×h , where v represents the number of nodes and h represents the feature dimensions. The first graph convolution layer maps X to a higher-dimensional space, and the second graph convolution layer maps X back to its original dimensions. The two consecutive graph convolutions can be represented as follows:

where Drop means the dropout layer, Wi is a learnable hyperparameter, d is the higher dimension than h, and σ is ELU, which is defined as:

Then, a graph pooling layer (gPool) is used to perform downsampling on the graph data. This reduces the size of the graph while retaining information on important nodes. In gPool, we rank the nodes based on importance and select k to generate a subgraph. The computation of gPool can be represented as follows: First, each node's feature vector Xi is projected using a trainable projection vector P to calculate a scalar projection value as an importance score ic:

where v represents the number of nodes, h represents the feature dimension. Then, sort the ic and select k largest values:

where idx is the number of selected top-k nodes, and ˜ ick is their weights. Feature vectors of the top-k nodes are used to update the adjacency matrix ˜ A and feature matrix ˜X:

where ⊙ denotes element-wise multiplication, and 1h represents an h-dimensional constant vector of 1 used to expand the dimensions of ˜ ic. In equation 11, the reason for multiplying the two extracted adjacency matrices is to avoid the nodes in the pooled graph becoming isolated as much as possible.

Spatial Decoder Block Sde is used to recover the compressed subgraph obtained by Sen into a feature supergraph. It consists of four network layers: a convolution layer, a graph unpooling layer (gUnpool), and two graph convolution layers.

Sde takes the adjacency matrix A and the feature matrix X as input. First, a CNN layer is used to extract local information from the feature matrix X and reduce the number of channels in X. Next, the data is upsampled using a gUnpool layer. The gUnpool layer corresponds to a gPool layer and uses the node selection information recorded during the downsampling process of the gPool layer. It places the existing nodes and their feature matrices back into their original positions in the graph, while the features of other nodes are set to zero. To avoid significant information loss, a skip connection is used to add the adjacency matrix of the upsampled graph to the corresponding adjacency matrix before downsampling in Sde. Finally, similar to Sen, two consecutive graph convolution layers are used: the first maps the feature matrix to a hidden high-dimensional space, and the second restores it to its original dimensions. Through these four layers, the Sde block ultimately recovers the compressed feature subgraph into a feature supergraph.

Spatial-Temporal Block Fusion

USleepNet comprises an alternating combination of the temporal prominence network and the spatial prominence network. It includes l layers of spatial-temporal encoders (STen), l layers of spatial-temporal decoders (STde), a bottom block (STb), and a segment classifier at the end. Specifically, each STen consists of a temporal encoder (Ten) followed by a spatial encoder (Sen), while each STde consists of a spatial decoder (Sde) followed by a temporal decoder (Tde). The bottom block (STb) is similar to STen but does not include any downsampling. In the encoder phase, it is essential to process with Ten first, as PSG data typically has high resolutions and contains noise. Convolution and pooling operations not only reduce dimensions but also remove noise, thereby extracting more representative features. These lowerdimensional features significantly reduce the computational complexity of subsequent graph processing.

4 Experiments

4.1 Datasets

We evaluate ST-USleepNet on three public datasets. ISRUCS1 [Khalighi et al., 2016] contains PSG data from 100 patients (55 male and 45 female), with a total of 87,187 sleep segments. The PSG data includes 10 channels of EEG, EOG, and EMG. ISRUC-S3 [Khalighi et al., 2016] contains PSG data from 10 patients (9 male and 1 female), with a total of 8,589 sleep segments. The PSG data includes 10 channels of EEG, EOG, and EMG. MASS-SS3 [O'reilly et al., 2014] contains PSG data from 62 healthy people (28 male and 34 female), with a total of 59,304 sleep segments. The PSG data includes 25 channels of EEG, EOG, EMG, and ECG. For each dataset, the PSG data follows the AASM standards [Berry et al., 2012] and is classified into five categories: Wake, N1, N2, N3, and REM.

4.2 Experiment Settings

In this study, we use a cross-subject dataset division and perform 10-fold cross-validation. The evaluation metrics used are accuracy (Acc) and F1-score. All experiments are performed using NVIDIA A100 Tensor Core GPU. We configure the batch size to 64, utilize the Adam optimizer with a learning rate of 1e-3, and train the model for 40 epochs.

4.3 Baseline Methods

To comprehensively evaluate the performance of STUSleepNet, we select 7 models as baselines. These include traditional Random Forest (RF) method [Dong et al., 2017]; classic CNN-based methods DeepSleepNet [Supratak et al., 2017] and TinySleepNet [Supratak and Guo, 2020]; sleep networks based on U-shaped structures U-time [Perslev et al., 2019] and SalientSleepNet [Jia et al., 2021b]; spatialtemporal sleep network GraphSleepNet [Jia et al., 2020]; and the latest general multivariate time series model FCSTGCN [Wang et al., 2024b].

4.4 Experiment Results

The results are shown in Table 1, where the best scores are highlighted in bold. Our model achieves state-of-theart performance across all metrics on all datasets. The largest performance gains are observed on the ISRUC S3 dataset. Specifically, ST-USleepNet improves accuracy by +1.14 (1.46%) compared to the strongest baseline [Jia et al., 2020], and F1-score by +1.02 (1.31%) compared to the bestperforming baseline [Jia et al., 2021b]. These results underscore the importance of extracting critical features for effective sleep staging.

4.5 Ablation Study

To verify the contribution of the module in ST-USleepNet, we design three variant models by removing the corresponding modules: 1) Without temporal prominence network (- T). 2) Without spatial prominence network (-S). 3) Without spatial-temporal graph construction (-ST). In the '-ST' variant model, we construct a spatial graph using the raw EEG signals in the whole segment based on the cosine similarity.

Table 2 shows the comparison results of the ablation experiments between the variant models and the original model on the three datasets. It can be seen that removing either the temporal prominence network or the spatial prominence network adversely affects the model's performance, and the impact is similar. This indicates that extracting characteristic sleep waveforms and salient spatial brain networks is important. Additionally, removing the temporal-spatial graph construction module also affects the performance, suggesting that we should also focus on the spatial-temporal coupling patterns contained throughout the entire sleep signal, rather than just temporal or spatial features.

4.6 Hyperparameter Analysis

We explore the impact of various hyperparameters and identify two with the most significant influence.

First, ST-USleepNet implements a multi-scale strategy in the temporal prominence network by dynamically varying convolutional kernel sizes and in the spatial prominence network by using different hidden layer dimensions, allowing the network to capture features at multiple scales effectively. To validate that the model adopts the optimal strategy, we experimented with the following combinations: 1)increasing then decreasing (our model's strategy); 2)consistently using small kernels and hidden layers (small); 3)consistently using large kernels (large); 4)decreasing then increasing (reverse). We found that using decreasing convolutional kernels and hidden layer dimensions in the encoder and increasing them in the decoder achieves the best performance, as shown in Table 3.

Second, we investigate the impact of depth in the U-shaped network on the performance and determine that a four-layer network is the optimal structure, as shown in Table 4.

4.7 Model Visualization

In ST-USleepNet, each module can extract one kind of sleep feature: the temporal prominence network can capture characteristic sleep waveforms from multiple channels simultaneously, revealing temporal patterns across different sleep stages; the spatial prominence network can capture the salient spatial brain networks, highlighting key brain regions involved in sleep; and the spatial-temporal graph construction module can capture the spatial-temporal coupling patterns, revealing how brain regions interact over time and space. In this section, we show the visualizations for these three types of sleep features, which not only provide insights into the underlying data patterns in different domains but also offer a level of interpretability for the predictive model.

Characteristic Sleep Waveforms

We visualize the model's point-wise labels for each channel, showcasing that ST-USleepNet effectively captures characteristic sleep waveforms across multiple channels for various sleep stages. As shown in Figure 4 (a), the model identifies typical waveforms for the five sleep stages, with results displayed for two channels per stage. Notably, it captures α

waves during the wake period, vertex waves in N1, spindles in N2, δ waves in N3, and during REM, EEG sawtooth waves and EOG rapid eye movements.

Salient Spatial Brain Network

We combine the spatial-temporal graphs generated by the model into a single spatial graph. From the bottom block's output, we select channels with the largest sums of degrees across different time steps within the same channel. Using this information, we extract corresponding subgraphs from the concatenated spatial graph to represent the salient spatial brain network for each sleep stage. The visualization results are shown in Figure 4 (b). In the figure, the left panel shows the complete brain network, while the right panels display the salient brain networks for each sleep stage. These networks vary significantly: during Wake, brain regions are highly active; in N1, brain activity increases as eye movement decreases; in N2, eye movement further weakens; in N3, brain activity decreases significantly; and during REM, brain activity remains low while eye movement becomes strong.

Spatial-Temporal Coupling Patterns

We visualize the spatial-temporal graph output by the model as shown in Figure 4 (c). In the figure, to highlight the spatial-temporal coupling patterns, the spatial layout at each time step is omitted. Different types of physiological signal channels are represented by distinct colors, and physiological signals from the same time step are arranged in a single column. It is evident that ST-USleepNet effectively captures spatial-temporal coupling patterns, and the patterns reflected are consistent with the salient spatial brain network.

5 Conclusion

We propose ST-USleepNet, a novel spatial-temporal coupling prominence network for sleep staging. Our model can simultaneously capture characteristic sleep waveforms and salient spatial brain networks across multiple channels. Also, by constructing the raw signals to graph, we capture the spatialtemporal coupling patterns. Experimental results show that ST-USleepNet reaches state-of-the-art performance. Visualization proves that our model has a degree of interpretability.

Acknowledgments

This research is supported by A*STAR, CISCO Systems (USA) Pte. Ltd., and the National University of Singapore under its Cisco-NUS Accelerated Digital Economy Corporate Laboratory (Award I21001E0002), the National University of Singapore President's Graduate Fellowship, and the National Natural Science Foundation of China (Grant No. 62306317).