Knowledge-augmented Pre-trained Language Models for Biomedical Relation Extraction

TL;DR

Biomedical relation extraction hinges on efficiently identifying interactions among genes, diseases, and chemicals in literature. This work rigorously evaluates knowledge-augmented pre-trained language models under a unified framework across five datasets and four relation scenarios, first optimizing three baseline PLMs and then augmenting the best model with textual descriptions, knowledge-base and literature embeddings, and molecular structure encodings. BioLinkBERT-large emerges as the strongest performer, while external context provides only modest improvements overall; however, smaller PLMs show notable gains from augmentation, suggesting that large models partly encode supervisory signals from added information. The findings emphasize the primacy of model choice and hyperparameter tuning for biomedical RE, and they offer guidance on when and how to integrate external knowledge to maximize transferability and efficiency.

Abstract

Automatic relationship extraction (RE) from biomedical literature is critical for managing the vast amount of scientific knowledge produced each year. In recent years, utilizing pre-trained language models (PLMs) has become the prevalent approach in RE. Several studies report improved performance when incorporating additional context information while fine-tuning PLMs for RE. However, variations in the PLMs applied, the databases used for augmentation, hyper-parameter optimization, and evaluation methods complicate direct comparisons between studies and raise questions about the generalizability of these findings. Our study addresses this research gap by evaluating PLMs enhanced with contextual information on five datasets spanning four relation scenarios within a consistent evaluation framework. We evaluate three baseline PLMs and first conduct extensive hyperparameter optimization. After selecting the top-performing model, we enhance it with additional data, including textual entity descriptions, relational information from knowledge graphs, and molecular structure encodings. Our findings illustrate the importance of i) the choice of the underlying language model and ii) a comprehensive hyperparameter optimization for achieving strong extraction performance. Although inclusion of context information yield only minor overall improvements, an ablation study reveals substantial benefits for smaller PLMs when such external data was included during fine-tuning.

Figures (2)

• Figure 1: Overview of the baseline model and all extensions we evaluate in our study. (a) Baseline model: We build one input example per entity pair, i.e., chemical-protein pair, in each sentence and mark the pair under investigation with special tokens. The sentence is embedded using a pretrained language model and the [CLS] token embedding is passed through an output layer performing the relation classification. (b) Model enhanced with additional embedded entity information: First, we lookup the pre-trained KB embeddings for the head and tail entity under investigation and input the concatenation of both to a multi-layer perceptron (MLP). The resulting embedding is concatenated to the contextualized input text embedding as additional input to the classification layer. (c) Model augmented with additional textual information: We conduct KB lookups to retrieve textual descriptions for the head entity, the tail entity or both using entity identifiers. We include the textual descriptions by appending them to the input text separated by the [SEP] token. (d) Model extended with molecular information of chemicals: For a given chemical $c$, we first retrieve its respective SMILES string $SM_c$, prepend the [CLS] token to $SM_c$ and feed it into the compound language model. We use the representation of the [CLS] token in the last hidden layer as encoding for the structure of $c$ and concatenate it with the language model's output to form the classification layer's input.
• Figure 2: Evaluation results of the models in scenarios with a reduced amount of training data. We perform six runs for each training set size and language model using different random seeds and report mean performance of all runs. For three of the five data sets tested, i.e., BC5CDR, CPI, and DDI, the models achieved promising results even with just 25 training documents. With ChemProt and ChemDisGene, the models benefit substantially from having access to more training data.