MTS-LOF: Medical Time-Series Representation Learning via Occlusion-Invariant Features

Joint-Embedding SSL is a prominent paradigm in representation learning, dedicated to discerning similarities and dissimilarities among data point pairs. It primarily involves maximizing the similarity between augmented views from identical inputs (referred to as ”positive” samples) and minimizing the similarity between augmented views from distinct inputs (referred to as ”negative” samples). A primary challenge in Joint-Embedding SSL is the potential risk of representation collapse. In such cases, the model consistently generates identical outputs, making the learned representations uninformative. In the realm of contrastive learning, the methodologies frequently incorporate positive samples and negative samples. The overarching objective is to bring positive samples closer to the anchor sample within an embedding space while simultaneously pushing negative samples farther apart. Contrastive learning techniques, including Contrastive Predictive Coding (CPC) [4], Simple Contrastive Learning Framework (SimCLR) [5], and Momentum Contrast (MoCo) [6], have demonstrated substantial progress. Unlike contrastive learning methods, non-contrastive approaches in Joint-Embedding SSL depart from the conventional utilization of negative samples. These techniques employ diverse strategies to mitigate the risk of representation collapse. For example, Bootstrap Your Own Latent (BYOL) [20] and Simple Siamese representation learning (SimSiam) [23] employ gradient stopping techniques and predictor modules to avert representation collapse. Furthermore, Variance-Invariance-Covariance Regularization (VICReg) [21] introduces covariance regularization as a means to penalize learned representations, thus mitigating the risk of collapse and enhancing the stability of embeddings.

Within the domain of SSL, masked language modeling (MLM) has emerged as a prominent method in Natural Language Processing (NLP) [8, 11, 9]. MLM and its auto-regressive variations [9] have transformed NLP, enabling the training of extensive language models that excel in a wide range of language comprehension and generation tasks. These methodologies entail predicting masked tokens within sentences or sentence pairs/triplets, harnessing extensive training data to attain remarkable performance. Concurrently, Masked Image Modeling (MIM) has advanced in parallel with MLM. These techniques encompass the deliberate masking of particular regions within images or feature spaces, directing models to discern intricate patterns, imputing missing information, and ultimately extracting valuable insights regarding the underlying structure of visual content. A seminal contribution in this domain is the context encoder approach [34], which masks regions within original images and predicts the missing pixels, laying the foundation for subsequent advancements. Recent studies [35, 26, 28] have delved into the utilization of MIM for pretraining vision Transformers [26, 36]. Additionally, Masked Autoencoder (MAE) [27] introduced a straightforward and scalable self-supervised learning approach for computer vision. This method involves masking random patches of the input image and subsequently reconstructing the missing pixels. Recent work proposed by Kong et al. [37] reveals the insight behind MIM is encouraging the network to learn the learning occlusion invariant feature.

SSL techniques have significantly impacted the field of time-series data analysis. Researchers have earnestly explored the adaptation of SSL approaches to capture the intricate temporal dependencies, patterns, and representations inherent in sequential data. Before the emergence of time-series contrastive learning, the predominant method utilized multitask learning [24, 14] to obtain time-series representations. These methods typically involved applying various transformations to the original time series to construct pretext tasks and pseudo labels. Subsequently, models were trained to recognize the specific transformations that were applied. However, motivated by the remarkable successes observed in other domains, contrastive learning has made inroads into the domain of time-series representation learning. For instance, SimCLR [5] has been extended to support EEG representation learning [16]. Contrastive Predictive Coding (CPC) [4] has employed predictive modeling in latent spaces to acquire time-series representations, showcasing competence in speech recognition tasks. The TS-TCC framework [13] has introduced innovative temporal and contextual contrasting modules founded on transformer blocks [25], enhancing the acquisition of discriminative representations. Additionally, TS2Vec [15] has presented a hierarchical contrasting framework with the capability to capture multi-scale temporal and instance dependencies within time series data. In addition to temporal contrastive learning, the direct application of MAE in the context of time series is put forth in [29].

