ConCAD: Contrastive Learning-based Cross Attention for Sleep Apnea Detection

The standard approach to diagnosis sleep apnea requires the patient to sleep overnight at a clinic setup and record the polysomnography (PSG) by various physiological sensors, and then the outputs of PSG are visually inspected by a clinical expert to give a diagnosis [6, 19, 23]. This process is always inconvenient and uncomfortable. Thus, some studies have begun to simplify the procedure of diagnosing sleep apnea by only using a single physiological data, such as ECG [3, 10], EEG [2], and the respiration signal [31]. Among these physiological data, ECG is a less intrusive option and also strongly related to sleep apnea.

To this end, several studies [18, 32] manually extract hand-crafted features and feed them to classifiers (e.g., random forest, support vector machine) for sleep apnea detection. Recently, with the development of deep learning methods, some studies [35, 3] extract RR interval (RRI) and the R-peak envelope (RPE) and build deep learning model to automatically learn representation and detect sleep apnea. Furthermore, several studies [10, 30] develop deep learning models to directly learn features from the raw ECG data and detect sleep apnea in an end-to-end style.

Another line of related work is feature fusion, which aims at combining different features to obtain a more effective representation. The frequently-used feature fusion techniques are concatenation [29], summation [11], and multiplication [39]. However, these operations evenly combine all the features together without considering the importance of each feature. Some of the features gathered will help the model make the right decision, while others can lead to significant misjudgment  [21].

Recently, the usage of attention learning mechanism has shown remarkable performance improvement for different tasks, such as natural language processing [37], image classification [33], and object tracking [8]. The attention mechanism highlights the effective discriminant parts of features while suppressing the redundant parts to a certain degree. To further take advantage of the features extracted from multi-modality inputs, a cross attention mechanism has been proposed to derive an attention mask from different inputs mutually. In [22], the authors use one modality (LiDAR) to generate an attention mask that controls the spatial features of a different modality (HSI). In [16], the authors derive cross attention maps for each pair of class features and query sample features to highlight specific regions and make the extracted features more discriminative.

All of the deep learning methods mentioned above are trained by the cross-entropy loss. The cross-entropy loss is the most commonly-used one in the classification tasks, which calculates the difference between the actual probability distribution of the data and the predicted probability distribution of the model [28]. As we previously introduced, the cross-entropy loss has some limitations. Thus, a supervised contrastive loss is added as an auxiliary regularization to alleviate problems in our proposed framework.

The contrastive loss has recently been widely used in self-surprised learning [14, 24, 5], which aims at clustering the similar data and pushing apart the dissimilar data. A supervised version of the contrastive loss is proposed by [20] to leverage the label information. Their proposed supervised contrastive learning contains two steps: First, the supervised contrastive loss is used to learn a representation to cluster the data from the same class and separate the data from different classes; Second, they froze the model and add a multi-layer perceptron (MLP) as a classifier on its top for the classification task. Recently, supervised contrasting learning has been used for different applications, such as image classification [20], few-shot classification[25], and semantic segmentation [36].

The expert knowledge is summarized by previous researches over the last centuries. In the field of detecting sleep apnea using ECG data, several previous studies [1, 9] have shown that the RR intervals (RRI) and R-peak envelope (RPE) are effective. To prepare the RRI and RPE data, we first detect the locations of the R-peaks by the Hamilton algorithm [13]. Then, we calculate the distance between R-peaks as the RRI and use the amplitudes of the R-peaks as the RPE. Since the RRI can be easily disturbed by unexpected ECG spikes, a median filter is used to eliminate the disturbance as suggested by [7]. Besides, since the number of the RRI or RPE is not always the same by giving a fixed time duration (e.g., 1 min), cubic interpolation was used to resample them to the same length [35].

We are augmenting the ECG, RRI and RPE by two simple approaches: random time shift and reversion. Given the data $\textbf{x}=[x_{0},x_{1},x_{2},\ldots,x_{n}]$, the random time shift will obtain $\textbf{x}_{shift}=[x_{t},x_{1+t},x_{2+t},\ldots,x_{n+t}]$, where $t$ is a randomly-generated number and represents the number of data points to shift. The revision will generate $\textbf{x}_{reverse}=[x_{n},x_{n-1},\ldots,x_{2},x_{1},x_{0}]$. The augmentation is conducted in each batch to provide more positives (i.e., instances with the same label) during batch training, which benefits a more robust clustering of the projection space. The augmentation process is presented as $\mathcal{A}{ug}(\textbf{x})$.

