BSTT: A Bayesian Spatial-Temporal Transformer for Sleep Staging
Authors: Yuchen Liu
Ziyu Jia
Venue: ICLR 2023 poster (2023)
Keywords: Spatial-Temporal Transformer
Sleep Staging
Bayesian Deep Learning

Abstract: 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-of-the-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.


Published as a conference paper at ICLR 2023
BSTT: A BAYESIAN FORMER FOR SLEEP STAGING
TRANS-
Yuchen Liu
Ziyu Jia"
Institute of Automation, CAS University of Chinese Academy of Sciences Institute of Computing Technology, CAS yuchen. www.....iootttttooooooo.o.oooooo com.
Institute of Automation, CAS
University of Chinese Academy of Sciences j jia z yuy out.ook..........o com
ABSTRACT
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 challenge. In particular, existing methods cannot adap- tively 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-templl ing. Our model is able to adaptively infer brain spatial- relations during sleep for spatial-temporal feature modeling through well-designed Bayesian re- lation inference component. Meanwhile, our model also includes spatial trans- former for extracting brain spatial features and a temporal transformer for cap- turing temporal features. Experiments show that our BSTT outperforms state-of- the-art baselines on ISRUC and MASS datasets. In addition, the visual analysis shows that the spatial- -emporll relations obtained by BSTT inference have certain interpretability for sleep staging.
transformer (BSTT), for sleep stag-
INTRODUCTION
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 known as sleep transition rules, to help sleep specialists identify sleep stages when sleep transitions occur. However, artificial sleep staging takes a long time, and the classification results are greatly affected by professional level and subjectivity (Supratak 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 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 al., 2019; Jia et al., 2022a;b). 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) 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; Jia al., 2020a). 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, is difficult for these methods explicitly demonstrate the relation of
"Corresponding author.

---

Published as a conference paper at ICLR 2023
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, propose a novel model called Bayesian spatial-temporal trans- former (BSTT) for sleep staging. The proposed model integrates the transformer and Bayesian relation inference in 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 re- lations of objects and generate the relation intensity graphs. Specifically, the main contributions of our BSTT are summarized as follows:
We design Bayesian relational inference component which can adaptively infer spatial- temporal relations of brain during sleep in the service of capturing spatial-temporal rela- tions.
We apply the spatial-temporal transformer architecture to simultaneously model spatial- temporal relations. 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-ar in multiple sleep staging datasets. The visual analysis shows that our model has a certain degree of interpretability for sleep staging.
2
RELATED WORK
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 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 al., 2016; Chambon et al., 20. For example, Chambon et al. propose a convolutional neural network that can extract temporal-invarian 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., 2021). 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 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 al., 2017). For example, Suratak 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 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 characteris- tics of sleep data. For example, Jia et al. propose 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 multi-view spatial-temport graph convolutional network based on domain generalization, which models the multi-view-baan based spatial characteristics of the brain (Jia et al., 2021a).
Although the above models achieve good classification performance, these models do not ade- quately model spatial-temporal properties effectively reason and capture spatial-temporal rela-
2

---