We begin by detailing the backbone network utilized in our study. Specifically, when dealing with a collection of multidimensional multivariate time series samples, each training sample denoted as $X\in\mathbb{R}^{t\times m}$ constitutes a sequence of $m$ univariate series represented as $x\in\mathbb{R}^{t}$, where $t$ signifies the time step. As illustrated in Figure. 1, our approach diverges from prior work [29], which employed a linear projection of each time step from the original time series into a $d$-dimensional vector space, treating them as tokens, using $\hat{X}=WX+b$, with $W\in\mathbb{R}^{m\times d}$ and $b\in\mathbb{R}^{d}$. Instead, we adopt a strategy where we initially generate patches as tokens by processing the time series through several convolutional layers. To streamline the discussion, we implement a modified version of the CNN1D architecture from TS-TCC [13]. The configuration of the initial convolutional layer varies depending on the dataset, with distinct kernel sizes ($k$) and strides ($s$). Subsequently, we stack three convolutional layers with fixed kernel size $8$ and stride $2$. These layers efficiently downsample the original time series by a factor of $8\times s$. Finally, a single convolutional layer with kernel size $1$ and stride $1$ reduces the feature channel dimensions to match the $d$-dimensional input required by the transformer encoder. Each convolutional layer is followed by a batch normalization (BN) [38] and Gaussian Error Linear Units (GELU) [39] beside the last one. Several factors motivate this choice. Firstly, akin to the concept of patching in vision transformers (ViT) [26], we consider a single time step as equivalent to a single pixel, lacking semantic meaning akin to words in a sentence. Thus, it becomes imperative to extract local semantic information for analyzing their connections. Secondly, as exemplified in [40], early convolutional layers can enhance the stability and convergence of the transformer during training. Lastly, this design facilitates further modifications and extensions of the backbone network, enabling the replacement of the simple CNN1D with more complex architectures [41].

After generating the patches through the convolutional layers, the tokens $\hat{X}\in\mathbb{R}^{p\times d}$, where $p=\frac{t}{8s}$ represents the number of patches, are encoded by adding a fixed 1D cosine positional embedding denoted as P, defined as follows:

Here, $\text{P}_{i,j}$ signifies the value of the embedding for the $i$-th position and the $j$-th dimension. Typically, $d$ matches the dimension of the Transformer encoder’s input embeddings. This positional encoding method enhances the model’s capability to capture sequential order and relationships within the time series data, which are crucial for various time-series analysis tasks. It ensures that the model can differentiate between different positions in the sequence, and when combined with the patch-based tokenization, it allows the transformer encoder to effectively process the data for downstream tasks. Subsequently, we feed these encoded tokens into the Transformer encoder. The Transformer encoder employs multi-head attention, where each attention head, denoted as $h=1,\dots,H$, independently processes the input tokens $\hat{X}$. Each head transforms the input into query matrices $Q_{h}$, key matrices $K_{h}$, and value matrices $V_{h}$, achieved through weight matrices $W_{Q}$, $W_{K}$, and $W_{V}$, respectively. Once the query, key, and value matrices are obtained, a scaled dot-product attention mechanism computes the attention output $O_{h}$. The softmax function is utilized to calculate attention weights, which determine the level of attention each patch pays to others. The calculation for $O_{h}$ is defined as:

where the $\sqrt{d_{k}}$ term refers to the dimension of the key vectors and is employed to scale the dot product to prevent it from becoming excessively large. The multi-head attention block incorporates Layer Normalization layers [42] to normalize the data and stabilize training. Additionally, a feed-forward network with residual connections is employed to capture intricate patterns within the data. Consequently, the outcome of the Transformer encoder’s operations is a representation denoted as $Z\in\mathbb{R}^{p\times d}$. Finally, a global average pooling layer is applied to obtain the final representation $\hat{Z}\in\mathbb{R}^{d}$, followed by a linear head for predicting results $Y\in\mathbb{R}^{c}$, where $c$ is the number of classes.