The feature extractor should be designed case by case. In this study, we have three feature extractors, which are used for learning features from the ECG, RRI and RPE separately, and named $\mathcal{F}_{ECG}$, $\mathcal{F}_{RRI}$ and $\mathcal{F}_{RPE}$.

$\mathcal{F}$ consists of 4 convolution blocks. The first three blocks are made of one convolutional layer, one batch normalization layer, one ReLU activation layer, one maxpooling layer, and one dropout layer. The feature map size of the convolutional layer in the first block is chosen to cover data points of two contiguous beats in case that the patterns between beats get missed. The last convolution block does not contain the maxpooling and dropout layer. The module can be represented as $\textbf{x}^{\prime}=\mathcal{F}(\textbf{x};\theta_{{F}})$, where x represents the input data and $\theta_{F}$ denotes the parameters of the module.

Since there are three kinds of data, we have three corresponding extractors, which are $\textbf{x}^{\prime}_{ECG}=\mathcal{F}_{ECG}(\textbf{x}_{ECG};\theta_{{F}_{ECG}})$ for ECG data, $\textbf{x}^{\prime}_{RRI}=\mathcal{F}_{RRI}(\textbf{x}_{RRI};\theta_{{F}_{RRI}})$ and $\textbf{x}^{\prime}_{RPE}=\mathcal{F}_{ECG}(\textbf{x}_{RPE};\theta_{{F}_{RPE}})$ for the expert knowledge-based features. More details of the extractors are described in Appendix 0.A.

Not all the deep features and expert features contribute equally to the classification task. Thus, we design a cross-attention module, $\mathcal{A}_{cross}$, to collaboratively learn their importance and concentrate more on the important ones. The cross-attention is designed to ask the model to concentrate on the particular features, which contribute more to distinguish the instances from different classes. Before computing the cross attention, since the outputs of feature extractors are likely to have different dimensions, we need to project them to the same space by a linear transformation. Given $\textbf{x}^{\prime}\in\mathbb{R}^{m\times{n}}$, the transformation is

where $\textbf{u}\in\mathbb{R}^{{m}}$ and $\textbf{V}\in\mathbb{R}^{n\times{k}}$ are trainable parameters.

After it, $\textbf{x}^{\prime\prime}_{ECG}$, $\textbf{x}^{\prime\prime}_{RRI}$ and $\textbf{x}^{\prime\prime}_{RPE}$ have the same dimension $k$. Then, we are going to compute the attention weights. Specifically,

where $i\in\mathcal{S}=\{ECG,RRI,RPE\}$ and $\boldsymbol{\alpha}\in\mathbb{R}^{3}$ is the attention weights. $\textbf{w}\in\mathbb{R}^{k}$ and $b\in\mathbb{R}$ are trainable parameters. The transformed $\textbf{x}^{\prime\prime}$ is passed through an one-layer MLP to learn the importance of different types of features synergistically, and the importance is normalized by a softmax function.

Lastly, we compute the context vector c by

The context vector is the fused feature vector that is the weighted sum of features from different inputs based on the learned importance. The cross-attention module can be represented as $\mathcal{A}_{cross}([\textbf{x}^{\prime}_{ECG}$, $\textbf{x}^{\prime}_{RRI}$, $\textbf{x}^{\prime}_{RPE}];\theta_{{A}_{cross}})$.

For most of the conventional classification tasks, cross-entropy (CE) loss is commonly used to adjust model weights during training. However, CE loss may be impaired by noisy labels [42] and induce representations with excessive discrimination towards training data [25]. In order words, CE loss is likely to result in sub-optimal generalization.

As a remedy, contrastive learning is adopted to assist the model to learn more general and robust features by maximizing intra-class similarity while minimizing inter-class similarity. Concretely, we propose a novel hybrid loss, which utilizes the supervised contrastive (SC) loss [20] as an auxiliary regularization to the standard CE loss.

In our experiments, two datasets, i.e., Apnea-ECG [26] and MIT-BIH Polysomnographic [17] obtained from Physionet [12], are used for performance evaluation and comparison. Both datasets are publicly available and have been used to study sleep apnea detection methods in previous researches.

