BSTT: A B AYESIAN S PATIAL -T EMPORAL T RANSFORMER FOR S LEEP S TAGING
Anonymous authors
Paper under double-blind review
A BSTRACT
Sleep staging is helpful in assessing sleep quality and diagnosing sleep disorders.
However, how to adequately capture the temporal and spatial relations of the brain during sleep remains a challenge.
In particular, existing methods cannot adaptively infer spatial-temporal relations of the brain under different sleep stages.
In this paper, we propose a novel Bayesian spatial-temporal relation inference neural network, named Bayesian spatial-temporal transformer (BSTT) for sleep staging.
Our model is able to adaptively infer brain spatial-temporal relations during sleep for spatial-temporal feature modeling through a well-designed Bayesian relation inference component.
Meanwhile, our model also includes a spatial transformer for extracting brain spatial features and a temporal transformer for capturing temporal features.
Experiments show that our BSTT outperforms state-ofthe-art baselines on ISRUC and MASS datasets.
In addition, the visual analysis shows that the spatial-temporal relations obtained by BSTT inference have certain interpretability for sleep staging.
I NTRODUCTION
Sleep staging is essential for assessing sleep quality and diagnosing sleep disorders.
Sleep specialists typically classify sleep stages based on the AASM sleep standard and polysomnography (PSG) recordings to aid in diagnosis.
The AASM standard not only provides criteria for determining each sleep period, but also documents conversion rules between different sleep stages, which is known as sleep transition rules, to help sleep specialists identify sleep stages when sleep transitions occur (Berry et al., 2012).
However, artificial sleep staging takes a long time, and the classification results are greatly affected by professional level and subjectivity (Supratak et al., 2017).
Therefore, automatic classification methods are applied into sleep staging to improve efficiency.
Traditional machine learning methods use artificially designed features for sleep staging, which improves the efficiency of staging to a certain extent (Fraiwan et al., 2012).
However, the accuracy of traditional machine learning methods relies heavily on feature engineering and feature selection, which still requires a lot of expert knowledge.
To address the above problems, deep learning methods have been applied to sleep staging and achieved satisfactory classification performance (Phan et al., 2019; Jain & Ganesan, 2021; Eldele et al., 2021b).
Most of the early deep learning methods focus on the temporal information of the sleep data, utilizing convolutional neural networks (CNN) and recurrent neural networks (RNN) to capture temporal features for sleep staging (Jain & Ganesan, 2021;
Perslev et al., 2019).
In addition, some studies have shown that the spatial topology of the brain behave differently in different sleep stages (Khanal, 2019), which means that both the temporal and spatial relations of the brain are both important during sleep.
Therefore, some researchers try to use the spatial and temporal characteristics of the brain for sleep staging (Jia et al., 2020b; Phan et al., 2022).
Although the above methods achieve good classification performance, it is challenging to model spatial and temporal relations.
Specifically, for the modeling of temporal relations, some approaches attempt to capture sleep transition rules in sleep to serve the identification of specific sleep stages.
However, it is difficult for these methods to explicitly demonstrate the relation of different sleep time slices in accordance with the AASM sleep standard.
Besides, for the modeling of spatial relations, spatial convolution operation is employed to extract the spatial features of the brain, which is insufficient that it may ignore the spatial topology of the brain by most methods (Zhou et al., 2021a;
Perslev et al., 2019).
A few researches utilize spatial topology and temporal relation information of brain for sleep staging by graph convolutional networks, but the constructed brain networks still lack interpretability to a certain extent (Jia et al., 2020b).
To address the above challenges, we propose a novel model called Bayesian spatial-temporal transformer (BSTT) for sleep staging.
The proposed model integrates the transformer and Bayesian relation inference in a unified framework.
Specifically, we design the spatial-temporal transformer architecture, which can capture the temporal and spatial features of the brain.
Besides, we propose the Bayesian relational inference component which comes in two forms, Bayesian temporal relation inference and Bayesian spatial relation inference.
Wherefore, it can infer the spatial-temporal relation of objects and generate the relation intensity graphs.
Specifically, the main contributions of our
BSTT are summarized as follows:
• We propose Bayesian relational inference component which can adaptively infer spatialtemporal relations of brain during sleep in the service of capturing spatial-temporal relations.
• We design the spatial-temporal transformer architecture to simultaneously model spatialtemporal relations.
It can effectively capture the spatial-temporal features of the brain and enhance the model’s ability to model spatial-temporal relations.
• Experimental results show that the proposed BSTT achieves the state-of-the-art in multiple sleep staging datasets.
The visual analysis shows that our model has a certain degree of interpretability for sleep staging.
R ELATED W ORK
Identifying sleep stages plays an important role in diagnosing and treating sleep disorders.
Earlier, the support vector machine (SVM) and random forest (RF) are used for sleep staging (Fraiwan et al., 2012).
However, these methods need hand-crafted features, which require a lot of prior knowledge.
Currently, deep learning methods have become the primary method for sleep staging.
Early deep learning methods extract temporal features of sleep signals for classification.
The earliest methods are based on the CNN models (Tsinalis et al., 2016; Chambon et al., 2018).
For example,
Chambon et al.
propose a convolutional neural network that can extract temporal-invariant features from sleep signals (Chambon et al., 2018).
Furthermore, Eldele et al. develope a multi-resolution
CNN with adaptive feature recalibration to extract representative features (Eldele et al., 2021a).
In addition, RNN models have been gradually used for sleep staging (Phan et al., 2019;
Perslev et al., 2019; Phan et al., 2018).
For example, Phan et al. propose a deep bidirectional RNN model with attention mechanism for single-channel EEG (Phan et al., 2018).
They then design an end-to-end hierarchical RNN architecture for capturing different levels of EEG signal features (Phan et al., 2019).
Some studies combine CNN with RNN (Supratak & Guo, 2020; Guillot & Thorey, 2021;
Dong et al., 2017).
For example, Suratak et al.
propose a hybrid model combining CNN and RNN to extract rich temporal features (Supratak et al., 2017).
In addition, Phan et al.
introduce transformer into the sleep staging task to capture the temporal context features of sleep signals (Phan et al., 2022).
Jia et al.
design a fully convolutional model to capture the typical waveform of sleep signals (Jia et al., 2021b).
Further, several studies have shown the importance of brain spatial relations for sleep staging (Khanal, 2019; Sakkalis, 2011).
Some researchers try to model the spatial-temporal characteristics of sleep data.
For example, Jia et al.
propose an adaptive deep learning model for sleep staging.
The proposed spatial-temporal graph convolutional network is used to extract spatial features and capture transition rules (Jia et al., 2020b).
They also propose a multi-view spatial-temporal graph convolutional network based on domain generalization, which models the multi-view-based spatial characteristics of the brain (Jia et al., 2021a).
Although the above models achieve good classification performance, these models do not adequately model spatial-temporal properties or effectively reason and capture spatial-temporal relations.
Therefore, our method attempts to model spatial-temporal relations using Bayesian inference, combined with state-of-the-art transformer architectures for sleep staging.
P RELIMINARIES
The proposed model processes the data of successive T sleep epochs and predicts the label of the time slice.
Each sleep epoch is defined as x ∈ R C × N , where C represents the number of channels of the sleep epoch (i.e. the EEG channel in this work. Since EEG signals from different channels are extracted from different regions of the brain, spatial relations are contained among these channels.) and N represents the number of sampling points in a sleep epoch.
The input sequence of sleep epochs is defined as x = { x 1 , x 2 , . . .
, x T } , where x i denotes a sleep epoch ( i ∈ [1 , 2 , . . .
, T ] ) and T is the number of sleep epochs.
The sleep staging problem is defined as: learning an artificial neural network F based on Bayesian spatial-temporal transformer, which can infer the spatial-temporal relation of the input sleep epoch sequence x and map it to the corresponding sleep stage (cid:98) Y , where (cid:98) Y is the classification result of the middle time slice x m .
According to the AASM standard, each (cid:98) Y ∈ { 0 , 1 , 2 , 3 , 4 } is matched to five sleep stages W, N1, N2, N3, and REM, respectively.
B AYESIAN S PATIAL -T EMPORAL T RANSFORMER
We propose a novel model named Bayesian spatial-temporal transformer (BSTT) for sleep staging.
The core ideas of our model are summarized as follows:
• Infer the spatial-temporal relation of the brain based on Bayesian inference method.
• Design the Bayesian transformer architecture while capturing the spatial-temporal features of the brain.
• Integrate Bayesian relational inference components and transformer architecture into a unified framework which can stage sleep period effectively.
The overall model is carefully designed to accurately classify different sleep stages.
4. A RCHITECTURE
The overall architecture of the proposed BSTT is shown in Figure 1.
The EEG signals are first encoded by the embedding layer.
The spatial-temporal relations are then inferred and modeled by Bayesian spatial-temporal transformer model.
Specifically, the Bayesian spatial-temporal transformer includes a Bayesian spatial transformer and a Bayesian temporal transformer.
The Bayesian spatial transformer can reason about spatial relations in the brain and capture spatial features.
The
Bayesian temporal transformer can reason about the temporal relation of consecutive sleep epochs and capture temporal features.
Finally, the predictions of different sleep stages are performed by the classification layer.
4.1. B AYESIAN RELATION INFERENCE
Capturing the spatial-temporal relations of brain signals during sleep can better serve sleep staging task.
However, due to the difficulty in inferring the spatial-temporal relation of sleep, current researches are insufficient for modeling the spatial-temporal relation.
Inspired by deep graph random process (DGP) proposed in recent research (Huang et al., 2020), we propose the Bayesian relation inference method.
Bayesian relational inference is the core component of our model, which can infer relations between each pair of object nodes and build relation intensity graphs.
In this article, an object node represents the embedding of an EEG channel or an embedding of a certain time slice.
The construction of the relation intensity graph is mainly divided into the following steps:
Step 1: Edge embedding initialization.
The input of the Bayesian relational inference is the node embeddings of the object.
Building a sleep relation intensity graph starts with generating edge embeddings.
Specifically, the node embeddings of the object are spliced and generated into edge embeddings.
We first apply a linear neural network f θ to generate edge embeddings:
[Equation]
where E n ∈ R B × N v × V represents the input n node objects, E n [ i, j ] ∈ R B × N e × 2 V represents the splicing of two node features, and E e ∈ R B × N e × E represents the generated edge embedding.
[Figure]
Step 2: Edge embedding coupling.
Coupling is one of the key steps in relation inference, the purpose is to obtain a summary graph of the relation between object nodes.
We assume that the edges of the summary graph ¨ α i,j are a summary of an n → ∞ of λ → 0 Bernoulli distributions, which means ¨ α i,j ∼ B ( n, λ ) , due to the uncertainty of spatial-temporal relations of brain (Huang et al., 2020). Drawing inspiration from VRNN (Chung et al., 2015), the parameters of the approximate posteriors are estimated from an RNN encoding features. However, unlike VRNN, Bayesian relation inference component contains the approximate posterior q (¨ α | E e ) , whose inference and sampling cannot be directly solved in a computationally tractable manner due to the infinity n .
By De MoivreLaplace theorem (Sheynin, 1977) and Theorem 1 provided in Appendix 3.1, we can subject these edge embeddings to a coupling transformation as follows:
[Equation]
where ζ ( · ) is softplus function, E e is the edge embedding generated in the first step, ϵ is a very small constant, L mean ( · ) and L std ( · ) are the mean coding network and the variance coding network,
M ∈ R B × N e is the approximation of summary graph which strengthens the representation of real spatial-temporal relations.
Step 3: Sleep relation intensity calculation.
The final step is to strengthen edge information and generate a relation intensity graph for downstream tasks.
The brain temporal-spatial network under real sleep is sparse based on existing research (Razi et al., 2017).
Therefore, generating a sparse graph based on edge embedding can not only highlight the representation of key relations, but also be more in line with the actual situation.
We employ the Gaussian graph transformation approach which produces a sparse sleep relation intensity graph G .
The specific calculation is defined as follows:
[Equation]
where m i , j ∈ M is the approximation of the summary graph obtained in Step 2, ε and ε
′ are the standard Gaussian random variable of the same dimension as M , (cid:101) a ∈ R B × N e is the Gaussian edge representation, S ∈ R B × N e is the task-related Gaussian variable, std is the standard deviation, mean is the mean value, ¯ α is the Gaussian transformation map, α ∈ R B × N e is the final sleep relation intensity graph, ζ ( · ) is softplus function and L( · ) is linear function.
Afterwards, we utilize an attention mechanism-based method to convert the node embeddings into feature embeddings based on the sleep relation intensity graph as the output of Bayesian relation inference.
The specific calculation is as follows:
[Equation]
[Equation] where f ATT ( · ) is attention layer, E n is the node embeddings of the input object, α is the sleep relation intensity graph, and E out ∈ R B × N v × V out is the output node embeddings.
We adopt variational inference to jointly optimise Baysian relation inference component.
Inspired by VRNN (Chung et al., 2015), we can use the evidence lower bound (ELBO) for joint learning and inference.
The details of how variational inference fits into our model are shown in the Appendix 3.2.
Specifically, we use two random variables need to be optimised to describe the same random process data. The resulting objective is to maximize the ELBO:
[Equation]
where S is the task-related Gaussian variable, ˜ A is the Gaussian edge embedding, q (cid:16) ˜
A , S | X 0:i (cid:17) is the prior distribution, and p (cid:16) ˜
A , S | X 0:i (cid:17) is the posterior distribution.
Since every variable in S are affected by ˜ A in Eq.(6), the KL term can be further written as:
[Equation]
Obviously the second term can be calculated, while the first term is tough to calculate because n →∞ .
According to Theorem 2 (The details are provided in Appendix 3.3), we can convert it into an easy-to-solve value to approximate the calculation.
We use the sum of the result of Eq.(11) in inferring the temporal and spatial relation and cross-entropy loss as the final loss.
4.1. B AYESIAN TRANSFORMER MODULE
Transformer shows convincing results in various sequence modeling tasks (Li et al., 2021; Luo et al., 2021; Zhou et al., 2021b).
However, traditional transformer does not have the ability to reason the relation between each pair of object nodes.
This results in a lack of interpretability of the attention graph generated by transformer.
Besides, the accuracy of traditional transformer is not good enough under some medical scenarios.
The Bayesian relation inference component we proposed can infer spatial-temporal relation efficiently.
Hence, we integrate the Bayesian relation inference component mentioned in Section 4.1.1 with the transformer in a unified framework to simultaneously reason and model the spatial-temporal features of sleep EEG data.
Bayesian spatial transformer.
To better construct the spatial functional connectivity of the brain and capture spatial features, we design the Bayesian spatial transformer.
It contains two components: a Bayesian spatial relation inference component and a position feed-forward network, of which the core component is a Bayesian relation inference component.
Specifically, the input of the Bayesian spatial transformer S is the embeddings of n spatial nodes.
First we add position encoding to the input to introduce position information:
[Equation]
(12) where S ∈ R ( B × N t ) × N s × V is the input spatial node embedding, P ep ∈ R ( B × N t ) × N s × V is the position encoding matrix, and ˜ S ∈ R ( B × N t ) × N s × V is the position-encoded spatial node embeddings.
For the position encoding matrix, we follow the groundbreaking work (Vaswani et al., 2017) and use the sine and cosine functions to calculate, the details of the calculation are shown in the Appendix 1.2.
We design a multi-head spatial Bayesian relation inference component which can reson about spatial relations to improve the representation learning ability of the model.
The details of Bayesian relation inference component have been described in Section 4.1.1.
The node embeddings with positional encoding are encoded as embeddings with spatial features after passing through the multi-head spatial Bayesian relation inference component and the feed-forward neural network layer:
[Equation]
(13) where ˜ S is the input node embedding, f BSRI ( · ) is the Bayesian spatial relation inference module which inferences the spatial relation, f FNN ( · ) is the feed-forward neural network layer, and
˜ S ′ ∈ R B × N t × V s is the spatial relation intensity graph based features.
Bayesian temporal transformer.
Similar to the Bayesian spatial transformer, in order to better capture the sleep transition rules, we design a Bayesian temporal transformer module to reason and model temporal features.
The calculations of e temporal relation intensity graph based features are as follows:
[Equation]
where f BTRI ( · ) is the Bayesian temporal relation inference component, f FNN ( · ) is the feed-forward neural network layer, and ˜ T ′ ∈ R B × N st is the temporal relation intensity graph based features.
˜ T ′ generates the classification results after passing through a linear classification layer f C : [Equation]
[Equation]
where (cid:98) Y is the classification result of BSTT.
E XPERIMENTS
To verify the effectiveness of the Bayesian spatial-temporal transformer, we evaluate it on the Institute of Systems and Robotics, University of Coimbra (ISRUC) and Montreal Archives of Sleep
Studies-SS3 (MASS-SS3) dataset.
5. D ATASET
ISRUC dataset contains the PSG recordings from 100 adult subjects.
Each PSG recording contains
EEG channels, 6 EOG channels, 3 EMG channels, and 1 ECG channel.
MASS-SS3 dataset contains the PSG recordings from 62 adult subjects.
Each PSG recording contains 20 EEG channels,
EOG channels, 3 EMG channels, and 1 ECG channel.
The recordings are divided into time slices according to a sleep epoch every 30s.
Sleep spacialists divide these time slices into five distinct sleep stages (W, N1, N2, N3, and REM) according to the AASM standard.
There are also motion artifacts at the beginning and end of each subject’s recording which are marked as unknown.
We follow the previous study and remove these recordings (Supratak et al., 2017).
5. E XPERIMENT S ETTINGS
We evaluate our model using k -fold cross-subject validation to ensure that the experiment results are correct and reliable.
We set k = 5 in order to test all recordings efficiently.
For optimizer, the
Adam and Adadelta optimizer are deployed in MASS-SS3 and ISRUC dataset.
We use multi-channel EEG data for sleep staging to better capture brain network structure.
Specifically, on the ISRUC dataset, we use all 6 EEG channels for experiments.
In the MASS-SS3 dataset we use 19 channels of EEG signals.
To comprehensively evaluate the Bayesian spatial-temporal transformer model and all baseline methods, we use accuracy (ACC), F1 Score and KAPPA to evaluate the models.
Specific information of the evaluation indicators and baseline models are shown in Appendix 2.
5. E XPERIMENT A NALYSIS
Table 1 and 2 indicate that the proposed model achieves the best performance compared to other baseline methods on both datasets.
Specifically, the MCNN and MMCNN utilize the CNN model to automatically extract sleep features, while RNN based methods such as DeepSleepNet and
TinySleepNet focus on the temporal context in sleep data, and model the multi-level temporal characteristics of the sleep process for sleep staging.
Further, GraphSleepNet and ST-Transformer simultaneously model the spatial-temporal relation during sleep and achieve satisfactory results.
However, GraphSleepNet and ST-Transformer cannot adequately reason the spatial-temporal relation, which limits the classification performance to a certain extent.
Our Bayesian ST-Transformer uses the multi-head Bayesian relation inference component to infer spatial-temporal relations to better model spatial and temporal relations.
Therefore, the proposed model achieves best classification performance on different datasets.
[Table]
[Table]
5. A BLATION EXPERIMENTS
To verify the effectiveness of each component in the Bayesian spatial-temporal transformer
[Figure]
we conduct ablation experiments to determine the impact of each module on the model’s performance.
Specifically
[Figure]
we design three variants of the Bayesian spatial-temporal transformer, including:
• Bayesian Spatial Transformer (BST), which removes the Bayesian temporal transformer module to determine the impact of modeling temporal relations on model performance. • Bayesian Temporal Transformer (BTT), which removes the Bayesian spatial transformer module to determine the impact of modeling spatial relations on model performance.
• Spatial-Temporal Transformer (STT), which removes the relational inference component to determine the impact of Bayesian relational inference on model performance.
Figure 2 demonstrates that the performance of the variant models degrades after removing certain component or module.
Among them, the removal of the relational inference component has the greatest impact on performance [Figure] which shows the importance of introducing relational inference to the sleep staging task.
In addition, only modeling the spatial relation or temporal relation of the data also lead to a decrease in performance.
It can be seen that modeling the spatial and temporal relation is helpful for the sleep staging task
[Figure] [Figure]
5. V ISUAL A NALYSIS
To verify the proposed Bayesian relational inference module can infer the spatial-temporal relations during sleep, we visualize and analysis the generated relational inference graphs
. 5.5. V ISUAL A NALYSIS OF S PATIAL R ELATION I NFERENCE
Some researches have shown that the functional connectivity of the brain varies during different sleep stages (Nguyen et al., 2018).
In order to analyze the role of the Bayesian spatial inference component of our model, we visualize the spatial relation intensity graph between EEG signal channels at different sleep periods, as shown in Figure 3.
The position of the nodes in the figure represents the position of the electrodes that output the EEG signals, and the edge is the relation intensity between the each two electrodes.
We notice that during the NREM period, brain connectivity is significantly stronger in light sleep (N1, N2) than that during deep sleep (N3).
It has been revealed that during light sleep, cerebral blood flow (CBF) and cerebral metabolic rate (CMR) are only about 3% to 10% lower than those of wakefulness while during deep sleep, these indexes have a significant decrease of 25% to 44% by previous study (Madsen & Vorstrup, 1991).
Synaptic connection activity is directly correlated with CBF and CMR, which is consistent with our connection intensity graph.
Madsen also reported that the level of brain synaptic activity during REM period is similar to that of the wake period, which matches our experimental findings.
high low
N
N
WAKE REM N tt- t t+ t+ N tt- t t+ t+
Other to N tt- t t+ t+
N tt- t t+ t+
Other to N tt- t t+ t+
N tt- t t+ t+
Other to N3 tt- t t+ t+
Wake tt- t t+ t+
Other to Wake tt- t t+ t+tREM tt- t t+ t+
Other to REM
. 5.5. V ISUAL A NALYSIS OF T EMPORAL R ELATION I NFERENCE
To analyze the contribution of the Bayesian temporal inference component of our model for sleep staging, we visualize the time slice relation intensity graphs as shown in Figure 4.
Figure 4 shows the temporal intensity graphs for time during unchanging period and when performing sleep transition.
In each graph, the nodes represent time slices while edges represent the strength of the relation between time slices.
The left graph of each graph pair (e.g. N1 marked green) represents the five time slices are in the same sleep period; the one on the right (e.g. Other to N1 marked blue) represents that sleep period transition occurs between the t − 1 and t time slices.
Figuredemonstrates that the temporal intensity learned by the model is stronger during the unchanging period than that when sleep transitions occur.
Previous studies have shown that the stability of the unchanging period is stronger, and sleep instability is the basis of sleep transition (Bassi et al., 2009), which is consistent with our experimental results.
Figure 4 also reports that when sleep transition occurs, the relation intensity between the t and t + 1 time slices are usually stronger.
Similarly, temporal intensity between the t − 2 and the last three time slices are usually weaker than that during the unchanging period, which is conducive to the interpretation of the sleeping transition.
C ONCLUSION
We propose a novel Bayesian spatial-temporal transformer model for sleep staging.
To our best knowledge, this is the first attempt to combine Bayesian relational inference with spatial-temporal transformer for sleep staging.
Our BSTT constructs spatial-temporal relations through Bayesian relational inference and applies the transformer architecture to capture brain spatial-temporal features for sleep staging.
The results show that BSTT can effectively improve the model performance and achieve state-of-the-art results.
In addition, visual analysis presents that the relation intensity graphs generated by the Bayesian relation inference have certain interpretability, which is consistent with the existing research and helps to reveal the potential working mechanism of our model.
Besides, the proposed BSTT is a general-framework which can inference the spatial-temporal relation of EEG data and perform satisfactory data forecasting.
In the future, the proposed method can be used for other EEG tasks, such as emotion recognition or motor imagery classification.
R EPRODUCIBILITY S TATEMENT AND E THICS S TATEMENT
We provide an open-source implementation of our BSTT and other baseline models.
Please check the Appendix 6 for the link.
The authors do not foresee any negative social impacts of this work.
All authors disclosed no relevant relationships
C ALCULATION DETAILS A.1. C ALCULATION OF EVALUATION INDICATORS
We use accuracy (ACC), F1 Score, and KAPPA to evaluate the our BSTT model and baseline models, the specific calculations are follows:
[Equation]
where TP C is the true positive for class C , N is the total number of recordings, F 1 C is the F1 score of per class, C is the number of sleep stages, P o is the observed agreement, and P e is the agreement by chance.
A.1. C ALCULATION OF P OSITION EMBEDDING MATRIX
We use the position embedding function in groundbreaking work(Vaswani et al., 2017) to calculate the position embedding of Bayesian transformer module.
The calculation of position embedding matrix is defined as:
[Equation]
A. B ASELINE D ETAILS
We compare the Bayesian spatial-temporal transformer with the following baselines:
• MCNN (Chambon et al., 2018): a deep learning sleep staging method utilizing multivariate multimodal PSG signals.
• MMCNN (Jia et al., 2020a): a multi-scale convolutional neural network for EEG signal classification.
• MLP+LSTM (Dong et al., 2017): a mixed neural network, which combines multilayer perceptron (MLP) and LSTM.
• DeepSleepNet
(Supratak et al., 2017): a mixed sleep staging method utilizing CNN and Bi-LSTM.
• TinySleepNet (Supratak & Guo, 2020): an efficient and lightweight EEG sleep staging network.
• U-Time
(Perslev et al., 2019): a temporal fully convolutional network based on U-Net architecture for sleep staging.
• GraphSleepNet
(Jia et al., 2020b): a spatial-temporal graph convolutional neural network that can adaptively learn spatial-temporal features.
• Spatial-Temporal Transformer
(Phan et al., 2022): a sleep staging method using transformer to extract spatial-temporal features.
A. K EY T HEOREMS A.3. T HEOREM
Let N (cid:0) µ, σ 2 (cid:1) denotes a Gaussian distribution with µ < 1 / 2 , and let B ( n, λ ) denotes a Binomial distribution with n → + ∞ and λ → 0 , where n is increasing while λ is decreasing.
There exists a real constant m such that if m =
nλ
and if we define:
[Equation]
2 ( x ) , where x ∈ (0 , 1) we have that: f 1 ( x ) attains its minimum on the interval (0, 1) and f 2 ( x ) − f ∗ 2 is bounded on the interval (0 , √ 2 / 2 − 1 / 2) , with:
[Equation]
, where l = 2 σ1 − 2 µ
Let m i,j = nλ i,j , with Theorem 1, direct parameterization of both the infinite parameter n and the near-zero parameter λ i,j can be avoided, which in turn allows for the re-parametrization trick of (Kingma & Welling, 2013) to be used. This trick draws samples from such Binomial distribution via its Gaussian proxy N ( m i,j , m i,j (1 − m i,j )) . Proof of Theorem 1.
The derivative of the function f 1 ( x ) over x can be written as:
[Equation]
We set it as 0 and solve for x , giving
[Equation]
where l =
[Equation]
[Equation]
Let x = nλ , the function f 2 ( x ) can be written as:
[Equation]
Let g ( nλ ) =
lim λ
→ 0 f 2 ( nλ ) , we have
[Equation]
Let z = √ 1 − nλ , we have
[Equation]
The derivative of function g ( z ) over z can be written as:
[Equation]
Given that z ∈ (0 , 1) , we have g ′ ( z ) < 0 .
Then g ( z ) attains its minimum 1 when z approaches 1.
Equivalently, f 2 ( n ) attains its minimum 1 when nλ approaches 0.
Considering Eq.(22), we find that nλ is bounded on (0, 1/2) if µ < 1 / 2 , We then calculate the difference between f 2 ( n ) and its minimum.
It can be written as
[Equation]
Let m =
nλ , the derivative of function ∆ f 2 ( m ) over m can be written as:
[Equation]
Then ∆ f 2 ( m ) is monotonically increasing over (0, 1/2).
Therefore ∆ f 2 ( m ) is bounded on (0 , √ 2 / 2 − 1 / 2) .
A.3. V ARIATIONAL INFERENCE IN B AYESIAN RELATION INFERENCE
Inspired by the variational RNN (VRNN) (Chung et al., 2015), we introduce variational inference to optimise our Bayesian relation inference component. VRNN introduces a latent variable z i,t to capture the uncertainty of input features at time t .
Such a latent variable is assumed to have a
Gaussian prior distribution p ( z i,t | h i,t − 1 ) dependent on the previous RNN hidden state h
i,t − 1 .
Its posterior distribution is approximated by a variational distribution q
( z i,t | x i,t , h i,t − 1 ) , allowing the use of the evidence lower bound (ELBO) for joint learning and inference.
It can be written as:
[Equation]
where Z is latent variables.
In Bayesian relation inference component, random variables ˜
A and
S can completely determines the random process in Bayesian relation inference as a whole.
The
Eq.(23) in VRNN can be rewritten as follow:
[Equation] [Equation] A.3. T HEOREM Suppose we are given two Binomial distributions, B ( n, λ ) and B (cid:0) n, λ 0 (cid:1) with n → + ∞ , λ
→and λ → 0 , where n is increasing while λ and λ 0 are decreasing.
There exists a real constant m and another real constant m (0) , such that if m = nλ and m (0) =
nλ (0) and if λ > λ (0) , we have:
[Equation]
By Theorem 2, we have a closed-form solution that is irrelevant to n for the ELBO.
The specific proof is as follows:
Let m = nλ and m (0)
= nλ (0) , we have:
[Equation]
We then take the right part,
[Equation]
By Taylor series’ theorem with Lagrange remainder, g can be written as: [Equation] =
[Equation]
− λ
[Equation]
There exists a θ ∈ (0 , 1) such that,
[Equation]
Given that n → ∞ and x ∈ (0 , 1) , we have R ′ j =2 ( x ) > 0 .
Therefore, R j =2 ( x ) is monotonically increasing over (1 , 1) .
Since λ > λ 0 , we have:
[Equation] [Equation]
We then seek to prove:
[Equation]
Let f ( x ) = n ( n − 1) x 2 / 2 − nx + 1 , we have:
[Equation]
Here, f ( x ) is an U-shaped parabola with axis x = 1 / ( n − 1) .
Inspired by previous work (Huang et al., 2020) and Theorem 1, we have nλ < 1 / 2 , then we have λ (0) <
λ < 1 / ( n − 1) , then f ( x ) is monotonically increasing over the support of λ (0) 0and , namely:
[Equation]
[Equation]
With Eq.(26) and Eq.(27), g can be written as:
[Equation]
Similarly, let h ( x ) = 1 − x + x 2 / 2 .
It is an U-shaped parabola with axis x = 1 such that:
[Equation]
Then we have:
[Equation]
Combining Eq.(25) and Eq.(28) concludes the proof.
A. P IPELINE GRAPH OF B AYESIAN SPATIAL - TEMPORAL TRANSFORMER
Figure 5 shows the overall of our Bayesian spatial-temporal transformer(BSTT).
The input of BSTT is multi-channel EEG signals with time context.
The embedded features first pass through the
Bayesian spatial transformer module to infer the spatial relations among channels and model their spatial features. The Bayesian temporal transformer module infer the temporal relations and model the temporal features. E m b e dd i n g L
[Equation]
er
B ay e s i a n
S p a t i a l T r a n s f o r m er
B ay e s i a n
T e m p o r a l T r a n s f o r m er
C l a ss i f
i c a t i o n
L ay er P re d i c t i o n Ŷ EE G S
i g n [Figure] l s
[Figure]
A. E XPERIMENT CONFIGURATION AND DETAILS A.5. T RAINING D ETAILS
The details of the resources for training and the versions of the software are provided in Table 4.
The details of the hyperparameters are provided in Table
[Table]
[Table]
[Table]
A. S UPPLEMENTARY E XPERIMENT R ESULTS A.6. S UPPLEMENTARY V ISUAL A NALYSIS
Figure 6 shows the intensity graphs of the temporal relation during some specific periods.
Specifically, we select the temporal intensity graphs when performing mutual transitions between REM period and N2 period and during REM period for visual analysis.
The proposed model pays more attention to the longer-interval relations when interpreting the duration of the REM period.
The
AASM standard states that the rule of continuous interpretation in REM period requires no rapid eye movement or significant waveform changes in one or more epochs after that (Berry et al., 2012).
Hence, sleep specialists tend to focus on the long-term relation intensity and the association of epochs following the target sleep epoch, which is consistent with our experimental results.
Furthermore, in our model, the process of mutual transition between REM period and N2 period tends to focus on the association between the last two sleep epochs.
The AASM standard mentions that when interpreting the mutual transition between REM period and N2 period, it is necessary to ensure that there is no additional transition of the subsequent EEG signals, otherwise it may be reinterpreted back to the original sleep period (Berry et al., 2012).
Our experimental results also confirm this result.
tt- t t+ t+ tt- t t+ t+
REM tt- t t+ t+
REM tt- t t+ t+ tt- t t+ t+
REM to N tt- t t+ t+
REM to N tt- t t+ t+ tt- t t+ t+
N2 to REM tt- t t+ t+ N2 to REM tt- t t+ t+
REM to N tt- t t+ t+
N2 to REM tt- t t+ t+
REM tt- t t+ t+
REM to N tt- t t+ t+
N2 to REM
A.6. A BLATION E XPERIMENT R ESULTS
Table 5 and 6 show the comparison of the results of BSTT ablation experiments on the two datasets.
Specifically, we design three variants of the Bayesian spatial-temporal transformer, including:
• Bayesian Spatial Transformer (BST), which removes the Bayesian temporal transformer module to determine the impact of modeling temporal relations on model performance. • Bayesian Temporal Transformer (BTT), which removes the Bayesian spatial transformer module to determine the impact of modeling spatial relations on model performance.
• Spatial-Temporal Transformer (STT), which removes the relational inference component to determine the impact of Bayesian relational inference on model performance. [Table] [Table] A
[Table] [Table]
C OMPARISON OF S OME M AIN B ASELINES WITH BSTT AT C LASS F1 S CORE
Table 7 and 8 show the comparison of class F1 score of BSTT and some main baseline models. [Table]
[Table]
ST-Transformer 91. BSTT (Our) 90. A.6. R ELATION I NTENSITY G RAPHS OF B AYESIAN R ELATION I NFERENCE C OMPONENT To further demonstrate the effectiveness of our proposed Bayesian spatial-temporal relation inference component, we display other spatial-temporal relationship intensity matrices generated by this component.
Evidently, brain spatial connectivity is very strong during Wake period and REM period, whereas brain spatial connectivity is gradually weakened during NREM periods and N3 period comes to the weakest.
Besides, the temporal relationship intensity matrices also have interpretability.
For example, when the sleep transition from N2 to N3 occurscomponent focus more on the relation between the second and third time slices and fourth time slices.
Same as sleep spacialists, the proposed component is more concerned with the persistence of the sleep signal after the transition has occurred during this period.
We replace the self-attention component in traditional transformer with Bayesian relation inference component we proposed. In traditional transformer, multi-head self-attention is proposed to generate
[Figure]
[Figure]
N1 to N N1 to N
N1 to REM
N1 to Wake
N2 to N N2 to N
N2 to Wake
N3 to N N3 to N
N3 to REM
N3 to Wake
REM to N
REM to N
REM to Wake
Wake to N
Wake to N Wake to N
Wake to REM high
[Figure]
attention graph from different angles and improve the performance.
Inspired by multi-head selfattention component, we propose multi-head Bayesian relation inference component to better infer spatial-temporal relation and generate relation graph.
The calculation is defined as follows:
[Equation]
Each head represents a single Bayesian relation inference component.
We have tried some values of n , and finally chose n = 3 to train the model. Table 9 shows the result on the ISRUC dataset when we set n to different values.
[Table]
A. C ODE AND D ATASET
The code of BSTT is available at: https://github.com/AnonymousBSTT/BSTT/BSTT.
The open-source dataset ISRUC is available at: https://github.com/AnonymousBSTT/BSTT/ISRUC.
The list of open-source baseline model code used in the experiments is as follows:
• DeepSleepNet: https://github.com/akaraspt/deepsleepnet • TinySleepNet: https://github.com/akaraspt/tinysleepnet • GraphSleepNet: https://github.com/ziyujia/GraphSleepNet • U-Time: https://github.com/perslev/U-Time
For the baseline models which are not publicly available, we implement and release the source code of them in Github according to the corresponding papers.
The list of them is given as follows:
• MCNN: https://github.com/AnonymousBSTT/BSTT/Baselines/MCNN • MMCNN: https://github.com/AnonymousBSTT/BSTT/Baselines/MMCNN
• MLP+LSTM: https://github.com/AnonymousBSTT/BSTT/Baselines/MLP • ST-Transformer: https://github.com/AnonymousBSTT/BSTT/Baselines/STT