DROPS: D EEP R ETRIEVAL OF P HYSIOLOGICAL S IGNALS VIA A TTRIBUTE - SPECIFIC C LINICAL P ROTOTYPES
Anonymous authors
Paper under double-blind review
A BSTRACT
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-speciﬁc 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 efﬁciently used for large-scale retrieval and clustering of physiological signals.
We also show that CPs capture attribute-speciﬁc semantic relationships.
I NTRODUCTION
Physiological data are being collected at a burgeoning rate.
Such growth is driven by the digitization of previous patient records, the presence of novel monitoring 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.
In the presence of large-scale datasets, retrieving instances based on some user-deﬁned criteria has been a longstanding goal within the machine learning community.
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.
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 ﬁbrillation, 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, IR systems within the healthcare domain fail to incorporate such patient-speciﬁc attributes.
Several researchers have implemented IR systems to retrieve instances from the electronic health records (Wang et al., 2019; Chamberlin et al., 2019).
However, most of these approaches do not incorporate an attribute-speciﬁc search and do not trivially extend to physiological signals.
Instead, we propose to learn embeddings, referred to as clinical prototypes (CPs).
CPs are efﬁcient descriptors of a combination of patient-speciﬁc 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 ﬁrst to design a contrastive learning based large-scale retrieval system for physiological signals.
Contributions.
Our contributions are the following:
• Attribute-speciﬁc clinical prototypes - we propose a supervised contrastive learning framework to learn embeddings, referred to as clinical prototypes (CPs), that are efﬁcient 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 speciﬁc patient-attribute combination and assign instances to various clusters.
R ELATED W ORK
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.
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-speciﬁc 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 predeﬁned tree hierarchy.
In contrast, we learn attribute-speciﬁc 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 ﬁrst 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.
M ETHODS 3. A TTRIBUTE - SPECIFIC C LINICAL P ROTOTYPES
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 speciﬁcity 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 efﬁcient descriptor of a set of attributes A
= { α 1 , α 2 , α 3 } where α 1 = disease class , α 2 = sex and α 3 = age .
Given this attribute-speciﬁc interpretation, we refer to such embeddings as attributespeciﬁc clinical prototypes or CPs.
We propose to learn such CPs, in an end-to-end manner, via contrastive learning, as explained next.
3. L EARNING A TTRIBUTE - SPECIFIC C LINICAL P ROTOTYPES
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 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 speciﬁcally, we attract each representation of an instance associated with a set of attributes, A , to the clinical prototype, p A , that describes that same set.
This can be achieved in two ways.
Hard assignment.
We encourage each representation, v i = f θ ( x i ) , of an instance, x i , associated with a particular set of attributes, A , to be similar to the single clinical prototype, p A , that describes the exact same set of attributes, and dissimilar to the remaining clinical prototypes, p j , where j (cid:54) =
A .
We quantify this similarity, s ( v i , p A ) , 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 :
[Equation]
The hard assignment approach assumes a bijective relationship between each representation, v i , and each clinical prototype, p A .
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, v i , is attracted to a subset of disease class-speciﬁc clinical prototypes, L ⊂ P .
We opted for this class-speciﬁc 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 ﬁve different disease classes.
To avoid such collapse, we modulate the attraction between each representation, v i , associated with the attribute set, A i , and each clinical prototype, p k , associated with the attribute set, A k .
More speciﬁcally, we introduce a weight, w ik , that is dependent on the discrepancy between the respective attribute sets, d ( A i , A k ) .
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. [Equation]
if disease class i
[Equation]
disease class k otherwise (
[Equation]
(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 reﬂect 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. [Figure]
To arrive at this arrangement, we propose a distance-based regularization term, derived as follows.
First, we use the L 2 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 ∈ R M × 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 ∈ R M/C × M/C with c ∈ [1 , . . .
, C ] where C is the number of class-speciﬁc clusters. Each class’ corresponding ground-truth sub-matrix, D c , is populated with distances, d E , that reﬂect the semantic relationships between CPs.
We choose d E = β × S where β ∈ R is a user-deﬁned distance and S ∈ Z
+ is the integer distance between attributes.
For example, a pair of CPs that differ according to sex with A 1 = { AFIB , M , < 25 } and A 2 = { AFIB , F , < 25 } , results in S = 1 and d E = β × 1 . We then calculate the mean-squared error between ˆ D c and D c ∀ c ∈ [1 , . . .
, C ] . Our ﬁnal objective function thus includes the soft contrastive loss (Eq. 3) and the regression loss.
[Equation]
E XPERIMENTAL D ESIGN 4. D ATASETS
We conduct our experiments using PyTorch (Paszke et al., 2019) on two large-scale datasets that consist of physiological time-series, such as the electrocardiogram (ECG), alongside cardiac arrhythmia labels.
Chapman (Zheng et al., 2020) consists of 12-lead ECG recordings from 10,646 patients alongside labels which we group into 4 major classes.
PTB-XL (Wagner et al., 2020) consists of 12-lead ECG recordings from 18,885 patients alongside labels which we group into 5 major classes (Strodthoff et al., 2020).
Further details can be found in Appendix A.
4. D EPLOYING C LINICAL P ROTOTYPES FOR R ETRIEVAL
We aim to exploit clinical prototypes to retrieve physiological instances, that satisfy certain criteria, from a large database.
One could argue that a simple text-based query, instead of CPs, in modern databases will allow for the retrieval of appropriate instances.
However, this only holds if all instances in the database are labelled.
Although we have used, as an exemplar, readily-available attributes such as sex and age, our method extends to other attributes which cannot be trivially searched in a database. In other words, the maximal utility of retrieval via CPs is realized when the troves of instances in the database are unlabelled . To reiterate, the clinical motivation for such retrieval is multi-fold. First, it allows researchers to discover patients that may be suitable for clinical trial recruitment.
This, in turn, may allow pharmaceutical companies to trial their drugs on a more diverse patient cohort that is more representative of the population.
Second, CPs can be deployed across datasets from different clinical institutions, for example, to retrieve patient sub-cohorts that satisfy a set of attributes.
These patient sub-cohorts can now be further analysed to identify sub-cohort variability, which in turn can guide future clinical treatment.
With that in mind, to reliably and quantitatively evaluate the retrieval capabilities of the CPs, we must have access to the ground-truth attributes associated with the instances in the held-out set.
To perform retrieval, we treat the CPs as the query set and retrieve the closest instances ( ↓ Euclidean distance) from a held-out dataset.
In evaluating this task of retrieval, we deploy a commonly used metric known as Precision at K (P@ K ).
This metric quantiﬁes whether at least one of the K retrieved instances matches the query.
In our context, however, a match can occur according to any of the three attributes: class , sex , and age .
Therefore, we extract a single attribute, α m , from the set of attributes, A , for each of the M clinical prototypes, and the same attribute, α i , for each of the K retrieved instances, and deﬁne the following attribute-speciﬁc P@ K .
[Equation]
4. D EPLOYING C LINICAL P ROTOTYPES FOR C LUSTERING
We also exploit clinical prototypes to cluster instances.
In this setting, and in contrast to the retrieval setting, we treat each representation of an instance as a query and search across the CPs to assign each representation a set of attributes.
From this perspective, CPs can be thought of as multiple centroids of a cluster which are used to label representations.
To evaluate these clusters, we ﬁrst calculate the pairwise Euclidean distance between each representation and CP before assigning each representation the speciﬁc attribute, α predi , of the CP to which it is closest.
Given the ground-truth attribute for each representation, α truei , we can calculate the accuracy, Acc ( α ) , of the assignments.
We can also quantify the agreement between the attribute assignments of the ground-truth, (cid:126)α true = { α 1 true , α 2 true , . . .
, α Ntrue } , and those obtained via clustering, (cid:126)α pred =
{ α 1 pred , α 2 pred , . . .
, α Npred } , by calculating the adjusted mutual information AMI ( α ) ∈ [0 , 1] .
[Equation]
where MI ( (cid:126)α true , (cid:126)α pred ) represents the mutual information between the ground-truth and the predicted set of attributes, and H ( (cid:126)α ) represents the entropy of the set of attributes.
4. B ASELINES
In the clustering scenario, we compare our method to the following approaches: 1) K-Means implements Expectation Maximization to arrive at class-speciﬁc centroids.
Instances are then assigned to the class of their nearest centroid.
We perform K-Means a) on the input instances ( K-Means Raw ) and b) on representations of instances learned using L contrastive − soft + L regression ( K-Means
Combined ).
2) DeepCluster (Caron et al., 2018) is an iterative method that performs K-Means on representations, pseudo-labels instances according to their assigned cluster, and then exploits such labels for supervised training.
The ﬁnal set of instance labels is taken from the epoch with the lowest validation loss.
3) Information Invariant Clustering (IIC)
(Ji et al., 2019) maximizes the mutual information between the posterior class probabilities assigned to an instance and its perturbed counterpart.
We adapt IIC to the time-series domain by perturbing instances with additive
Gaussian noise, (cid:15) ∼ N (0 , σ ) .
We also incorporate auxiliary over-clustering as it was shown to signiﬁcantly improve performance.
The ﬁnal set of instance labels is chosen by taking the argmax of the output probabilities.
4) SeLA (Asano et al., 2020) pseudo-labels instances by implementing the
Sinkhorn-Knopp algorithm to allow for supervised training.
We pseudo-label instances after each epoch of training and use prior class information to determine the number of clusters, which should boost performance.
5) DTCluster (Han et al., 2019) calculates the distance between representations and cluster prototypes to generate a probability distribution over classes.
The KL divergence between this distribution and a target distribution is minimized.
4. H YPERPARAMETERS
We chose the temperature parameters, τ = 0 .
1 , as per (Kiyasseh et al., 2020), and τ w = 1 .
The regression regularization term (Eq. 6) requires the speciﬁcation of β .
We chose β = 0 .
2 due to our choice of distance metric (squared Euclidean distance) and the number of attribute groups.
If β were too small in magnitude, the intra-cluster separability of clinical prototypes would be too small (see
Appendix ?? ).
If β were too large, then the clinical prototypes from different classes would begin to overlap with one another and class separability diminishes.
For both datasets, sex ∈ { M , F } , and age is converted to quartiles.
For the Chapman and PTB-XL datasets, | class | = 4 and 5 , respectively, and thus M = | class | × | sex | × | age | = 32 and 40 , respectively.
We use a network with 1D convolutional modules to generate representations.
Further network and implementation details can be found in Appendix B.
E XPERIMENTAL R ESULTS 5. D EPLOYING C LINICAL P ROTOTYPES FOR R ETRIEVAL
We exploit CPs to retrieve attribute-speciﬁc physiological signals from a database.
In this section, we task the CPs with retrieving the closest K =
[1 , 5 , 10] instances in a held-out set and, in Table 1, illustrate the resulting P@ K .
We also stratify the results according to the number of matched attributes between the query (CPs) and retrieved instances.
For example, # of matched attributes =implies that a perfect match has occurred for all attributes: class, sex, and age. [Table]
In Table 1, we show that clinical prototypes perform well in retrieving instances with at least one attribute match.
This can be seen for # of matched attributes ≥ 1 where [-R@
-]P@1 = 76 . 3 and 70 for the Chapman and PTB-XL datasets, respectively.
In the former case, this implies that 76 .
3% of the
CPs retrieved an instance that had at least one match.
We also ﬁnd that as the number of retrieved instances increases ( ↑ K ), so does the likelihood of obtaining one that has at least one attribute match.
For example, as K = 1 −→ 10 , the P@ K = 76 .
3 −→ 100 on the Chapman dataset.
Such a ﬁnding illustrates the potential utility of clinical prototypes as tools for retrieval.
On the other hand, we ﬁnd that, when deployed to retrieve a single instance ( K = 1 ), clinical prototypes struggle to retrieve one whose attribute sets match more than once.
This is shown by
[-R@
-]P@1 = 26 . 9 and 19 . 5 for the Chapman and PTB-XL datasets, respectively when the # of matched attributes ≥ 2 .
However, this struggle is signiﬁcantly alleviated as more instances are retrieved.
For example, for the Chapman dataset, as K = 1 −→ 10 , the P@ K = 26 .
9 −→ 93 .
In other words, when retrieving ten instances, 93.8% of the CPs ﬁnd one whose attribute set matches at least twice.
On a similar note, we ﬁnd that in searching for a larger # of attributes that match, the retrieval performance of the clinical prototypes decreases (P@1 = 76 .
3 −→ 2 . 5 ).
We hypothesize that this behaviour is due to our training paradigm’s preference for the discrimination of classes over that of the remaining attributes.
Therefore, perfect matches are more likely to occur for the class attribute.
Nonetheless, the task of discriminating ECG recordings according to sex, for example, is notoriously difﬁcult (Attia et al., 2019).
To qualitatively evaluate the retrieval capabilities of our clinical prototypes, we choose CPs at random to form the query set, calculate their Euclidean distance to the representations of instances in a held-out dataset, and retrieve the 5 which are closest.
In Fig.
3, we visualize the ECG signals that correspond to these closest representations and colour their borders green if their class label matches that of the query, and red otherwise.
We can see that at least 60% of the retrieved instances match the class label.
[Figure]
5. D EPLOYING C LINICAL P ROTOTYPES FOR C LUSTERING
So far, we have illustrated the utility of exploiting clinical prototypes for retrieval purposes.
In this section, we evaluate the utility of CPs for clustering.
In Table.
2, we illustrate the clustering performance of DROPS relative to that of state-of-the-art clustering methods.
There are several takeaways from Table 2.
First, our supervised contrastive learning paradigm generates rich representations.
This can be seen by the superiority of K - Means Combined over
K - Means Raw where the Acc(class) = 81 .
7 and 29 . 2 , respectively.
Note that this approach does not yet exploit clinical prototypes and thus allows for the reliable evaluation of the representations alone.
We provide further evidence in support of the richness of these representations in Sec.
Second, we ﬁnd that leveraging clinical prototypes, as is done with DROPS , builds upon the rich representations and further improves clustering performance.
This can be seen by the performance of
DROPS Combined which achieves an Acc(class) = 90 . 3 and 76 . 0 on the Chapman and PTB-XL datasets, respectively which outperforms not only its K - Means counterpart but also recent state-ofthe-art methods such as DeepCluster , IIC , and SeLA .
Such a ﬁnding underscores the vital role clinical prototypes play in clustering.
We also note that our method can be trivially incorporated into
DeepCluster , for example, to generate pseudo-labels.
Lastly, in Table 2b, we show that DROPS [Table] is ﬂexible enough to simultaneously cluster according to multiple attributes, a feature that does not trivially extend to other methods.
Such ﬂexibility can provide researchers with improved control over sensitive attributes [Table]
5. C APTURING S EMANTIC R ELATIONSHIPS B ETWEEN C LINICAL P ROTOTYPES
Up until now, we have shown that clinical prototypes can be successfully deployed for retrieval and clustering.
We encouraged such prototypes to exhibit inter and intra-cluster separability, with the aim of capturing semantic relationships between them.
In this section, we look to conﬁrm whether such semantic relationships are indeed captured.
In Fig.
4 (top row), we illustrate the t-SNE projection of representations in the training set alongside the average class-speciﬁc CP.
In Fig.
4 (bottom row), we illustrate the t-SNE projection of the CPs colour-coded and shaded according to sex and age groups, respectively.
For clarity, we follow the same coding scheme presented in Fig.
2 (left).
We show that our training paradigm leads to clinical prototypes that can satisfy the semantic relationships between the class, sex, and age attributes.
This can be seen by the high separability between, and ordering of, the sex and age attributes in Fig.
4 (Bottom Row).
This observation holds regardless of the class attribute or the dataset that is experimented with.
We further note the high degree of similarity between these projections, empirically-derived, and expected (see Fig.
2 left).
We claim that the adoption of this formation by the CPs is driven primarily by our weighting mechanism (Eq. 3) and distance-based regularization term (Eq. 6).
Empirical support for this claim can be found in
Appendix ?? .
D ISCUSSION AND F UTURE W ORK
In this paper, we propose a supervised contrastive learning framework, DROPS, for the retrieval and clustering of physiological signals.
In the process, we learn representations, entitled clinical prototypes, that efﬁciently describe the state of a patient associated with a set of attributes, e.g., disease, sex, and age.
We show that clinical prototypes can capture semantic relationships and allow for the learning of rich representations advantageous for both the retrieval and clustering of physiological signals.
We now elucidate several avenues worth exploring.
Learning disentangled clinical prototypes.
CPs are attribute-speciﬁc representations.
Disentangling these representations into their constituent attributes might facilitate the interpretability of
CPs and their application to fair ML.
Preliminary exploratory work is performed in Appendix C.
Incorporate continuous attributes.
CPs are currently limited to discrete attributes.
Scaling CPs to a continuous set and large number of attributes will allow for more ﬁne-grained retrieval and further increase their utility.