• •

  Apnea-ECG: The apnea-ECG database is provided by Philipps University, which is the most commonly-used dataset for ECG-based sleep apnea studies. It contains 70 single-lead ECG recordings of varying lengths between 7 hours to 10 hours, sampled at the rate of 100 Hz. Each segment of 1 min ECG data is annotated by the expert as either apnea or normal event. The datasets are officially split into two sets by the provider: a released set of 35 recordings and a withheld set of 35 recordings. After removing the data with an unreasonable heartbeat rate, the released set contains 16,888 segments and the withheld set contains 17,120 segments.

• •

  MIT-BIH Polysomnographic: The MIT-BIH Polysomnographic database is collected by Boston’s Beth Israel Hospital Sleep Laboratory. It contains over 80 hours of polysomnographic (PSG) recordings during sleep. Each recording includes a single channel of ECG annotated beat-by-beat and EEG and respiration signals. Each segment of 30 seconds of data is annotated with respect to sleep stages and apnea. After removing the data with an unreasonable heartbeat rate, the final dataset contains 9,717 segments.

To valid the performance of our framework, we use several state-of-the-art methods as our benchmark methods:

• •

  Support Vector Machine (SVM), Random Forest (RF), K-Nearest Neighbor (KNN), and Multi-Layer Perception (MLP) is adopted with 10 popular hand-crafted features from ECG data (e.g., RMSSD, NN50, etc.) as benchmark methods according to the work by [18].

• •

  LeNet-5: In [35], a LeNet-5 convolutional neural network is used to learn features from RRI and RPE for sleep apnea detection.

• •

  CNN+LSTM: In [3], three different deep learning architectures are proposed. We adopt their best performing architecture, i.e., CNN+LSTM, as one of our benchmark methods.

• •

  ResNet: In [38], a strong baseline model with the ResNet structure is proposed for time series classification, including ECG classification. So, we also use it for comparison.

• •

  CNN-4: In [10], a four-layer CNN-based model with a novel pooling layer is proposed to detect sleep apnea from ECG data directly, and we compare it with our proposed method as well.

• •

  CNN-6: In [30], several models with a different number of convolutional layers are designed to predict sleep apnea with ECG data. We adopt their best performing one, which contains 6 convolutional layers for our comparison.

We also compare the following approaches to show the improvements of the proposed framework step by step:

• •

  $\mathcal{F}_{ECG}+\mathcal{C}{lf}$: It employs a CNN-based feature extractor to learn the deep features from the raw ECG data and then classifies the targets by several fully-connected layers. It can be considered as a standard architecture of a naive deep learning model.

• •

  $\mathcal{F}_{ECG}+\mathcal{F}_{RRI}+\mathcal{F}_{RPE}+\mathcal{C}{lf}$: Besides the raw ECG data, RRI and RPE are used as expert knowledge inputs. A simple concatenation is used to combine the features from ECG, RRI, and RPE. Then the concatenated features are sent for classification.

• •

  $\mathcal{F}_{ECG}+\mathcal{F}_{RRI}+\mathcal{F}_{RPE}+\mathcal{A}_{cross}+\mathcal{C}{lf}$: Instead of using the simple concatenation, a cross-attention mechanism is proposed to collaboratively fuse the features from ECG, RRI, and RPE.

• •

  $\mathcal{F}_{ECG}+\mathcal{F}_{RRI}+\mathcal{F}_{RPE}+\mathcal{A}_{cross}+\mathcal{P}{roj}_{SC}+\mathcal{C}{lf}$: The proposed hybrid loss is used to update the model’s parameter by adding an auxiliary projection during training to learn more general and useful features.

• •

  ConCAD ($\mathcal{A}{ug}+\mathcal{F}_{ECG}+\mathcal{F}_{RRI}+\mathcal{F}_{RPE}+\mathcal{A}_{cross}+\mathcal{P}{roj}_{SC}+\mathcal{C}{lf}$): Data augmentation is used with the architecture mentioned above to help learn more general and robust features to boost performance.