Published as a conference paper at ICLR 2023
tions. Therefore, our method attempts to model spatial-temporal relations using Bayesian inference, combined with state-of-thee teefthhertt transformer architectures for sleep staging.
3 PRELIMINARIES
The proposed model processes data from successive T sleep epochs and predicts the label of the epoch in the middle. Each sleep epoch is defined as I E RCXN, 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 sleep epoch. The input sequence of sleep epochs defined as a {x1, ,2...., ) where z; denotes a sleep epoch (i [1. (1......,T] T])
and 7 is the number of sleep epochs.
The sleep staging problem defined as: learning an artificial neural network F based on Bayesian spatial-temporal transformer, which infer the spatial-temporal relations of the input sleep epoch sequence x and map to the corresponding sleep stage Y. where Y the classification result of the middle epoch Im- According to the AASM standard, each Y E 2, is matched to five sleep stages W, NI, N2, N3, and REM, respectively.
BAYESIAN SPATIAL-TEMPORAA TRANSFORMER
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-tempont relations of the brain based on Bayesian inference method.
Design the Bayesian transformer architecture while capturing the tentiall features of the brain.
Integrate Bayesian relational inference components and transformer architecture into a uni- fied framework which can stage sleep period effectively.
The overall model is carefully designed to accurately classify different sleep stages.
4.1 ARCHITECTURE
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 module. Specifically, the Bayesian spatial-temporal trans- former includes a Bayesian spatial transformer and Bayesian temporal transformer. The Bayesian spatial transformer reason about spatial relations in the brain and capture spatial features. The Bayesian temporal transformer can reason about the temporal relations of consecutive sleep epochs and capture temporal features. Finally, the predictions of different sleep stages performed by the classification layer.
4.1.1 BAYESIAN RELATION INFERENCE
Capturing the spatial- tempoaltemmporl relations of brain signals during sleep can better serve sleep staging task. However, due to the difficulty in inferring the ppatial-temporal relations of sleep, current re- searches are insufficient for modeling the spatial-temporal relations. Inspired by deep graph random process (DGP) proposed 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 in- fer relations between each pair of object nodes and build relation intensity graphs. In this article, an object node represents the embedding of EEG channel the embedding of 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 sleep relation intensity graph starts with generating edge embeddings. Specifically, the node embeddings of the object are spliced and generated into edge
3

---