We now delve into the concept of Occlusion-Invariant Feature Learning and motivation behind the proposed MTS-LOF framework. To establish a comprehensive understanding, it is essential to revisit the MAE objective, a pivotal component of our methodology. The primary purpose of this objective is to predict the original pixel values of the masked patches/tokens $\hat{X}\odot M$. Let it be known that we represent $M\in\mathbb{R}^{p}$, where each $m_{i}$ within $M$ is assigned the value of $1$ for masked tokens and $0$ for unmasked tokens, effectively delineating the regions of the time series data that are concealed and revealed. This MAE objective, working in tandem with a transformer decoder denoted as $d_{\phi}(\cdot)$ and the latent representation $Z^{m}$ obtained from the unmasked patches $\hat{X}\odot(1-M)$, is formally articulated as follows:

In this equation, $\mathcal{L}_{MAE}$ quantifies the mean absolute error between the predicted values obtained from the decoder and the actual values of the masked tokens. Fundamentally, the MAE objective encapsulates the idea of reconstructing the entirety of the input data, even when faced with incomplete or occluded information. It is worth noting that in an ideal scenario, with an over-parameterized neural network, the network can memorize all seen training samples. This implies that with the latent representation $Z$ of the complete input tokens $\hat{X}$, we have:

Consequently, we can reformulate the MAE objective with a distance function $\mathcal{D}$:

which signifies that the primary objective of MAE training is to acquire the Occlusion-Invariant Feature.

Moreover, it is essential to recall the fundamental concept of Joint-Embedding SSL, which is centered on the objective of minimizing the distance between two latent representations derived from augmented views of the same input. This objective can be generally formulated as follows:

where we denote the two augmentations as $\mathcal{A}_{1}$ and $\mathcal{A}_{2}$. Importantly, it should be noted that the distance function can be both parametric and non-parametric. It means it can encompass simple functions such as mean square error and cosine similarity, or it can incorporate projection heads or predictor networks, as is common in Joint-Embedding SSL settings. We can establish a connection between Joint-Embedding SSL and the MAE framework by assuming that $\mathcal{A}{1}$ preserves the original input while $\mathcal{A}{2}$ involves random masking operations. Consequently, we can define the primary similarity objective of the proposed MTS-LOF Framework:

This objective captures the essence of the MTS-LOF framework, aiming to minimize the distance between the original representation $Z$ and the representation $Z^{m}$ obtained after masking a portion of the input data. This formulation bears a resemblance to non-contrastive learning algorithms, which are prone to the risk of representation collapse. To mitigate this risk, a covariance regularization technique is employed to prevent representation collapse. Specifically, the latent representation $Z^{m}$ of the masked patches is penalized using Total Coding Rate (TCR), which is formulated as:

where $b$ represents the batch size, $d$ signifies the dimensionality of the representation, and $\epsilon>0$ is a chosen size of distortion.

In practical implementation, the algorithm is supposed to effectively capture features that are consistently occlusion-invariant. To achieve this, we apply multiple mask operations to the patches denoted as $\hat{X}$. This operation is defined as follows:

where $N$ refers to the number of masks. Subsequently, the similarity objective is employed to compare the representations of the unmasked patches with each representation of the masked patches:

where we use the negative cosine similarity to measure the distance between two representation vectors in this context. It’s important to note that we calculate $\mathcal{L}_{sim}$ using the globally pooled representations $\hat{Z}$ and $\hat{Z}_{i}^{m}$, as well as the covariance regularization $\mathcal{L}_{TCR}$. A transformer decoder is employed to decode the unmasked patches encoded in the previous step, along with all the masked tokens. To provide positional awareness, position embeddings are added so that each token can determine its respective position. It’s worth mentioning that each masked token is shared and can learn a vector. Visual demonstration of the entire workflow of the MTS-LOF framework is demonstrated in Figure. 2, where we set a hyperparameter $\lambda$ to balance the $\mathcal{L}_{sim}$ and $\mathcal{L}_{TCR}$.

Our experiments encompass various medical time series datasets, including those related to human activity recognition, sleep stage classification, and epileptic seizure prediction. Additionally, to assess the generalizability and transferability of the proposed MTS-LOF framework beyond medical time series, we utilize a fault diagnosis dataset. The datasets are summarized in Table. I.

