[Paper Review] CROCS: Clustering and Retrieval of Cardiac Signals Based on Patient Disease Class, Sex, and Age

Dani Kiyasseh, Tingting Zhu+1 more|arXiv (Cornell University)|Nov 28, 2020

Machine Learning in Healthcare4 citations

TL;DR

CROCS is a supervised contrastive learning framework that clusters and retrieves cardiac signals using patient attributes like disease class, sex, and age. By learning clinical prototypes that attract representations of similar patient profiles, CROCS outperforms state-of-the-art methods in clustering (90.3% accuracy on Chapman, 72.8% AMI) and retrieval while enabling semantically interpretable, attribute-based organization of ECG data.

ABSTRACT

The process of manually searching for relevant instances in, and extracting information from, clinical databases underpin a multitude of clinical tasks. Such tasks include disease diagnosis, clinical trial recruitment, and continuing medical education. This manual search-and-extract process, however, has been hampered by the growth of large-scale clinical databases and the increased prevalence of unlabelled instances. To address this challenge, we propose a supervised contrastive learning framework, CROCS, where representations of cardiac signals associated with a set of patient-specific attributes (e.g., disease class, sex, age) are attracted to learnable embeddings entitled clinical prototypes. We exploit such prototypes for both the clustering and retrieval of unlabelled cardiac signals based on multiple patient attributes. We show that CROCS outperforms the state-of-the-art method, DTC, when clustering and also retrieves relevant cardiac signals from a large database. We also show that clinical prototypes adopt a semantically meaningful arrangement based on patient attributes and thus confer a high degree of interpretability.

Motivation & Objective

• To address the challenge of manually searching large, unlabelled clinical databases for relevant cardiac signals.
• To develop a unified framework that supports both clustering and retrieval of ECG signals based on multiple patient attributes.
• To improve interpretability by organizing learned clinical prototypes in a semantically meaningful way according to patient characteristics.
• To outperform existing methods in clustering and retrieval performance on real-world ECG databases.

Proposed method

• CROCS uses a neural network to embed cardiac signals into dense representations based on input ECG data.
• It introduces learnable clinical prototypes for each unique combination of patient attributes (e.g., disease class, sex, age), acting as attractors for similar signal representations.
• The framework employs a soft-assignment contrastive loss (NCE-soft) to encourage representations to cluster around prototypes sharing the same attributes.
• A regularization term is applied to ensure prototypes maintain semantic relationships between attributes, improving interpretability.
• The model is trained end-to-end using a combined loss: NCE-soft loss for representation learning and a regularization loss for prototype semantics.
• After training, prototypes enable both clustering (by assigning signals to nearest prototype) and retrieval (by querying with attribute combinations).

Experimental results

Research questions

• RQ1 Can a unified deep learning framework effectively perform both clustering and retrieval of unlabelled cardiac signals using multiple patient attributes?
• RQ2 How does the use of learnable clinical prototypes improve clustering and retrieval performance compared to existing methods?
• RQ3 To what extent do the learned prototypes reflect meaningful semantic relationships between patient attributes like disease class, sex, and age?
• RQ4 Does soft assignment of representations to prototypes yield better performance than hard assignment in this multi-attribute setting?
• RQ5 Can the framework maintain high interpretability while achieving state-of-the-art performance in both clustering and retrieval tasks?

Key findings

• CROCS achieved 90.3% accuracy and 72.8% AMI on the Chapman ECG database for clustering based on disease class, outperforming the state-of-the-art DTC method.
• On the PTB-XL dataset, CROCS achieved 76.0% accuracy and 76.1% AMI for disease class clustering, demonstrating robustness across datasets.
• The framework with soft assignment and regularization (NCE-soft + reg) consistently outperformed variants, achieving 57.4% accuracy and 38.0% AMI for sex and age attributes on Chapman.
• Clinical prototypes were found to be semantically arranged, with similar attribute combinations forming coherent clusters, enhancing interpretability.
• Ablation studies confirmed that soft assignment and the regularization term significantly improved performance and prototype interpretability.
• CROCS successfully retrieved relevant cardiac signals based on multi-attribute queries, with performance superior to DTC in both clustering and retrieval tasks.