CE-SSL: Computation-Efficient Semi-Supervised Learning for ECG-based Cardiovascular Diseases Detection

TL;DR

A computation-efficient semi-supervised learning paradigm (CE-SSL) is proposed for robust and computation-efficient CVDs detection using ECG that not only outperforms the state-of-the-art methods in multi-label CVDs detection but also consumes fewer GPU footprints, training time, and parameter storage space.

Abstract

The label scarcity problem is the main challenge that hinders the wide application of deep learning systems in automatic cardiovascular diseases (CVDs) detection using electrocardiography (ECG). Tuning pre-trained models alleviates this problem by transferring knowledge learned from large datasets to downstream small datasets. However, bottlenecks in computational efficiency and detection performance limit its clinical applications. It is difficult to improve the detection performance without significantly sacrificing the computational efficiency during model training. Here, we propose a computation-efficient semi-supervised learning paradigm (CE-SSL) for robust and computation-efficient CVDs detection using ECG. It enables a robust adaptation of pre-trained models on downstream datasets with limited supervision and high computational efficiency. First, a random-deactivation technique is developed to achieve robust and fast low-rank adaptation of pre-trained weights. Subsequently, we propose a one-shot rank allocation module to determine the optimal ranks for the update matrices of the pre-trained weights. Finally, a lightweight semi-supervised learning pipeline is introduced to enhance model performance by leveraging labeled and unlabeled data with high computational efficiency. Extensive experiments on four downstream datasets demonstrate that CE-SSL not only outperforms the state-of-the-art methods in multi-label CVDs detection but also consumes fewer GPU footprints, training time, and parameter storage space. As such, this paradigm provides an effective solution for achieving high computational efficiency and robust detection performance in the clinical applications of pre-trained models under limited supervision. Code and Supplementary Materials are available at https://github.com/KAZABANA/CE-SSL

Figures (7)

• Figure 1: Overview of CE-SSL.a. The architecture of the pre-trained backbone model. It consists of three convolution blocks, eight self-attention blocks, and one classification block in the base backbone. It is pre-trained on a public 12-lead ECG dataset using the supervised multi-label binary cross-entropy loss. b. Random-deactivation low-rank adaptation. On the downstream datasets, the pre-trained weights $\{W_{0}^{i}\}_{i=1}^{n}$ in the backbone are updated by the proposed random-deactivation low-rank adaptation. It randomly activates or deactivates the low-rank matrices ($\{B^{i}\}_{i=1}^{n}$ and $\{A^{i}\}_{i=1}^{n}$) in each trainable layer with a given probability $p$. Note that all the pre-trained weights are frozen during model training. All the low-rank matrices are activated and merged into the pre-trained weights in the testing stage. The merge process generates an ensemble network combining all the sub-networks produced by the random deactivation operation. c. One-shot rank allocation. The ranks of the low-rank matrices are determined by the proposed one-shot rank allocation method using only one gradient backward on the labeled samples. The matrices with high importance are allocated with a higher rank than those with low importance. d. Lightweight semi-supervised learning. During the low-rank adaptation process, unlabeled samples from the downstream datasets are combined with the labeled samples to estimate the statistics in batch-normalization layers. Subsequently, only the labeled data is forwarded to the self-attention blocks for CVDs detection, and the unlabeled data is released in GPU memory. This lightweight semi-supervised pipeline improves the model performance in a computational-efficient way.
• Figure 2: Comparison between CE-SSL and parameter-efficient semi-supervised methods on the base backbone. Circles with various colors denote different models, and their size represents the number of trainable parameters. The training time for each optimization iteration (Time/iter) of different methods is also reported. The gray dotted lines represent the performance of the FixMatch baseline without parameter-efficient training (approximately 9.505M trainable parameters). The first row of the figure presents the performance of different models with sufficient parameter budgets ($r=16$), while the second row reports their performance under limited parameter budgets ($r=4$).
• Figure 3: Effect of the ratio of labeled samples for model training. We adjust the ratio of the labeled samples in the dataset from 0.05 to 0.15 and report the averaged performance of different models across four datasets and six random seeds.
• Figure 4: Effect of the rank initialization. Averaged performance of the CE-SSL across four datasets and six random seeds under different initial ranks $r$.
• Figure 5: Effect of the ratio of important weight matrices. We adjust the ratio of the important weight matrices to the total number of weight matrices and report the averaged performance across four datasets and six random seeds. Important weights are adapted with rank $r$ while the remaining weights are adapted with rank $\frac{1}{2}r$.