DROPS: Deep Retrieval of Physiological Signals via Attribute-specific Clinical Prototypes

1 INTRODUCTION . Physiological data are being collected at a burgeoning rate . Such growth is driven by the digitization of previous patient records , the presence of novel health monitoring and recording systems , and the recent recommendation to facilitate the exchange of health records ( European Commission , 2019 ) . This engenders large-scale datasets from which the manual extraction of clinically-useful insight is non-trivial . Such insight can include , but is not limited to , medical diagnoses , prognoses , or treatment . In the presence of large-scale datasets , retrieving instances based on some user-defined criteria has been a longstanding goal within the machine learning community ( Manning et al. , 2008 ) . This information retrieval ( IR ) process typically consists of a query that is used to search through a large database and retrieve matched instances . Within healthcare , the importance of an IR system is threefold ( Hersh & Hickam , 1998 ; Hersh , 2008 ) . First , it provides researchers with greater control over which patients to choose for clinical trial recruitment . Second , IR systems can serve as an educational and diagnostic tool , allowing physicians to identify seemingly similar patients who exhibit different clinical parameters and vice versa . Lastly , if the query were to consist of sensitive attributes such as sex , age , and race , then such a system would allow researchers to more reliably evaluate the individual and counterfactual fairness of a particular model ( Verma & Rubin , 2018 ) . To illustrate this point , let us assume the presence of a query instance that corresponds to a patient with an abnormality of the heart , atrial fibrillation , who is male and under the age of 25 . To reliably determine the sensitivity of a model with respect to sex , one would observe its response when exposed to a counterfactual instance , namely the exact same instance but with a different sex label ( Kusner et al. , 2017 ) . At present , deep-learning based IR systems within the healthcare domain fail to incorporate such patient-specific attributes . Existing IR systems which retrieve instances from the electronic health records ( Wang et al. , 2019 ; Chamberlin et al. , 2019 ) do not incorporate an attribute-specific search and do not trivially extend to physiological signals . In this paper , we propose to learn embeddings , referred to as clinical prototypes ( CPs ) . CPs are efficient descriptors of a combination of patient-specific attributes , such as disease , sex , and age . We learn these embeddings via contrastive learning whereby representations of instances are encouraged to be similar to their corresponding clinical prototype and dissimilar to the others . To the best of our knowledge , we are the first to design a supervised contrastive learning based large-scale retrieval system for electrocardiogram ( ECG ) signals . Contributions . Our contributions are the following : • Attribute-specific clinical prototypes - we propose a supervised contrastive learning framework to learn embeddings , referred to as clinical prototypes ( CPs ) , that are efficient descriptors of a set of patient attributes , e.g. , disease , sex , and age . • Deep retrieval and clustering - we exploit CPs to retrieve instances corresponding to a specific patient-attribute combination and assign instances to various clusters . 2 RELATED WORK . Clinical representation learning involves meaningfully representing clinical data for solving tasks . Most research attempts to learn representations of electronic health records ( EHRs ) ( Miotto et al. , 2016 ; Gee et al. , 2019 ; Liu et al. , 2019 ; Li et al. , 2020b ; Biswal et al. , 2020 ; Darabi et al. , 2020 ) in a generative manner . For example , Landi et al . ( 2020 ) and Huang et al . ( 2019 ) implement an autoencoder to learn patient representations . These representations are then clustered either in a hierarchical manner or via K-means . Other methods involve learning prototypes . For example , Li et al . ( 2020a ) propose to do so via the ProtoNCE loss . Our approach , unlike theirs , exploits readily-available patient-attribute data and is not dependent upon the K-means algorithm . Moreover , Van Looveren & Klaise ( 2019 ) learn to perturb prototypes to derive interpretable counterfactual instances . Most similar to our work is that of Kiyasseh et al . ( 2020 ) , which learns patient-specific representations while pre-training via contrastive learning , and Garnot & Landrieu ( 2020 ) where the distance between class prototypes learned in an end-to-end manner is regularized based on a predefined tree hierarchy . In contrast , we learn attribute-specific prototypes via supervised contrastive learning and capture their semantic relationships via distance-based regularization . Clinical information retrieval whereby instances similar to a query are retrieved was first introduced in 1990 ( Hersh & Greenes , 1990 ) . Most research in this domain revolves around text ( Gurulingappa et al. , 2016 ; Wang et al. , 2017 ; Rhine , 2017 ; Wallace et al. , 2016 ) . For example , D ’ Avolio et al . ( 2010 ) map text to SNOMED concepts to retrieve clinical documents . More recently , IR has been performed with biomedical images , and is referred to as content-based image retrieval ( Saritha et al. , 2019 ; Chittajallu et al. , 2019 ) . Others have extended this concept to EHR data ( Goodwin & Harabagiu , 2018 ; Wang et al. , 2019 ; Chamberlin et al. , 2019 ) . For example , Chamberlin et al . ( 2019 ) implement rudimentary IR methods such as divergence from randomness on the UPMC and MIMIC III ( Johnson et al. , 2016 ) datasets with the aim of discovering patient cohorts . In contrast to such methods , we implement a deep-learning based clinical information retrieval system for physiological signals . 3 METHODS . 3.1 ATTRIBUTE-SPECIFIC CLINICAL PROTOTYPES . Information retrieval systems typically necessitate a query of some sort that is exploited to search through a large database and retrieve instances that satisfy criteria outlined by the initial query . Such a query can take on a multitude of forms ( e.g. , text , image , audio , etc . ) depending on the modality of instances in the database . As we are primarily interested in large databases comprising physiological signals , we design a query that is based on such signals . Moreover , the type and specificity of instances that are retrieved highly depend on the criteria outlined by a query . In our context , these criteria comprise patient attribute information such as disease class , sex , and age . As a result , our query should be capable of retrieving physiological instances that are associated with the aforementioned patient attributes . To achieve this , we propose to learn a set of query embeddings , P , analogous to word embeddings in natural language processing , where |P | = M , representing each of the M possible patient-attribute combinations within a dataset . Each embedding , pA ∈ P , is an efficient descriptor of a set of attributes A = { α1 , α2 , α3 } where α1 = disease class , α2 = sex and α3 = age . Given this attribute-specific interpretation , we refer to such embeddings as attributespecific clinical prototypes or CPs . We propose to learn such CPs , in an end-to-end manner , via contrastive learning , as explained next . 3.2 LEARNING ATTRIBUTE-SPECIFIC CLINICAL PROTOTYPES . We assume the presence of a learner , fθ : x ∈ RD −→ v ∈ RE parameterized by θ , that maps a D-dimensional input , x , to an E-dimensional representation , v. In information retrieval systems , a reliable query should accurately retrieve instances that satisfy certain criteria . When dealing with representations , such reliability can be designed for by ensuring that the query embedding is in a similar subspace , and thus in proximity , to the representations of instances that satisfy said criteria . To achieve this proximity , we can attract query embeddings to representations in the database which share patient attribute information and repel them from those which do not . More specifically , we attract each representation of an instance associated with a set of attributes , A , to the clinical prototype , pA , that describes that same set . This can be achieved in two ways . Hard assignment . We encourage each representation , vi = fθ ( xi ) , of an instance , xi , associated with a particular set of attributes , A , to be similar to the single clinical prototype , pA , that describes the exact same set of attributes , and dissimilar to the remaining clinical prototypes , pj , where j 6= A . We quantify this similarity , s ( vi , pA ) , using the cosine similarity , with a temperature parameter , τs . By mapping each representation to a single CP , we refer to this as a hard assignment . To encourage this behaviour , one would optimize the following objective function for a mini-batch of size , B : Lcontra−hard = − 1 B B∑ i=1 log ( es ( vi , pA ) ∑M j e s ( vi , pj ) ) ( 1 ) s ( vi , pA ) = vi · pA ‖vi‖‖pA‖ · 1 τs ( 2 ) The hard assignment approach assumes a bijective relationship between each representation , vi , and each clinical prototype , pA . This many-to-one mapping implies that CPs without a perfect attribute match ( i.e. , a near miss ) do not leverage potentially useful information from representations that exhibit some , albeit imperfect , overlap in patient attributes . Soft assignment . To overcome the limitations of a hard assignment , we propose a soft assignment approach whereby each representation , vi , is attracted to a subset of disease class-specific clinical prototypes , L ⊂ P . We opted for this class-specific setup to avoid erroneously attracting representations to CPs from a different class . This attraction would hinder retrievals based on disease class . Recall , though , that the CPs in L still describe different sets of non-class attributes . By attracting each representation to these CPs uniformly , CPs within a class will collapse to a single point , and thus be unable to distinguish between various patient attributes . We support this claim in Fig . 1a where we illustrate the t-SNE projection of the CPs for five different disease classes . To avoid such collapse , we modulate the attraction between each representation , vi , associated with the attribute set , Ai , and each clinical prototype , pk , associated with the attribute set , Ak . More specifically , we introduce a weight , wik , that is dependent on the discrepancy between the respective attribute sets , d ( Ai , Ak ) . This formulation , characterized by greater attraction between CPs and representations with more similar attribute sets , results in the CPs shown in Fig . 1b . Formally , we optimize the following objective function . Lcontra−soft = − 1 B B∑ i=1 [ M∑ k=1 ωik log ( es ( vi , pk ) ∑M j e s ( vi , pj ) ) ] ( 3 ) ωik = { ed ( Ai , Ak ) ∑|L| j e d ( Ai , Aj ) if disease classi = disease classk 0 otherwise ( 4 ) d ( Ai , Ak ) = [ δ ( disease classi = disease classk ) + δ ( sexi = sexk ) + δ ( agei = agek ) ] · 1 τω ( 5 ) where δ is the Kronecker delta function that evaluates to one if the argument is true and zero otherwise and τω is a temperature parameter that determines how soft the assignment is . For example , as τω −→ 0 , the loss-term approaches the hard assignment formulation . Arrangement of clinical prototypes . Clinical prototypes , learned in an end-to-end manner , as presented will exhibit a high and desirable degree of inter-class separability . However , prototypes within a class are still at risk of collapsing to a select few points . This would decrease their utility for attribute-based querying of instances . We would prefer to learn CPs that not only distinguish between different sets of attributes , but also reflect the semantic relationships between one another , as is done with word embeddings in natural language processing . Capturing semantic relationships is desirable as it allows for improved interpretability of the retrieval process . For example , representations that are equidistant from two CPs can indicate a graded difference in a particular attribute such as age . In Fig . 2 ( left ) , we illustrate the desired arrangement of 16 CPs for two arbitrary classes . To arrive at this arrangement , we propose a distance-based regularization term , derived as follows . First , we use the L2 norm to normalize all CPs in the set , P , and calculate the empirical pairwise Euclidean distance between them . As a result , we generate the matrix D̂ ∈ RM×M ( Fig . 2 right ) . As we are only interested in the pairwise distances for prototypes within the same class , we focus on the sub-matrices , D̂c ∈ RM/C×M/C with c ∈ [ 1 , . . . , C ] where C is the number of class-specific clusters . Each class ’ corresponding ground-truth sub-matrix , Dc , is populated with distances , dE , that reflect the semantic relationships between CPs . We choose dE = β × S where β ∈ R is a user-defined distance and S ∈ Z+ is the integer distance between attributes . For example , a pair of CPs that differ according to sex with A1 = { AFIB , M , < 25 } and A2 = { AFIB , F , < 25 } , results in S = 1 and dE = β × 1 . We then calculate the mean-squared error between D̂c and Dc ∀c ∈ [ 1 , . . . , C ] . Our final objective function thus includes the soft contrastive loss ( Eq . 3 ) and the regression loss . Lreg = C M2 C∑ c=1 ( D̂c −Dc ) 2 ( 6 ) Lcomb = Lcontra−soft + Lreg ( 7 )

Summary: This paper proposes to tackle the problem of retrieving and clustering physiological signals by learning clinical prototypes via supervised contrastive learning. Three readily available patient attributes, disease, age, and sex, are used to assist the learning. A hard assignment of samples to prototype is proposed, and further relaxed to a soft one to utilize the samples that do not have exactly matched attribute set. A regularization term is also proposed to encourage intra-cluster distances. Two ECG datasets are used to evaluate the proposed model.