DROPS: DEEP RETRIEVAL OF PHYSIOLOGICAL SIG-NALS VIA ATTRIBUTE-SPECIFIC CLINICAL PROTO-TYPES

Abstract

The ongoing digitization of health records within the healthcare industry results in large-scale datasets. Manually extracting clinically-useful insight from such datasets is non-trivial. However, doing so at scale while simultaneously leveraging patient-specific attributes such as sex and age can assist with clinical-trial enrollment, medical school educational endeavours, and the evaluation of the fairness of neural networks. To facilitate the reliable extraction of clinical information, we propose to learn embeddings, known as clinical prototypes (CPs), via supervised contrastive learning. We show that CPs can be efficiently used for large-scale retrieval and clustering of physiological signals based on multiple patient attributes. We also show that CPs capture attribute-specific semantic relationships.

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. (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.

2. RELATED WORK

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.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, p A ∈ 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 ∈ R D -→ v ∈ R E parameterized by θ, that maps a D-dimensional input, x, to an E-dimensional representation, v. In information retrieval systems, a



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 ofKiyasseh et al. (2020), which learns patient-specific representations while pre-training via contrastive learning, and Garnot & Landrieu

