[Paper Review] Improving BERT Model Using Contrastive Learning for Biomedical Relation Extraction

Peng Su, Yifan Peng+1 more|arXiv (Cornell University)|Apr 28, 2021

Topic Modeling33 references4 citations

TL;DR

This paper proposes a contrastive learning framework to enhance BERT's representation for biomedical relation extraction by integrating linguistic knowledge through shortest dependency path (SDP)-aware data augmentation and leveraging external knowledge bases for large-scale pre-training. The method achieves state-of-the-art performance on ChemProt, DDI, and PPI benchmarks and improves model faithfulness by reducing prediction shifts, indicating stronger reliance on rationales.

ABSTRACT

Contrastive learning has been used to learn a high-quality representation of the image in computer vision. However, contrastive learning is not widely utilized in natural language processing due to the lack of a general method of data augmentation for text data. In this work, we explore the method of employing contrastive learning to improve the text representation from the BERT model for relation extraction. The key knob of our framework is a unique contrastive pre-training step tailored for the relation extraction tasks by seamlessly integrating linguistic knowledge into the data augmentation. Furthermore, we investigate how large-scale data constructed from the external knowledge bases can enhance the generality of contrastive pre-training of BERT. The experimental results on three relation extraction benchmark datasets demonstrate that our method can improve the BERT model representation and achieve state-of-the-art performance. In addition, we explore the interpretability of models by showing that BERT with contrastive pre-training relies more on rationales for prediction. Our code and data are publicly available at: https://github.com/udel-biotm-lab/BERT-CLRE.

Motivation & Objective

• To improve BERT's text representation for biomedical relation extraction using contrastive learning.
• To address the challenge of limited text data augmentation in NLP by incorporating linguistic knowledge via shortest dependency paths (SDP).
• To enhance model generalization by constructing large-scale training data from external knowledge bases.
• To improve model interpretability by measuring faithfulness to input rationales using a novel prediction shift metric.
• To demonstrate state-of-the-art performance and robustness on standard biomedical relation extraction benchmarks.

Proposed method

• Introduce a contrastive pre-training step that generates positive pairs by applying synonym replacement to words outside the shortest dependency path (SDP) between entity pairs.
• Preserve all words on the SDP during data augmentation to maintain relation-relevant information and ensure semantic consistency.
• Leverage external knowledge bases (e.g., from distant supervision) to generate additional training instances for contrastive pre-training.
• Train a BERT model with contrastive learning objective to maximize agreement between augmented views while minimizing similarity to negative pairs.
• Use a prediction shift metric to evaluate model faithfulness by measuring changes in predictions after SDP-agnostic input perturbations.
• Apply the method to three biomedical relation extraction tasks: ChemProt, DDI, and PPI, using BioBERT and PubMedBERT as backbones.

Experimental results

Research questions

• RQ1 Can contrastive learning improve BERT's representation for biomedical relation extraction when combined with linguistically informed data augmentation?
• RQ2 How does incorporating external knowledge bases affect the generalization of contrastively pre-trained BERT models in relation extraction?
• RQ3 Does contrastive pre-training enhance the faithfulness of model predictions to rationales, particularly the shortest dependency path?
• RQ4 To what extent does the model’s robustness to input perturbations increase after contrastive pre-training?
• RQ5 Can a novel prediction shift metric effectively measure the faithfulness of model rationales in biomedical relation extraction?

Key findings

• The proposed method achieves state-of-the-art performance on all three benchmark datasets: ChemProt, DDI, and PPI, outperforming standard BERT and PubMedBERT.
• On the ChemProt dataset, BioBERT+CLEK reduced prediction shifts from 246 to 191 (22% decrease), and PubMedBERT+CLEK from 248 to 189 (24% decrease).
• On the DDI dataset, prediction shifts dropped from 111 to 89 (20% decrease) for BioBERT+CLEK and from 90 to 75 (17% decrease) for PubMedBERT+CLEK.
• On the PPI dataset, prediction shifts were reduced from 51 to 33 (35% decrease) for BioBERT+CLEK and from 49 to 34 (31% decrease) for PubMedBERT+CLEK.
• The reduction in prediction shifts indicates that contrastively pre-trained models rely more on the shortest dependency path (SDP) as rationales, enhancing interpretability and faithfulness.
• The model demonstrates increased robustness to small, non-SDP input changes, confirming improved generalization and stability.