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

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 . 1 Introduction . Clinical databases can comprise instances that are either unlabelled or labelled with patient attribute information , such as disease class , sex , and age . The process of manually searching for relevant instances in , and extracting information from , such databases underpin a multitude of tasks [ 1 ] . For example , clinicians extract a disease diagnosis from patient data , researchers involved in clinical trials search for and recruit patients satisfying specific inclusion criteria [ 2 ] , and educators retrieve relevant information as part of the continuing medical education scheme [ 3 ] . This manual search-and-extract process , however , is hampered by the rapid growth of large-scale clinical databases and the increased prevalence of unlabelled instances ; those for which patient attribute information is unavailable . Given such a setting , in this paper , we are interested in addressing two questions : given a large , unlabelled clinical database , ( 1 ) how do we extract attribute information from such unlabelled instances ? and ( 2 ) how do we reliably search for and retrieve relevant instances ? To address the former , the task of clustering holds value . In this setting , a centroid groups together instances that share some similarities . Recent research has focused on exploiting existing clustering algorithms , such as k-means , to group similar patients from electronic heath record ( EHR ) data [ 4 , 5 ] . Such methods , however , are exclusively unsupervised ; they do not exploit patient attribute information . To address the second question , the task of information retrieval holds promise . In this setting , a query associated with a set of desired attributes is exploited to retrieve a relevant instance . Recent research has focused predominantly on retrieving medical images [ 6 ] , clinical text [ 7 ] , and EHR data [ 8 ] , with minimal emphasis on medical time-series data [ 9 ] . These methods do not extend to cardiac 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . time-series data nor do they account for search based on multiple patient attributes . Most notably , previous work performs either clustering or retrieval , and not both . In this work , we address both questions while exploiting large-scale electrocardiogram ( ECG ) databases comprising patient attribute information . Our contributions are the following : ( 1 ) we propose a supervised contrastive learning framework , CROCS , in which we attract representations of cardiac signals associated with a unique set of patient attributes to embeddings , entitled clinical prototypes . Such attribute-specific prototypes , which create “ islands ” of similar representations [ 10 ] , allow for both the clustering and retrieval of cardiac signals based on multiple patient attributes . ( 2 ) We show that CROCS outperforms the state-of-the-art method , DTC , in the clustering setting and retrieves relevant cardiac signals from a large database . At the same time , clinical prototypes adopt a semantically meaningful arrangement and thus confer a high degree of interpretability . 2 Related work . Clinical representation learning and clustering Learning meaningful representations of clinical data is an ongoing research endeavour . Recent research has focused on learning representations from EHR data [ 11 , 12 , 13 , 14 , 15 ] and via auto-encoders , which are then clustered using existing methods , such as k-means [ 4 , 5 ] . As for time-series data , auto-encoders are learned with [ 16 ] or without [ 17 ] an auxiliary clustering objective , salient features ( shapelets ) are identified in an unsupervised manner [ 18 , 19 ] , and patient-specific representations are learned via contrastive learning [ 20 ] . Li et al . [ 21 ] learn prototypes , or representative embeddings , via the ProtoNCE loss and cluster instances using k-means . Their work builds upon recent research in the contrastive learning literature [ 22 , 23 , 24 ] . Similar to our work is that of Gee et al . [ 11 ] where prototypes are learned for the clustering of time-series signals . Their prototypes , however , can not cluster instances based on multiple patient attributes and do not extend to the retrieval setting . Clinical information retrieval ( IR ) Retrieving clinical data from a large database has been a longstanding goal of researchers within healthcare [ 25 ] . Such research has involved the retrieval of clinical documents [ 26 , 7 , 27 , 28 ] where , for example , Avolio et al . [ 29 ] map text queries to an ontology known as SNOMED , before retrieving relevant clinical documents . Recent research has focused on the retrieval of biomedical images [ 6 , 30 ] , and EHR data [ 31 , 32 ] to discover patient cohorts in a clinical database [ 8 ] . Goodwin et al . [ 9 ] implement an unsupervised patient cohort retrieval system by exploiting clinical text and time-series data . These approaches , however , do not explore cardiac signals , can not account for multiple patient attributes , and are unable to also cluster instances . To the best of our knowledge , we are the first to design a learning framework that allows for both the clustering and retrieval of cardiac signals based on multiple patient attributes . 3 Background . Supervised clustering We learn a function , fθ : x ∈ RD −→ v ∈ RE , parameterized by θ , that maps a D-dimensional input , x , to an E-dimensional representation , v. We also have a labelled dataset , Dl = { xi , Ai } Nli=1 , where each instance , xi , is associated with a set of discrete patient attributes , Ai = { αic , αis , αia } where αc = disease class , αs = sex and αa = age . Supervised clustering can involve learning M Nl centroids with each representing a unique set of attributes , { αjc , αjs , αja } ∈ { Aj } Mj=1 , and grouping similar instances together . Given unlabelled instances , { xu } Nu=1 , the centroid closest to each representation , vu = fθ ( xu ) , is used to infer the latter ’ s attributes . In this work , we learn cluster centroids which are more formally introduced in Sec . 4.1 . Information retrieval IR involves searching through a large , unlabelled dataset , { xu } Nu=1 , and retrieving a relevant instance , xu . However , relevance , defined based on whether an instance satisfies some criteria , is difficult to ascertain when instances are unlabelled . Typically , a query in the form of an embedding which represents a desired set of attributes , Aj , retrieves its closest ( and most relevant ) representation , vu = fθ ( xu ) , and infers the latter ’ s attributes . In this work , we learn a set of query embeddings . As will become apparent in Sec . 5 , these embeddings can also be treated as centroids , like those outlined in supervised clustering , and will thus serve a dual purpose . 4 Methods . 4.1 Attribute-specific clinical prototypes . We aim to learn embeddings , referred to as clinical prototypes , that can be exploited for both the clustering and retrieval of cardiac signals based on multiple patient attributes . In the clustering setting , the goal is to annotate unlabelled instances with a set of patient attributes . To that end , we exploit clinical prototypes as centroids of clusters to which such unlabelled instances are assigned ( see Fig . 1 left ) . In the retrieval setting , the goal is to retrieve unlabelled instances based on a set of patient attributes . To that end , we exploit each clinical prototype as a query to search through an unlabelled database and retrieve instances to which it is most similar ( see Fig . 1 right ) . In designing clinical prototypes , we take inspiration from the field of natural language processing ( NLP ) where a learnable word embedding represents a unique word . In our case , each prototype represents a unique combination of discrete patient attributes . Formally , given the attributes , αc , αs , and αa , we would have M = |αc| × |αs| × |αa| such unique combinations denoted by A = { αjc , αjs , αja } Mj=1 . We associate each combination , Aj ∈ A , with a learnable prototype , pAj ∈ RE , for a set of M prototypes , P = { pAj } Mj=1 . Note that this framework extends to any number of attributes . In the next section , we outline how to learn these clinical prototypes . 4.2 Learning attribute-specific clinical prototypes . Clinical prototypes will serve a dual purpose of attribute-specific clustering and retrieval . As such , prototypes will need to be in proximity to a subset of representations of instances associated with a specific set of patient attributes . To achieve this proximity , we leverage the contrastive learning framework which involves a sequence of attractions and repulsions , as explained next . Hard assignment We encourage the representation , vi = fθ ( xi ) , of an instance , xi , associated with a set of attributes , Ai ∈ A , to be similar to the single clinical prototype , pAj , that shares the exact same set of attributes ( i.e. , Ai = Aj ) , and dissimilar to the remaining clinical prototypes , { pAk } k 6=j . To achieve this , we optimize LNCE−hard for a mini-batch of B instances ( Eq . 1 ) . Intuitively , it heavily penalizes the learner if less probability mass is placed on the similarity of vi and pAj than on the similarity of other representation-and-prototype pairs . We choose to quantify the cosine similarity of pairs , s ( vi , pAj ) , alongside a temperature parameter , τs , as is done by Kiyasseh et al . [ 20 ] . LNCE−hard = − 1 B B∑ i=1 log ( es ( vi , pAj ) ∑M k e s ( vi , pAk ) ) s ( vi , pAj ) = vi.pAj ‖vi‖‖pAj‖ . 1 τs ( 1 ) We refer to this many-to-one mapping from representations to clinical prototype as a hard assignment . Such an assignment , however , implies that a prototype is unlikely to extract potentially useful information from a representation whose attributes are not exactly the same as those of the prototype . We quantify this limitation in Sec . 6.4 and propose an alternative assignment next . Soft assignment To overcome the limitation of a hard assignment , we encourage the representation , vi , to be similar to a subset of clinical prototypes , L ⊂ P ( see Fig . 2 ) . We must take caution , however , to avoid erroneously attracting representations to clinical prototypes from a different class . Doing so would reduce class-specific margins and thus hinder the downstream clustering and retrieval tasks . Uniform attraction . We chose the subset to include prototypes , L = { pAl } |L| l=1 , that share the same disease class , αic , with the representation , vi , implying that Al ∈ { αic , αls , αla } |L| l=1 ⊂ A . Note that the clinical prototypes in the subset , L , continue to represent attributes that vary along the dimensions of patient sex and age ( αs , αa ) . Therefore , attracting vi to prototypes in L uniformly will likely cause the latter to become minimally distinguishable across sex and age . This is an undesired outcome in light of our goal of learning attribute-specific prototypes . We avoid this behaviour by modulating the degree of attraction between vi and all prototypes in the set , P , as outlined next . Modulated attraction . The attractive force between vi and pAj is reflected by the corresponding LNCE−hard term ( Eq . 1 ) . By placing less probability mass on s ( vi , pAj ) ( i.e. , less similarity ) than on s ( vi , pAk ) ∀k , k 6= j , the learner incurs a higher loss and thus attracts the pair . We extend this logic to all prototypes to obtainM LNCE−hard terms per representation . To modulate these attractions , we introduce a weight , wij ∈ { wij } Mj=1 , as a coefficient of the j-th loss term ( Eq . 2 ) . Each weight , wij , quantifies the degree of matching between attributes of the representation , Ai = { αic , αis , αia } , and those of the clinical prototype , Aj = { αjc , αjs , αja } , as reflected by q ( Ai , Aj ) ∈ R. We define 1 as the indicator function and τω as a temperature parameter that determines how soft the representation-andprototype attraction is . For example , as τω −→∞ , this approach reverts to the uniform attraction setup . The intuition is that a stronger attraction ( ↑ wij ) should exist for a representation-and-prototype pair that shares more attributes . We also avoid the erroneous attraction of pairs from different classes ( i.e. , αic 6= αjc ) by setting ωij = 0 . When visualizing the UMAP projection [ 33 ] of prototypes learned with a uniform attraction ( τω =∞ ) ( Fig . 3 left ) and those learned with a modulated attraction ( τω 6=∞ ) ( Fig . 3 centre ) , we show that the latter become more linearly separable across sex . LNCE−soft = − 1 B B∑ i=1 M∑ j=1 ωij log ( es ( vi , pAj ) ∑M k e s ( vi , pAk ) ) ( 2 ) ωij = { eq ( Ai , Aj ) ∑|L| l e q ( Ai , Aj ) if αic = α j c 0 otherwise q ( Ai , Aj ) = 1 τω . [ 1 ( αic = α j c ) + 1 ( α i s = α j s ) + 1 ( α i a = α j a ) ] Arrangement of clinical prototypes Clinical prototypes would confer a high degree of interpretability if they also captured the semantic relationships between attributes . Concretely , prototypes representing similar attribute sets ( e.g. , adjacent age groups ) should be similar to one another . This is analogous to the high similarity of word embeddings representing semantically similar words [ 34 ] . To capture these semantic relationships , we encourage class-specific prototypes to maintain some desired distance between one another . As such , each pair of M clinical prototypes , pAj , pAk ∀j , k ∈ [ 1 , M ] is associated with an empirical and ground-truth ( desired ) distance . For the former , we normalize the prototypes ( L2 norm ) and calculate their Euclidean distance , d̂jk = ∥∥pAj − pAk∥∥2 ∀j , k ∈ [ 1 , M ] . For the latter , we define the ground-truth distance as djk = β × dH ∈ R , where dH ( Aj , Ak ) ∈ Z+ is the Hamming distance between a pair of discrete attribute sets . Intuitively , the Hamming distance counts the number of attribute mismatches and β ∈ R penalizes each mismatch . Therefore , we can generate an empirical set , { d̂jk } Mj , k=1 and a ground-truth set , { djk } Mj , k=1 , of distance values . By minimizing the mean-squared error between these two sets , we learn clinical prototypes that adopt a semantically meaningful arrangement ( see Fig . 3 right ) . Since we are only interested in adopting this arrangement for prototypes of the same class ( i.e. , αjc = α k c ) , we incorporate the regularization term , Lreg , into the final objective function , Ltot . Lreg = M∑ j , k=1 ( d̂jk − djk ) 2 ⇔ αjc = αkc Ltot = LNCE−soft + Lreg ( 3 )

Summary: This paper proposes a clustering and information retrieval algorithm for unannotated large databases of physiological signals (in this case ECG signals). The method is based on a contrastive learning approach, which has recently been suggested for semi-supervised learning on ECG data. The proposed approaches are evaluated on two large ECG datasets and compared with state-of-the-art approaches.