Hierarchical Pretraining for Biomedical Term Embeddings

HiPrBERT leverages two main sources of data. The first is the UMLS, a knowledge graph that encodes relations across many different medical vocabularies. These terms have no inherent order to them, and there are many different types of relations between pairs of terms. In addition to the UMLS knowledge graph, we have a collection of various hierarchies that we can leverage. Specifically, PheCODE is a hierarchy containing ICD codes that can be represented as a forest of trees. The root of each tree is a separate concept, and children of a node will represent a more specific concept. LOINC is another hierarchy representing laboratory observations, containing 171,191 nodes from 27 trees, whose depth varies from 2 to 13; Similarly, RxNorm and CPT are also represented as forests focusing on medication and procedure terms, respectively. PheCode contains 1,601 nodes, RxNorm contains 192,683 nodes; and CPT contains 10,360 nodes. In these hierarchies, the structure contains more information than UMLS on the “closeness” between various biomedical terms, which can be used to accomplish a fine embedding better discriminating closed related terms from moderately related terms.

It is worth noting that although the number of terms in each hierarchy is significantly lower than the number of terms in the UMLS, we expect that we can obtain enough high-quality training pairs from the hierarchy to enhance the embeddings in most relevant regions of the embedding space. In practical terms, each hierarchy consists of two mappings: one from parents to children and one from codes to the biomedical term strings.

HiPrBERT takes in an input term $s$ and outputs a corresponding embedding $\textbf{e}_{s}\in R^{d}$. Specifically, the input $s$ is first converted into a series of tokens, which are then encoded by HiPrBERT into a series of $d$ dimensional hidden state vectors

$$[\text{CLS}],\mathbf{t}_{0},\mathbf{t}_{1},...,\mathbf{t}_{n},[\text{SEP}]\xrightarrow{\textsc{HiPrBERT}}\mathbf{h}_{[\text{CLS}]},\mathbf{h}_{0},\mathbf{h}_{1},...,\mathbf{h}_{n},\mathbf{h}_{[\text{SEP}]}.$$

The embedding of $s$ is defined to be the latent vector corresponding to the [CLS] token

$$s\rightarrow\mathbf{e}_{s}=\mathbf{h}_{[\text{CLS}]}\in R^{d}.$$

Similar to SapBERT, our approach learns term representations by maximizing the embedding similarity between term-term pairs that are “close” and minimizing embedding similarities between term-term pairs that are “far”. We define the embedding similarity between terms $s_{i}$ and $s_{j}$ as $S_{ij}=\cos(\mathbf{e}_{i},\mathbf{e}_{j}).$ We also define following distances to quantify the resemblance between terms $s_{i}$ and $s_{j}.$ These particular choices of the numerical value are not important and only their order matters in training embeddings.

1. 1.

   If $s_{i}$ and $s_{j}$ are from the UMLS, $d(s_{i},s_{j})=\begin{cases}0&\text{$s_{i}$ and $s_{j}$ are synonyms;}\\ 3&\text{otherwise.}\end{cases}$

2. 2.

   If $s_{i}$ and $s_{j}$ are from a hierarchy, $d(s_{i},s_{j})=\begin{cases}0&\text{$s_{i}$ and $s_{j}$ are synonyms;}\\ 1&\text{$s_{i}$ and $s_{j}$ have the same parent (a sibling pair);}\\ 2&\text{$s_{i}$ and $s_{j}$ are a parent-child pair;}\\ 3&\text{otherwise.}\end{cases}$

When sampling UMLS term data, we use an online triplet miner to select negative pairs. Specifically, among all triplets of terms $(s_{a},s_{p},s_{n})$, where $(s_{a},s_{p})$ are synonymous and $(s_{a},s_{n})$ are non-synonymous, based on initial embeddings, $(s_{a},s_{p},s_{n})\rightarrow(\mathbf{e}_{a},\mathbf{e}_{p},\mathbf{e}_{n})$, we consider the difference between $\mbox{cos}(\mathbf{e}_{a},\mathbf{e}_{p})$ and $\mbox{cos}(\mathbf{e}_{a},\mathbf{e}_{n}),$ and select the triplets with this difference $>0.25$ to be included in our minibatch for further training. We do the same for UMLS relational data.

For hierarchical data, we leverage the structure of the tree to construct minibatches. For example, we use distance 0 pairs as positive samples, and distance $>0$ pairs as negative samples. We do this with every distance to encourage separation between varying levels of similarity.

Given an anchor term $s_{i}$ and a set of terms $\Omega_{i}$, we can define the sets

