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: Cardiac signals when unlabeled are difficult to query for given a set of attributes of interest such as disease class, age and sex.  The paper proposes a cluster-based retrieval system for this purpose. The algorithm proposed is of two-fold - (1) Learn a function to map a cardiac signal to an embedding closer in representation to a clinical prototype (embedding representative of a unique attribute set)   - this is done using a contrastive-learning framework assuming an initial-set of labeled data where cardiac-signals are annotated with class, age and sex and (2) Given a query attribute set - use its clinical prototype embedding to rank the  cardiac signals. Here the function from step 1 is used to get an embedding representation for cardiac signals and then their Euclidean distance from clinical prototypes is considered for ranking in the retrieval. The paper evaluates clustering and retrieval tasks on two datasets Chapman and PTB-XL.