We adopt the data splitting settings from TS-TCC [13] for Epilepsy, Sleep-EDF, FD, and HAR. The data is divided into training (60%), validation (20%), and testing (20%) sets. For the Sleep-EDF dataset, we additionally perform subject-wise splitting. In the case of SHHS, we initially divide the records into two equal parts subject-wisely to obtain EEG data from two distinct channels. Each part is then subdivided into training (60%), validation (20%), and testing (20%) sets. We conducted experiments five times using different random seeds ($\{2019,2020,2021,2022,2023\}$), and the results are reported as averages across all runs. Model training, including SSL pretraining, linear probing, and fine-tuning, was carried out using the AdamW optimizer. The training duration was 40 epochs with a learning rate of 5e-4, weight decay set to 0.05, and a batch size of 128. The initial convolutional layers have specific kernel sizes and strides for each dataset: ${k=8,s=1}$ for Epilepsy, ${k=25,s=6}$ for Sleep-EDF, ${k=32,s=4}$ for FD, ${k=8,s=1}$ for HAR, and ${k=25,s=6}$ for SHHS. Hyperparameters are set as follows: $\lambda=100$, mask ratio=0.8, and the number of masks=20. The transformer decoder configuration matches that of the encoder while having 4 layers. We implemented the algorithms and models using PyTorch (Code available at https://github.com/HuayuLiArizona/MST-LOF.git) and conducted all experiments on an NVIDIA RTX 3090 GPU.

In the pursuit of assessing the performance of our proposed approach, we conducted a comprehensive comparative analysis against several baseline methods: SSL-ECG [14], CPC [4], SimCLR [5], TS-TCC [13], TST [29], and TS2Vec [15]. This evaluation encompassed the domains of Human Activity Recognition (HAR), Sleep-EDF, and Epilepsy, and focused on standard linear probing results, employing linear classifiers on top of frozen representations from the SSL-pretrained models. Table. II showcases the results of these comparisons, including accuracy and the Macro F1 score (MF1). Notably, our approach demonstrates a significant improvement in performance across all three domains. In the domain of Human Activity Recognition (HAR), our approach achieves an accuracy of 93.05% and an MF1 score of 93.09, surpassing all baseline methods. In comparison, SSL-ECG, CPC, and SimCLR show lower accuracy and MF1 scores, highlighting the superior capability of our method in recognizing human activities from time series data. For Sleep-EDF, our approach attains an accuracy of 84.35% and an MF1 score of 73.52, demonstrating a substantial enhancement compared to the baseline methods. TS-TCC, while showing competitive results, falls short of our approach in terms of MF1 score, emphasizing the effectiveness of our method in sleep stage classification. In the challenging domain of Epilepsy, our approach outperforms all baseline methods, achieving an impressive accuracy of 98.33% and an MF1 score of 97.41. This showcases the robustness and high predictive power of our approach in distinguishing between subjects with and without epileptic seizures.

These results validate the effectiveness of our proposed method, which leverages surprising self-supervised representation learning capability, in enhancing the predictive performance across various medical time series domains. Our approach consistently exhibits superior accuracy and F1 scores, highlighting its potential to make a significant impact in medical diagnostics and related applications. The success of our method can be attributed to the inherent ability of the occlusion invariant feature learning to capture intricate temporal and local patterns within the time series data, which is vital for distinguishing subtle differences in medical conditions. Moreover, our approach showcases the adaptability of self-supervised learning in a diverse range of medical domains, including human activity recognition, sleep stage classification, and epileptic seizure prediction. This versatility is a testament to the broad utility and applicability of our method in real-world medical scenarios.

To further assess the adaptability and performance of the MTS-LOF framework, we conducted experiments under semi-supervised learning conditions, specifically on the Sleep-EDF dataset. In this scenario, we fine-tuned the pre-trained model using varying percentages of partially labeled data, specifically 1%, 5%, 10%, 50%, and 100% of randomly selected subsets from the complete dataset. The results of these semi-supervised experiments are depicted in Figure. 3, and they are compared to the results obtained under fully supervised learning conditions with fully labeled data, which yielded an F1 score of 75.55. The outcomes of these semi-supervised experiments are insightful. We observe that the MTS-LOF framework demonstrates its robustness and adaptability across different levels of labeled data. Even with just 1% of the data labeled, the model achieves a competitive F1 score of 72.77. This suggests that the framework can effectively leverage minimal labeled data to produce meaningful results, which is especially relevant in real-world scenarios where labeling medical data can be a time-consuming and costly process.

As the percentage of labeled data increases, we notice a gradual improvement in the F1 score. This trend highlights the framework’s ability to capitalize on additional labeled data for enhancing predictive performance. Notably, when 100% of the data is labeled, the F1 score reaches 75.95, surpassing the fully supervised learning result. This finding is promising, as it implies that SSL pretraining benefits the performance of supervised learning. Comparing these semi-supervised results with the fully supervised counterpart, we can infer that the MTS-LOF framework offers a powerful solution in scenarios where labeling medical data comprehensively is challenging or impractical. The ability to achieve competitive performance even with minimal labeled data reaffirms the framework’s applicability in scenarios where obtaining fully labeled datasets may be a limiting factor.

The MTS-LOF framework not only enhances performance but also improves the transferability of learned representations across diverse data distributions. We conducted experiments using the FD and SHHS datasets to explore this aspect. Baseline models, trained with supervised learning, utilized fully labeled data from each dataset. In contrast, SSL pretraining utilized unlabeled data to acquire generalized representations, which were later fine-tuned with fully labeled target domain data through linear probing. It’s essential to note that these models were exclusively trained in one domain (the source domain) and subsequently tested in various domains (the target domains).

In the SHHS dataset, we evaluated the SSL-pretrained model’s ability to transfer learned representations across different domains, namely labeled domains C4-A1 and C3-A2. The results consistently demonstrate high cross-domain accuracy when assessing accuracy (Table. III). A similar pattern is observed when examining F1 scores (Table. IV). The SSL-pretrained model consistently outperforms in cross-domain performance regarding F1 scores, with particularly significant improvements in domain C3-A2. These results underscore the effectiveness of SSL pretraining in enhancing the transferability of learned representations in the SHHS dataset. Tables. V and VI provide accuracy and F1 score results for the FD dataset. These tables demonstrate strong transferability, with consistently high accuracy and F1 scores, even when shifting between different domains. These results showcase that MTS-LOF’s exceptional generalizability extends beyond medical time series data.

To comprehensively assess the impact of two critical hyperparameters, specifically the number of masks and mask ratio, we performed an ablation study using the HAR dataset. Figure. 4 illustrates the F1 scores obtained across various combinations of these hyperparameters. In the initial part of the ablation study, we explored the effect of different numbers of masks while keeping the mask ratio constant at 0.8. Clearly, the F1 score shows significant sensitivity to the number of masks, with a significant increase as we move from one mask to 20 masks. However, increasing the number of masks to forty does not lead to a substantial improvement. In the subsequent part of the ablation study, we investigated the effect of different mask ratios while maintaining a constant number of masks at 20. The results indicate significant sensitivity of the F1 score to changes in the mask ratio. Notably, with an increase in the mask ratio from 0.5 to 0.9, the F1 score steadily improves, with the most significant gain occurring with a mask ratio of 0.8.

In this section, we provide visual insights into the learned representations on three distinct datasets: Epilepsy, Sleep-EDF, and HAR. For visual exploration, we utilize t-SNE (t-Distributed Stochastic Neighbor Embedding) [45] to produce two-dimensional representations of the latent features obtained through different training methodologies. Figure 5 presents a side-by-side comparison of t-SNE visualizations for three distinct training approaches: SSL, supervised training, and 5% fine-tuning. These visualizations aim to clarify how each training paradigm influences the distribution and clustering of data points within the latent space. This information is essential for assessing the quality and discriminative capabilities of the learned representations. Our observations from the visualizations indicate that the features learned through SSL are already linearly separable, and a 5% fine-tuning can yield results comparable to supervised training.