$$\Omega_{i}^{(0)}(d_{0})=\left\{j\in\Omega_{i}\mid d(s_{i},s_{j})\leq d_{0}\right\}\leavevmode\nobreak\ \leavevmode\nobreak\ \mbox{and}\leavevmode\nobreak\ \leavevmode\nobreak\ \Omega_{i}^{(1)}(d_{0})=\left\{j\in\Omega_{i}\mid d(s_{i},s_{j})>d_{0}\right\}.$$

In other words, $\Omega_{i}^{(0)}(d_{0})$ contains all terms that are at most distance $d_{0}$ from $s_{i}$, $\Omega_{i}^{(1)}(d_{0})$ contains all terms that are further than $d_{0}$ away. Our goal is to create embeddings such that the similarity between $s_{i}$ and terms in $\Omega_{i}^{(0)}(d_{0})$ is greater than that between $s_{i}$ and terms in $\Omega_{i}^{(1)}(d_{0})$. We use the multi-similarity loss [17]. For UMLS data, we have the standard MS loss function.

$$\sum_{i=1}^{k}\left[\alpha^{-1}\log\left(1+\sum_{j\in\Omega_{i}^{(0)}(0)}e^{-\alpha\left(S_{ij}-\lambda\right)}\right)+\beta^{-1}\log\left(1+\sum_{j\in\Omega_{i}^{(1)}(0)}e^{\beta\left(S_{ij}-\lambda\right)}\right)\right],$$

where $\alpha=2,\beta=2,\lambda=.5$. Note that the terms in $\Omega_{i}^{(1)}(0)$ come from the triplet mining procedure. For hierarchical data, we use a modified loss:

$$\sum_{d_{0}=0}^{2}\sum_{i=1}^{k}\left[\alpha^{-1}\log\left(1+\sum_{j\in\Omega_{i}^{(0)}(d_{0})}e^{-\alpha\left(S_{ij}-\lambda\right)}\right)+\beta^{-1}\log\left(1+\sum_{j\in\Omega_{i}^{(1)}(d_{0})}e^{\beta\left(S_{ij}-\lambda\right)}\right)\right],$$

with the same set of tuning parameters.

Our training process is similar to that of SapBERT, with the main key difference being the loss functions that were used. Using PyTorch [18] and the transformers library [19], our model was initialized from PubMedBERT [20] and trained using AdamW [21] with a learning rate of $2\times 10^{-5}$, a weight decay rate of 0.01, and linear learning rate scheduler. We use a training batch size of 256, and train on the preprocessed UMLS synonym data, UMLS relation data, and hierachical data for one epoch. This equates to about 120 thousand iterations, and takes less than 10 hours on a single GPU machine.

To objectively evaluate our models, we randomly selected evaluation pairs from hierarchies that were not used in model training. For each evaluation pair, we calculated the cosine similarity between the respective embeddings to determine their relatedness. The quality of the embedding was measured using the AUC under the ROC curve for discriminating between distance $i$ pairs and distance $j$ pairs, where $0\leq i<j\leq 3$. In addition, we have also evaluated the embedding performance via Spearman’s correlation and precision-recall curve.

For relatedness tasks, we used pairs of terms in our holdout set for various relations to test the models. There are many different types of relationships, and we report three of clinical importance, as well as the average of the 28 most common relations. We also included performance on the Cadec term normalization task.

We compare HiPrBERT with a set of competitors including SapBERT, CODER, PubMedBERT, BioBERT, BioGPT and DistilBERT, where the SapBERT is retrained without using testing data for generating fair comparisons.

The AUC values for discriminating pairs of different distances are reported in Table 3. HiPrBERT, fine-tuned on hierarchical datasets, outperforms all its competitors in every category, except for 1 vs 3, where it’s performance is very close to CODER. The most noteworthy improvement is in the 0 vs 1 task, where models have to distinguish synonyms from very closely related pairs, such as “Type 1 Diabetes” and “Type 2 Diabetes”. We have also reported the results using Spearman’s rank correlation coefficient in Table 4, and the conclusions are similar.

We also see significant improvements in all relatedness tasks (Table 5). For example, the AUC in the “Causative” category improves from 91.9% to 98.1% in comparison with the second best embedding generated by CODER. Similar improvement has been also observed in detecting “May Cause/Treat” and “Method of” relations. Overall, the average performance of the model in detecting the 28 most common relationships improved from 88.6% to 93.7% in comparison with the next best embedding. This demonstrates a substantial improvement in our ability to capture more nuanced information. It is worth noting that HiPrBERT’s performance in Cadec is on par with other existing models, indicating that our model does not compromise on performance in similarity tasks while achieving improvements in other areas. Lastly, the comparison results based on Spearman’s correlation (Table 4) and precision-recall curve (not reported) are similar.