SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sample-level contrastive loss. The sample-level contrastive loss  $\mathcal{L}_{\mathrm{S}}$  [18, 43] for the input sample  $x_{i}$  is defined as:
$$
\mathcal {L} _ {\mathrm {S}} = \mathbb {E} _ {\boldsymbol {x} _ {i} \in \mathcal {D}} \left[ - \log \frac {\exp \left(\boldsymbol {h} _ {i} \cdot \widetilde {\boldsymbol {h}} _ {i}\right)}{\sum_ {j = 1} ^ {| \mathcal {D} |} \left(\exp \left(\boldsymbol {h} _ {i} \cdot \widetilde {\boldsymbol {h}} _ {j}\right) + \mathbb {1} _ {[ i \neq j ]} \exp \left(\boldsymbol {h} _ {i} \cdot \boldsymbol {h} _ {j}\right)\right)} \right] \tag {2}
$$
where  $|\mathcal{D}|$  represents the total number of samples in the dataset  $\mathcal{D}$  and  $\cdot$  denotes dot product. The  $\mathbb{1}_{[i\neq j]}$  is an indicator function that equals 0 when  $i = j$  and 1 otherwise.
4.4 Trial-Level Contrastive Block
For an input sample  $\pmb{x} \in \mathcal{R}_i$ , where  $\mathcal{R}_i$  is a collection of all samples segmented from trial  $\pmb{r}_i$ , we feed it into the contrastive encoder  $G$  to generate a sample-level representation  $\pmb{h} = G(\pmb{x})$ . To seize trial-level data consistency, we assume that the representation of the anchor sample  $\pmb{x} \in \mathcal{R}_i$  is close to the representation of sample  $\pmb{x}^+$  that also come from the trial  $\pmb{r}_i$ . In contrast, the representation of the anchor sample  $\pmb{x}$  is far away from the representation of sample  $\pmb{x}^-$  that come from a different trial  $\pmb{r}_j$ , where  $\pmb{x}^- \in \mathcal{R}_j$ . In other words, we have positive pair  $(\pmb{x}, \pmb{x}^+)$  and negative pair  $(\pmb{x}, \pmb{x}^-)$ .
Trial-level contrastive loss. The trial-level contrastive loss  $\mathcal{L}_{\mathbb{R}}$  [15, 18] for the input sample  $x$  is defined as:
$$
\mathcal {L} _ {\mathrm {R}} = \mathbb {E} _ {\boldsymbol {x} \in \mathcal {D}} \left[ \mathbb {E} _ {\boldsymbol {x} ^ {+} \in \mathcal {R} _ {i}} \left[ - \log \frac {\exp \left(\sin \left(\boldsymbol {h} , G \left(\boldsymbol {x} ^ {+}\right)\right) / \tau\right)}{\sum_ {j = 1} ^ {J} \sum_ {\boldsymbol {x} ^ {-} \in \mathcal {R} _ {j}} \left(\exp \left(\sin \left(\boldsymbol {h} , G \left(\boldsymbol {x} ^ {-}\right)\right) / \tau\right)\right)} \right] \right] \tag {3}
$$
where  $J$  is the total number of trials in the dataset  $\mathcal{D}$ . The  $\mathrm{sim}(\pmb{u},\pmb{v}) = \pmb{u}^T\pmb{v} / \| \pmb{u}\| \| \pmb{v}\|$  denotes the cosine similarity, and  $\tau$  is a temperature parameter to adjust the scale. The  $G(\pmb{x}^{+})$  and  $G(\pmb{x}^{-})$  are learned representations of samples  $\pmb{x}^{+}\in \mathcal{R}_{i}$  and  $\pmb{x}^{-}\in \mathcal{R}_{j}$ , respectively. To measure the trial-level loss for sample  $\pmb{x}$ , we iterate all the  $\pmb{x}^{+}$  in  $\mathcal{R}_i$ , and averaging across  $|\mathcal{R}_i| - 1$  positive pairs.
In this block, we do NOT learn a trial-level embedding representing the entire trial. Instead, we learn a representation for each sample within the trial while considering trial-level data consistency. Similarly, we follow this protocol for the patient-level contrastive block.
4.5 Patient-Level Contrastive Block
For an input sample  $\pmb{x} \in \mathcal{P}_i$ , where  $\mathcal{P}_i$  denotes all samples from patient  $\pmb{p}_i$ , we feed it into the contrastive encoder  $G$  to generate a sample-level representation  $h = G(\pmb{x})$ . Similar to the above
trial-level contrastive block, we have positive pair  $(\pmb{x},\pmb{x}^{+})$  and negative pair  $(\pmb{x},\pmb{x}^{-})$ , in which  $\pmb{x}^{+}$  come from the same patient while  $\pmb{x}^{-}$  come from a different patient.
Patient-level contrastive loss. The patient-level contrastive loss  $\mathcal{L}_{\mathrm{P}}$  [15] for the input sample  $x$  is defined as:
$$
\mathcal {L} _ {\mathrm {P}} = \mathbb {E} _ {\boldsymbol {x} \in \mathcal {D}} \left[ \mathbb {E} _ {\boldsymbol {x} ^ {+} \in \mathcal {P} _ {i}} \left[ - \log \frac {\exp \left(\sin \left(\boldsymbol {h} , G \left(\boldsymbol {x} ^ {+}\right)\right) / \tau\right)}{\sum_ {j = 1} ^ {M} \sum_ {\boldsymbol {x} ^ {-} \in \mathcal {P} _ {j}} \left(\exp \left(\sin \left(\boldsymbol {h} , G \left(\boldsymbol {x} ^ {-}\right)\right) / \tau\right)\right)} \right] \right] \tag {4}
$$
where  $M$  is the total number of patients in the dataset  $\mathcal{D}$ . In this block, the  $G(\pmb{x}^{+})$  and  $G(\pmb{x}^{-})$  are learned representations of samples  $\pmb{x}^{+} \in \mathcal{P}_{i}$  and  $\pmb{x}^{-} \in \mathcal{P}_{j}$ , respectively.
4.6 Overall Loss Function
The overall loss function  $\mathcal{L}$  consists of four loss terms. The observation-level loss  $\mathcal{L}_{\mathrm{O}}$  and sample-level loss  $\mathcal{L}_{\mathrm{S}}$  encourage the encoder to learn robust representations that are invariant to perturbations. The trial-level loss  $\mathcal{L}_{\mathrm{R}}$  and patient-level loss  $\mathcal{L}_{\mathrm{P}}$  compel the encoder to learn cross-sample features within a trial or a patient. In summary, the overall loss function of the proposed COMET model is:
$$
\mathcal {L} = \lambda_ {1} \mathcal {L} _ {\mathrm {O}} + \lambda_ {2} \mathcal {L} _ {\mathrm {S}} + \lambda_ {3} \mathcal {L} _ {\mathrm {R}} + \lambda_ {4} \mathcal {L} _ {\mathrm {P}} \tag {5}
$$
where  $\lambda_1, \lambda_2, \lambda_3, \lambda_4 \in [0, 1]$  are hyper-coefficients that control the relative importance and adjust the scales of each level's loss. Users can simply turn off specific data levels by setting  $\lambda$  of those levels to 0. We set  $\lambda_1 + \lambda_2 + \lambda_3 + \lambda_4 = 1$ . We calculate the total loss by taking the expectation of  $\mathcal{L}$  across all samples  $\boldsymbol{x} \in \mathcal{D}$ . In practice, the contrastive losses are calculated within a mini-batch.
We compare the COMET model with six baselines on three datasets. Following the setup in SimCLR [18], we use unlabeled data to pre-train encoder  $G$  and evaluate it in two downstream settings (Appendix E): partial fine-tuning (P-FT; i.e., linear evaluation [18]) and full fine-tuning (F-FT). All datasets are split into training, validation, and test sets in patient-independent setting (Figure 3), which is very challenging due to patient variations (More explanations in Appendix F.1).
Datasets. (1) AD [44] has 23 patients, 663 trials, and 5967 multivariate EEG samples. There are 4329, 891, and 747 samples in training, validation, and test sets. The sampling rate (frequency) is  $256\mathrm{Hz}$ . Each sample is a one-second interval with 256 observations. A binary label based on whether the patient has Alzheimer's disease is assigned to each sample. (2) PTB [45] has 198 patients, 6237 trials, and 62370 multivariate ECG samples. There are 53950, 3400,
Figure 3: Patient-dependent/independent Setting. In the patient-dependent setting, samples from the same patient can appear in both the training and test sets. In contrast, in the patient-independent setting, samples from the same patient are exclusively included in the training or test set.
and 5020 samples in training, validation, and test sets. The sampling rate (frequency) is  $250\mathrm{Hz}$ . Each sample is a heartbeat with 300 observations. A binary label based on whether the patient has Myocardial infarction is assigned to each sample. (3) TDBRAIN [46] has 72 patients, 624 trials, and 11856 multivariate EEG samples. There are 8208, 1824, and 1824 samples in training, validation, and test sets. The sampling rate (frequency) is  $256\mathrm{Hz}$ . Each sample is a one-second interval with 256 observations. A binary label based on whether the patient has Parkinson's disease is assigned to each sample. See appendix D for more details about data statistics, train-test split, and data preprocessing.
Baselines. We compare with 6 state-of-the-art methods: TS2vec [13], Mixing-up [16], TS-TCC [27], SimCLR [43], CLOCS [15] and TF-C [26]. Since TF-C is designed for transfer learning, we implement its downstream tasks the same as ours for a fair comparison. The evaluation metrics are accuracy, precision (macro-averaged), recall(macro-averaged), F1 score(macro-averaged, AUROC(macro-averaged), and AUPRC(macro-averaged).
Table 2: Partial fine-tuning results. A logistic regression(LG) classifier  $L$  is added on top of a frozen encoder  $G$  on dataset AD. Only fine-tuning the classifier  $L$  .
Models	Accuracy	Precision	Recall	F1 score	AUROC	AUPRC
TS2vec	66.48±5.33	67.72±5.09	67.40±5.20	66.32±5.46	74.12±6.88	72.96±7.21
TF-C	77.03±2.80	75.79±5.07	64.27±5.03	64.85±5.56	80.71±4.03	79.27±4.15
Mixing-up	46.16±1.38	52.62±4.90	50.81±1.32	37.37±1.98	64.42±6.49	62.85±6.07
TS-TCC	59.71±8.63	61.66±8.63	60.33±8.26	58.66±8.39	67.53±10.04	68.33±9.37
SimCLR	57.16±2.05	56.67±3.91	53.57±2.21	49.11±4.26	56.67±3.91	52.10±1.41
CLOCS	66.99±2.84	67.17±2.96	67.33±2.99	66.91±2.88	73.71±3.62	72.58±3.54
COMET (Ours)	76.09±4.21	77.36±3.97	74.68±4.62	74.80±4.83	81.30±4.97	80.50±5.31
Implementation. We use a ten residual blocks dilated CNN module [47] as backbones for encoders  $G$ . To preserve the magnitude feature of time series, which is often an important feature in time series, we utilize timestamp masking as proposed in [13] as an augmentation method. While many existing works apply data augmentation directly on raw time series, we use a fully connected projection layer to map the raw input into an embedding with a higher feature dimension. This strategy helps prevent situations where some parts of the raw data are zero, which could result in ineffective augmentation if timestamp masking were directly applied. Additionally, to ensure stability during training, we incorporate a hierarchical loss [13], which encapsulates the observation and sample-level losses. To guarantee there are samples with the same trial and patient IDs in a batch for trial and patient levels contrasting, we design two specific shuffle functions. See Appendix C for further details. For datasets whose trial level is absent or limited to one trial per patient, we provide a solution in Appendix F.2.
We conduct experiments with five seeds(41-45) based on the same data split to account for variance and report the mean and standard deviation. All experiments expect baseline SimCLR run on NVIDIA RTX 4090. Baseline SimCLR runs on Google Colab NVIDIA A100. See further implementation details in Appendix E.
5.1 Results on Partial Fine-tuning
Setup. Linear evaluation is a widely used P-FT setup to evaluate the quality of representation, where a linear classifier is trained on top of a frozen network for sample classification [18, 15, 48, 49]. This evaluation method allows for a quick assessment of representation quality. In linear evaluation, we add a logistic regression classifier  $L$  on top of the pre-trained encoder  $G$ . The training process utilizes  $100\%$  labeled training data for sample classification. Notably, the parameters of the encoder  $G$  are frozen during training, ensuring that only the classifier  $L$  is fine-tuned. See further details in E.1.
Results. We report the experimental results of the P-FT setup on AD in Table 2. On average, our COMET model claims a large margin of more than  $7\%$  over all baselines on the F1 score on AD.
5.2 Results on Full Fine-tuning
Setup. In an F-FT setup, we add a two-layer fully connected network classifier  $P$  to the pre-trained encoder  $G$ . The training process utilizes  $100\%$ ,  $10\%$ , and  $1\%$  of the labeled training data for sample classification, respectively. Unlike the P-FT approach, both the encoder  $G$  and classifier  $P$  are trainable now, allowing for fine-tuning of the entire network structure. See further details in E.2.
Results. The results for F-FT on the three datasets are presented in Table 3 and Table 4. In general, COMET demonstrates success in 46 out of 54 tests conducted across the three datasets, considering six different evaluation metrics. With  $100\%$  labeled data, the F1 score of COMET outperforms the best baseline, TS2Vec, by  $2\%$  on the AD dataset and surpasses the best baseline, Mixing-up, by  $4\%$  on the TDBrain dataset. Furthermore, it achieves a result comparable to the best baseline, TF-C, on the PTB dataset.
Notably, the F1 score of COMET surpasses the best baseline, TF-C, by  $14\%$  and  $13\%$  with label fractions of  $10\%$  and  $1\%$ , respectively, on the AD dataset. Additionally, on the TDBrain dataset, the F1 score of COMET outperforms the best baseline, Mixing-up, by  $2\%$  and  $8\%$  with label fractions of  $10\%$  and  $1\%$ , respectively. Similarly, the F1 score of COMET outperforms the best baseline, CLOCS, by  $0.17\%$  and  $2.66\%$  with label fractions of  $10\%$  and  $1\%$ , respectively, on the PTB dataset. It is interesting to observe that the F1 score of COMET with label fractions of  $10\%$  and  $1\%$  outperforms a label fraction of  $100\%$  on the AD and PTB datasets. This suggests a potential overfitting of COMET to the training data. A similar phenomenon is observed with SimCLR, TF-C, CLOCS, and TS2vec, where a higher fraction of labeled data did not necessarily lead to improved performance.
As a result, COMET demonstrates its superiority and stability across all the datasets. Furthermore, COMET outperforms SOTAs methods significantly with  $10\%$  and  $1\%$  label fractions, highlighting the effectiveness of our contrastive pre-training approach in reducing the reliance on labeled data.
5.3 Ablation Study, Visualization, Additional Downstream Tasks, and Heavy Duty Baseline
Ablation study. We verify the effectiveness of each contrastive block and other COMET variants. Besides, we study the impact of hyperparameter  $\lambda$ . (Appendix G.1)
Visualization. We visualize the embedding space and show our model can learn more distinctive and robust representations. (Appendix G.2)
Additional downstream tasks. Apart from the classification tasks in Section 5.1-5.2, we show the proposed COMET outperforms baselines in a wide range of downstream tasks, including clustering and anomaly detection. (Appendix G.3)
Heavy duty baseline. We show the superiority of our model does NOT result from the newly added contrastive blocks (i.e., increased model parameters). COMET outperforms the heavy-duty SimCLR and TS2Vec with four contrastive blocks and four times pre-training epochs. (Appendix G.4)
Table 4: Full fine-tuning results of ECG datasets. A two-layer fully connected networks classifier  $P$  is trained on top of pre-trained encoder  $G$  on ECG dataset PTB. Fine-tuning the whole network includes  $G$  and  $P$ . We use  ${100}\% ,{10}\%$  ,and  $1\%$  of labeled training data in PTB.
Datasets	Fraction	Models	Accuracy	Precision	Recall	F1 score	AUROC	AUPRC
PTB	100%	TS2vec	85.14±1.66	87.82±2.21	76.84±3.99	79.66±3.63	90.50±1.59	90.07±1.73
TF-C	87.50±2.43	85.50±3.04	82.68±4.04	83.77±3.50	77.59±19.22	80.62±15.10
Mixing-up	87.61±1.48	89.56±2.80	79.30±1.67	82.56±2.00	89.29±1.26	88.94±1.04
TS-TCC	82.24±3.55	85.28±5.08	69.46±5.85	72.08±6.85	87.60±2.51	86.26±3.00
SimCLR	84.19±1.32	85.85±1.89	73.89±2.95	76.84±2.80	85.85±1.89	70.70±2.36
CLOCS	86.39±2.76	87.06±2.81	77.95±4.79	80.71±4.78	90.41±2.07	87.35±3.36
COMET (Ours)	87.84±1.98	87.67±1.72	81.14±3.68	83.45±3.22	92.95±1.56	87.47±2.82
10%	TS2vec	82.49±4.71	80.39±5.04	83.35±0.91	80.18±4.04	93.03±1.03	92.81±0.97
TF-C	85.37±1.23	82.80±2.35	79.94±0.71	81.09±1.14	81.57±15.60	84.57±8.12
Mixing-up	87.05±1.41	86.56±3.24	80.61±2.68	82.62±1.99	89.28±1.43	87.22±2.76
TS-TCC	83.38±1.53	85.11±2.11	72.24±2.45	75.27±2.64	86.06±1.76	84.34±2.08
SimCLR	83.84±2.15	87.19±1.34	72.51±4.63	75.44±4.77	87.19±1.34	69.99±3.84
CLOCS	88.25±2.48	88.64±2.12	81.40±4.64	83.84±4.03	91.91±2.40	89.76±3.94
COMET (Ours)	88.49±3.28	88.98±2.60	81.65±6.00	84.01±5.61	94.83±1.08	92.48±2.22
1%	TS2vec	81.95±1.85	78.78±0.98	83.83±0.36	79.77±1.44	90.99±1.34	88.69±1.44
TF-C	85.82±2.34	82.79±3.20	81.86±2.36	82.21±2.62	89.14±2.89	89.47±2.92
Mixing-up	84.71±4.11	82.33±6.07	78.17±4.89	79.81±5.33	89.04±3.57	84.67±5.16
TS-TCC	78.61±1.19	78.27±1.79	64.56±2.53	66.23±3.16	79.97±2.17	76.98±2.00
SimCLR	84.19±1.49	87.15±1.27	73.21±3.43	76.31±3.26	87.15±1.27	70.58±2.71
CLOCS	88.80±2.82	88.11±4.37	84.47±3.41	85.57±3.41	94.73±1.59	94.04±2.14
COMET (Ours)	90.52±1.90	88.58±2.93	88.23±1.98	88.23±2.10	95.08±1.50	93.78±1.98
This paper introduces COMET, a hierarchical contrastive representation learning framework tailored for medical time series. COMET leverages all data levels of medical time series, including patient, trial, sample, and observation levels, to capture the intricate complexities of the data. Through extensive experiments on three diverse datasets, we demonstrate that our method surpasses existing state-of-the-art methods in medical time series analysis. Our framework also shows its effectiveness in patient-independent downstream tasks, highlighting its potential to advance medical time series analysis and improve patient care and diagnosis.
One limitation of our work is the presence of label conflicts between different levels. In patient-level consistency, we assume all samples belonging to the same patient are positive while others are negative. In trial-level consistency, we consider samples from the same trial as positive samples and others as negative. This means a positive pair at the patient level may be considered negative at the trial level, as they do not belong to the same trial.
In future research, we aim to investigate the efficacy of our method across a wider range of datasets, with a particular focus on ECG datasets. Additionally, we intend to explore approaches to integrate disease-level consistency within our self-supervised contrastive framework.
We discuss the broader impacts with potential negative social impacts in Appendix H.
This work is partially supported by the National Science Foundation under Grant No. 2245894. Any opinions, findings, conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the funders.
Appendix A A Medical Time Series Example
Here we use EEG data of Alzheimer's as an example. An EEG dataset has many patients with or without Alzheimer's (healthy control). Data collectors take multiple trials on a specific patient. These trials could be collected continuously within a short time or across a long period but follow the same experiment manner. Usually, the timestamps are too long for the deep learning pipeline to learn. For example, a 5 minutes trial with a sampling rate  $256\mathrm{Hz}$  has 76800 timestamps. Researchers generally use data preprocessing to split a trial into many samples, such as 1-second and 3-second short samples, for further representation learning. Each observation denotes a scalar or a vector of real value at a specific timestamp. An experiment with a sampling rate  $256\mathrm{Hz}$  has 256 observations in one second.
Appendix B Data Augmentation Banks
Binomial masking: Generate a mask following a binomial distribution that masks some timestamps of a sample, setting all channels at the masked timestamps to zero.
Channel binomial masking: Generate a mask following a binomial distribution that masks some channels of some timestamps of a sample, setting only a subset of channels at the masked timestamps to zero.
Continuous masking: Mask some continuous sequences of timestamps of a sample, setting all channels at the masked timestamps to zero.
Channel continuous masking: Mask some continuous sequences of timestamps of a sample, setting only a random half of the channels at the masked timestamps to zero.
All true: Do not apply any masking to the sample. Output the raw sample.
Appendix C Shuffle Function Banks
In real-world scenarios, ensuring the presence of samples from the same trial or patient within a training batch becomes increasingly low probability as the number of patients grows. This situation can hinder learning meaningful representations at the trial and patient levels. To address this situation, we have designed two distinct shuffle functions that serve to rearrange samples while also upholding the requirement to include samples originating from the same trial and patient. These functions are called the "trial shuffle" and the "batch shuffle".
Trial shuffle: This function shuffles samples originating from the same trial and subsequently shuffles the trial order. Initially, we arrange the samples by sorting them based on their trial IDs. Next, samples from the same trial are grouped into sets, and the order of samples within each trial set is shuffled. Finally, we sort the trial sets themselves while preserving the order of samples within each respective trial set.
Batch shuffle: This function shuffles samples in a batch and subsequently shuffles the order of batches. The logic of trial and batch shuffle are similar. Initially, we arrange the samples by sorting them based on their trial IDs. Next, we group samples into batch sets, and the order of samples within each batch set is shuffled. Finally, we sort the batch sets themselves while preserving the order of samples within each respective batch set.
Random shuffle: Besides the two specifically designed shuffle functions, this random shuffle function shuffles all the samples in the dataset.
All the shuffle functions mentioned here are designed to shuffle samples within the dataset before training. During training, it is also essential to shuffle samples each epoch to prevent the model from memorizing the dataset and encourage it to learn more useful representations. To address this, we implemented a specially crafted BatchSampler class in PyTorch, following the "batch shuffle" approach. This BatchSampler shuffles the samples locally within each epoch, ensuring that the pre-shuffled sample order is not disrupted significantly. This approach guarantees that each batch contains samples from the same trial. It's worth noting that when a batch consists of samples from the same trial, it must have samples from the same patient.
Appendix D Data Preprocessing
D.1 AD Data Preprocessing
The AD dataset [44] comprises EEG recordings from 12 patients with Alzheimer's disease and 11 healthy controls. Each patient has an average of  $30.0 \pm 12.5$  trials. Each trial corresponds to a 5-second interval with 1280 timestamps (sampled at  $256\mathrm{Hz}$ ) and includes 16 channels. Prior to further processing, each trial is scaled using a standardScaler. To facilitate analysis, we segment each trial into nine half-overlapping samples, where each sample has a duration of 256 timestamps (equivalent to 1 second). Additionally, we assign a unique trial ID and patient ID to each sample based on its origin in terms of the patient and trial. We split training, validation, and test sets in a patient-independent way. We use samples from patient IDs 17 and 18 as the validation set and samples from IDs 19 and 20 as the test set. The rest of the samples are all put into the training set.
D.2 PTB Data Preprocessing