For the Apnea-ECG dataset, we train all the models, including the proposed ConCAD method and benchmark methods, on the released dataset and test them on the withheld dataset. For the MIT-BIH PSG dataset, 10-fold cross-validation is used to examine the performance as there is no predefined training and test set. Moreover, since some existing studies [35, 29, 40] have shown that adjacent segment information helps analyze the sleep-related problems, the labeled segment with its surrounding $\pm 2$ segments of the ECG data is also included in our study. Thus, we will examine segments of 1 and 5 min on the Apnea-ECG dataset and test segments of 0.5 and 2.5 min on the MIT-BIH PSG dataset. In addition, some of the deep learning-based benchmark methods (i.e., [38, 10, 30]) are modified by increasing the pooling size and replacing flatten layer with GlobalAveragePooling for the input of 5 min and 2.5 min as they do no have a version to handle data with adjacent segments, and processing very long vector with their original structures exceeds our hardware memory limitation.

The proposed ConCAD model is trained by the AMSGrad optimizer [27], and all its parameters are initialized using HeNormal initializer [15]. An initial learning rate of 0.005 is chosen and it decreases to 0.001 after certain epochs (e.g. 200 epochs). Moreover, the L2 regularization is added to the feature extractor $\mathcal{F}$ to prevent the model from overfitting into the noise or artifacts.

We compare the performance of ConCAD with other state-of-the-art benchmark methods, and the results are listed in Table 1. Our proposed framework achieves an accuracy of $88.75\%$ with the 1 min segment input and $91.22\%$ with the 5 min segment input on the Apnea-ECG dataset, and $82.50\%$ with the 1 min segment input and $83.47\%$ with the 5 min segment input on the MIT-BIH PSG dataset, which outperforms other benchmark methods. Besides, we can see deep learning methods can adaptively learn features from a different length of input while the machine learning methods with hand-crafted features are more sensitive to the change of the input length. With the adjacent segments, the deep learning model can learn more effective features for classification tasks.

We also examine the proposed framework step by step to show the effectiveness of each step in Table 2. We can see that the performance can be worse if we simply concatenate the deep features with the expert features as some of the features will help the model make the better judgment possible, while others are likely to act as noise and thereby lead to more errors. With the cross attention module $\mathcal{A}_{cross}$, we can see that the model learns a better-fused feature representation by learning an attention mask synergistically from each other. The new fused feature representation maintains the effective discriminant parts of features while suppressing the irrelevant parts. Specifically, the accuracy improves from $83.41\%$ to $85.35\%$ with the 1 min segment input and from $85.48\%$ to $89.43\%$ with the 5 min segment input on the Apnea-ECG dataset. On the MIT-BIH PSG dataset, the accuracy improves from $78.83\%$ to $80.22\%$ with the 0.5 min segment input and from $80.11\%$ to $81.77\%$ with the 2.5 min segment input.

In Table 2, we can also see a further improvement by using the proposed hybrid loss, which takes advantage of both CE loss and SC loss. The accuracy increases to $87.16\%$ with the 1 min segment input and $89.43\%$ with the 5 min segment input on the Apnea-ECG dataset. On the MIT-BIH PSG dataset, the accuracy increases to $81.83\%$ with the 1 min segment input and $81.77\%$ with the 5 min segment input. The hybrid loss boosts the performance by promoting the model to learn more general and discriminant feature representations in case that the model overfits into the training data by learning features with excessive discrimination.

In Fig. 2, the t-SNE plots show the learned feature representation with the CE loss and the proposed hybrid loss. We can see that the hybrid loss promotes to more compact clustering of the instances from the same class while the representation with CE loss is more scattered. We also think the attention module benefits contrastive learning by focusing on parts of features when increasing the agreement among instances in positive pairs and encouraging the difference among instances in negative pairs. It is similar to human behavior that human usually tends to recognize an unseen data by comparing the most relevant parts with known ones. By using data augmentation, the proposed model can learn more general feature representation with a more clear boundary and achieve better performance.

In addition, the hybrid loss enables the classification tasks with limited training labeled data. The limitation of labeled data is a prevalent and critical problem in the healthcare field. We train the model by using a fraction of the training set and test it on the entire test set. The results are shown in Fig. 3 in terms of the macro F1 score. F1 score can clearly show the quality of the model when the dataset is imbalanced. With only $1\%$ of the training data, the proposed ConCAD model can still achieve an F1 score of $0.67$ on the Apnea-ECG dataset and $0.59$ on the MIT-BIH PSG dataset. However, the CE loss performs poorly and skews into the majority class. Furthermore, the proposed model only requires $10\%$ of the training data to make a reasonable classification while the naive deep learning model needs more than $50\%$ to get decent performance. Hence, the proposed model with hybrid loss largely outperforms a naive deep learning approach with CE loss on smaller datasets.