Published as a conference paper at ICLR 2023
EEG Signals
Embedding Layer
Bayesian
Bayesian temporal transformer
Postioo
spatial transformer
Posiiio
Finhedling
LayerNorm
NoyerNorm
Multi-head Bayesian
Multt-heed Bayeslan
Spatial relational inference
Reidmal
temporal relational inference/ T Temporal relation intensity graph based features
Resihale
Spatial relation intensity
or graph based features
LayerNorm
LayerNorm
imear
Linear
Rel. U ropout
ReL,U,Dropon
inear
inear
Classification Layer
Figure 1: The overall architecture of the proposed Bayesian spatial-temporal transformer for sleep staging. BSTT includes two Bayesian transformer modules, spatial Bayesian transformer and temporal Bayesian transformer. For each transformer module, the input features are passed through the position embedding and layernorm layer. Then the multi-head Bayesian relation inference com- ponent infers the object's 's spatial or temporal relation and captures the spatial-tempor features. The residual connection used prevent overfitting and gradient disappearance. Here, means add.
embeddings. We first apply linear neural network fo generate edge embeddings:
E. = fo E, (i,j]) (1,j E [1, && i!=
(1)
where E, E BB represents the input n node objects, Ja E 2B represents the splicing of two node embeddings, and E. E Bx.., x represents the generated edge embeddings.
Step 2: Edge embedding coupling. Coupling is one of the key steps in relation inference, the purpose to obtain a summary graph of the relation between object nodes (Huang et al., 2020). We assume that the edges of the summary graph Ma, are a summary of an n > Do A - Binomial distributions, which means m.j.j B(n, due to the uncertainty of spatial-temporal relations of brain. Drawing from the Virtual Recurrent Neural Network (VRNN) model, the parameters of the approximate posteriors However, in contrast to VRNN, Bayesian relation inference component contains an approximate posterior (m Ex), whose inference and sampling cannot be solved in computationally feasible manner due to its infinite n. By De SoivreLapplace theorem (Sheynin, 1977) and DGP (Huang et al., 2020), can subject these edge embeddings a coupling transformation follows:
estimated using Recurrent Neural Network (RNN) encode features.
nij LLmean (Ee,s) +
ō = List Eo,,))
(3)
+ 2m; 0
4n 3.
mij
(4)
2
where ((-) is softplus function, E. is the edge embedding generated in the first step. is very small constant, (.)ean.) and Lsta (-) are implemented by neural networks for estimating the mean and standard deviation respectively, mi,5 M is the approximation of Binomial edge variable in the summary graph, and M IR N. is the approximation of summary graph which strengthens the representation real spatial or 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 relation intensities of brain temporal- spatial network during sleep is sparse based existing research (Razi et al., 2017). Therefore,
4

---

Published as a conference paper at ICLR 2023
generating sparse graph based edge embedding 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 sparse sleep relation intensity graph G. The specific calculation defined as follows:
0ijj
= [m]]] x Eiji + miji fij
(5)
Si.j ( m mean
āi,j + [ā (o..)
(6)
ØU SL x āij
(7)
0ij = (L (ā,,))
(8)
where mij M is the approximation of the edges of the summary graph obtained in Step 2, and e' are the standard Gaussian random variable of the same dimension M, ā E BBxN. is the Gaussian edge representation, S E TR BxNo is the task-related Gaussian variable, øj,j is calculated in (3), std is the standard deviation, mean is the mean value, ã is the Gaussian transformation map, a E R BxxNe is the final sleep relation intensity graph, <(.) is softplus function, and L(-) is linear function. Afterwards, we utilize attention mechanism-baser method 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:
Eont fGAL (E,
(9)
where (-) is the graph attention layer, En is the node embeddings of the input object, is the sleep relation intensity graph, and Eout E fRxxxN N, Xout is the output node embeddings.
4.1.2 LEARNING OF BAYESIAN RELATION INFERENCE
We adopt variational inference 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. 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:
Σ [KTL KT. ( (A,S)) Xo:) lP (A,S) Xo:)) FAS log (Y, A.S]] }
M
(10)
where S is the task-related Gaussian variable, A is the Gaussian graph embedding, (A,S Xoa) is the prior distribution, and S Xo: is the posterior distribution. Since every variable in S are affected by À in (6), the KL term can further written as:
Σ
KL B(n, B| n, (x.
(LIEÉ
(11)
+ +Bai KL (ān i Pi.j ārg 1ootilll
(0) pig.aa āng * x (0)
Obviously the second term be calculated, while the first term is tough to calculate because n + Do. According to Theorem provided in DGP (Huang et al., 2020), we can convert into an easy-to-solve value approximate the calculation.
4.1.3 BAYESIAN TRANSFORMER MODULE
Transformer shows convincing results in various sequence modeling tasks (Li et al., 2021; Luo et al., 2021: Zhou et al., 202 However, traditional transformer does not have the ability to reason the relation between each pair of object nodes. This results in 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 proposed can infer
5

---

Published as a conference paper at ICLR 2023
peatial-temporre relations 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 -emiraall features of sleep EEG data.
Bayesian spatial transformer. To better construct the spatial functional connectivity of the brain and capture spatial features, design the Bayesian spatial transformer. It contains components: a Bayesian spatial relation inference component and a position feed-forward network, which the core component is 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:
S PPP
(12)
where E RB( the input spatial node embedding, Pep R(BXNN)XNXX the po- sition encoding matrix, and Š E (B( 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.
We design a multi-head spatial Bayesian relation inference component which 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 1.1.1. The node embeddings with positional encoding are encoded as embeddings with spatial features after passing through the multi-head spa- tial Bayesian relation inference component and the feed-forward neural network layer:
Š' = fFNN (FASRI ($) (s))
(13)
where S is the input node embedding, f(SSR((-) is the Bayesian spatial relation inference module which inferences the spatial relation, fFNN (-) is the feed-forward neural network layer, and S' BBXN 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 the temporal relation intensity graph based features as follows:
T= S' PPP
(14)
T' JENN (FETRI
(T))
(15)
where fETRI(-) is the Bayesian temporal relation inference component, fFNN (-) is the forward neural network layer, and T' IR Bx., N.. is the temporal relation intensity graph based features. The classification results are generated by T" after passing through linear classification layer fc:
8 T''
(16)
where Y is the classification result of BSTT.
5 EXPERIMENTS
To verify the effectiveness of the Bayesian spatial-temporal transformer, we evaluate it on the In- stitute of Systems and Robotics, University of Coimbra (ISRUC) and Montreal Archives of Sleep Studies-SSt SS3 (MASS-SS3) dataset.
5.1 DATASET
ISRUC dataset contains the PSG recordings from 100 adult subjects. Each PSG recording contains EEG channels, 6 EOG channels, 3 EMG channels, and 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 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 also motion artifacts at the beginning and end each subject's 's recording which are marked as unknown. We follow the previous study and remove these recordings (Supratak et al., 2017).

---

Published as a conference paper at ICLR 2023
5.2 EXPERIMENT SETTINGS
We evaluate our model using k-fold cross-subject validation to ensure that the experiment results are correct and reliable. We k order to 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. Specifi- cally, on the ISRUC dataset, we use all 6 EEG channels for experiments. In the MASS-SS3 dataset we use channels of EEG signals. To comprehensively evaluate the Bayesian spatial-temporal transformer model and all baseline methods, we accuracy (ACC), Fl Score, and KAPPA to eval- uate the models. Specific information of the evaluation indicators and baseline models are shown in Appendix 2.
5.3 EXPERIMENT ANALYSIS
Table 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 automatically extract sleep features, while RNN based methods such DeepSleepNet and TinySleepNet focus on the temporal context in sleep data, and model the multi-level temporal char- acteristics of the sleep process for sleep staging. Further, GraphSleepNet and ST-Transformer simul- taneously model the spatial-temporal relations during sleep and achieve satisfactory results. How- ever, GraphSleepNet and ST-arosffrrerr cannot adequately reason the ppatial-temporal relations, which limits the classification performance to certain extent. Our Bayesian ST- TT-Transforme uses 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 1: Comparison of Bayesian spatial-tempord transformer and baselines on ISRUC dataset
Method
ACC(%)
F1 Scroe %)
KAPPA (%)
MCNN
78.23
76.37
71.79
MMCNN
76.83
78.93
74.17
MLP+LSTM
76.08
72.87
69.16
DeepSleepNet
75.71
73.16
68,82
TinySleepNet
76.92
75.15
70.26
U-Time
75.52
71.04
67.92
GraphSleepNet
80.18
78.19
74.54
ST-Transformer
80.35
78.05
74.71
BSTT (Our)
81.96"
80.30*
76.78*
"indicates the significant differences between our model and other models (p < 0.05).
Table 2: Comparison of Bayesian spatial-temporal transformer and baselines on MASS dataset
Method
ACC(%)
F1 Score( %)
KAPPA(%)
MCNN
86.31
80.79
81.26
MMCNN
80.15
71.03
69.42
MLP+LSTM
86.31
81.55
80.12
DeepSleepNet
85,92
79.81
79.09
SleeSNe
83.30
77.40
79,95
U-Time
85,23
78.36
77.94
GraphSleepNet
88.36
83.25
83.30
ST-Transformer
88.64
84.53
83.16
BSTT (Our)
89.50*
85.00*
84.37'
"indicates the significant differences between our model and other models (p < 0.05).

---

Published as a conference paper at ICLR 2023
5.4 ABLATION EXPERIMENTS
To verify the effectiveness of each component in the Bayesian spatial-temporal transformer, we conduct ablation experiments determine the impact of each module on the model's performance. Specifically, we design three variants of the Bayesian spatial- transformer, including:
Bayesian Spatial Transformer (BST), which removes the Bayesian temporal transformer module to determine the impact of modeling temporal relations 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 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 performance, 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 decrease in performance. It can be that modeling the spatial and temporal relation helpful for the sleep staging task.
0.84
0,800
0.8050.800
0.003
0.80
STT(No Bepesian inference)
0.771
0.78
0.763
BTTINo spattal)
0.750
BST(No tamporal)
0.74
BSTT[Our)
ACC
Score
KAPPA
Figure 2: Ablation experiment results of Bayesian spatial-tempont transformer on ISRUC dataset.
5.5
VISUAL. ANALYSIS
To verify the proposed Bayesian relational inference module can infer the seatial-temporal relations during sleep, we visualize and analysis the generated relational inference graphs.
5.5.1 VISUAL ANALYSIS OF SPATIAL RELATION INFERENCE
Some researches have shown that the functional connectivity of the brain varies during different sleep stages (Nguyen 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 chan- nels different sleep periods, as shown Figure 3. The position of the nodes in the figure repre- sents 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 significant decrease 25% to 44% by previous study (Madsen Vorstrup, 1991). Synaptic connection activity directly correlated with CBF and CMR, which is consistent with connection intensity graph. Madsen also reported that the level of brain synaptic activity during REM period similar to that of the wake period, which matches our experimental findings.
8

---

Published as a conference paper at ICLR 2023
Fpt
Fl
Fp
F4
F
Fpb
Fp2
Fl
Pp
FL
R
Fo
Fol
F4
F8
03
P4
Plps
g48
Plps
Plpy
P4P8
10
0%
p2
00
06
N1
N2
N3
REM
WAKE
Figure 3: The graph shows the average of the brain spatial intensity over time. The spatial relation during the WAKE, REM and N1 period is strong, while the that during the N2 and N3 period is weak.
5P2
Other
[-]
t+2
H-]
612
1-1
t12
t-1
11
BBM
Otber
to RMM
Figure 4: Intensity graphs of the temporal relation during different sleep periods and when sleep transitions occur.
5.5.2 VISUAL ANALYSIS OF TEMPORAL RELATION INFERENCE
To analyze the contribution the Bayesian temporal inference component of our model for sleep staging. visualize the time slice relation intensity graphs as shown in Figure 4. In each graph, the nodes represent time slices while edges represent the strength of the relation between time slices (Here, time slice represents the manually divided EEG signals of 30s duration.) The left graph of each graph pair (e.g. N1 marked green) represents the five time slices (nodes) are in the same sleep period. The one on the right (e.g. Other to N1 marked blue) represents the and time slices are in one sleep period while the other three are in another sleep period (Sleep period transition occurs between the -- and time slices.) In Figure 4, for each pair of graphs, the edges the left graph are, average, darker than those in the right one. This means that model tends to think that EEG signals are more closely related to each other during a single sleep stage. 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 consistent with our experimental results. Figure also reports that when sleep transition occurs, the relation intensity between the and + 1 time slices are usually stronger. Similarly, temporal intensity between the f 2 and the last three time slices usually weaker than that during the unchanging period, which is conducive the interpretation of the sleeping transition.
6 CONCLUSION
We propose a novel Bayesian knowledge, this is the first attempt to combine Bayesian relational inference with spatial-tempores transformer for sleep staging. Our BSTT constructs spatial-temporal relations through Bayesian re- lational 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, which can inference the ppatial-temporal relations of
transformer model for sleep staging. To our best
the proposed BSTT is EEG data and perform satisfactory data forecasting. the future, the proposed method can be used for other EEG tasks, such emotion recognition motor imagery classification.

---

Published as a conference paper at ICLR 2023
7 REPRODUCIBILITY STATEMENT AND ETHICS STATEMENT
We provide an open-source implementation of our BSTT and other baseline models. The code of BSTT is available at: https: / BSTT. Please check the Appendix 7 for links of the baseline methods.
t.ub.
The authors do not foresee any negative social impacts of this work. All authors disclosed no relevant relationships.
REFERENCES
Alejandro Bassi, Ennio A Vivaldi, and Adrián Ocampo-Gareés. The time course of the probability of transition into and out of rem sleep. Sleep, 32(5):655-669 2009.
Stanislas Chambon, Mathieu N Galtier, Pierrick Arnal, Gilles Wainrib, and Alexandre Gramfort. A deep learning architecture for temporal sleep stage classification using multivariate and multi- modal time series. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 26(4): 758 7898769,22 2018.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Ben- gio. A recurrent latent variable model for sequential data. Advances in neural information pro- cessing systems, 28, 2015.
Hao Dong, Akara Supratak, Wei Pan, Chao Wu, Paul M Matthews, and Yike Guo. Mixed neural network approach for temporal sleep stage classification. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 26(2):324 -33, 2017.
Emadeldeen Eldele, Zhenghua Chen, Chengyu Liu, Min Wu, Chee-Keong Kwoh, Xiaoli Li, and Cuntai Guan. An attention-basee deep learning approach for sleep stage classification with single- channel eeg. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 29:809-818 -18, 2021.
Luay Fraiwan, Khaldon Lweesy, Natheer Khasawneh, Heinrich Wenz, and Hartmut Dickhaus. Auto- mated sleep stage identification system based time- frequenyy analysis of a single eeg channel and random forest classifier. Computer methods and programs in biomedicine, 108(1):10-19) 2012.
Antoine Guillot and Valentin Thorey, Robustsleepnet: Transfer learning for automated sleep staging at scale. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 29: 19:1441-1451, 2021.
Hengguan Huang, Fuzhao Xue, Hao Wang, and Ye Wang. Deep graph random process for relational- thinking based speech recognition. In International Conference on Machine Learning, pp. 4531- 4541. PMLR, 2020.
Ritika Jain and Ramakrishnan Angarai Ganesan. An efficient sleep scoring method using visibility graph and temporal features of single-channel eeg. In 2021 43rd Annual International Conference of the IEEE Engineering in Medicine & Biology Society (EMBC), pp. 6306-6309. IEEE, 2021.
Ziyu Jia, Xiyang Cai, Gaoxing Zheng, Jing Wang, and Youfang Lin. Sleepprintnet: multivariate multimodal neural network based on physiological time-series for automatic sleep staging. IEEE Transactions on Artificial Intelligence, ( 13)):248-257 2020a.
Ziyu Jia, Youfang Lin, Jing Wang, Ronghao Zhou, Xiaojun Ning, Yuanlai He, and Yaoshuai Zhao. Graphsleepnet: Adaptive spatial-temporal graph convolutional networks for sleep stage classifi- cation. In IJCAI, pp. 1324 1330, 2020b.
Ziyu Jia, Youfang Lin, Jing Wang, Xiaojun Ning, Yuanlai He, Ronghao Zhou, Yuhan Zhou, and H Lehman Li-wei. Multi-view spatial-temppral graph convolutional networks with domain gen- cralization for sleep stage classification. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 29:1977-1986, 2021a.
10

---

Published as a conference paper at ICLR 2023
Ziyu Jia, Youfang Lin, Jing Wang, Xuehui Wang, Peiyi Xie, and Yingbin Zhang. Salientsleepnet: Multimodal salient wave detection network for sleep staging. arXiv preprint #rXiv:2105.13866 2021b.
Ziyu Jia, Xiyang Cai, and Zehui Jiao. Multi-modal physiological signals based squeeze-and- excitation network with domain adversarial learning for sleep staging. IEEE Sensors Journal, 2022a.
Ziyu Jia, Junyu Ji, Xinliang Zhou, and Yuhan Zhou. Hybrid spiking neural network for sleep elec- troencephalogram signals. Science China Information Sciences, 65(4): 140403, 2022b.
Rajani Khanal. Brain Connectivity During Different Sleep Stages Using EEG and NIRS. PhD thesis, Flinders University, College of Science and Engineering. 2019.
Pengfei Li, Peixiang Zhong, Kezhi Mao, Dongzhe Wang, Xuefeng Yang. Yunfeng Liu, Jianxiong Yin, and Simon See. Act: an attentive convolutional transformer for efficient text classification. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 13261-13269, 2021.
Yunpeng Luo. Jiayi Ji, Xiaoshuai Sun, Liujuan Cao. Yongjian Wu, Feiyue Huang, Chia-Wen Lin, and Rongrong Ji. Dual-level collaborative transformer for image captioning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 2286 2293, 2021.
PL Madsen and S Vorstrup. Cerebral blood flow and metabolism during sleep. Cerebrovascular and brain metabolism reviews, 3(4):281-296, 1991.
Thien Nguyen, Olajide Babawale, Tae Kim, Hang Joon Jo, Hanli Liu, and Jae Gwan Kim. Exploring brain functional connectivity in rest and sleep states: a fnirs study. Scientific reports, 8(1):1-10, 2018.
Mathias Perslev, Michael Jensen, Sune Darkner, Poul Jørgen Jennum, and Christian lgel. U-time: A fully convolutional network for time series segmentation applied to sleep staging. Advances in Neural Information Processing Systems, 32, 2019.
Huy Phan, Fernando Andreotti, Navin Cooray, Oliver Y Chén, and Maarten De Vos. Automatic sleep stage classification using single-channel eeg: Learning sequential features with attention-based recurrent neural networks. In 2018 40th annual international conference of the IEEE engineering in medicine and biology society (EMBC), pp. 1452-1455. IEEE, 2018
Huy Phan, Fernando Andreotti, Navin Cooray, Oliver Y Chén, and Maarten De Vos. Seqsleep- net: end-to-end hierarchical recurrent neural network for sequence-to-sequenee automatic sleep staging. IEEE Transactions Neural Systems and Rehabilitation Engineering, 27(3):400-410, 2019.
Huy Phan, Kaare B Mikkelsen, Oliver Chen, Philipp Koch, Alfred Mertins, and Maarten De Vos. Sleeptransformer: Automatic sleep staging with interpretability and uncertainty quantification. IEEE Transactions on Biomedical Engineering, 2022.
Adeel Razi, Mohamed L Seghier, Yuan Zhou, Peter McColgan, Peter Zeidman, Hae-Jeong Park, Olaf Sporns, Geraint Rees, and Karl J Friston. Large-scale dems for resting-state fmri. Network Neuroscience, (3), 2017.
Vangelis Sakkalis. Review of advanced techniques for the estimation of brain connectivity measured with eeg/meg. Computers in biology and medicine. 41(12):1110-11 2011.
Oscar B Sheynin. Laplace's theory of errors. Archive for history of exact sciences, 17(1):1-61, 1977.
Akara Supratak and Yike Guo. Tinysleepnet: An efficient deep learning model for sleep stage scoring based on raw single-channel eeg. In 2020 42nd Anmual International Conference of the IEEE Engineering in Medicine Biology Society (EMBC), pp. 641 6444. IEFE, 2020.
11

---

Published as a conference paper at ICLR 2023
Akara Supratak, Hao Dong, Chao Wu, and Yike Guo. Deepsleepnet: A model for automatic sleep stage scoring based on raw single-channel eeg. IEEE Transactions on Neural Systems and Reha- bilitation Engineering. 25(11):1998 11119998-000,,, 2017.
Orestis Tsinalis, Paul M Matthews, Yike Guo, and Stefanos Zafeiriou. Automatic sleep stage scoring with single-channel eeg using convolutional neural networks. arXiv preprini arXiv: 16X.0:6600.01683, 2016.
A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, Jones, A. N. Gomez, L. Kaiser, and Polo- sukhin. Attention is all you need. In arXiv, 2017.
Dongdong Zhou, Qi Xu, Jian Wang, Jiacheng Zhang, Guoqiang Hu, Lauri Kettunen, Zheng Chang, and Fengyu Cong. Lightsleepnet: A lightweight deep model for rapid sleep stage classification with spectrograms. In 2021 43rd Annual International Conference of the IEEE Engineering in Medicine & Biology Society (EMBC), pp. 43-46. IEEE, 2021a.
Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang Informer: Beyond efficient transformer for long sequence time-series forecasting, In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 11106-10 111 11115, 2021b.
12


Reviews:
[{"review_id": "svMbXSyYO2", "reviewer": "Reviewer_USPW", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a novel model named BSTT for sleep staging. BSST integrates the transformer and Bayesian relation inference to simultaneously model spatial-temporal relations and effectively capture the spatial-temporal features of the brain. Bayesian relation inference, the core component of BBST, comes in two forms, Bayesian temporal relation inference and Bayesian spatial relation inference and is proposed to infer spatial-temporal relations of brain during sleep in the service of capturing spatial-temporal relations. Experimental results show that BSTT achieves the state-of-the-art in multiple sleep staging datasets.\n", "summary_of_the_review": "It is interesting to infer the space-time relationship during sleep. This work provides some meaningful insights for sleep staging based on deep learning. Some details of the presentation need minor modification.", "strength_and_weaknesses": "*Strength\n\n- The problem of modeling temporal and spatial relations in an interpretable way for sleep staging is well motivated and explained.\n\n- This paper introduces the Bayesian relational inference in detail in terms of graphics, figures and formulas. How data and output flow are explained in the Bayesian relational inference step by step.\n\n- This paper compares some recently proposed baselines, and obtains the results of SOTA on two datasets, and the accuracy is improved by two percentage points.\n\n- The ablation experiments were valuable in demonstrating the impact of the proposed methods.\n\n\n*Weaknesses\n\n- The authors should add a pipline diagram so that readers can understand the overall process.\n\n- In the problem definition of preliminaries, the input data only reflects temporal information and doesn't very reflect spatial information. The problem defined is more like a regular time series problem, which is difficult to directly contact with the spatial-temporal relations emphasized above.\n\n- Some components used in the paper aren't introduced or cited in detail, such as the multi-head spatial Bayesian relation inference component.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "UfXg-g3zJzD", "reviewer": "Reviewer_dJC8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method for sleep staging based on Bayesian spatial-temporal transformer.", "summary_of_the_review": "The paper has some novel contributions but does not outperform existing methods by any significant margin. Moreover, the writeup needs a lot of work. The terms are not explained. One of the two interpretations is either not very convincing or not explained well.", "strength_and_weaknesses": "Strengths:\n- Surpasses state-of-the-art for sleep staging\n- Shows some spacial interpretability\n\n\nWeaknesses:\n- Improvements are minor (<1% for MASS dataset)\n- Figure 1 both transformers are identical.\n- Section 4, the method is not well written. The terminology is not explained.\n- Figure 4 is not explained well and it is unclear how it provides interpretation.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "STTbRoo5meL", "reviewer": "Reviewer_sN5h", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In EEG signals to be used in sleep stage classification, there are multi-channels due to the location of brains (spatial), and previous sleep stages can affect the current sleep stage. This paper considers the Bayesian spatial-temporal transformer to classify the sleep stage. The authors proposed the Bayesian transformer module to be applied to capture the spatial-temporal characteristics of EEG. Experiments with the other algorithms were conducted, showing better performance. Also, visual analysis shows the relationship between channels or previous stages. The interpretation is reasonable in the area of the sleep study.", "summary_of_the_review": "This paper is a good one in the application area. The contribution is somewhat low in the machine learning algorithm and learning representation. Also, I have concerns about more explanation, re-organization, and polishing the proofs.\nThis paper is a good one in the application area. In the machine learning algorithm and learning representation, the contribution is somewhat low. Also, I have concerns of more explanation, re-organization, and polishing the proofs.\n", "strength_and_weaknesses": "Pros:\nThe performance is better than other algorithms, and the experiment is conducted compared with various algorithms. Also, the usage of the Bayesian transformer is interesting. This paper is strong in the application aspects.     \n\nCons:\nIn algorithm aspects, there are many issues. At first, the authors used the Bayesian transformer module. Many people know the transformer model, but some readers cannot be familiar with Bayesian transformers. It is better to add a more comprehensive explanation for the Bayesian transformer and the motivation for using this module. Next, there are some theoretical statements in the paper and Appendix. However, the organization of these materials is not well presented. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9xcXEQq7HAa", "reviewer": "Reviewer_s23K", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a combination of Bayesian relational inference with spatial-temporal transformer for automatic sleep scoring. In addition to testing the new deep learning model's accuracy in sleep scoring, it provides ablation experiment results and visualizations of relation intensity graphs that come from this model as a by product.", "summary_of_the_review": "The idea of Bayesian spatial-temporal transformer is interesting and seems appropriate for sleep scoring. Also the experimental results are very strong. But the paper should be written more clearly.", "strength_and_weaknesses": "Strength: The experimental results are strong, and there are enough details to make the results reproducible with github code. The idea of spatial-temporal transformer (and some level of interpretability with the graphs) is very promising for this application domain.\n\nWeakness: The presentation of this paper can be improved by a lot, especially in section 4.1.1. It's unclear what each variable means, and what their dimensions are. Also in Figure 1, should one of the two transformers be spatial? \nIn results, what were the standard deviation from cross validation? Finally, in Table 8, the results in the last row are bolded even when it's not the highest F1 score compared to the baselines.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster