Rewire-then-Probe: A Contrastive Recipe for Probing Biomedical Knowledge of Pre-trained Language Models (2022)

Abstract: Knowledge probing is crucial for understanding the knowledge transfer mechanism behind the pre-trained language models (PLMs). Despite the growing progress of probing knowledge for PLMs in the general domain, specialised areas such as the biomedical domain are vastly under-explored. To facilitate this, we release a well-curated biomedical knowledge probing benchmark, MedLAMA, constructed based on the Unified Medical Language System (UMLS) Metathesaurus. We test a wide spectrum of state-of-the-art PLMs and probing approaches on our benchmark, reaching at most 3% of acc@10. While highlighting various sources of domain-specific challenges that amount to this underwhelming performance, we illustrate that the underlying PLMs have a higher potential for probing tasks. To achieve this, we propose Contrastive-Probe, a novel self-supervised contrastive probing approach, that adjusts the underlying PLMs without using any probing data. While Contrastive-Probe pushes the acc@10 to 28%, the performance gap still remains notable. Our human expert evaluation suggests that the probing performance of our Contrastive-Probe is still under-estimated as UMLS still does not include the full spectrum of factual knowledge. We hope MedLAMA and Contrastive-Probe facilitate further developments of more suited probing techniques for this domain. Our code and dataset are publicly available at https://github.com/cambridgeltl/medlama .

## introduction pre-trained language models (plms; @xcite ) have orchestrated incredible progress on myriads of few-or zero-shot language understanding tasks, by pre-training model parameters in a task-agnostic way and transferring knowledge to specific downstream tasks via finetuning @xcite . 1 the data and code implementation are available at https://github.com/cambridgeltl/medlama . ## medlama to facilitate research of knowledge probing in the biomedical domain, we create the medlama benchmark based on the largest biomedical knowledge graph umls @xcite . umls foot_0 is a comprehensive metathesaurus containing 3.6 million entities and more than 35.2 million knowledge triples over 818 relation types which are integrated from various ontologies, including snomed ct, mesh and the ncbi taxonomy. creating a lama-style @xcite probing benchmark from such a knowledge graph poses its own challenges: (1) umls is a collection of knowledge graphs with more than 150 ontologies constructed by different organisations with very different schemata and emphasis; (2) a significant amount of entity names (from certain vocabularies) are unnatural language (e.g., t(8;21)(q22;q22) denoting an observed karyotypic abnormality) which can hardly be understood by the existing plms, with tokenisation tailored for natural language; (3) some queries (constructed from knowledge triples) can have up to hundreds of answers (i.e., 1-to-n relations), complicating the interpretation of probing performance; and (4) some queries may expose answers in themselves (e.g., answer within queries), making it challenging to interpret relative accuracy scores. selection of relationship types. in order to obtain high-quality knowledge queries, we conducted multiple rounds of manual filtering on the relation level to exclude uninformative relations or relations that are only important in the ontological context but do not contain interesting semantics as a natural language (e.g, taxonomy and measurement relations). we also excluded relations with insufficient triples/entities. then, we manually checked the knowledge triples for each relation to filter out those that contain unnatural language entities and ensure that their queries are semantically meaningful. additionally, in the cases of 1-to-n relations where there are multiple gold answers for the same query, we constrained all the queries to contain at most 10 gold answers. these steps resulted in 19 relations with each containing 1k randomly sampled knowledge queries. ## existing multi-token knowledge probing approaches while the pioneer works in plm knowledge probing mainly focused on the single-token entities, many recent works have started exploring the solutions for the multi-token scenario @xcite @xcite . these knowledge probing approaches can be categorised, based on answer search space and reliance on mlm head, into three categories: mask predict, generation-based, and retrieval-based. hiv + in + ##fect + ##ions 0.45 0.2 0.6 0.35 shock (d)contrastive probe (b) generative plms (a)mask predict dextran 40 may treat [mask] . elvitegravir may prevent [mask][mask][mask][mask] . single-token multi-token epistaxis bert bart/t5 elvitegravir may prevent [mask] <@xmath0. bert elvitegravir may prevent [mask]. bert mlm mlm mlm nasal mass pain epistaxis rhinorrhea … entities autoregressive decoder (c)mask average nasal mass pain epistaxis rhinorrhea … entities epistaxis query embedding entity embeddings … [cls] multi-token decoding nearest neighbor search used approaches to probe knowledge for masked plms (e.g. bert). the mask predict approach uses the mlm head to fill a single mask token for a cloze-style query, and the output token is subjected to the plm vocabulary @xcite . since many real-world entity names are encoded with multiple tokens, the mask predict approach has also been extended to predict multitoken answers using the conditional masked language model @xcite . figure 2 (a) shows the prediction process. specifically, given a query, the probing task is formulated as: 1) filling masks in parallel independently (independent); 2) filling masks from left to right autoregressively (order); 3) filling tokens sorted by the maximum confidence greedily (confidence). after all mask tokens are replaced with the initial predictions, the predictions can be further refined by iteratively modifying one token at a time until convergence or until the maximum number of iterations is reached @xcite . for example, order+order represents that the answers are initially predicted by order and then refined by order. in this paper we examined two of these approaches, i.e. independent and order+order, based on our initial exploration. ## contrastive-probe: cloze-style task as a self-retrieving game to better transform the plm encoders for the cloze-style probing task, we propose contrastive-probe which pre-trains on a small number of sentences sampled from the plm's original pre-training corpora with a contrastive selfsupervising objective, inspired by the mirror-bert @xcite . our contrastive pretraining does not require the mlm head or any additional external knowledge, and can be completed in less than one minute on 2 × 2080ti gpus. self-supervised contrastive rewiring. we randomly sample a small set of sentences (e.g. 10k, see §5.2 for stability analysis of contrastive-probe on several randomly sampled sets), and replace their tail tokens (e.g. the last 50% excluding the full stop) with a [mask] token. then these transformed sentences are taken as the queries of the cloze-style self-retrieving game. in the following we show an example of transforming a sentence into a cloze-style query: sentence: social-distancing largely reduces coronavirus infections. query: social-distancing largely @xcite . ## experiments in this section we conduct extensive experiments to verify whether contrastive-probe is effective for probing biomedical plms. first, we experiment with contrastive-probe and existing probing approaches on medlama benchmark ( §5.1). then, we conduct in-depth analysis of the stability and applicability of contrastive-probe in probing biomedical plms ( §5.2). finally, we report an evaluation of a biomedical expert on the probing predictions and highlight our findings ( §5.3). ## conclusion in this work, we created a carefully curated biomedical probing benchmark, medlama, from the umls knowledge graph. we illustrated that state-of-the-art probing techniques and biomedical pre-trained languages models (plms) struggle to cope with the challenging nature (e.g. multitoken answers) of this specialised domain, reaching only an underwhelming 3% of acc@10. to reduce the gap, we further proposed a novel contrastive recipe which rewires the underlying plms without using any probing-specific data and illustrated that with a lightweight pre-training their accuracies could be pushed to 24%. our experiments also revealed that different layers of transformers encode different types of information, reflected by their individual success at handling certain types of prompts. additionally, using a human expert, we showed that the existing evaluation criteria could overpenalise the models as many valid responses that plms produce are not in the ground truth umls knowledge graph. this further highlights the importance of having a human in the loop to better understand the potentials and limitations of plms in encoding domain specific factual knowledge. our findings indicate that the real lower bound on the amount of factual knowledge encoded by plms is higher than we estimated, since such bound can be continuously improved by optimising both the encoding space (e.g. using our selfsupervised contrastive learning technique) and the input space (e.g. using the prompt optimising techniques @xcite ). we leave further exploration of integrating the two possibilities to future work.